1 /*
2  * Copyright (C) 2023 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 package com.android.tradefed.util.image;
17 
18 import com.android.annotations.VisibleForTesting;
19 import com.android.tradefed.log.LogUtil.CLog;
20 import com.android.tradefed.util.FileUtil;
21 
22 import com.google.common.cache.CacheBuilder;
23 import com.google.common.cache.CacheLoader;
24 import com.google.common.cache.LoadingCache;
25 import com.google.common.cache.RemovalListener;
26 import com.google.common.cache.RemovalNotification;
27 
28 import java.io.File;
29 import java.io.IOException;
30 import java.util.concurrent.TimeUnit;
31 
32 /**
33  * For some of the incremental device update, we need the baseline files to compute diffs. This
34  * utility helps keeping track of them.
35  */
36 public class DeviceImageTracker {
37 
38     private static DeviceImageTracker sDefaultInstance;
39 
40     private final LoadingCache<String, FileCacheTracker> mImageCache;
41     private final File mCacheDir;
42 
43     /** Track information of the device image cached and its metadata */
44     public class FileCacheTracker {
45         public File zippedDeviceImage;
46         public File zippedBootloaderImage;
47         public File zippedBasebandImage;
48         public String buildId;
49         public String branch;
50         public String flavor;
51 
FileCacheTracker( File zippedDeviceImage, File zippedBootloaderImage, File zippedBasebandImage, String buildId, String branch, String flavor)52         FileCacheTracker(
53                 File zippedDeviceImage,
54                 File zippedBootloaderImage,
55                 File zippedBasebandImage,
56                 String buildId,
57                 String branch,
58                 String flavor) {
59             this.zippedDeviceImage = zippedDeviceImage;
60             this.zippedBootloaderImage = zippedBootloaderImage;
61             this.zippedBasebandImage = zippedBasebandImage;
62             this.buildId = buildId;
63             this.branch = branch;
64             this.flavor = flavor;
65         }
66     }
67 
getDefaultCache()68     public static DeviceImageTracker getDefaultCache() {
69         if (sDefaultInstance == null) {
70             sDefaultInstance = new DeviceImageTracker();
71         }
72         return sDefaultInstance;
73     }
74 
75     @VisibleForTesting
DeviceImageTracker()76     protected DeviceImageTracker() {
77         try {
78             mCacheDir = FileUtil.createTempDir("image_file_cache_dir");
79         } catch (IOException e) {
80             throw new RuntimeException(e);
81         }
82         RemovalListener<String, FileCacheTracker> listener =
83                 new RemovalListener<String, FileCacheTracker>() {
84                     @Override
85                     public void onRemoval(RemovalNotification<String, FileCacheTracker> n) {
86                         if (n.wasEvicted()) {
87                             FileUtil.recursiveDelete(n.getValue().zippedDeviceImage);
88                             FileUtil.deleteFile(n.getValue().zippedBootloaderImage);
89                             FileUtil.deleteFile(n.getValue().zippedBasebandImage);
90                         }
91                     }
92                 };
93         mImageCache =
94                 CacheBuilder.newBuilder()
95                         .maximumSize(20)
96                         .expireAfterAccess(1, TimeUnit.DAYS)
97                         .removalListener(listener)
98                         .build(
99                                 new CacheLoader<String, FileCacheTracker>() {
100                                     @Override
101                                     public FileCacheTracker load(String key) throws IOException {
102                                         // We manually seed and manage the cache
103                                         // no need to load.
104                                         return null;
105                                     }
106                                 });
107         Runtime.getRuntime()
108                 .addShutdownHook(
109                         new Thread() {
110                             @Override
111                             public void run() {
112                                 cleanUp();
113                             }
114                         });
115     }
116 
117     /**
118      * Tracks a given device image to the device serial that was flashed with it
119      *
120      * @param serial The device that was flashed with the image.
121      * @param deviceImage The image flashed onto the device.
122      * @param buildId The build id associated with the device image.
123      * @param branch The branch associated with the device image.
124      * @param flavor The build flavor associated with the device image.
125      */
trackUpdatedDeviceImage( String serial, File deviceImage, File bootloader, File baseband, String buildId, String branch, String flavor)126     public void trackUpdatedDeviceImage(
127             String serial,
128             File deviceImage,
129             File bootloader,
130             File baseband,
131             String buildId,
132             String branch,
133             String flavor) {
134         if (bootloader == null) {
135             CLog.d("Skip tracking image, bootloader is null.");
136             return;
137         }
138         if (deviceImage == null) {
139             CLog.d("Skip tracking image, device image is null.");
140             return;
141         }
142         File copyInCacheDeviceImage = new File(mCacheDir, serial + "_device_image");
143         FileUtil.recursiveDelete(copyInCacheDeviceImage);
144         File copyInCacheBootloader = new File(mCacheDir, serial + "_bootloader");
145         FileUtil.deleteFile(copyInCacheBootloader);
146         File copyInCacheBaseband = null;
147         if (baseband != null) { // Baseband is optional on some devices
148             copyInCacheBaseband = new File(mCacheDir, serial + "_baseband");
149             FileUtil.deleteFile(copyInCacheBaseband);
150         }
151         try {
152             if (deviceImage.isDirectory()) {
153                 CLog.d("Tracking device image as directory.");
154                 FileUtil.recursiveHardlink(deviceImage, copyInCacheDeviceImage);
155             } else {
156                 FileUtil.hardlinkFile(deviceImage, copyInCacheDeviceImage);
157             }
158             FileUtil.hardlinkFile(bootloader, copyInCacheBootloader);
159             if (copyInCacheBaseband != null) {
160                 FileUtil.hardlinkFile(baseband, copyInCacheBaseband);
161             }
162             mImageCache.put(
163                     serial,
164                     new FileCacheTracker(
165                             copyInCacheDeviceImage,
166                             copyInCacheBootloader,
167                             copyInCacheBaseband,
168                             buildId,
169                             branch,
170                             flavor));
171         } catch (IOException e) {
172             invalidateTracking(serial);
173             CLog.e(e);
174         }
175     }
176 
invalidateTracking(String serial)177     public void invalidateTracking(String serial) {
178         mImageCache.invalidate(serial);
179     }
180 
181     @VisibleForTesting
cleanUp()182     protected void cleanUp() {
183         mImageCache.invalidateAll();
184         FileUtil.recursiveDelete(mCacheDir);
185     }
186 
187     /** Returns the device image that was tracked for the device. Null if none was tracked. */
getBaselineDeviceImage(String serial)188     public FileCacheTracker getBaselineDeviceImage(String serial) {
189         return mImageCache.getIfPresent(serial);
190     }
191 }
192