From 0a1c1ef5d55601c4b2fc263e75462a11f2e3914c Mon Sep 17 00:00:00 2001
From: Brennan Stehling <stehlib@amazon.com>
Date: Wed, 19 Oct 2022 13:57:18 -0700
Subject: [PATCH] handles race condition and coordinates with tests for
 handling behavior

---
 .../StorageBackgroundEventsRegistry.swift     | 18 +++++++-
 ...StorageBackgroundEventsRegistryTests.swift | 45 ++++++++++++++++---
 2 files changed, 54 insertions(+), 9 deletions(-)

diff --git a/AmplifyPlugins/Storage/Sources/AWSS3StoragePlugin/Support/Internal/StorageBackgroundEventsRegistry.swift b/AmplifyPlugins/Storage/Sources/AWSS3StoragePlugin/Support/Internal/StorageBackgroundEventsRegistry.swift
index a726ed617b..84f94eb2af 100644
--- a/AmplifyPlugins/Storage/Sources/AWSS3StoragePlugin/Support/Internal/StorageBackgroundEventsRegistry.swift
+++ b/AmplifyPlugins/Storage/Sources/AWSS3StoragePlugin/Support/Internal/StorageBackgroundEventsRegistry.swift
@@ -7,16 +7,23 @@
 
 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?
 
+    // override for use with unit tests
+    static var notificationCenter: NotificationCenter?
+
     /// Handles background events for URLSession on iOS.
     /// - Parameters:
     ///   - identifier: session identifier
@@ -27,10 +34,17 @@ class StorageBackgroundEventsRegistry {
 
         return await withCheckedContinuation { (continuation: CheckedContinuation<Bool, Never>) in
             self.continuation = continuation
+            notifyWaiting(for: identifier)
         }
     }
 
-    // MARK: - Internal -
+    /// Notifies observes when waiting for continuation to be resumed.
+    /// - Parameters:
+    ///   - identifier: session identifier
+    private static func notifyWaiting(for identifier: String) {
+        let notificationCenter = notificationCenter ?? NotificationCenter.default
+        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) {
diff --git a/AmplifyPlugins/Storage/Tests/AWSS3StoragePluginTests/Support/Internal/StorageBackgroundEventsRegistryTests.swift b/AmplifyPlugins/Storage/Tests/AWSS3StoragePluginTests/Support/Internal/StorageBackgroundEventsRegistryTests.swift
index d2fe6f430a..92f9b8baa1 100644
--- a/AmplifyPlugins/Storage/Tests/AWSS3StoragePluginTests/Support/Internal/StorageBackgroundEventsRegistryTests.swift
+++ b/AmplifyPlugins/Storage/Tests/AWSS3StoragePluginTests/Support/Internal/StorageBackgroundEventsRegistryTests.swift
@@ -18,7 +18,25 @@ class StorageBackgroundEventsRegistryTests: XCTestCase {
         let otherIdentifier = UUID().uuidString
         StorageBackgroundEventsRegistry.register(identifier: identifier)
 
-        let done = asyncExpectation(description: "done", expectedFulfillmentCount: 2)
+        let notificationCenter = NotificationCenter()
+        StorageBackgroundEventsRegistry.notificationCenter = notificationCenter
+        defer {
+            StorageBackgroundEventsRegistry.notificationCenter = nil
+        }
+
+        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()
+            }
+        }
 
         Task {
             let handled = await StorageBackgroundEventsRegistry.handleEventsForBackgroundURLSession(identifier: identifier)
@@ -26,16 +44,23 @@ class StorageBackgroundEventsRegistryTests: XCTestCase {
             XCTAssertTrue(handled)
         }
 
+        await waitForExpectations([waiting])
+
+        let didContinue = await handleEvents(for: identifier)
+        XCTAssertTrue(didContinue)
+        await waitForExpectations([done])
+
+        let otherDone = asyncExpectation(description: "other done")
+
         Task {
             let otherHandled = await StorageBackgroundEventsRegistry.handleEventsForBackgroundURLSession(identifier: otherIdentifier)
-            await done.fulfill()
+            await otherDone.fulfill()
             XCTAssertFalse(otherHandled)
         }
 
-        handleEvents(for: identifier)
-        handleEvents(for: otherIdentifier)
-
-        await waitForExpectations([done])
+        let didNotContinue = await handleEvents(for: otherIdentifier)
+        XCTAssertFalse(didNotContinue)
+        await waitForExpectations([otherDone])
     }
 
     func testHandlingUnregisteredIdentifier() async throws {
@@ -55,10 +80,16 @@ class StorageBackgroundEventsRegistryTests: XCTestCase {
     }
 
     // Simulates URLSessionDelegate behavior
-    func handleEvents(for identifier: String) {
+    func handleEvents(for identifier: String) async -> Bool {
+        await Task.yield()
+
         if let continuation = StorageBackgroundEventsRegistry.getContinuation(for: identifier) {
             continuation.resume(returning: true)
             StorageBackgroundEventsRegistry.removeContinuation(for: identifier)
+            return true
+        } else {
+            print("No continuation for identifier: \(identifier)")
+            return false
         }
     }