Add singleton-ish instance methods for model downloader (public API). (#7124)

* Restores ML Pods after M77.

* Fix Package.swift

* Refactor model info to be in-memory.

* Add improved model download functionality w/ single set of download handlers per task.

* Rename model download to better match URLSessionTask.

* Rename model download to better match URLSessionTask.

* Update unit test to read/write to user defaults.

* Configure downloader with Firebase app.

* Minor updates to model downloader and tests.

* Singleton for model downloader instance, for default and custom app.

* Manually add app lifecycle handling + TODOs for Firebase Components.

* Replace options with app name as a property in ModelDownloader.

* Update unit test to check model downloader instance creation + minor rename.

Author: manjanac

FirebaseMLModelDownloader/Sources/ModelDownloadTask.swift

@@ -37,7 +37,7 @@ class DownloadHandlers {
 
 /// Manager to handle model downloading device and storing downloaded model info to persistent storage.
 class ModelDownloadTask: NSObject {
-  let app: FirebaseApp
+  private let appName: String
   private(set) var modelInfo: ModelInfo
   private var downloadTask: URLSessionDownloadTask?
   private let downloadHandlers: DownloadHandlers
@@ -48,11 +48,11 @@ class ModelDownloadTask: NSObject {
                                                 delegate: self,
                                                 delegateQueue: nil)
 
-  init(app: FirebaseApp, modelInfo: ModelInfo,
+  init(modelInfo: ModelInfo, appName: String,
        progressHandler: DownloadHandlers.ProgressHandler? = nil,
        completion: @escaping DownloadHandlers.Completion) {
-    self.app = app
     self.modelInfo = modelInfo
+    self.appName = appName
     downloadHandlers = DownloadHandlers(
       progressHandler: progressHandler,
       completion: completion
@@ -90,7 +90,7 @@ extension ModelDownloadTask: URLSessionDownloadDelegate {
     modelInfo.path = savedURL.absoluteString
     /// Write model to user defaults.
     do {
-      try modelInfo.writeToDefaults(app: app, defaults: .firebaseMLDefaults)
+      try modelInfo.writeToDefaults(.firebaseMLDefaults, appName: appName)
     } catch {
       downloadHandlers
         .completion(.failure(.internalError(description: error.localizedDescription)))
@@ -123,7 +123,7 @@ extension ModelDownloadTask: URLSessionDownloadDelegate {
 /// Extension to handle post-download operations.
 extension ModelDownloadTask {
   var downloadedModelFileName: String {
-    return "fbml_model__\(app.name)__\(modelInfo.name)"
+    return "fbml_model__\(appName)__\(modelInfo.name)"
   }
 
   /// Build custom model object from model info.

FirebaseMLModelDownloader/Sources/ModelDownloader.swift

@@ -13,6 +13,7 @@
 // limitations under the License.
 
 import Foundation
+import FirebaseCore
 
 /// Possible errors with model downloading.
 public enum DownloadError: Error, Equatable {
@@ -49,7 +50,54 @@ public enum ModelDownloadType {
 }
 
 /// Downloader to manage custom model downloads.
-public struct ModelDownloader {
+public class ModelDownloader {
+  /// Name of the app associated with this instance of ModelDownloader.
+  private let appName: String
+
+  /// Shared dictionary mapping app name to a specific instance of model downloader.
+  // TODO: Switch to using Firebase components.
+  private static var modelDownloaderDictionary: [String: ModelDownloader] = [:]
+
+  /// Private init for downloader.
+  private init(app: FirebaseApp) {
+    appName = app.name
+    NotificationCenter.default.addObserver(
+      self,
+      selector: #selector(deleteModelDownloader),
+      name: Notification.Name("FIRAppDeleteNotification"),
+      object: nil
+    )
+  }
+
+  /// Handles app deletion notification.
+  @objc private func deleteModelDownloader(notification: Notification) {
+    if let userInfo = notification.userInfo,
+      let appName = userInfo["FIRAppNameKey"] as? String {
+      ModelDownloader.modelDownloaderDictionary.removeValue(forKey: appName)
+      // TODO: Clean up user defaults
+      // TODO: Clean up local instances of app
+    }
+  }
+
+  /// Model downloader with default app.
+  public static func modelDownloader() -> ModelDownloader {
+    guard let defaultApp = FirebaseApp.app() else {
+      fatalError("Default Firebase app not configured.")
+    }
+    return modelDownloader(app: defaultApp)
+  }
+
+  /// Model Downloader with custom app.
+  public static func modelDownloader(app: FirebaseApp) -> ModelDownloader {
+    if let downloader = modelDownloaderDictionary[app.name] {
+      return downloader
+    } else {
+      let downloader = ModelDownloader(app: app)
+      modelDownloaderDictionary[app.name] = downloader
+      return downloader
+    }
+  }
+
   /// Downloads a custom model to device or gets a custom model already on device, w/ optional handler for progress.
   public func getModel(name modelName: String, downloadType: ModelDownloadType,
                        conditions: ModelDownloadConditions,

FirebaseMLModelDownloader/Sources/ModelInfo.swift

@@ -41,9 +41,15 @@ struct ModelInfo {
     self.size = size
   }
 
-  init?(fromDefaults defaults: UserDefaults, name: String, app: FirebaseApp) {
+  /// Get user defaults key prefix.
+  private static func getUserDefaultsKeyPrefix(appName: String, modelName: String) -> String {
     let bundleID = Bundle.main.bundleIdentifier ?? ""
-    let defaultsPrefix = "\(bundleID).\(app.name).\(name)"
+    return "\(bundleID).\(appName).\(modelName)"
+  }
+
+  // TODO: Move reading and writing to user defaults to a new file.
+  init?(fromDefaults defaults: UserDefaults, modelName: String, appName: String) {
+    let defaultsPrefix = ModelInfo.getUserDefaultsKeyPrefix(appName: appName, modelName: modelName)
     guard let downloadURL = defaults
       .value(forKey: "\(defaultsPrefix).model-download-url") as? String,
       let url = URL(string: downloadURL),
@@ -52,20 +58,19 @@ struct ModelInfo {
       let path = defaults.value(forKey: "\(defaultsPrefix).model-path") as? String else {
       return nil
     }
-    self.name = name
+    name = modelName
     self.downloadURL = url
     self.modelHash = modelHash
     self.size = size
     self.path = path
   }
 
-  func writeToDefaults(app: FirebaseApp, defaults: UserDefaults) throws {
+  func writeToDefaults(_ defaults: UserDefaults, appName: String) throws {
     guard let modelPath = path else {
       throw DownloadedModelError
         .fileIOError(description: "Could not save model info to user defaults.")
     }
-    let bundleID = Bundle.main.bundleIdentifier ?? ""
-    let defaultsPrefix = "\(bundleID).\(app.name).\(name)"
+    let defaultsPrefix = ModelInfo.getUserDefaultsKeyPrefix(appName: appName, modelName: name)
     defaults.setValue(downloadURL.absoluteString, forKey: "\(defaultsPrefix).model-download-url")
     defaults.setValue(modelHash, forKey: "\(defaultsPrefix).model-hash")
     defaults.setValue(size, forKey: "\(defaultsPrefix).model-size")

FirebaseMLModelDownloader/Sources/ModelInfoRetriever.swift

@@ -34,8 +34,8 @@ extension ModelInfoResponse {
 
 /// Model info retriever for a model from local user defaults or server.
 class ModelInfoRetriever: NSObject {
-  /// Current Firebase app.
-  var app: FirebaseApp
+  /// Current Firebase app options
+  private var options: FirebaseOptions
   /// Model info associated with model.
   var modelInfo: ModelInfo?
   /// Model name.
@@ -44,10 +44,10 @@ class ModelInfoRetriever: NSObject {
   var installations: Installations
 
   /// Associate model info retriever with current Firebase app, and model name.
-  init(app: FirebaseApp, modelName: String) {
-    self.app = app
+  init(modelName: String, options: FirebaseOptions, installations: Installations) {
     self.modelName = modelName
-    installations = Installations.installations(app: app)
+    self.options = options
+    self.installations = installations
   }
 }
 
@@ -70,8 +70,8 @@ extension ModelInfoRetriever {
 
   /// Construct model fetch base URL.
   var modelInfoFetchURL: URL {
-    let projectID = app.options.projectID ?? ""
-    let apiKey = app.options.apiKey
+    let projectID = options.projectID ?? ""
+    let apiKey = options.apiKey
     var components = URLComponents()
     components.scheme = "https"
     components.host = "firebaseml.googleapis.com"

FirebaseMLModelDownloader/Tests/Integration/ModelDownloaderIntegrationTests.swift

@@ -14,6 +14,7 @@
 
 import XCTest
 @testable import FirebaseCore
+@testable import FirebaseInstallations
 @testable import FirebaseMLModelDownloader
 
 extension UserDefaults {
@@ -47,8 +48,9 @@ final class ModelDownloaderIntegrationTests: XCTestCase {
     }
     let testModelName = "image-classification"
     let modelInfoRetriever = ModelInfoRetriever(
-      app: testApp,
-      modelName: testModelName
+      modelName: testModelName,
+      options: testApp.options,
+      installations: Installations.installations(app: testApp)
     )
     let expectation = self.expectation(description: "Wait for FIS auth token.")
     modelInfoRetriever.getAuthToken(completion: { result in
@@ -72,8 +74,9 @@ final class ModelDownloaderIntegrationTests: XCTestCase {
     }
     let testModelName = "pose-detection"
     let modelInfoRetriever = ModelInfoRetriever(
-      app: testApp,
-      modelName: testModelName
+      modelName: testModelName,
+      options: testApp.options,
+      installations: Installations.installations(app: testApp)
     )
     let downloadExpectation = expectation(description: "Wait for model info to download.")
     modelInfoRetriever.downloadModelInfo(completion: { error in
@@ -111,10 +114,10 @@ final class ModelDownloaderIntegrationTests: XCTestCase {
     let functionName = #function.dropLast(2)
     let testModelName = "\(functionName)-test-model"
     let modelInfoRetriever = ModelInfoRetriever(
-      app: testApp,
-      modelName: testModelName
+      modelName: testModelName,
+      options: testApp.options,
+      installations: Installations.installations(app: testApp)
     )
-
     let urlString =
       "https://tfhub.dev/tensorflow/lite-model/ssd_mobilenet_v1/1/metadata/1?lite-format=tflite"
     let url = URL(string: urlString)!
@@ -127,8 +130,7 @@ final class ModelDownloaderIntegrationTests: XCTestCase {
     )
     let expectation = self.expectation(description: "Wait for model to download.")
     let modelDownloadManager = ModelDownloadTask(
-      app: testApp,
-      modelInfo: modelInfoRetriever.modelInfo!,
+      modelInfo: modelInfoRetriever.modelInfo!, appName: testApp.name,
       progressHandler: { progress in
         XCTAssertNotNil(progress)
       }

FirebaseMLModelDownloader/Tests/Unit/ModelDownloaderUnitTests.swift

@@ -14,6 +14,7 @@
 
 import XCTest
 @testable import FirebaseCore
+@testable import FirebaseInstallations
 @testable import FirebaseMLModelDownloader
 
 /// Mock options to configure default Firebase app.
@@ -73,21 +74,21 @@ final class ModelDownloaderUnitTests: XCTestCase {
     )
     // This fails because there is no model path.
     do {
-      try modelInfo.writeToDefaults(app: testApp, defaults: .getTestInstance())
+      try modelInfo.writeToDefaults(.getTestInstance(), appName: testApp.name)
     } catch {
       XCTAssertNotNil(error)
     }
     modelInfo.path = testModelPath
     // This shouldn't fail because model info object is now complete.
     do {
-      try modelInfo.writeToDefaults(app: testApp, defaults: .getTestInstance())
+      try modelInfo.writeToDefaults(.getTestInstance(), appName: testApp.name)
     } catch {
       XCTFail(error.localizedDescription)
     }
     guard let savedModelInfo = ModelInfo(
       fromDefaults: .getTestInstance(cleared: false),
-      name: testModelName,
-      app: testApp
+      modelName: testModelName,
+      appName: testApp.name
     ) else {
       XCTFail("Model info not saved to user defaults.")
       return
@@ -106,8 +107,9 @@ final class ModelDownloaderUnitTests: XCTestCase {
     let functionName = #function
     let testModelName = "\(functionName)-test-model"
     let modelInfoRetriever = ModelInfoRetriever(
-      app: testApp,
-      modelName: testModelName
+      modelName: testModelName,
+      options: testApp.options,
+      installations: Installations.installations(app: testApp)
     )
     let sampleResponse: String = """
     {
@@ -137,7 +139,18 @@ final class ModelDownloaderUnitTests: XCTestCase {
     // This is an example of a functional test case.
     // Use XCTAssert and related functions to verify your tests produce the correct
     // results.
-    let modelDownloader = ModelDownloader()
+    guard let testApp = FirebaseApp.app() else {
+      XCTFail("Default app was not configured.")
+      return
+    }
+
+    let modelDownloader = ModelDownloader.modelDownloader()
+
+    let modelDownloaderWithApp = ModelDownloader.modelDownloader(app: testApp)
+
+    /// These should point to the same instance.
+    XCTAssert(modelDownloader === modelDownloaderWithApp)
+
     let conditions = ModelDownloadConditions()
 
     // Download model w/ progress handler
@@ -161,7 +174,7 @@ final class ModelDownloaderUnitTests: XCTestCase {
     }
 
     // Access array of downloaded models
-    modelDownloader.listDownloadedModels { result in
+    modelDownloaderWithApp.listDownloadedModels { result in
       switch result {
       case .success:
         // Pick model(s) for further use