[storage] Simplify callback implementation

FirebaseStorage/Sources/Internal/StorageTokenAuthorizer.swift

@@ -37,7 +37,6 @@ class StorageTokenAuthorizer: NSObject, GTMSessionFetcherAuthorizer {
     request?.setValue(googleAppID, forHTTPHeaderField: "x-firebase-gmpid")
 
     var tokenError: NSError?
-    let callbackQueue = fetcherService.callbackQueue ?? DispatchQueue.main
     let fetchTokenGroup = DispatchGroup()
     if let auth {
       fetchTokenGroup.enter()
@@ -101,19 +100,19 @@ class StorageTokenAuthorizer: NSObject, GTMSessionFetcherAuthorizer {
 
   var userEmail: String?
 
-  let fetcherService: GTMSessionFetcherService
+  let callbackQueue: DispatchQueue
   private let googleAppID: String
   private let auth: AuthInterop?
   private let appCheck: AppCheckInterop?
 
   private let serialAuthArgsQueue = DispatchQueue(label: "com.google.firebasestorage.authorizer")
 
   init(googleAppID: String,
-       fetcherService: GTMSessionFetcherService,
+       callbackQueue: DispatchQueue = DispatchQueue.main,
        authProvider: AuthInterop?,
        appCheck: AppCheckInterop?) {
     self.googleAppID = googleAppID
-    self.fetcherService = fetcherService
+    self.callbackQueue = callbackQueue
     auth = authProvider
     self.appCheck = appCheck
   }

FirebaseStorage/Sources/Storage.swift

@@ -123,19 +123,7 @@ import FirebaseCore
   @objc public var uploadChunkSizeBytes: Int64 = .max
 
   /// A `DispatchQueue` that all developer callbacks are fired on. Defaults to the main queue.
-  @objc public var callbackQueue: DispatchQueue {
-    get {
-      ensureConfigured()
-      guard let queue = fetcherService?.callbackQueue else {
-        fatalError("Internal error: Failed to initialize fetcherService callbackQueue")
-      }
-      return queue
-    }
-    set(newValue) {
-      ensureConfigured()
-      fetcherService?.callbackQueue = newValue
-    }
-  }
+  @objc public var callbackQueue: DispatchQueue = .main
 
   /// Creates a `StorageReference` initialized at the root Firebase Storage location.
   /// - Returns: An instance of `StorageReference` referencing the root of the storage bucket.
@@ -317,7 +305,8 @@ import FirebaseCore
   private static func initFetcherServiceForApp(_ app: FirebaseApp,
                                                _ bucket: String,
                                                _ auth: AuthInterop?,
-                                               _ appCheck: AppCheckInterop?)
+                                               _ appCheck: AppCheckInterop?,
+                                               _ callbackQueue: DispatchQueue)
     -> GTMSessionFetcherService {
     objc_sync_enter(fetcherServiceLock)
     defer { objc_sync_exit(fetcherServiceLock) }
@@ -334,7 +323,7 @@ import FirebaseCore
       fetcherService?.allowLocalhostRequest = true
       let authorizer = StorageTokenAuthorizer(
         googleAppID: app.options.googleAppID,
-        fetcherService: fetcherService!,
+        callbackQueue: callbackQueue,
         authProvider: auth,
         appCheck: appCheck
       )
@@ -390,7 +379,8 @@ import FirebaseCore
     guard fetcherService == nil else {
       return
     }
-    fetcherService = Storage.initFetcherServiceForApp(app, storageBucket, auth, appCheck)
+    fetcherService = Storage.initFetcherServiceForApp(app, storageBucket, auth, appCheck,
+                                                      callbackQueue)
     if usesEmulator {
       fetcherService?.allowLocalhostRequest = true
       fetcherService?.allowedInsecureSchemes = ["http"]

FirebaseStorage/Sources/StorageDownloadTask.swift

@@ -85,7 +85,6 @@ open class StorageDownloadTask: StorageObservableTask, StorageTaskManagement {
   }
 
   private var fetcher: GTMSessionFetcher?
-  private var fetcherCompletion: ((Data?, NSError?) -> Void)?
   var downloadData: Data?
   // Hold completion in object to force it to be retained until completion block is called.
   var completionData: ((Data?, Error?) -> Void)?
@@ -104,7 +103,7 @@ open class StorageDownloadTask: StorageObservableTask, StorageTaskManagement {
     self.fetcher?.stopFetching()
   }
 
-  func enqueueImplementation(resumeWith resumeData: Data? = nil) {
+  private func enqueueImplementation(resumeWith resumeData: Data? = nil) {
     dispatchQueue.async { [weak self] in
       guard let self = self else { return }
       self.state = .queueing
@@ -153,32 +152,29 @@ open class StorageDownloadTask: StorageObservableTask, StorageTaskManagement {
         }
       }
       self.fetcher = fetcher
-
-      // Capture self here to retain until completion.
-      self.fetcherCompletion = { [self] (data: Data?, error: NSError?) in
-        defer {
-          self.removeAllObservers()
-          self.fetcherCompletion = nil
-        }
-        self.fire(for: .progress, snapshot: self.snapshot)
-
-        // Handle potential issues with download
-        if let error {
+      self.state = .running
+      Task {
+        do {
+          let data = try await self.fetcher?.beginFetch()
+          // Fire last progress updates
+          self.fire(for: .progress, snapshot: self.snapshot)
+
+          // Download completed successfully, fire completion callbacks
+          self.state = .success
+          if let data {
+            self.downloadData = data
+          }
+          self.fire(for: .success, snapshot: self.snapshot)
+        } catch {
+          self.fire(for: .progress, snapshot: self.snapshot)
           self.state = .failed
-          self.error = StorageErrorCode.error(withServerError: error, ref: self.reference)
+          self.error = StorageErrorCode.error(
+            withServerError: error as NSError,
+            ref: self.reference
+          )
           self.fire(for: .failure, snapshot: self.snapshot)
-          return
         }
-        // Download completed successfully, fire completion callbacks
-        self.state = .success
-        if let data {
-          self.downloadData = data
-        }
-        self.fire(for: .success, snapshot: self.snapshot)
-      }
-      self.state = .running
-      self.fetcher?.beginFetch { [self] data, error in
-        self.fetcherCompletion?(data, error as? NSError)
+        self.removeAllObservers()
       }
     }
   }

FirebaseStorage/Sources/StorageObservableTask.swift

@@ -168,12 +168,11 @@ import Foundation
 
   func fire(handlers: [String: (StorageTaskSnapshot) -> Void],
             snapshot: StorageTaskSnapshot) {
-    let callbackQueue = fetcherService.callbackQueue ?? DispatchQueue.main
     objc_sync_enter(StorageObservableTask.self)
     let enumeration = handlers.enumerated()
     objc_sync_exit(StorageObservableTask.self)
     for (_, handler) in enumeration {
-      callbackQueue.async {
+      reference.storage.callbackQueue.async {
         handler.value(snapshot)
       }
     }

FirebaseStorage/Sources/StorageReference.swift

@@ -206,7 +206,7 @@ import Foundation
                                    file: nil)
 
     task.completionData = completion
-    let callbackQueue = fetcherService.callbackQueue ?? DispatchQueue.main
+    let callbackQueue = storage.callbackQueue
 
     task.observe(.success) { snapshot in
       let error = self.checkSizeOverflow(task: snapshot.task, maxSize: maxSize)
@@ -288,7 +288,7 @@ import Foundation
 
     if let completion {
       task.completionURL = completion
-      let callbackQueue = fetcherService.callbackQueue ?? DispatchQueue.main
+      let callbackQueue = storage.callbackQueue
 
       task.observe(.success) { snapshot in
         callbackQueue.async {

FirebaseStorage/Sources/StorageUploadTask.swift

@@ -115,38 +115,36 @@ import Foundation
 
       // Process fetches
       self.state = .running
+      Task {
+        do {
+          let data = try await self.uploadFetcher?.beginFetch()
+          // Fire last progress updates
+          self.fire(for: .progress, snapshot: self.snapshot)
+
+          // Upload completed successfully, fire completion callbacks
+          self.state = .success
 
-      self.fetcherCompletion = { [self] (data: Data?, error: NSError?) in
-        // Fire last progress updates
-        self.fire(for: .progress, snapshot: self.snapshot)
+          guard let data = data else {
+            fatalError("Internal Error: uploadFetcher returned with nil data and no error")
+          }
 
-        // Handle potential issues with upload
-        if let error {
+          if let responseDictionary = try? JSONSerialization
+            .jsonObject(with: data) as? [String: AnyHashable] {
+            let metadata = StorageMetadata(dictionary: responseDictionary)
+            metadata.fileType = .file
+            self.metadata = metadata
+          } else {
+            self.error = StorageErrorCode.error(withInvalidRequest: data)
+          }
+          self.finishTaskWithStatus(status: .success, snapshot: self.snapshot)
+        } catch {
+          self.fire(for: .progress, snapshot: self.snapshot)
           self.state = .failed
-          self.error = StorageErrorCode.error(withServerError: error, ref: self.reference)
+          self.error = StorageErrorCode.error(withServerError: error as NSError,
+                                              ref: self.reference)
           self.metadata = self.uploadMetadata
           self.finishTaskWithStatus(status: .failure, snapshot: self.snapshot)
-          return
-        }
-        // Upload completed successfully, fire completion callbacks
-        self.state = .success
-
-        guard let data = data else {
-          fatalError("Internal Error: fetcherCompletion returned with nil data and nil error")
-        }
-
-        if let responseDictionary = try? JSONSerialization
-          .jsonObject(with: data) as? [String: AnyHashable] {
-          let metadata = StorageMetadata(dictionary: responseDictionary)
-          metadata.fileType = .file
-          self.metadata = metadata
-        } else {
-          self.error = StorageErrorCode.error(withInvalidRequest: data)
         }
-        self.finishTaskWithStatus(status: .success, snapshot: self.snapshot)
-      }
-      self.uploadFetcher?.beginFetch { [weak self] (data: Data?, error: Error?) in
-        self?.fetcherCompletion?(data, error as NSError?)
       }
     }
   }
@@ -202,7 +200,6 @@ import Foundation
   }
 
   private var uploadFetcher: GTMSessionUploadFetcher?
-  private var fetcherCompletion: ((Data?, NSError?) -> Void)?
   private var uploadMetadata: StorageMetadata
   private var uploadData: Data?
   // Hold completion in object to force it to be retained until completion block is called.
@@ -247,7 +244,6 @@ import Foundation
   func finishTaskWithStatus(status: StorageTaskStatus, snapshot: StorageTaskSnapshot) {
     fire(for: status, snapshot: snapshot)
     removeAllObservers()
-    fetcherCompletion = nil
   }
 
   private func GCSEscapedString(_ input: String?) -> String? {