Add async to Websocket internal completion handlers

Author: AnthonyMDev

apollo-ios/Sources/ApolloWebSocket/OperationMessage.swift

@@ -57,24 +57,24 @@ final class OperationMessage {
     self.serialized = serialized
   }
 
-  func parse(handler: (ParseHandler) -> Void) {
+  func parse(handler: (ParseHandler) async -> Void) async {
     guard let serialized = self.serialized else {
-      handler(ParseHandler(nil,
-                           nil,
-                           nil,
-                           WebSocketError(payload: nil,
-                                          error: nil,
-                                          kind: .serializedMessageError)))
+      await handler(ParseHandler(nil,
+                                 nil,
+                                 nil,
+                                 WebSocketError(payload: nil,
+                                                error: nil,
+                                                kind: .serializedMessageError)))
       return
     }
 
     guard let data = self.serialized?.data(using: (.utf8) ) else {
-      handler(ParseHandler(nil,
-                           nil,
-                           nil,
-                           WebSocketError(payload: nil,
-                                          error: nil,
-                                          kind: .unprocessedMessage(serialized))))
+      await handler(ParseHandler(nil,
+                                 nil,
+                                 nil,
+                                 WebSocketError(payload: nil,
+                                                error: nil,
+                                                kind: .unprocessedMessage(serialized))))
       return
     }
 
@@ -89,18 +89,18 @@ final class OperationMessage {
       type = json["type"] as? String
       payload = json["payload"] as? JSONObject
 
-      handler(ParseHandler(type,
-                           id,
-                           payload,
-                           nil))
+      await handler(ParseHandler(type,
+                                 id,
+                                 payload,
+                                 nil))
     }
     catch {
-      handler(ParseHandler(type,
-                           id,
-                           payload,
-                           WebSocketError(payload: payload,
-                                          error: error,
-                                          kind: .unprocessedMessage(serialized))))
+      await handler(ParseHandler(type,
+                                 id,
+                                 payload,
+                                 WebSocketError(payload: payload,
+                                                error: error,
+                                                kind: .unprocessedMessage(serialized))))
     }
   }
 }

apollo-ios/Sources/ApolloWebSocket/WebSocketTask.swift

