Handle expired url file paths with retries. (#2276)

annzimmer · schmidt-sebastian · commit c858d4ce599c · 2021-01-08T09:57:31.000-07:00
diff --git a/firebase-ml-modeldownloader/src/main/java/com/google/firebase/ml/modeldownloader/FirebaseModelDownloader.java b/firebase-ml-modeldownloader/src/main/java/com/google/firebase/ml/modeldownloader/FirebaseModelDownloader.java
@@ -172,14 +172,27 @@ private Task<CustomModel> getCustomModelTask(
 
             // if modelHash matches current local model just return local model.
             // Should be handled by above case but just in case.
-            if (currentModel != null
-                && currentModel
-                    .getModelHash()
-                    .equals(incomingModelDetails.getResult().getModelHash())) {
-              if (!currentModel.getLocalFilePath().isEmpty()
+            if (currentModel != null) {
+              // is this the same model?
+              if (currentModel
+                      .getModelHash()
+                      .equals(incomingModelDetails.getResult().getModelHash())
+                  && currentModel.getLocalFilePath() != null
+                  && !currentModel.getLocalFilePath().isEmpty()
                   && new File(currentModel.getLocalFilePath()).exists()) {
                 return Tasks.forResult(currentModel);
               }
+
+              // is download already in progress for this hash?
+              if (currentModel.getDownloadId() != 0) {
+                CustomModel downloadingModel =
+                    sharedPreferencesUtil.getDownloadingCustomModelDetails(modelName);
+                if (downloadingModel != null
+                    && downloadingModel
+                        .getModelHash()
+                        .equals(incomingModelDetails.getResult().getModelHash()))
+                  return Tasks.forResult(downloadingModel);
+              }
               // todo(annzimmer) this shouldn't happen unless they are calling the sdk with multiple
               // sets of download types/conditions.
               //  this should be a download in progress - add appropriate handling.
@@ -193,19 +206,77 @@ && new File(currentModel.getLocalFilePath()).exists()) {
                     downloadTask -> {
                       if (downloadTask.isSuccessful()) {
                         // read the updated model
-                        CustomModel downloadedModel =
-                            sharedPreferencesUtil.getCustomModelDetails(modelName);
-                        // trigger the file to be moved to permanent location.
-                        fileDownloadService.loadNewlyDownloadedModelFile(downloadedModel);
-                        return Tasks.forResult(downloadedModel);
+                        CustomModel updatedModel =
+                            sharedPreferencesUtil.getDownloadingCustomModelDetails(modelName);
+                        if (updatedModel == null) {
+                          // either download failed or it completed really fast.
+                          return Tasks.forResult(
+                              sharedPreferencesUtil.getCustomModelDetails(modelName));
+                        }
+                        // trigger the file to be moved to permanent location
+                        // This handles immediate download and completion.
+                        fileDownloadService.loadNewlyDownloadedModelFile(updatedModel);
+                        updatedModel =
+                            sharedPreferencesUtil.getDownloadingCustomModelDetails(modelName);
+                        // download complete - get current model.
+                        if (updatedModel == null) {
+                          updatedModel = sharedPreferencesUtil.getCustomModelDetails(modelName);
+                        }
+                        return Tasks.forResult(updatedModel);
+                      } else {
+                        return retryExpiredUrlDownload(modelName, conditions, downloadTask, 2);
                       }
-                      return Tasks.forException(new Exception("File download failed."));
                     });
           }
           return Tasks.forException(incomingModelDetailTask.getException());
         });
   }
 
