fix(API): Memory leak in GraphQLOperation (#1562)

* Fix memory leaks

* chore(api): prevent unsafe usage of removePair

Co-authored-by: Diego Costantino <[email protected]>

Author: dnys1

AmplifyPlugins/API/AWSAPICategoryPlugin/Operation/AWSAPIOperation+APIOperation.swift

@@ -45,6 +45,8 @@ extension AWSRESTOperation: APIOperation {
             return
         }
 
+        mapper.removePair(for: self)
+
         let apiOperationResponse = APIOperationResponse(error: error, response: response)
         do {
             try apiOperationResponse.validate()

AmplifyPlugins/API/AWSAPICategoryPlugin/Operation/AWSGraphQLOperation+APIOperation.swift

@@ -45,6 +45,8 @@ extension AWSGraphQLOperation: APIOperation {
             return
         }
 
+        mapper.removePair(for: self)
+
         let apiOperationResponse = APIOperationResponse(error: error, response: response)
         do {
             try apiOperationResponse.validate()

AmplifyPlugins/API/AWSAPICategoryPlugin/URLSessionBehavior/OperationTaskMapper.swift

@@ -74,15 +74,14 @@ class OperationTaskMapper {
 
     /// Not inherently thread safe--this must be called from `concurrencyQueue`
     private func removePair(operationId: UUID?, taskId: Int?) {
-        OperationTaskMapper.concurrencyQueue.sync {
-            if let operationId = operationId {
-                operations[operationId] = nil
-                taskIdsByOperationId[operationId] = nil
-            }
-            if let taskId = taskId {
-                tasks[taskId] = nil
-                operationIdsByTaskId[taskId] = nil
-            }
+        dispatchPrecondition(condition: .onQueue(Self.concurrencyQueue))
+        if let operationId = operationId {
+            operations[operationId] = nil
+            taskIdsByOperationId[operationId] = nil
+        }
+        if let taskId = taskId {
+            tasks[taskId] = nil
+            operationIdsByTaskId[taskId] = nil
         }
     }
 }

AmplifyPlugins/API/AWSAPICategoryPluginTests/Operation/AWSGraphQLOperationTests.swift

@@ -10,33 +10,32 @@ import XCTest
 @testable import AmplifyTestCommon
 @testable import AWSAPICategoryPlugin
 
-class AWSGraphQLOperationTests: XCTestCase {
-
-    // TODO: Complete the implementation below
-
-    //    override func setUp() {
-    //    }
-    //
-    //    override func tearDown() {
-    //    }
-    //
-    //    func testGraphQLOperationSuccess() {
-    //        XCTFail("Not yet implemented.")
-    //    }
-    //
-    //    func testGraphQLOperationValidationError() {
-    //        XCTFail("Not yet implemented.")
-    //    }
-    //
-    //    func testGraphQLOperationEndpointConfigurationError() {
-    //        XCTFail("Not yet implemented.")
-    //    }
-    //
-    //    func testGraphQLOperationJSONSerializationError() {
-    //        XCTFail("Not yet implemented.")
-    //    }
-    //
-    //    func testGraphQLOperationInterceptorError() {
-    //        XCTFail("Not yet implemented.")
-    //    }
+@available(iOS 13.0, *)
+class AWSGraphQLOperationTests: AWSAPICategoryPluginTestBase {
+
+    /// Tests that upon completion, the operation is removed from the task mapper.
+    func testOperationCleanup() {
+        let request = GraphQLRequest(apiName: apiName,
+                                     document: testDocument,
+                                     variables: nil,
+                                     responseType: JSONValue.self)
+
+        let operation = apiPlugin.query(request: request, listener: nil)
+
+        guard let operation = operation as? AWSGraphQLOperation else {
+            XCTFail("Operation is not an AWSGraphQLOperation")
+            return
+        }
+
+        let receivedCompletion = expectation(description: "Received completion")
+        let sink = operation.resultPublisher.sink { _ in
+            receivedCompletion.fulfill()
+        } receiveValue: { _ in }
+        defer { sink.cancel() }
+
+        wait(for: [receivedCompletion], timeout: 1)
+        let task = operation.mapper.task(for: operation)
+        XCTAssertNil(task)
+    }
+
 }