Skip to content

Commit 0399724

Browse files
annzimmerannzimmer
authored
Add model file clean-up on initial load for all models found in share… (#2353)
* Add model file clean-up on initial load for all models found in shared preferences. * fix merge error * Address reviewer comments. * formatting * change to tasks and base deletion on file index instead of name.
1 parent 2de6a3e commit 0399724

File tree

4 files changed

+305
-55
lines changed

4 files changed

+305
-55
lines changed

‎firebase-ml-modeldownloader/src/main/java/com/google/firebase/ml/modeldownloader/internal/ModelFileDownloadService.java

Lines changed: 26 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -67,6 +67,8 @@ public class ModelFileDownloadService {
6767
private final SharedPreferencesUtil sharedPreferencesUtil;
6868
private final FirebaseMlLogger eventLogger;
6969

70+
private boolean isInitialLoad;
71+
7072
@GuardedBy("this")
7173
// Mapping from download id to broadcast receiver. Because models can update, we cannot just keep
7274
// one instance of DownloadBroadcastReceiver per RemoteModelDownloadManager object.
@@ -87,6 +89,7 @@ public ModelFileDownloadService(
8789
downloadManager = (DownloadManager) context.getSystemService(Context.DOWNLOAD_SERVICE);
8890
this.fileManager = ModelFileManager.getInstance();
8991
this.sharedPreferencesUtil = new SharedPreferencesUtil(firebaseApp);
92+
this.isInitialLoad = true;
9093
this.eventLogger = FirebaseMlLogger.getInstance();
9194
}
9295

@@ -96,12 +99,14 @@ public ModelFileDownloadService(
9699
DownloadManager downloadManager,
97100
ModelFileManager fileManager,
98101
SharedPreferencesUtil sharedPreferencesUtil,
99-
FirebaseMlLogger eventLogger) {
102+
FirebaseMlLogger eventLogger,
103+
boolean isInitialLoad) {
100104
this.context = firebaseApp.getApplicationContext();
101105
this.downloadManager = downloadManager;
102106
this.fileManager = fileManager;
103107
this.sharedPreferencesUtil = sharedPreferencesUtil;
104108
this.eventLogger = eventLogger;
109+
this.isInitialLoad = isInitialLoad;
105110
}
106111

107112
/**
@@ -404,7 +409,8 @@ public File loadNewlyDownloadedModelFile(CustomModel model) {
404409
new CustomModel(
405410
model.getName(), model.getModelHash(), model.getSize(), 0, newModelFile.getPath()));
406411

407-
// todo(annzimmer) Cleans up the old files if it is the initial creation.
412+
maybeCleanUpOldModels();
413+
408414
return newModelFile;
409415
} else if (statusCode == DownloadManager.STATUS_FAILED) {
410416
Log.d(TAG, "Model downloaded failed.");
@@ -418,6 +424,24 @@ public File loadNewlyDownloadedModelFile(CustomModel model) {
418424
return null;
419425
}
420426

427+
private Task<Void> maybeCleanUpOldModels() {
428+
if (!isInitialLoad) {
429+
return Tasks.forResult(null);
430+
}
431+
432+
// only do once per initialization.
433+
isInitialLoad = false;
434+
435+
// for each custom model directory, find out the latest model and delete the other files.
436+
// If no corresponding model, clean up the full directory.
437+
try {
438+
fileManager.deleteNonLatestCustomModels();
439+
} catch (FirebaseMlException fex) {
440+
Log.d(TAG, "Failed to clean up old models.");
441+
}
442+
return Tasks.forResult(null);
443+
}
444+
421445
private FirebaseMlException getExceptionAccordingToDownloadManager(Long downloadId) {
422446
int errorCode = FirebaseMlException.INTERNAL;
423447
String errorMessage = "Model downloading failed";

‎firebase-ml-modeldownloader/src/main/java/com/google/firebase/ml/modeldownloader/internal/ModelFileManager.java

Lines changed: 54 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -23,7 +23,6 @@
2323
import androidx.annotation.Nullable;
2424
import androidx.annotation.VisibleForTesting;
2525
import androidx.annotation.WorkerThread;
26-
import com.google.android.gms.common.internal.Preconditions;
2726
import com.google.firebase.FirebaseApp;
2827
import com.google.firebase.ml.modeldownloader.CustomModel;
2928
import com.google.firebase.ml.modeldownloader.FirebaseMlException;
@@ -44,10 +43,12 @@ public class ModelFileManager {
4443
private static final int INVALID_INDEX = -1;
4544
private final Context context;
4645
private final FirebaseApp firebaseApp;
46+
private final SharedPreferencesUtil sharedPreferencesUtil;
4747

4848
public ModelFileManager(@NonNull FirebaseApp firebaseApp) {
4949
this.context = firebaseApp.getApplicationContext();
5050
this.firebaseApp = firebaseApp;
51+
this.sharedPreferencesUtil = new SharedPreferencesUtil(firebaseApp);
5152
}
5253

5354
/**
@@ -61,6 +62,23 @@ public static ModelFileManager getInstance() {
6162
return FirebaseApp.getInstance().get(ModelFileManager.class);
6263
}
6364

65+
void deleteNonLatestCustomModels() throws FirebaseMlException {
66+
File root = getDirImpl("");
67+
68+
boolean ret = true;
69+
if (root.isDirectory()) {
70+
for (File f : root.listFiles()) {
71+
// for each custom model sub directory - extract customModelName and clean up old models.
72+
String modelName = f.getName();
73+
74+
CustomModel model = sharedPreferencesUtil.getCustomModelDetails(modelName);
75+
if (model != null) {
76+
deleteOldModels(modelName, model.getLocalFilePath());
77+
}
78+
}
79+
}
80+
}
81+
6482
/**
6583
* Get the directory where the model is supposed to reside. This method does not ensure that the
6684
* directory specified does exist. If you need to ensure its existence, you should call
@@ -178,6 +196,40 @@ public synchronized File moveModelToDestinationFolder(
178196
return modelFileDestination;
179197
}
180198

199+
/**
200+
* Deletes old models in the custom model directory, except the {@code latestModelFilePath}. This
201+
* should only be called when no files are in use or more specifically when the first
202+
* initialization, otherwise it may remove a model that is in use.
203+
*
204+
* @param latestModelFilePath The file path to the latest custom model.
205+
*/
206+
@WorkerThread
207+
public synchronized void deleteOldModels(
208+
@NonNull String modelName, @NonNull String latestModelFilePath) {
209+
File modelFolder = getModelDirUnsafe(modelName);
210+
if (!modelFolder.exists()) {
211+
return;
212+
}
213+
214+
File latestFile = new File(latestModelFilePath);
215+
int latestIndex = Integer.parseInt(latestFile.getName());
216+
File[] modelFiles = modelFolder.listFiles();
217+
218+
boolean isAllDeleted = true;
219+
int fileInt;
220+
for (File modelFile : modelFiles) {
221+
try {
222+
fileInt = Integer.parseInt(modelFile.getName());
223+
} catch (NumberFormatException ex) {
224+
// unexpected file - ignore
225+
fileInt = Integer.MAX_VALUE;
226+
}
227+
if (fileInt < latestIndex) {
228+
isAllDeleted = isAllDeleted && modelFile.delete();
229+
}
230+
}
231+
}
232+
181233
/**
182234
* Deletes all previously cached Model File(s) and the model root folder.
183235
*
@@ -203,7 +255,7 @@ boolean deleteRecursively(@Nullable File root) {
203255

204256
boolean ret = true;
205257
if (root.isDirectory()) {
206-
for (File f : Preconditions.checkNotNull(root.listFiles())) {
258+
for (File f : root.listFiles()) {
207259
ret = ret && deleteRecursively(f);
208260
}
209261
}

‎firebase-ml-modeldownloader/src/test/java/com/google/firebase/ml/modeldownloader/internal/ModelFileDownloadServiceTest.java

Lines changed: 60 additions & 7 deletions
Original file line numberDiff line numberDiff line change
@@ -95,6 +95,8 @@ public class ModelFileDownloadServiceTest {
9595
File testAppModelFile;
9696

9797
private ModelFileDownloadService modelFileDownloadService;
98+
99+
private ModelFileDownloadService modelFileDownloadServiceInitialLoad;
98100
private SharedPreferencesUtil sharedPreferencesUtil;
99101
@Mock DownloadManager mockDownloadManager;
100102
@Mock ModelFileManager mockFileManager;
@@ -116,7 +118,21 @@ public void setUp() throws IOException {
116118

117119
modelFileDownloadService =
118120
new ModelFileDownloadService(
119-
app, mockDownloadManager, mockFileManager, sharedPreferencesUtil, mockStatsLogger);
121+
app,
122+
mockDownloadManager,
123+
mockFileManager,
124+
sharedPreferencesUtil,
125+
mockStatsLogger,
126+
false);
127+
128+
modelFileDownloadServiceInitialLoad =
129+
new ModelFileDownloadService(
130+
app,
131+
mockDownloadManager,
132+
mockFileManager,
133+
sharedPreferencesUtil,
134+
mockStatsLogger,
135+
true);
120136

121137
matrixCursor = new MatrixCursor(new String[] {DownloadManager.COLUMN_STATUS});
122138
testTempModelFile = File.createTempFile("fakeTempFile", ".tflite");
@@ -519,7 +535,6 @@ public void ensureModelDownloaded_alreadyInProgess_UrlExpired() throws Exception
519535
when(mockDownloadManager.enqueue(any())).thenReturn(DOWNLOAD_ID).thenReturn(downloadId2);
520536

521537
// first download will get cancelled and cleaned - up before intent is sent.
522-
523538
// Complete the second download
524539
Intent downloadCompleteIntent = new Intent(DownloadManager.ACTION_DOWNLOAD_COMPLETE);
525540
downloadCompleteIntent.putExtra(DownloadManager.EXTRA_DOWNLOAD_ID, downloadId2);
@@ -761,6 +776,7 @@ public void loadNewlyDownloadedModelFile_successFilePresent()
761776
CustomModel retrievedModel = sharedPreferencesUtil.getCustomModelDetails(MODEL_NAME);
762777
assertEquals(retrievedModel, customModelDownloadComplete);
763778
verify(mockDownloadManager, times(1)).remove(anyLong());
779+
verify(mockFileManager, never()).deleteNonLatestCustomModels();
764780
verify(mockStatsLogger, times(1))
765781
.logDownloadEventWithErrorCode(
766782
eq(CUSTOM_MODEL_DOWNLOADING),
@@ -770,7 +786,8 @@ public void loadNewlyDownloadedModelFile_successFilePresent()
770786
}
771787

772788
@Test
773-
public void loadNewlyDownloadedModelFile_successNoFile() throws FileNotFoundException {
789+
public void loadNewlyDownloadedModelFile_successNoFile()
790+
throws FileNotFoundException, FirebaseMlException {
774791
// Not found
775792
assertNull(modelFileDownloadService.getDownloadingModelStatusCode(0L));
776793
matrixCursor.addRow(new Integer[] {DownloadManager.STATUS_SUCCESSFUL});
@@ -783,6 +800,7 @@ public void loadNewlyDownloadedModelFile_successNoFile() throws FileNotFoundExce
783800
assertNull(modelFileDownloadService.loadNewlyDownloadedModelFile(CUSTOM_MODEL_DOWNLOADING));
784801
assertNull(sharedPreferencesUtil.getDownloadingCustomModelDetails(MODEL_NAME));
785802
verify(mockDownloadManager, times(1)).remove(anyLong());
803+
verify(mockFileManager, never()).deleteNonLatestCustomModels();
786804
verify(mockStatsLogger, times(1))
787805
.logDownloadEventWithErrorCode(
788806
eq(CUSTOM_MODEL_DOWNLOADING),
@@ -792,20 +810,21 @@ public void loadNewlyDownloadedModelFile_successNoFile() throws FileNotFoundExce
792810
}
793811

794812
@Test
795-
public void loadNewlyDownloadedModelFile_Running() {
813+
public void loadNewlyDownloadedModelFile_Running() throws FirebaseMlException {
796814
// Not found
797815
assertNull(modelFileDownloadService.getDownloadingModelStatusCode(0L));
798816
matrixCursor.addRow(new Integer[] {DownloadManager.STATUS_RUNNING});
799817
when(mockDownloadManager.query(any())).thenReturn(matrixCursor);
800818
assertNull(modelFileDownloadService.loadNewlyDownloadedModelFile(CUSTOM_MODEL_DOWNLOADING));
801819
assertNull(sharedPreferencesUtil.getCustomModelDetails(MODEL_NAME));
802820
verify(mockDownloadManager, never()).remove(anyLong());
821+
verify(mockFileManager, never()).deleteNonLatestCustomModels();
803822
verify(mockStatsLogger, never())
804823
.logDownloadEventWithErrorCode(any(), anyBoolean(), any(), any());
805824
}
806825

807826
@Test
808-
public void loadNewlyDownloadedModelFile_Failed() {
827+
public void loadNewlyDownloadedModelFile_Failed() throws FirebaseMlException {
809828
// Not found
810829
assertNull(modelFileDownloadService.getDownloadingModelStatusCode(0L));
811830
matrixCursor =
@@ -820,10 +839,44 @@ public void loadNewlyDownloadedModelFile_Failed() {
820839
assertNull(modelFileDownloadService.loadNewlyDownloadedModelFile(CUSTOM_MODEL_DOWNLOADING));
821840
assertNull(sharedPreferencesUtil.getCustomModelDetails(MODEL_NAME));
822841
verify(mockDownloadManager, times(1)).remove(anyLong());
823-
verify(mockStatsLogger, never())
824-
.logDownloadEventWithErrorCode(any(), anyBoolean(), any(), any());
825842
verify(mockStatsLogger, times(1))
826843
.logDownloadFailureWithReason(
827844
eq(CUSTOM_MODEL_DOWNLOADING), eq(false), eq(DownloadManager.ERROR_INSUFFICIENT_SPACE));
845+
verify(mockFileManager, never()).deleteNonLatestCustomModels();
846+
}
847+
848+
@Test
849+
public void loadNewlyDownloadedModelFile_initialLoad_successFilePresent()
850+
throws FirebaseMlException, FileNotFoundException {
851+
// Not found
852+
assertNull(modelFileDownloadServiceInitialLoad.getDownloadingModelStatusCode(0L));
853+
matrixCursor.addRow(new Integer[] {DownloadManager.STATUS_SUCCESSFUL});
854+
when(mockDownloadManager.query(any())).thenReturn(matrixCursor);
855+
when(mockDownloadManager.openDownloadedFile(anyLong()))
856+
.thenReturn(
857+
ParcelFileDescriptor.open(testTempModelFile, ParcelFileDescriptor.MODE_READ_ONLY));
858+
when(mockDownloadManager.remove(anyLong())).thenReturn(1);
859+
860+
when(mockFileManager.moveModelToDestinationFolder(any(), any())).thenReturn(testAppModelFile);
861+
862+
assertEquals(
863+
modelFileDownloadServiceInitialLoad.loadNewlyDownloadedModelFile(CUSTOM_MODEL_DOWNLOADING),
864+
testAppModelFile);
865+
866+
// second attempt should not call deleteNonLatestCustomModels a second time.
867+
assertEquals(
868+
modelFileDownloadServiceInitialLoad.loadNewlyDownloadedModelFile(CUSTOM_MODEL_DOWNLOADING),
869+
testAppModelFile);
870+
871+
CustomModel retrievedModel = sharedPreferencesUtil.getCustomModelDetails(MODEL_NAME);
872+
assertEquals(retrievedModel, customModelDownloadComplete);
873+
verify(mockDownloadManager, times(2)).remove(anyLong());
874+
verify(mockFileManager, times(1)).deleteNonLatestCustomModels();
875+
verify(mockStatsLogger, times(2))
876+
.logDownloadEventWithErrorCode(
877+
eq(CUSTOM_MODEL_DOWNLOADING),
878+
eq(true),
879+
eq(DownloadStatus.SUCCEEDED),
880+
eq(ErrorCode.NO_ERROR));
828881
}
829882
}

0 commit comments

Comments
 (0)