fix(API): DataRace - subscription cancel and OperationTaskMapper on reset (#1684)

Author: lawmicha

AmplifyPlugins/API/APICategoryPlugin.xcodeproj/xcshareddata/xcschemes/AWSAPICategoryPlugin.xcscheme

@@ -90,6 +90,7 @@
       buildConfiguration = "Debug"
       selectedDebuggerIdentifier = "Xcode.DebuggerFoundation.Debugger.LLDB"
       selectedLauncherIdentifier = "Xcode.DebuggerFoundation.Launcher.LLDB"
+      enableThreadSanitizer = "YES"
       launchStyle = "0"
       useCustomWorkingDirectory = "NO"
       ignoresPersistentStateOnLaunch = "NO"

AmplifyPlugins/API/AWSAPICategoryPlugin/AWSAPIPlugin+Resettable.swift

@@ -13,8 +13,6 @@ extension AWSAPIPlugin: Resettable {
     public func reset(onComplete: @escaping BasicClosure) {
         mapper.reset()
 
-        mapper = OperationTaskMapper()
-
         let waitForReset = DispatchSemaphore(value: 0)
         session.reset { waitForReset.signal() }
         waitForReset.wait()

AmplifyPlugins/API/AWSAPICategoryPlugin/Operation/AWSGraphQLSubscriptionOperation.swift

@@ -21,6 +21,8 @@ final public class AWSGraphQLSubscriptionOperation<R: Decodable>: GraphQLSubscri
     var subscriptionItem: SubscriptionItem?
     var apiAuthProviderFactory: APIAuthProviderFactory
 
+    private let subscriptionQueue = DispatchQueue(label: "AWSGraphQLSubscriptionOperation.subscriptionQueue")
+
     init(request: GraphQLOperationRequest<R>,
          pluginConfig: AWSAPICategoryPluginConfiguration,
          subscriptionConnectionFactory: SubscriptionConnectionFactory,
@@ -42,11 +44,14 @@ final public class AWSGraphQLSubscriptionOperation<R: Decodable>: GraphQLSubscri
     }
 
     override public func cancel() {
-        if let subscriptionItem = subscriptionItem, let subscriptionConnection = subscriptionConnection {
-            subscriptionConnection.unsubscribe(item: subscriptionItem)
-            let subscriptionEvent = SubscriptionEvent<GraphQLResponse<R>>.connection(.disconnected)
-            dispatchInProcess(data: subscriptionEvent)
+        subscriptionQueue.sync {
+            if let subscriptionItem = subscriptionItem, let subscriptionConnection = subscriptionConnection {
+                subscriptionConnection.unsubscribe(item: subscriptionItem)
+                let subscriptionEvent = SubscriptionEvent<GraphQLResponse<R>>.connection(.disconnected)
+                dispatchInProcess(data: subscriptionEvent)
+            }
         }
+
         super.cancel()
         dispatch(result: .successfulVoid)
         finish()
@@ -90,27 +95,28 @@ final public class AWSGraphQLSubscriptionOperation<R: Decodable>: GraphQLSubscri
         let pluginOptions = request.options.pluginOptions as? AWSPluginOptions
 
         // Retrieve the subscription connection
-        do {
-            subscriptionConnection = try subscriptionConnectionFactory
-                .getOrCreateConnection(for: endpointConfig,
-                                       authService: authService,
-                                       authType: pluginOptions?.authType,
-                                       apiAuthProviderFactory: apiAuthProviderFactory)
-        } catch {
-            let error = APIError.operationError("Unable to get connection for api \(endpointConfig.name)", "", error)
-            dispatch(result: .failure(error))
-            finish()
-            return
+        subscriptionQueue.sync {
+            do {
+                subscriptionConnection = try subscriptionConnectionFactory
+                    .getOrCreateConnection(for: endpointConfig,
+                                              authService: authService,
+                                              authType: pluginOptions?.authType,
+                                              apiAuthProviderFactory: apiAuthProviderFactory)
+            } catch {
+                let error = APIError.operationError("Unable to get connection for api \(endpointConfig.name)", "", error)
+                dispatch(result: .failure(error))
+                finish()
+                return
+            }
+
+            // Create subscription
+
+            subscriptionItem = subscriptionConnection?.subscribe(requestString: request.document,
+                                                                 variables: request.variables,
+                                                                 eventHandler: { [weak self] event, _ in
+                self?.onAsyncSubscriptionEvent(event: event)
+            })
         }
-
-        // Create subscription
-
-        subscriptionItem = subscriptionConnection?.subscribe(requestString: request.document,
-                                                             variables: request.variables,
-                                                             eventHandler: { [weak self] event, _ in
-            self?.onAsyncSubscriptionEvent(event: event)
-        })
-
     }
 
     private func onAsyncSubscriptionEvent(event: SubscriptionItemEvent) {

AmplifyPlugins/API/AWSAPICategoryPlugin/URLSessionBehavior/OperationTaskMapper.swift

@@ -69,6 +69,10 @@ class OperationTaskMapper {
         OperationTaskMapper.concurrencyQueue.sync {
             operations.values.forEach { $0.cancelOperation() }
             tasks.values.forEach { $0.cancel() }
+            operations = [:]
+            tasks = [:]
+            operationIdsByTaskId = [:]
+            taskIdsByOperationId = [:]
         }
     }
 

AmplifyPlugins/API/AWSAPICategoryPluginTests/Operation/AWSGraphQLSubscriptionOperationCancelTests.swift

@@ -170,8 +170,13 @@ class AWSGraphQLSubscriptionOperationCancelTests: XCTestCase {
         let connectionCreation = expectation(description: "connection factory called")
         let mockSubscriptionConnectionFactory = MockSubscriptionConnectionFactory(onGetOrCreateConnection: { _, _, _, _ in
             connectionCreation.fulfill()
-            sleep(5)
-            throw APIError.invalidConfiguration("something went wrong", "", nil)
+            return MockSubscriptionConnection(onSubscribe: { (_, _, eventHandler) -> SubscriptionItem in
+                let item = SubscriptionItem(requestString: "", variables: nil, eventHandler: { _, _ in
+                })
+                eventHandler(.connection(.connecting), item)
+                return item
+            }, onUnsubscribe: {_ in
+            })
         })
 
         setUp(subscriptionConnectionFactory: mockSubscriptionConnectionFactory)
@@ -184,7 +189,7 @@ class AWSGraphQLSubscriptionOperationCancelTests: XCTestCase {
         let receivedFailure = expectation(description: "Received failure")
         receivedFailure.isInverted = true
         let receivedValue = expectation(description: "Received value for connecting")
-        receivedValue.isInverted = true
+        receivedValue.assertForOverFulfill = false
 
         let valueListener: GraphQLSubscriptionOperation<JSONValue>.InProcessListener = { _ in
             receivedValue.fulfill()

AmplifyPlugins/API/Podfile.lock

@@ -116,4 +116,4 @@ SPEC CHECKSUMS:
 
 PODFILE CHECKSUM: 54e87158e45936fe60756d4a417b89eb64ea1c51
 
-COCOAPODS: 1.10.1
+COCOAPODS: 1.11.2