From 2cdaf7fa7e975843e74b7b59cd4e5db00b01c267 Mon Sep 17 00:00:00 2001
From: Abhash Kumar Singh <thisisabhash@gmail.com>
Date: Mon, 9 Dec 2024 15:58:23 -0800
Subject: [PATCH] fix(predictions): use serial queue for WebSocketSession
 delegate queue

---
 .../Liveness/Service/FaceLivenessSession.swift       | 12 ++++++++----
 .../Liveness/Service/WebSocketSession.swift          |  8 +++++++-
 2 files changed, 15 insertions(+), 5 deletions(-)

diff --git a/AmplifyPlugins/Predictions/AWSPredictionsPlugin/Liveness/Service/FaceLivenessSession.swift b/AmplifyPlugins/Predictions/AWSPredictionsPlugin/Liveness/Service/FaceLivenessSession.swift
index 092532d8e3..1723f0c688 100644
--- a/AmplifyPlugins/Predictions/AWSPredictionsPlugin/Liveness/Service/FaceLivenessSession.swift
+++ b/AmplifyPlugins/Predictions/AWSPredictionsPlugin/Liveness/Service/FaceLivenessSession.swift
@@ -71,7 +71,9 @@ public final class FaceLivenessSession: LivenessService {
     }
 
     public func closeSocket(with code: URLSessionWebSocketTask.CloseCode) {
-        websocket.close(with: code)
+        livenessServiceDispatchQueue.async {
+            self.websocket.close(with: code)
+        }
     }
 
     public func initializeLivenessStream(withSessionID sessionID: String, userAgent: String = "") throws {
@@ -89,14 +91,16 @@ public final class FaceLivenessSession: LivenessService {
 
         savedURLForReconnect = url
         let signedConnectionURL = signer.sign(url: url)
-        websocket.open(url: signedConnectionURL)
+        livenessServiceDispatchQueue.async {
+            self.websocket.open(url: signedConnectionURL)
+        }
     }
 
     public func send<T>(
         _ event: LivenessEvent<T>,
         eventDate: @escaping () -> Date = Date.init
     ) {
-        livenessServiceDispatchQueue.sync {
+        livenessServiceDispatchQueue.async {
             let encodedPayload = self.eventStreamEncoder.encode(
                 payload: event.payload,
                 headers: [
@@ -107,7 +111,7 @@ public final class FaceLivenessSession: LivenessService {
             )
 
             let dateForSigning: Date
-            if let serverDate = serverDate {
+            if let serverDate = self.serverDate {
                 dateForSigning = serverDate
             } else {
                 dateForSigning = eventDate()
diff --git a/AmplifyPlugins/Predictions/AWSPredictionsPlugin/Liveness/Service/WebSocketSession.swift b/AmplifyPlugins/Predictions/AWSPredictionsPlugin/Liveness/Service/WebSocketSession.swift
index f900f3dfc3..5888cde6e4 100644
--- a/AmplifyPlugins/Predictions/AWSPredictionsPlugin/Liveness/Service/WebSocketSession.swift
+++ b/AmplifyPlugins/Predictions/AWSPredictionsPlugin/Liveness/Service/WebSocketSession.swift
@@ -15,13 +15,19 @@ final class WebSocketSession {
     private var receiveMessage: ((Result<URLSessionWebSocketTask.Message, Error>) -> WebSocketMessageResult)?
     private var onSocketClosed: ((URLSessionWebSocketTask.CloseCode) -> Void)?
     private var onServerDateReceived: ((Date?) -> Void)?
+    private let delegateQueue: OperationQueue
 
     init() {
+        self.delegateQueue = OperationQueue()
+        self.delegateQueue.maxConcurrentOperationCount = 1
+        self.delegateQueue.qualityOfService = .userInteractive
+        
         self.urlSessionWebSocketDelegate = Delegate()
+
         self.session = URLSession(
             configuration: .default,
             delegate: urlSessionWebSocketDelegate,
-            delegateQueue: .init()
+            delegateQueue: delegateQueue
         )
     }