+  private Task<CustomModel> retryExpiredUrlDownload(
+      @NonNull String modelName,
+      @Nullable CustomModelDownloadConditions conditions,
+      Task<Void> downloadTask,
+      int retryCounter)
+      throws Exception {
+    if (downloadTask.getException().getMessage().contains("Retry: Expired URL")) {
+      // this is likely an expired url - retry once.
+      Task<CustomModel> retryModelDetails =
+          modelDownloadService.getCustomModelDetails(
+              firebaseOptions.getProjectId(), modelName, null);
+      // no local model - start download.
+      return retryModelDetails.continueWithTask(
+          executor,
+          retryModelDetailTask -> {
+            if (retryModelDetailTask.isSuccessful()) {
+              // start download
+              return fileDownloadService
+                  .download(retryModelDetailTask.getResult(), conditions)
+                  .continueWithTask(
+                      executor,
+                      retryDownloadTask -> {
+                        if (retryDownloadTask.isSuccessful()) {
+                          // read the updated model
+                          CustomModel downloadedModel =
+                              sharedPreferencesUtil.getCustomModelDetails(modelName);
+                          // TODO(annz) trigger file move here as well... right
+                          // now it's temp
+                          // call loadNewlyDownloadedModelFile
+                          return Tasks.forResult(downloadedModel);
+                        }
+                        if (retryCounter > 1) {
+                          return retryExpiredUrlDownload(
+                              modelName, conditions, downloadTask, retryCounter - 1);
+                        }
+                        return Tasks.forException(
+                            new Exception("File download failed. Too many attempts."));
+                      });
+            }
+            return Tasks.forException(retryModelDetailTask.getException());
+          });
+    }
+    return Tasks.forException(new Exception("File download failed."));
+  }
+
   /**
    * Triggers the move to permanent storage of successful model downloads and lists all models
    * downloaded to device.
diff --git a/firebase-ml-modeldownloader/src/main/java/com/google/firebase/ml/modeldownloader/internal/ModelFileDownloadService.java b/firebase-ml-modeldownloader/src/main/java/com/google/firebase/ml/modeldownloader/internal/ModelFileDownloadService.java
@@ -40,6 +40,7 @@
 import com.google.firebase.ml.modeldownloader.CustomModelDownloadConditions;
 import java.io.File;
 import java.io.FileNotFoundException;
+import java.util.Date;
 import java.util.regex.Matcher;
 import java.util.regex.Pattern;
 
@@ -120,21 +121,23 @@ Task<Void> ensureModelDownloaded(CustomModel customModel) {
       return Tasks.forException(new Exception("Failed to schedule the download task"));
     }
 
-    return registerReceiverForDownloadId(newDownloadId);
+    return registerReceiverForDownloadId(newDownloadId, customModel.getName());
   }
 
-  private synchronized DownloadBroadcastReceiver getReceiverInstance(long downloadId) {
+  private synchronized DownloadBroadcastReceiver getReceiverInstance(
+      long downloadId, String modelName) {
     DownloadBroadcastReceiver receiver = receiverMaps.get(downloadId);
     if (receiver == null) {
       receiver =
-          new DownloadBroadcastReceiver(downloadId, getTaskCompletionSourceInstance(downloadId));
+          new DownloadBroadcastReceiver(
+              downloadId, modelName, getTaskCompletionSourceInstance(downloadId));
       receiverMaps.put(downloadId, receiver);
     }
     return receiver;
   }
 
-  private Task<Void> registerReceiverForDownloadId(long downloadId) {
-    BroadcastReceiver broadcastReceiver = getReceiverInstance(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
     // for the same download id, the same broadcast receiver will be notified only once.
     context.registerReceiver(
@@ -249,10 +252,13 @@ public void maybeCheckDownloadingComplete() throws Exception {
       if (matcher.find()) {
         String modelName = matcher.group(matcher.groupCount());
         CustomModel downloadingModel = sharedPreferencesUtil.getCustomModelDetails(modelName);
-        Integer statusCode = getDownloadingModelStatusCode(downloadingModel.getDownloadId());
-        if (statusCode == DownloadManager.STATUS_SUCCESSFUL
-            || statusCode == DownloadManager.STATUS_FAILED) {
-          loadNewlyDownloadedModelFile(downloadingModel);
+        if (downloadingModel != null) {
+          Integer statusCode = getDownloadingModelStatusCode(downloadingModel.getDownloadId());
+          if (statusCode != null
+              && (statusCode == DownloadManager.STATUS_SUCCESSFUL
+                  || statusCode == DownloadManager.STATUS_FAILED)) {
+            loadNewlyDownloadedModelFile(downloadingModel);
+          }
         }
       }
     }
@@ -264,7 +270,7 @@ public File loadNewlyDownloadedModelFile(CustomModel model) throws Exception {
     Long downloadingId = model.getDownloadId();
     String downloadingModelHash = model.getModelHash();
 
-    if (downloadingId == null || downloadingModelHash == null) {
+    if (downloadingId == 0 || downloadingModelHash.isEmpty()) {
       // no downloading model file or incomplete info.
       return null;
     }
@@ -318,11 +324,13 @@ private class DownloadBroadcastReceiver extends BroadcastReceiver {
     // Download Id is captured inside this class in memory. So there is no concern of inconsistency
     // with the persisted download id in shared preferences.
     private final long downloadId;
+    private final String modelName;
     private final TaskCompletionSource<Void> taskCompletionSource;
 
     private DownloadBroadcastReceiver(
-        long downloadId, TaskCompletionSource<Void> taskCompletionSource) {
+        long downloadId, String modelName, TaskCompletionSource<Void> taskCompletionSource) {
       this.downloadId = downloadId;
+      this.modelName = modelName;
       this.taskCompletionSource = taskCompletionSource;
     }
 
@@ -351,6 +359,12 @@ public void onReceive(Context context, Intent intent) {
 
       if (statusCode != null) {
         if (statusCode == DownloadManager.STATUS_FAILED) {
+          if (checkErrorCausedByExpiry(id, modelName)) {
+            // retry as a new download
+            // todo change to FirebaseMlException retry error.
+            taskCompletionSource.setException(new Exception("Retry: Expired URL"));
+            return;
+          }
           // todo add failure reason and logging
           System.out.println("Download Failed for id: " + id);
           taskCompletionSource.setException(new Exception("Failed"));
@@ -367,5 +381,28 @@ public void onReceive(Context context, Intent intent) {
       // Status code is null or not one of success or fail.
       taskCompletionSource.setException(new Exception("Model downloading failed"));
     }
+
+    private boolean checkErrorCausedByExpiry(Long downloadId, String modelName) {
+      CustomModel model = sharedPreferencesUtil.getCustomModelDetails(modelName);
+
+      if (model == null) {
+        return false;
+      }
+
+      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);
+        }
+      }
+      return false;
+    }
   }
 }
diff --git a/firebase-ml-modeldownloader/src/test/java/com/google/firebase/ml/modeldownloader/FirebaseModelDownloaderTest.java b/firebase-ml-modeldownloader/src/test/java/com/google/firebase/ml/modeldownloader/FirebaseModelDownloaderTest.java
@@ -61,6 +61,7 @@ public class FirebaseModelDownloaderTest {
   private static final String MODEL_NAME = "MODEL_NAME_1";
   private static final String MODEL_URL = "https://project.firebase.com/modelName/23424.jpg";
   private static final long URL_EXPIRATION = 604800L;
+  private static final long DOWNLOAD_ID = 99;
 
   private static final CustomModelDownloadConditions DEFAULT_DOWNLOAD_CONDITIONS =
       new CustomModelDownloadConditions.Builder().build();
@@ -73,6 +74,8 @@ public class FirebaseModelDownloaderTest {
   private final CustomModel CUSTOM_MODEL = new CustomModel(MODEL_NAME, MODEL_HASH, 100, 0);
   private final CustomModel UPDATE_CUSTOM_MODEL_URL =
       new CustomModel(MODEL_NAME, UPDATE_MODEL_HASH, 100, MODEL_URL, URL_EXPIRATION + 10L);
+  private final CustomModel UPDATE_CUSTOM_MODEL_DOWNLOADING =
+      new CustomModel(MODEL_NAME, UPDATE_MODEL_HASH, 100, DOWNLOAD_ID);
   private CustomModel customModelUploaded;
   private CustomModel customModelLoaded;
 
@@ -262,6 +265,31 @@ public void getModel_latestModel_localExists_UpdateFound() throws Exception {
     assertEquals(customModel, customModelUploaded);
   }
 
+  @Test
+  public void getModel_latestModel_localExists_DownloadInProgress() throws Exception {
+    CustomModel customModelLoadedWithDownload =
+        new CustomModel(MODEL_NAME, MODEL_HASH, 100, 99, expectedDestinationFolder + "/0");
+
+    when(mockPrefs.getCustomModelDetails(eq(MODEL_NAME))).thenReturn(customModelLoadedWithDownload);
+    when(mockPrefs.getDownloadingCustomModelDetails(eq(MODEL_NAME)))
+        .thenReturn(UPDATE_CUSTOM_MODEL_DOWNLOADING);
+
+    when(mockModelDownloadService.getCustomModelDetails(
+            eq(TEST_PROJECT_ID), eq(MODEL_NAME), eq(MODEL_HASH)))
+        .thenReturn(Tasks.forResult(UPDATE_CUSTOM_MODEL_URL));
+
+    TestOnCompleteListener<CustomModel> onCompleteListener = new TestOnCompleteListener<>();
+    Task<CustomModel> task =
+        firebaseModelDownloader.getModel(
+            MODEL_NAME, DownloadType.LATEST_MODEL, DEFAULT_DOWNLOAD_CONDITIONS);
+    task.addOnCompleteListener(executor, onCompleteListener);
+    CustomModel customModel = onCompleteListener.await();
+
+    verify(mockPrefs, times(2)).getCustomModelDetails(eq(MODEL_NAME));
+    assertThat(task.isComplete()).isTrue();
+    assertEquals(customModel, UPDATE_CUSTOM_MODEL_DOWNLOADING);
+  }
+
   @Test
   public void getModel_latestModel_noLocalModel() throws Exception {
     when(mockPrefs.getCustomModelDetails(eq(MODEL_NAME)))
@@ -459,6 +487,58 @@ public void getModel_local_noLocalModel() throws Exception {
     assertEquals(customModel, CUSTOM_MODEL);
   }
 
+  @Test
+  public void getModel_local_noLocalModel_urlRetry() throws Exception {
+    when(mockPrefs.getCustomModelDetails(eq(MODEL_NAME))).thenReturn(null).thenReturn(CUSTOM_MODEL);
+    when(mockModelDownloadService.getCustomModelDetails(
+            eq(TEST_PROJECT_ID), eq(MODEL_NAME), eq(null)))
+        .thenReturn(Tasks.forResult(CUSTOM_MODEL));
+    when(mockFileDownloadService.download(any(), eq(DOWNLOAD_CONDITIONS)))
+        .thenReturn(Tasks.forException(new Exception("Retry: Expired URL")))
+        .thenReturn(Tasks.forResult(null));
+    when(mockFileDownloadService.loadNewlyDownloadedModelFile(eq(customModelUploaded)))
+        .thenReturn(firstDeviceModelFile);
+    TestOnCompleteListener<CustomModel> onCompleteListener = new TestOnCompleteListener<>();
+    Task<CustomModel> task =
+        firebaseModelDownloader.getModel(MODEL_NAME, DownloadType.LOCAL_MODEL, DOWNLOAD_CONDITIONS);
+    task.addOnCompleteListener(executor, onCompleteListener);
+    CustomModel customModel = onCompleteListener.await();
+
+    verify(mockPrefs, times(3)).getCustomModelDetails(eq(MODEL_NAME));
+    verify(mockFileDownloadService, times(2)).download(any(), eq(DOWNLOAD_CONDITIONS));
+    assertThat(task.isComplete()).isTrue();
+    assertEquals(customModel, CUSTOM_MODEL);
+  }
+
+  @Test
+  public void getModel_local_noLocalModel_urlRetry_maxTries() throws Exception {
+    when(mockPrefs.getCustomModelDetails(eq(MODEL_NAME)))
+        .thenReturn(null)
+        .thenReturn(null)
+        .thenReturn(CUSTOM_MODEL);
+    when(mockModelDownloadService.getCustomModelDetails(
+            eq(TEST_PROJECT_ID), eq(MODEL_NAME), eq(null)))
+        .thenReturn(Tasks.forResult(CUSTOM_MODEL));
+    when(mockFileDownloadService.download(any(), eq(DOWNLOAD_CONDITIONS)))
+        .thenReturn(Tasks.forException(new Exception("Retry: Expired URL")));
+    TestOnCompleteListener<CustomModel> onCompleteListener = new TestOnCompleteListener<>();
+    Task<CustomModel> task =
+        firebaseModelDownloader.getModel(MODEL_NAME, DownloadType.LOCAL_MODEL, DOWNLOAD_CONDITIONS);
+    task.addOnCompleteListener(executor, onCompleteListener);
+    try {
+      onCompleteListener.await();
+    } catch (Exception ex) {
+      assertThat(ex.getMessage().contains("download failed")).isTrue();
+      assertThat(ex.getMessage().contains("Too many attempts")).isTrue();
+    }
+
+    verify(mockPrefs, times(2)).getCustomModelDetails(eq(MODEL_NAME));
+    verify(mockFileDownloadService, times(3)).download(any(), eq(DOWNLOAD_CONDITIONS));
+    verify(mockFileDownloadService, never()).loadNewlyDownloadedModelFile(any());
+    assertThat(task.isComplete()).isTrue();
+    assertThat(task.isSuccessful()).isFalse();
+  }
+
   @Test
   public void getModel_local_noLocalModel_error() throws Exception {
     when(mockPrefs.getCustomModelDetails(eq(MODEL_NAME)))
diff --git a/firebase-ml-modeldownloader/src/test/java/com/google/firebase/ml/modeldownloader/internal/ModelFileDownloadServiceTest.java b/firebase-ml-modeldownloader/src/test/java/com/google/firebase/ml/modeldownloader/internal/ModelFileDownloadServiceTest.java
