From ad538164c62c47670499a58e63552c14006c50c3 Mon Sep 17 00:00:00 2001
From: Di Wu <github@wudi.me>
Date: Tue, 4 Jun 2024 03:29:24 +0000
Subject: [PATCH] fix(datastore): reconnect the AppSyncRealTimeClient once
 connection is stale (#3734)

---
 .../AppSyncRealTimeClient.swift                      | 12 +++++++++++-
 1 file changed, 11 insertions(+), 1 deletion(-)

diff --git a/AmplifyPlugins/API/Sources/AWSAPIPlugin/AppSyncRealTimeClient/AppSyncRealTimeClient.swift b/AmplifyPlugins/API/Sources/AWSAPIPlugin/AppSyncRealTimeClient/AppSyncRealTimeClient.swift
index 25a695f9b5..9e70c658a5 100644
--- a/AmplifyPlugins/API/Sources/AWSAPIPlugin/AppSyncRealTimeClient/AppSyncRealTimeClient.swift
+++ b/AmplifyPlugins/API/Sources/AWSAPIPlugin/AppSyncRealTimeClient/AppSyncRealTimeClient.swift
@@ -298,6 +298,16 @@ actor AppSyncRealTimeClient: AppSyncRealTimeClientProtocol {
         .eraseToAnyPublisher()
     }
 
+    private func reconnect() async {
+        do {
+            log.debug("[AppSyncRealTimeClient] Reconnecting")
+            await disconnect()
+            try await connect()
+        } catch {
+            log.debug("[AppSyncRealTimeClient] Failed to reconnect, error: \(error)")
+        }
+    }
+
     private static func decodeAppSyncRealTimeResponseError(_ data: JSONValue?) -> [Error] {
         let knownAppSyncRealTimeRequestErorrs =
             Self.decodeAppSyncRealTimeRequestError(data)
@@ -414,7 +424,7 @@ extension AppSyncRealTimeClient {
             .sink(receiveValue: {
                 self.log.debug("[AppSyncRealTimeClient] KeepAlive timed out, disconnecting")
                 Task { [weak self] in
-                    await self?.disconnect()
+                    await self?.reconnect()
                 }.toAnyCancellable.store(in: &self.cancellables)
             })
             .store(in: &cancellablesBindToConnection)