@@ -16,7 +16,7 @@ final class WebSocketTask<Operation: GraphQLOperation>: Sendable, Cancellable {
   /// - Parameter completionHandler: A completion handler to fire when the operation has a result.
   init(_ ws: WebSocketTransport,
        _ operation: Operation,
-       _ completionHandler: @escaping @Sendable (_ result: Result<JSONObject, any Error>) -> Void) {
+       _ completionHandler: @escaping @Sendable (_ result: Result<JSONObject, any Error>) async -> Void) {
     sequenceNumber = ws.sendHelper(operation: operation, resultHandler: completionHandler)
     transport = ws
   }

apollo-ios/Sources/ApolloWebSocket/WebSocketTransport.swift

@@ -50,7 +50,7 @@ public class WebSocketTransport: @unchecked Sendable {
   private var queue: [Int: String] = [:]
 
   @Atomic
-  private var subscribers = [String: (Result<JSONObject, any Error>) -> Void]()
+  private var subscribers = [String: (Result<JSONObject, any Error>) async -> Void]()
   @Atomic
   private var subscriptions : [String: String] = [:]
   let processingQueue = DispatchQueue(label: "com.apollographql.WebSocketTransport")
@@ -187,8 +187,8 @@ public class WebSocketTransport: @unchecked Sendable {
     return websocket.write(ping: data, completion: completionHandler)
   }
 
-  private func processMessage(text: String) {
-    OperationMessage(serialized: text).parse { parseHandler in
+  private func processMessage(text: String) async {
+    await OperationMessage(serialized: text).parse { parseHandler in
       guard
         let type = parseHandler.type,
         let messageType = OperationMessage.Types(rawValue: type) else {
@@ -215,14 +215,14 @@ public class WebSocketTransport: @unchecked Sendable {
         // subscriber probably unsubscribed.
         if let responseHandler = subscribers[id] {
           if let payload = parseHandler.payload {
-            responseHandler(.success(payload))
+            await responseHandler(.success(payload))
           } else if let error = parseHandler.error {
-            responseHandler(.failure(error))
+            await responseHandler(.failure(error))
           } else {
             let websocketError = WebSocketError(payload: parseHandler.payload,
                                                 error: parseHandler.error,
                                                 kind: .neitherErrorNorPayloadReceived)
-            responseHandler(.failure(websocketError))
+            await responseHandler(.failure(websocketError))
           }
         }
       case .complete:
@@ -266,8 +266,10 @@ public class WebSocketTransport: @unchecked Sendable {
   }
 
   private func notifyErrorAllHandlers(_ error: any Error) {
-    for (_, handler) in subscribers {
-      handler(.failure(error))
+    Task {
+      for (_, handler) in subscribers {
+        await handler(.failure(error))
+      }
     }
   }
 
@@ -335,10 +337,13 @@ public class WebSocketTransport: @unchecked Sendable {
     self.websocket.delegate = nil
   }
 
-  func sendHelper<Operation: GraphQLOperation>(operation: Operation, resultHandler: @escaping @Sendable (_ result: Result<JSONObject, any Error>) -> Void) -> String? {
-    let body = config.requestBodyCreator.requestBody(for: operation,
-                                                     sendQueryDocument: true,
-                                                     autoPersistQuery: false)
+  func sendHelper<Operation: GraphQLOperation>(operation: Operation, resultHandler: @escaping @Sendable (_ result: Result<JSONObject, any Error>) async -> Void) -> String? {
+    let body = config.requestBodyCreator.requestBody(
+      for: operation,
+      sendQueryDocument: true,
+      autoPersistQuery: false,
+      clientAwarenessMetadata: config.clientAwarenessMetadata
+    )
     let identifier = config.operationMessageIdCreator.requestId()
 
     let messageType: OperationMessage.Types
@@ -460,36 +465,40 @@ extension WebSocketTransport: NetworkTransport {
 
     #warning("TODO: stream never finishes. WebSocketTask does not report subscription termination.")
     return AsyncThrowingStream { continuation in
-      let task = WebSocketTask(self, operation) { [weak store, contextIdentifier] result in
-        Task {
+      let wsTask = WebSocketTask(self, operation) { [weak store, contextIdentifier] result in
+        do {
           try Task.checkCancellation()
-
+          
           let jsonBody = try result.get()
           let response = GraphQLResponse(operation: operation, body: jsonBody)
-
+          
           if let store = store {
             let (graphQLResult, parsedRecords) = try await response.parseResult()
-
+            
             try Task.checkCancellation()
-
+            
             guard let records = parsedRecords else {
               continuation.yield(graphQLResult)
               return
             }
-
+            
             try await store.publish(records: records, identifier: contextIdentifier)
             continuation.yield(graphQLResult)
-
+            
           } else {
             let graphQLResult = try await response.parseResultFast()
             try Task.checkCancellation()
-
+            
             continuation.yield(graphQLResult)
           }
+        } catch {
+          continuation.finish(throwing: error)
         }
-      }
+      }      
 
-      continuation.onTermination = { _ in task.cancel() }
+      continuation.onTermination = {
+        _ in wsTask.cancel()
+      }
     }
   }
 }
@@ -592,7 +601,9 @@ extension WebSocketTransport: WebSocketClientDelegate {
   }
 
   public func websocketDidReceiveMessage(socket: any WebSocketClient, text: String) {
-    self.processMessage(text: text)
+    Task {
+      await self.processMessage(text: text)
+    }
   }
 
   public func websocketDidReceiveData(socket: any WebSocketClient, data: Data) {