Clean up/remove failed or completed download attempts. (#2298)

* Clean up/remove failed or completed download attempts.
Removes old download attempts from the download managers and resets the models in shared preferences.

* fix guarded by issues

* minor updates

Author: annzimmer

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

@@ -26,6 +26,7 @@
 import android.os.Build.VERSION;
 import android.os.Build.VERSION_CODES;
 import android.os.ParcelFileDescriptor;
+import android.util.Log;
 import android.util.LongSparseArray;
 import androidx.annotation.GuardedBy;
 import androidx.annotation.NonNull;
@@ -57,6 +58,9 @@
  */
 public class ModelFileDownloadService {
 
+  private static final String TAG = "ModelFileDownloadSer";
+  private static final int COMPLETION_BUFFER_IN_MS = 60 * 5 * 1000;
+
   private final DownloadManager downloadManager;
   private final Context context;
   private final ModelFileManager fileManager;
@@ -125,10 +129,30 @@ public Task<Void> download(
 
   @VisibleForTesting
   Task<Void> ensureModelDownloaded(CustomModel customModel) {
-    // todo check model not already in progress of being downloaded
-
-    // todo remove any failed download attempts
-
+    // todo add logging for explicitly requested
+    // check model download already in progress
+    CustomModel downloadingModel =
+        sharedPreferencesUtil.getDownloadingCustomModelDetails(customModel.getName());
+    if (downloadingModel != null && downloadingModel.getDownloadId() != 0) {
+      Integer statusCode = getDownloadingModelStatusCode(downloadingModel.getDownloadId());
+      Date now = new Date();
+
+      // check if download has completed or still has time to finish.
+      // Give a buffer above url expiry to continue if in progress.
+      if (statusCode != null
+          && (statusCode == DownloadManager.STATUS_SUCCESSFUL
+              || statusCode == DownloadManager.STATUS_FAILED
+              || (customModel.getDownloadUrlExpiry()
+                  > (now.getTime() - COMPLETION_BUFFER_IN_MS)))) {
+        // download in progress for this hash already - return this result.
+        return registerReceiverForDownloadId(
+            downloadingModel.getDownloadId(), customModel.getName());
+      }
+    }
+    if (downloadingModel != null) {
+      // remove failed download attempts
+      removeOrCancelDownload(downloadingModel.getName(), downloadingModel.getDownloadId());
+    }
     // schedule new download of model file
     Long newDownloadId = scheduleModelDownload(customModel);
     if (newDownloadId == null) {
@@ -138,18 +162,30 @@ Task<Void> ensureModelDownloaded(CustomModel customModel) {
     return registerReceiverForDownloadId(newDownloadId, customModel.getName());
   }
 
+  /** Removes or cancels the downloading model if exists. */
+  synchronized void removeOrCancelDownload(String modelName, Long downloadId) {
+    if (downloadManager != null && downloadId != 0) {
+      downloadManager.remove(downloadId);
+    }
+    sharedPreferencesUtil.setFailedUploadedCustomModelDetails(modelName);
+  }
+
   private synchronized DownloadBroadcastReceiver getReceiverInstance(
       long downloadId, String modelName) {
-    DownloadBroadcastReceiver receiver = receiverMaps.get(downloadId);
+    DownloadBroadcastReceiver receiver = this.receiverMaps.get(downloadId);
     if (receiver == null) {
       receiver =
           new DownloadBroadcastReceiver(
               downloadId, modelName, getTaskCompletionSourceInstance(downloadId));
-      receiverMaps.put(downloadId, receiver);
+      this.receiverMaps.put(downloadId, receiver);
     }
     return receiver;
   }
 
+  private synchronized void removeReceiverInstance(long downloadId) {
+    this.receiverMaps.remove(downloadId);
+  }
+
   private Task<Void> registerReceiverForDownloadId(long downloadId, String modelName) {
     BroadcastReceiver broadcastReceiver = getReceiverInstance(downloadId, modelName);
     // It is okay to always register here. Since the broadcast receiver is the same via the lookup
@@ -162,15 +198,25 @@ private Task<Void> registerReceiverForDownloadId(long downloadId, String modelNa
 
   @VisibleForTesting
   synchronized TaskCompletionSource<Void> getTaskCompletionSourceInstance(long downloadId) {
-    TaskCompletionSource<Void> taskCompletionSource = taskCompletionSourceMaps.get(downloadId);
+    TaskCompletionSource<Void> taskCompletionSource = this.taskCompletionSourceMaps.get(downloadId);
     if (taskCompletionSource == null) {
       taskCompletionSource = new TaskCompletionSource<>();
-      taskCompletionSourceMaps.put(downloadId, taskCompletionSource);
+      this.taskCompletionSourceMaps.put(downloadId, taskCompletionSource);
     }
 
     return taskCompletionSource;
   }
 
+  @VisibleForTesting
+  synchronized boolean existTaskCompletionSourceInstance(long downloadId) {
+    TaskCompletionSource<Void> taskCompletionSource = this.taskCompletionSourceMaps.get(downloadId);
+    return (taskCompletionSource != null);
+  }
+
+  private synchronized void removeTaskCompletionSourceInstance(long downloadId) {
+    this.taskCompletionSourceMaps.remove(downloadId);
+  }
+
   @VisibleForTesting
   synchronized Long scheduleModelDownload(@NonNull CustomModel customModel) {
     if (downloadManager == null) {
@@ -253,8 +299,8 @@ private synchronized ParcelFileDescriptor getDownloadedFile(Long downloadingId)
     try {
       fileDescriptor = downloadManager.openDownloadedFile(downloadingId);
     } catch (FileNotFoundException e) {
-      // todo(annz)
-      System.out.println("Downloaded file is not found");
+      // todo replace with FirebaseMlException
+      Log.d(TAG, "Downloaded file is not found: " + e);
     }
     return fileDescriptor;
   }
@@ -286,12 +332,18 @@ public File loadNewlyDownloadedModelFile(CustomModel model) throws FirebaseMlExc
     String downloadingModelHash = model.getModelHash();
 
     if (downloadingId == 0 || downloadingModelHash.isEmpty()) {
-      // no downloading model file or incomplete info.
+      // Clear the downloading info completely.
+      // It handles the case: developer clear the app cache but downloaded model file in
+      // DownloadManager's cache would not be cleared.
+      removeOrCancelDownload(model.getName(), model.getDownloadId());
       return null;
     }
 
     Integer statusCode = getDownloadingModelStatusCode(downloadingId);
     if (statusCode == null) {
+      Log.d(TAG, "Download failed - no download status available.");
+      // No status code, it may mean no such download or no download manager.
+      removeOrCancelDownload(model.getName(), model.getDownloadId());
       return null;
     }
 
@@ -301,16 +353,21 @@ public File loadNewlyDownloadedModelFile(CustomModel model) throws FirebaseMlExc
       if (fileDescriptor == null) {
         // reset original model - removing download id.
         sharedPreferencesUtil.setFailedUploadedCustomModelDetails(model.getName());
-        // todo call the download register?
+        removeOrCancelDownload(model.getName(), model.getDownloadId());
+        // todo log this?
         return null;
       }
 
       // Try to move it to destination folder.
-      File newModelFile = fileManager.moveModelToDestinationFolder(model, fileDescriptor);
+      File newModelFile;
+      try {
+        // TODO add logging
+        newModelFile = fileManager.moveModelToDestinationFolder(model, fileDescriptor);
+      } finally {
+        removeOrCancelDownload(model.getName(), model.getDownloadId());
+      }
 
       if (newModelFile == null) {
-        // reset original model - removing download id.
-        // todo call the download register?
         sharedPreferencesUtil.setFailedUploadedCustomModelDetails(model.getName());
         return null;
       }
@@ -323,9 +380,9 @@ public File loadNewlyDownloadedModelFile(CustomModel model) throws FirebaseMlExc
       // todo(annzimmer) Cleans up the old files if it is the initial creation.
       return newModelFile;
     } else if (statusCode == DownloadManager.STATUS_FAILED) {
-      // reset original model - removing download id.
+      // reset original model - removing downloading details.
       sharedPreferencesUtil.setFailedUploadedCustomModelDetails(model.getName());
-      // todo - determine if the temp files need to be clean up? Does one exist?
+      removeOrCancelDownload(model.getName(), model.getDownloadId());
     }
     // Other cases, return as null and wait for download finish.
     return null;
@@ -398,6 +455,13 @@ public void onReceive(Context context, Intent intent) {
       }
 
       Integer statusCode = getDownloadingModelStatusCode(downloadId);
+      // check to prevent DuplicateTaskCompletionException - this was already updated and removed.
+      // Just return.
+      if (!existTaskCompletionSourceInstance(downloadId)) {
+        removeReceiverInstance(downloadId);
+        return;
+      }
+
       synchronized (ModelFileDownloadService.this) {
         try {
           context.getApplicationContext().unregisterReceiver(this);
@@ -409,31 +473,45 @@ public void onReceive(Context context, Intent intent) {
           // move on.
         }
 
-        receiverMaps.remove(downloadId);
-        taskCompletionSourceMaps.remove(downloadId);
+        removeReceiverInstance(downloadId);
+        removeTaskCompletionSourceInstance(downloadId);
       }
 
       if (statusCode != null) {
         if (statusCode == DownloadManager.STATUS_FAILED) {
-          eventLogger.logDownloadFailureWithReason(
-              sharedPreferencesUtil.getDownloadingCustomModelDetails(modelName),
-              false,
-              getFailureReason(id));
-          if (checkErrorCausedByExpiry(id, modelName)) {
-            // retry as a new download
-            // todo change to FirebaseMlException retry error.
-            taskCompletionSource.setException(new Exception("Retry: Expired URL"));
-            return;
+          int failureReason = getFailureReason(id);
+          CustomModel downloadingModel =
+              sharedPreferencesUtil.getDownloadingCustomModelDetails(modelName);
+          if (downloadingModel != null) {
+            eventLogger.logDownloadFailureWithReason(downloadingModel, false, failureReason);
+            if (checkErrorCausedByExpiry(downloadingModel.getDownloadUrlExpiry(), failureReason)) {
+              // retry as a new download
+              // todo change to FirebaseMlException retry error - or whatever we decide is
+              // appropriate.
+              taskCompletionSource.setException(new Exception("Retry: Expired URL"));
+              return;
+            }
           }
           taskCompletionSource.setException(getExceptionAccordingToDownloadManager(id));
           return;
         }
 
         if (statusCode == DownloadManager.STATUS_SUCCESSFUL) {
+          CustomModel model = sharedPreferencesUtil.getDownloadingCustomModelDetails(modelName);
+          if (model == null) {
+            // model might have been updated already get the downloaded model.
+            model = sharedPreferencesUtil.getCustomModelDetails(modelName);
+            if (model == null) {
+              // todo add logging here.
+              taskCompletionSource.setException(
+                  new FirebaseMlException(
+                      "No model associated with name: " + modelName,
+                      FirebaseMlException.INVALID_ARGUMENT));
+              return;
+            }
+          }
           eventLogger.logDownloadEventWithExactDownloadTime(
-              sharedPreferencesUtil.getDownloadingCustomModelDetails(modelName),
-              ErrorCode.NO_ERROR,
-              DownloadStatus.SUCCEEDED);
+              model, ErrorCode.NO_ERROR, DownloadStatus.SUCCEEDED);
           taskCompletionSource.setResult(null);
           return;
         }
@@ -443,25 +521,11 @@ public void onReceive(Context context, Intent intent) {
       taskCompletionSource.setException(new Exception("Model downloading failed"));
     }
 
-    private boolean checkErrorCausedByExpiry(Long downloadId, String modelName) {
-      CustomModel model = sharedPreferencesUtil.getCustomModelDetails(modelName);
-
-      if (model == null) {
-        return false;
-      }
-
+    private boolean checkErrorCausedByExpiry(long downloadUrlExpiry, int failureReason) {
       final Date time = new Date();
 
-      if (model.getDownloadUrlExpiry() < time.getTime()) {
-        Cursor cursor =
-            (downloadManager == null || downloadId == null)
-                ? null
-                : downloadManager.query(new Query().setFilterById(downloadId));
-        if (cursor != null && cursor.moveToFirst()) {
-          int reason = cursor.getInt(cursor.getColumnIndex(DownloadManager.COLUMN_REASON));
-          // 400 implies possibility of url expiry
-          return (reason == 400);
-        }
+      if (failureReason == 400 && downloadUrlExpiry < time.getTime()) {
+        return true;
       }
       return false;
     }