1 /*
<lambda>null2  * Copyright (C) 2024 The Android Open Source Project
3  *
4  * Licensed under the Apache License, Version 2.0 (the "License");
5  * you may not use this file except in compliance with the License.
6  * You may obtain a copy of the License at
7  *
8  *      http://www.apache.org/licenses/LICENSE-2.0
9  *
10  * Unless required by applicable law or agreed to in writing, software
11  * distributed under the License is distributed on an "AS IS" BASIS,
12  * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
13  * See the License for the specific language governing permissions and
14  * limitations under the License.
15  */
16 
17 package android.tools.traces.io
18 
19 import android.tools.Scenario
20 import android.tools.io.BUFFER_SIZE
21 import android.tools.io.FLICKER_IO_TAG
22 import android.tools.io.ResultArtifactDescriptor
23 import android.tools.io.RunStatus
24 import android.tools.traces.deleteIfExists
25 import android.tools.withTracing
26 import android.util.Log
27 import java.io.BufferedInputStream
28 import java.io.BufferedOutputStream
29 import java.io.File
30 import java.io.FileInputStream
31 import java.io.FileOutputStream
32 import java.util.zip.ZipEntry
33 import java.util.zip.ZipOutputStream
34 
35 /**
36  * Creates artifacts avoiding duplication.
37  *
38  * If an artifact already exists, append a counter at the end of the filename
39  */
40 class ArtifactBuilder {
41     private var runStatus: RunStatus? = null
42     private var scenario: Scenario? = null
43     private var outputDir: File? = null
44     private var files: Map<ResultArtifactDescriptor, File> = emptyMap()
45     private var counter = 0
46 
47     fun withScenario(value: Scenario): ArtifactBuilder = apply { scenario = value }
48 
49     fun withOutputDir(value: File): ArtifactBuilder = apply { outputDir = value }
50 
51     fun withStatus(value: RunStatus): ArtifactBuilder = apply { runStatus = value }
52 
53     fun withFiles(value: Map<ResultArtifactDescriptor, File>): ArtifactBuilder = apply {
54         files = value
55     }
56 
57     fun build(): FileArtifact {
58         return withTracing("ArtifactBuilder#build") {
59             val scenario = scenario ?: error("Missing scenario")
60             require(!scenario.isEmpty) { "Scenario shouldn't be empty" }
61             val artifactFile = createArtifactFile()
62             Log.d(FLICKER_IO_TAG, "Creating artifact archive at $artifactFile")
63 
64             writeToZip(artifactFile, files)
65 
66             FileArtifact(scenario, artifactFile, counter)
67         }
68     }
69 
70     private fun createArtifactFile(): File {
71         val fileName = getArtifactFileName()
72 
73         val outputDir = outputDir ?: error("Missing output dir")
74         // Ensure output directory exists
75         outputDir.mkdirs()
76         return outputDir.resolve(fileName)
77     }
78 
79     private fun getArtifactFileName(): String {
80         val runStatus = runStatus ?: error("Missing run status")
81         val scenario = scenario ?: error("Missing scenario")
82         val outputDir = outputDir ?: error("Missing output dir")
83 
84         var artifactAlreadyExists = existsArchiveFor(outputDir, scenario, counter)
85         while (artifactAlreadyExists && counter < 100) {
86             artifactAlreadyExists = existsArchiveFor(outputDir, scenario, ++counter)
87         }
88 
89         require(!artifactAlreadyExists) {
90             val files =
91                 try {
92                     outputDir.listFiles()?.filterNot { it.isDirectory }?.map { it.absolutePath }
93                 } catch (e: Throwable) {
94                     null
95                 }
96             "An archive for $scenario already exists in ${outputDir.absolutePath}. " +
97                 "Directory contains ${files?.joinToString()?.ifEmpty { "no files" }}"
98         }
99 
100         return runStatus.generateArchiveNameFor(scenario, counter)
101     }
102 
103     private fun existsArchiveFor(outputDir: File, scenario: Scenario, counter: Int): Boolean {
104         return RunStatus.values().any {
105             outputDir.resolve(it.generateArchiveNameFor(scenario, counter)).exists()
106         }
107     }
108 
109     private fun addFile(zipOutputStream: ZipOutputStream, artifact: File, nameInArchive: String) {
110         Log.v(FLICKER_IO_TAG, "Adding $artifact with name $nameInArchive to zip")
111         val fi = FileInputStream(artifact)
112         val inputStream = BufferedInputStream(fi, BUFFER_SIZE)
113         inputStream.use {
114             val entry = ZipEntry(nameInArchive)
115             zipOutputStream.putNextEntry(entry)
116             val data = ByteArray(BUFFER_SIZE)
117             var count: Int = it.read(data, 0, BUFFER_SIZE)
118             while (count != -1) {
119                 zipOutputStream.write(data, 0, count)
120                 count = it.read(data, 0, BUFFER_SIZE)
121             }
122         }
123         zipOutputStream.closeEntry()
124         artifact.deleteIfExists()
125     }
126 
127     private fun writeToZip(file: File, files: Map<ResultArtifactDescriptor, File>) {
128         ZipOutputStream(BufferedOutputStream(FileOutputStream(file), BUFFER_SIZE)).use {
129             zipOutputStream ->
130             val writtenFileNames = HashSet<String>()
131             files.forEach { (descriptor, artifact) ->
132                 if (!writtenFileNames.contains(descriptor.fileNameInArtifact)) {
133                     addFile(
134                         zipOutputStream,
135                         artifact,
136                         nameInArchive = descriptor.fileNameInArtifact
137                     )
138                     writtenFileNames.add(descriptor.fileNameInArtifact)
139                 } else {
140                     Log.d(
141                         FLICKER_IO_TAG,
142                         "Not adding duplicated ${descriptor.fileNameInArtifact} to zip"
143                     )
144                 }
145             }
146         }
147     }
148 }
149