fix(Storage): Addressing issues with race conditions (#2452)

Author: sebaland

Amplify/Core/Support/AmplifyTask+OperationTaskAdapters.swift

@@ -69,8 +69,14 @@ public class AmplifyInProcessReportingOperationTaskAdapter<Request: AmplifyOpera
     public init(operation: AmplifyInProcessReportingOperation<Request, InProcess, Success, Failure>) {
         self.operation = operation
         self.childTask = ChildTask(parent: operation)
-        resultToken = operation.subscribe(resultListener: resultListener)
-        inProcessToken = operation.subscribe(inProcessListener: inProcessListener)
+        resultToken = operation.subscribe(resultListener: { [weak self] result in
+            guard let self = self else { return }
+            self.resultListener(result)
+        })
+        inProcessToken = operation.subscribe(inProcessListener: { [weak self] inProcess in
+            guard let self = self else { return }
+            self.inProcessListener(inProcess)
+        })
     }
 
     deinit {

Amplify/Core/Support/AtomicDictionary.swift

@@ -7,41 +7,43 @@
 
 import Foundation
 
-actor AtomicDictionary<Key: Hashable, Value> {
+final class AtomicDictionary<Key: Hashable, Value> {
+    private let lock: NSLocking
     private var value: [Key: Value]
 
     init(initialValue: [Key: Value] = [Key: Value]()) {
+        self.lock = NSLock()
         self.value = initialValue
     }
 
     var count: Int {
-        value.count
+        lock.execute { value.count }
     }
 
     var keys: [Key] {
-        Array(value.keys)
+        lock.execute { Array(value.keys) }
     }
 
     var values: [Value] {
-        Array(value.values)
+        lock.execute { Array(value.values) }
     }
 
     // MARK: - Functions
 
     func getValue(forKey key: Key) -> Value? {
-        value[key]
+        lock.execute { value[key] }
     }
 
     func removeAll() {
-        value = [:]
+        lock.execute { value = [:] }
     }
 
     @discardableResult
     func removeValue(forKey key: Key) -> Value? {
-        value.removeValue(forKey: key)
+        return lock.execute { value.removeValue(forKey: key) }
     }
 
     func set(value: Value, forKey key: Key) {
-        self.value[key] = value
+        lock.execute { self.value[key] = value }
     }
 }

Amplify/DefaultPlugins/AWSHubPlugin/AWSHubPlugin.swift

@@ -60,18 +60,14 @@ final public class AWSHubPlugin: HubCategoryPlugin {
                        isIncluded filter: HubFilter? = nil,
                        listener: @escaping HubListener) -> UnsubscribeToken {
         let filteredListener = FilteredListener(for: channel, filter: filter, listener: listener)
-        Task {
-            await dispatcher.insert(filteredListener)
-        }
+        dispatcher.insert(filteredListener)
 
         let unsubscribeToken = UnsubscribeToken(channel: channel, id: filteredListener.id)
         return unsubscribeToken
     }
 
     public func removeListener(_ token: UnsubscribeToken) {
-        Task {
-            await dispatcher.removeListener(withId: token.id)
-        }
+        dispatcher.removeListener(withId: token.id)
     }
 
     // MARK: - Custom Plugin methods
@@ -80,8 +76,8 @@ final public class AWSHubPlugin: HubCategoryPlugin {
     ///
     /// - Parameter token: The UnsubscribeToken of the listener to check
     /// - Returns: True if the dispatcher has a listener registered with `token`
-    public func hasListener(withToken token: UnsubscribeToken) async -> Bool {
-        return await dispatcher.hasListener(withId: token.id)
+    public func hasListener(withToken token: UnsubscribeToken) -> Bool {
+        return dispatcher.hasListener(withId: token.id)
     }
 
 }

Amplify/DefaultPlugins/AWSHubPlugin/Internal/HubChannelDispatcher.swift

@@ -25,22 +25,22 @@ final class HubChannelDispatcher {
     ///
     /// - Parameter id: The ID of the listener to check
     /// - Returns: True if the dispatcher has a listener registered with `id`
-    func hasListener(withId id: UUID) async -> Bool {
-        return await listenersById.getValue(forKey: id) != nil
+    func hasListener(withId id: UUID) -> Bool {
+        return listenersById.getValue(forKey: id) != nil
     }
 
     /// Inserts `listener` into the `listenersById` dictionary by its ID
     ///
     /// - Parameter listener: The listener to add
-    func insert(_ listener: FilteredListener) async {
-        await listenersById.set(value: listener, forKey: listener.id)
+    func insert(_ listener: FilteredListener) {
+        listenersById.set(value: listener, forKey: listener.id)
     }
 
     /// Removes the listener identified by `id` from the `listeners` dictionary
     ///
     /// - Parameter id: The ID of the listener to remove
-    func removeListener(withId id: UUID) async {
-        await listenersById.removeValue(forKey: id)
+    func removeListener(withId id: UUID) {
+        listenersById.removeValue(forKey: id)
     }
 
     /// Dispatches `payload` to all listeners on `channel`
@@ -63,7 +63,7 @@ final class HubChannelDispatcher {
     /// after you issue an `await Amplify.reset()`, you may wish to add additional sleep around your code
     /// that calls `await Amplify.reset()`.
     func destroy() async {
-        await listenersById.removeAll()
+        listenersById.removeAll()
         messageQueue.cancelAllOperations()
         await withCheckedContinuation { continuation in
             messageQueue.addBarrierBlock {
@@ -75,16 +75,14 @@ final class HubChannelDispatcher {
 
 extension HubChannelDispatcher: HubDispatchOperationDelegate {
     var listeners: [FilteredListener] {
-        get async {
-            return Array(await listenersById.values)
-        }
+        return Array(listenersById.values)
     }
 }
 
 protocol HubDispatchOperationDelegate: AnyObject {
     /// Used to let a dispatch operation retrieve the list of listeners at the time of invocation, rather than the time
     /// of queuing.
-    var listeners: [FilteredListener] { get async }
+    var listeners: [FilteredListener] { get }
 }
 
 final class HubDispatchOperation: Operation {
@@ -121,14 +119,12 @@ final class HubDispatchOperation: Operation {
             return
         }
 
-        Task {
-            guard let listeners = await delegate?.listeners else {
-                return
-            }
-
-            let dispatcher = SerialDispatcher(channel: channel, payload: payload)
-            dispatcher.dispatch(to: listeners)
+        guard let listeners = delegate?.listeners else {
+            return
         }
+        
+        let dispatcher = SerialDispatcher(channel: channel, payload: payload)
+        dispatcher.dispatch(to: listeners)
     }
 
 }

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

@@ -40,12 +40,10 @@ class StorageBackgroundEventsRegistryTests: XCTestCase {
             XCTAssertFalse(otherHandled)
         }
 
-        Task {
-            handleEvents(for: identifier)
-            handleEvents(for: otherIdentifier)
-        }
-
         await waitForExpectations([done])
+
+        handleEvents(for: identifier)
+        handleEvents(for: otherIdentifier)
     }
 
     func testHandlingUnregisteredIdentifier() async throws {