From a0e916f93122abe75e0f25aa940303b553ef9cc9 Mon Sep 17 00:00:00 2001
From: Di Wu <github@wudi.me>
Date: Thu, 18 Apr 2024 14:54:21 -0700
Subject: [PATCH] fix(core): potential memory leak in WebSocketClient

---
 .../Core/AWSPluginsCore/WebSocket/WebSocketClient.swift     | 6 +++---
 1 file changed, 3 insertions(+), 3 deletions(-)

diff --git a/AmplifyPlugins/Core/AWSPluginsCore/WebSocket/WebSocketClient.swift b/AmplifyPlugins/Core/AWSPluginsCore/WebSocket/WebSocketClient.swift
index bc677744f8..cc1149ac27 100644
--- a/AmplifyPlugins/Core/AWSPluginsCore/WebSocket/WebSocketClient.swift
+++ b/AmplifyPlugins/Core/AWSPluginsCore/WebSocket/WebSocketClient.swift
@@ -265,7 +265,7 @@ extension WebSocketClient: URLSessionWebSocketDelegate {
 extension WebSocketClient {
     /// Monitor network status. Disconnect or reconnect when the network drops or comes back online.
     private func startNetworkMonitor() {
-        networkMonitor.publisher.sink(receiveValue: { stateChange in
+        networkMonitor.publisher.sink(receiveValue: { [weak self] stateChange in
             Task { [weak self] in
                 await self?.onNetworkStateChange(stateChange)
             }
@@ -304,7 +304,7 @@ extension WebSocketClient {
             return closeCode
         }
         .compactMap { $0 }
-        .sink(receiveCompletion: { _ in }) { closeCode in
+        .sink(receiveCompletion: { _ in }) { [weak self] closeCode in
             Task { [weak self] in await self?.retryOnCloseCode(closeCode) }
         }
         .store(in: &cancelables)
@@ -319,7 +319,7 @@ extension WebSocketClient {
             }
             return false
         }
-        .sink(receiveCompletion: { _ in }) { _ in
+        .sink(receiveCompletion: { _ in }) { [weak self] _ in
             Task { [weak self] in
                 await self?.retryWithJitter.reset()
             }