Handle expired download URL with support for retries (#7329)

* Restores ML Pods after M77.

* Fix Package.swift

* Re-add catalyst to GHA workflow.

* Update error handling.

* Add logging for background download failure.

* Handle expired download URL.

* Remove empty file.

* Fix integration test.

* Minor refactor.

Author: manjanac

FirebaseMLModelDownloader/Sources/ModelDownloadTask.swift

@@ -55,23 +55,31 @@ class ModelDownloadTask: NSObject {
   private lazy var downloadSession = URLSession(configuration: .ephemeral,
                                                 delegate: self,
                                                 delegateQueue: nil)
+  /// Model info retriever in case of retries.
+  private let modelInfoRetriever: ModelInfoRetriever
+  /// Number of retries in case of model download URL expiry.
+  private var numberOfRetries: Int = 1
   /// Telemetry logger.
   private let telemetryLogger: TelemetryLogger?
 
   init(remoteModelInfo: RemoteModelInfo, appName: String, defaults: UserDefaults,
+       modelInfoRetriever: ModelInfoRetriever,
        telemetryLogger: TelemetryLogger? = nil,
        progressHandler: DownloadHandlers.ProgressHandler? = nil,
        completion: @escaping DownloadHandlers.Completion) {
     self.remoteModelInfo = remoteModelInfo
     self.appName = appName
+    self.modelInfoRetriever = modelInfoRetriever
     self.telemetryLogger = telemetryLogger
     self.defaults = defaults
     downloadHandlers = DownloadHandlers(
       progressHandler: progressHandler,
       completion: completion
     )
   }
+}
 
+extension ModelDownloadTask {
   /// Asynchronously download model file to device.
   func resumeModelDownload() {
     guard downloadStatus == .notStarted else { return }
@@ -116,7 +124,42 @@ extension ModelDownloadTask: URLSessionDownloadDelegate {
     }
 
     guard (200 ..< 299).contains(response.statusCode) else {
-      // TODO: Handle download url expiry + retries.
+      /// Possible failure due to download URL expiry.
+      if response.statusCode == 400 {
+        let currentDateTime = Date()
+        /// Retry download if allowed.
+        guard currentDateTime > remoteModelInfo.urlExpiryTime, numberOfRetries > 0 else {
+          downloadStatus = .failed
+          downloadHandlers
+            .completion(.failure(.internalError(description: ModelDownloadTask.ErrorDescription
+                .expiredModelInfo)))
+          return
+        }
+        numberOfRetries -= 1
+        modelInfoRetriever.downloadModelInfo { result in
+          switch result {
+          case let .success(downloadModelInfoResult):
+            switch downloadModelInfoResult {
+            /// New model info was downloaded from server.
+            case let .modelInfo(remoteModelInfo):
+              self.remoteModelInfo = remoteModelInfo
+              self.resumeModelDownload()
+            /// This should not ever be the case - model info cannot be unmodified within ModelDownloadTask.
+            case .notModified:
+              DispatchQueue.main.async {
+                self.downloadHandlers
+                  .completion(.failure(.internalError(description: ModelDownloadTask
+                      .ErrorDescription.expiredModelInfo)))
+              }
+            }
+          case let .failure(downloadError):
+            self.downloadStatus = .failed
+            DispatchQueue.main.async {
+              self.downloadHandlers.completion(.failure(downloadError))
+            }
+          }
+        }
+      }
       return
     }
 
@@ -196,5 +239,6 @@ extension ModelDownloadTask {
       "Could not get server response for model downloading."
     static let saveModel: StaticString =
       "Unable to save downloaded remote model file."
+    static let expiredModelInfo = "Unable to update expired model info."
   }
 }

FirebaseMLModelDownloader/Sources/ModelDownloader.swift

@@ -17,7 +17,7 @@ import FirebaseCore
 import FirebaseInstallations
 
 /// Possible errors with model downloading.
-public enum DownloadError: Error, Equatable {
+public enum DownloadError: Error {
   /// No model with this name found on server.
   case notFound
   /// Caller does not have necessary permissions for this operation.
@@ -33,7 +33,7 @@ public enum DownloadError: Error, Equatable {
 }
 
 /// Possible errors with locating model on device.
-public enum DownloadedModelError: Error, Equatable {
+public enum DownloadedModelError: Error {
   /// File system error.
   case fileIOError(description: String)
   /// Model not found on device.
@@ -282,6 +282,7 @@ extension ModelDownloader {
             remoteModelInfo: remoteModelInfo,
             appName: self.appName,
             defaults: self.userDefaults,
+            modelInfoRetriever: modelInfoRetriever,
             telemetryLogger: self.telemetryLogger,
             progressHandler: progressHandler,
             completion: completion

FirebaseMLModelDownloader/Tests/Integration/ModelDownloaderIntegrationTests.swift

@@ -157,11 +157,19 @@ final class ModelDownloaderIntegrationTests: XCTestCase {
       urlExpiryTime: Date()
     )
 
+    let modelInfoRetriever = ModelInfoRetriever(
+      modelName: testModelName,
+      options: testApp.options,
+      installations: Installations.installations(app: testApp),
+      appName: testApp.name
+    )
+
     let expectation = self.expectation(description: "Wait for model to download.")
     let modelDownloadManager = ModelDownloadTask(
       remoteModelInfo: remoteModelInfo,
       appName: testApp.name,
       defaults: .createTestInstance(testName: #function),
+      modelInfoRetriever: modelInfoRetriever,
       progressHandler: { progress in
         XCTAssertLessThanOrEqual(progress, 1)
         XCTAssertGreaterThanOrEqual(progress, 0)