implements new Storage API for handleEventsForBackgroundURLSession

Amplify/Categories/Storage/StorageCategory+ClientBehavior.swift

@@ -53,8 +53,9 @@ extension StorageCategory: StorageCategoryBehavior {
         try await plugin.list(options: options)
     }
 
-    public func handleBackgroundEvents(identifier: String) async -> Bool {
-        await plugin.handleBackgroundEvents(identifier: identifier)
+    @discardableResult
+    public func handleEventsForBackgroundURLSession(identifier: String) async -> Bool {
+        await plugin.handleEventsForBackgroundURLSession(identifier: identifier)
     }
 
 }

Amplify/Categories/Storage/StorageCategoryBehavior.swift

@@ -88,6 +88,7 @@ public protocol StorageCategoryBehavior {
     /// Handles background events which are related to URLSession
     /// - Parameter identifier: identifier
     /// - Returns: returns true if the identifier is handled by Amplify
-    func handleBackgroundEvents(identifier: String) async -> Bool
+    @discardableResult
+    func handleEventsForBackgroundURLSession(identifier: String) async -> Bool
 
 }

AmplifyPlugins/Storage/Sources/AWSS3StoragePlugin/AWSS3StoragePlugin+AsyncClientBehavior.swift

@@ -119,10 +119,8 @@ extension AWSS3StoragePlugin {
         return try await taskAdapter.value
     }
 
-    public func handleBackgroundEvents(identifier: String) async -> Bool {
-        await withCheckedContinuation { (continuation: CheckedContinuation<Bool, Never>) in
-            StorageBackgroundEventsRegistry.handleBackgroundEvents(identifier: identifier, continuation: continuation)
-        }
+    public func handleEventsForBackgroundURLSession(identifier: String) async -> Bool {
+        await StorageBackgroundEventsRegistry.shared.handleEventsForBackgroundURLSession(identifier: identifier)
     }
 
 }

AmplifyPlugins/Storage/Sources/AWSS3StoragePlugin/Service/Storage/AWSS3StorageService.swift

@@ -108,7 +108,9 @@ class AWSS3StorageService: AWSS3StorageServiceBehaviour, StorageServiceProxy {
         self.awsS3 = awsS3
         self.bucket = bucket
 
-        StorageBackgroundEventsRegistry.register(identifier: identifier)
+        Task {
+            await StorageBackgroundEventsRegistry.shared.register(identifier: identifier)
+        }
 
         delegate.storageService = self
 
@@ -125,7 +127,9 @@ class AWSS3StorageService: AWSS3StorageServiceBehaviour, StorageServiceProxy {
     }
 
     deinit {
-        StorageBackgroundEventsRegistry.unregister(identifier: identifier)
+        Task {
+            await StorageBackgroundEventsRegistry.shared.unregister(identifier: identifier)
+        }
     }
 
     func reset() {

AmplifyPlugins/Storage/Sources/AWSS3StoragePlugin/Support/Internal/StorageBackgroundEventsRegistry.swift

@@ -7,50 +7,71 @@
 
 import Foundation
 
+extension Notification.Name {
+    static let StorageBackgroundEventsRegistryWaiting = Notification.Name("StorageBackgroundEventsRegistryWaiting")
+}
+
 /// Background events registry.
 ///
 /// Discussion:
 ///    Multiple URLSession instances could be running background events with their own unique identifier. Those can be run
 ///    independently of the Amplify Storage plugin and this function will indiciate if it will handle the given identifier.
-class StorageBackgroundEventsRegistry {
+actor StorageBackgroundEventsRegistry {
     typealias StorageBackgroundEventsContinuation = CheckedContinuation<Bool, Never>
-    static var identifier: String?
-    static var continuation: StorageBackgroundEventsContinuation?
+
+    @MainActor
+    static let shared = StorageBackgroundEventsRegistry()
+
+    private var identifier: String?
+    private var continuation: StorageBackgroundEventsContinuation?
+
+    // override for use with unit tests
+    internal private(set) var notificationCenter: NotificationCenter?
+
+    func change(notificationCenter: NotificationCenter?) {
+        self.notificationCenter = notificationCenter
+    }
 
     /// Handles background events for URLSession on iOS.
     /// - Parameters:
     ///   - identifier: session identifier
     ///   - completionHandler: completion handler
     /// - Returns: indicates if the identifier was registered and will be handled
-    static func handleBackgroundEvents(identifier: String, continuation: StorageBackgroundEventsContinuation) {
-        if self.identifier == identifier {
+    func handleEventsForBackgroundURLSession(identifier: String) async -> Bool {
+        guard self.identifier == identifier else { return false }
+
+        return await withCheckedContinuation { (continuation: CheckedContinuation<Bool, Never>) in
             self.continuation = continuation
-        } else {
-            continuation.resume(returning: false)
+            notifyWaiting(for: identifier)
         }
     }
 
-    // MARK: - Internal -
+    /// Notifies observes when waiting for continuation to be resumed.
+    /// - Parameters:
+    ///   - identifier: session identifier
+    private func notifyWaiting(for identifier: String) {
+        notificationCenter?.post(name: Notification.Name.StorageBackgroundEventsRegistryWaiting, object: identifier)
+    }
 
     // The storage plugin will register the session identifier when it is configured.
-    static func register(identifier: String) {
+    func register(identifier: String) {
         self.identifier = identifier
     }
 
     // When the storage function is deinitialized it will unregister the session identifier.
-    static func unregister(identifier: String) {
+    func unregister(identifier: String) {
         if self.identifier == identifier {
             self.identifier = nil
         }
     }
 
     // When URLSession is done processing background events it will use this function to get the completion handler.
-    static func getContinuation(for identifier: String) -> StorageBackgroundEventsContinuation? {
+    func getContinuation(for identifier: String) -> StorageBackgroundEventsContinuation? {
         self.identifier == identifier ? continuation : nil
     }
 
     // Once the background event completion handler is used it can be cleared.
-    static func removeContinuation(for identifier: String) {
+    func removeContinuation(for identifier: String) {
         if self.identifier == identifier {
             self.continuation = nil
         }

AmplifyPlugins/Storage/Sources/AWSS3StoragePlugin/Support/Internal/StorageServiceSessionDelegate.swift

@@ -53,13 +53,15 @@ extension StorageServiceSessionDelegate: URLSessionDelegate {
     func urlSessionDidFinishEvents(forBackgroundURLSession session: URLSession) {
         logURLSessionActivity("Session did finish background events")
 
-        if let identifier = storageService?.identifier,
-           let continuation = StorageBackgroundEventsRegistry.getContinuation(for: identifier) {
-            // Must be run on main thread as covered by Apple Developer docs.
-            Task { @MainActor in
-                continuation.resume(returning: true)
+        Task {
+            if let identifier = session.configuration.identifier,
+               let continuation = await StorageBackgroundEventsRegistry.shared.getContinuation(for: identifier) {
+                // Must be run on main thread as covered by Apple Developer docs.
+                Task { @MainActor in
+                    continuation.resume(returning: true)
+                }
+                await StorageBackgroundEventsRegistry.shared.removeContinuation(for: identifier)
             }
-            StorageBackgroundEventsRegistry.removeContinuation(for: identifier)
         }
     }
 

AmplifyPlugins/Storage/Tests/AWSS3StoragePluginTests/Support/Internal/StorageBackgroundEventsRegistryTests.swift

@@ -16,60 +16,82 @@ class StorageBackgroundEventsRegistryTests: XCTestCase {
     func testRegisteringAndUnregister() async throws {
         let identifier = UUID().uuidString
         let otherIdentifier = UUID().uuidString
-        StorageBackgroundEventsRegistry.register(identifier: identifier)
+        await StorageBackgroundEventsRegistry.shared.register(identifier: identifier)
 
-        let done = asyncExpectation(description: "done", expectedFulfillmentCount: 2)
+        let notificationCenter = NotificationCenter()
+        await StorageBackgroundEventsRegistry.shared.change(notificationCenter: notificationCenter)
+        defer {
+            Task {
+                await StorageBackgroundEventsRegistry.shared.change(notificationCenter: nil)
+            }
+        }
 
-        Task {
-            let handled = await withCheckedContinuation { (continuation: CheckedContinuation<Bool, Never>) in
-                StorageBackgroundEventsRegistry.handleBackgroundEvents(identifier: identifier, continuation: continuation)
-                Task {
-                    await done.fulfill()
-                }
+        let done = asyncExpectation(description: "done")
+        let waiting = asyncExpectation(description: "waiting")
+
+        notificationCenter.addObserver(forName: Notification.Name.StorageBackgroundEventsRegistryWaiting, object: nil, queue: nil) { notification in
+            guard let notificationIdentifier = notification.object as? String else {
+                XCTFail("Identifier not defined")
+                return
+            }
+            XCTAssertEqual(notificationIdentifier, identifier)
+            Task {
+                await waiting.fulfill()
             }
-            XCTAssertTrue(handled)
         }
 
         Task {
-            let otherHandled = await withCheckedContinuation { (continuation: CheckedContinuation<Bool, Never>) in
-                StorageBackgroundEventsRegistry.handleBackgroundEvents(identifier: otherIdentifier, continuation: continuation)
-                Task {
-                    await done.fulfill()
-                }
-            }
-            XCTAssertFalse(otherHandled)
+            let handled = await StorageBackgroundEventsRegistry.shared.handleEventsForBackgroundURLSession(identifier: identifier)
+            await done.fulfill()
+            XCTAssertTrue(handled)
         }
 
+        await waitForExpectations([waiting])
+
+        let didContinue = await handleEvents(for: identifier)
+        XCTAssertTrue(didContinue)
         await waitForExpectations([done])
 
-        handleEvents(for: identifier)
-        handleEvents(for: otherIdentifier)
+        let otherDone = asyncExpectation(description: "other done")
+
+        Task {
+            let otherHandled = await StorageBackgroundEventsRegistry.shared.handleEventsForBackgroundURLSession(identifier: otherIdentifier)
+            await otherDone.fulfill()
+            XCTAssertFalse(otherHandled)
+        }
+
+        let didNotContinue = await handleEvents(for: otherIdentifier)
+        XCTAssertFalse(didNotContinue)
+        await waitForExpectations([otherDone])
     }
 
     func testHandlingUnregisteredIdentifier() async throws {
         let identifier = UUID().uuidString
         let otherIdentifier = UUID().uuidString
-        StorageBackgroundEventsRegistry.register(identifier: otherIdentifier)
+        await StorageBackgroundEventsRegistry.shared.register(identifier: otherIdentifier)
 
         let done = asyncExpectation(description: "done")
 
         Task {
-            let handled = await withCheckedContinuation { (continuation: CheckedContinuation<Bool, Never>) in
-                StorageBackgroundEventsRegistry.handleBackgroundEvents(identifier: identifier, continuation: continuation)
-                Task {
-                    await done.fulfill()
-                }
-            }
+            let handled = await StorageBackgroundEventsRegistry.shared.handleEventsForBackgroundURLSession(identifier: identifier)
+            await done.fulfill()
             XCTAssertFalse(handled)
         }
 
         await waitForExpectations([done])
     }
 
     // Simulates URLSessionDelegate behavior
-    func handleEvents(for identifier: String) {
-        if let continuation = StorageBackgroundEventsRegistry.getContinuation(for: identifier) {
+    func handleEvents(for identifier: String) async -> Bool {
+        await Task.yield()
+
+        if let continuation = await StorageBackgroundEventsRegistry.shared.getContinuation(for: identifier) {
             continuation.resume(returning: true)
+            await StorageBackgroundEventsRegistry.shared.removeContinuation(for: identifier)
+            return true
+        } else {
+            print("No continuation for identifier: \(identifier)")
+            return false
         }
     }
 

AmplifyTestCommon/Mocks/MockStorageCategoryPlugin.swift

@@ -176,7 +176,7 @@ class MockStorageCategoryPlugin: MessageReporter, StorageCategoryPlugin {
         return try await taskAdapter.value
     }
 
-    func handleBackgroundEvents(identifier: String) async -> Bool {
+    func handleEventsForBackgroundURLSession(identifier: String) async -> Bool {
         false
     }
 

AmplifyTests/CategoryTests/Hub/AmplifyOperationHubTests.swift

@@ -217,7 +217,7 @@ class MockDispatchingStoragePlugin: StorageCategoryPlugin {
         return operation
     }
 
-    func handleBackgroundEvents(identifier: String) async -> Bool {
+    func handleEventsForBackgroundURLSession(identifier: String) async -> Bool {
         false
     }