From 73c1505cf741d0f585156c4934cff1ed04ce9256 Mon Sep 17 00:00:00 2001
From: Nicholas Barnett <n-barnett@hotmail.com>
Date: Fri, 26 Apr 2024 10:47:55 -0500
Subject: [PATCH] fix(socket-errors): fixed (#1598)

* fix(socket-errors): fixed
---
 .../Sockets/use-stacks-api-socket-client.ts   | 47 +++++++++++--------
 .../BlockList/Sockets/useSubscribeBlocks.ts   | 13 ++---
 2 files changed, 34 insertions(+), 26 deletions(-)

diff --git a/src/app/_components/BlockList/Sockets/use-stacks-api-socket-client.ts b/src/app/_components/BlockList/Sockets/use-stacks-api-socket-client.ts
index fdb5f231e..5143c7313 100644
--- a/src/app/_components/BlockList/Sockets/use-stacks-api-socket-client.ts
+++ b/src/app/_components/BlockList/Sockets/use-stacks-api-socket-client.ts
@@ -4,53 +4,60 @@ import { StacksApiSocketClient } from '@stacks/blockchain-api-client';
 
 export interface StacksApiSocketClientInfo {
   client: StacksApiSocketClient | null;
-  connect: (handleOnConnect?: (client: StacksApiSocketClient) => void) => void;
+  connect: (
+    handleOnConnect?: (client: StacksApiSocketClient) => void,
+    handleError?: (client: StacksApiSocketClient | null) => void
+  ) => void;
   disconnect: () => void;
 }
 
 export function useStacksApiSocketClient(apiUrl: string): StacksApiSocketClientInfo {
   const [socketClient, setSocketClient] = useState<StacksApiSocketClient | null>(null);
-  const socketUrlTracker = useRef<string | null>(null);
   const isSocketClientConnecting = useRef(false);
 
+  const disconnect = useCallback(() => {
+    if (socketClient?.socket.connected) {
+      socketClient?.socket.removeAllListeners();
+      socketClient?.socket.close();
+      setSocketClient(null);
+      isSocketClientConnecting.current = false;
+    }
+  }, [socketClient]);
+
   const connect = useCallback(
-    async (handleOnConnect?: (client: StacksApiSocketClient) => void) => {
+    async (
+      handleOnConnect?: (client: StacksApiSocketClient) => void,
+      handleError?: (client: StacksApiSocketClient | null) => void
+    ) => {
       if (!apiUrl) return;
       if (socketClient?.socket.connected || isSocketClientConnecting.current) {
         return;
       }
       try {
         isSocketClientConnecting.current = true;
-        const socketUrl = apiUrl;
-        socketUrlTracker.current = socketUrl;
-        const client = await StacksApiSocketClient.connect({ url: socketUrl });
+        const client = await StacksApiSocketClient.connect({ url: apiUrl });
         client.socket.on('connect', () => {
           setSocketClient(client);
           handleOnConnect?.(client);
-          isSocketClientConnecting.current = false;
         });
         client.socket.on('disconnect', () => {
-          setSocketClient(null);
-          isSocketClientConnecting.current = false;
+          client.socket.removeAllListeners();
+          client.socket.close();
+          disconnect();
         });
         client.socket.on('connect_error', error => {
-          setSocketClient(null);
-          isSocketClientConnecting.current = false;
+          client.socket.removeAllListeners();
+          client.socket.close();
+          disconnect();
+          handleError?.(client);
         });
       } catch (error) {
-        setSocketClient(null);
-        isSocketClientConnecting.current = false;
+        disconnect();
       }
     },
-    [apiUrl, socketClient]
+    [apiUrl, socketClient, disconnect]
   );
 
-  const disconnect = useCallback(() => {
-    if (socketClient?.socket.connected) {
-      socketClient?.socket.close();
-    }
-  }, [socketClient]);
-
   return {
     client: socketClient,
     connect,
diff --git a/src/app/_components/BlockList/Sockets/useSubscribeBlocks.ts b/src/app/_components/BlockList/Sockets/useSubscribeBlocks.ts
index 4f7dddcad..d4e50a212 100644
--- a/src/app/_components/BlockList/Sockets/useSubscribeBlocks.ts
+++ b/src/app/_components/BlockList/Sockets/useSubscribeBlocks.ts
@@ -11,11 +11,12 @@ interface Subscription {
 
 export function useSubscribeBlocks(
   liveUpdates: boolean,
-  handleBlock: (block: NakamotoBlock | Block) => void
+  handleBlock: (block: NakamotoBlock | Block) => void,
+  handleError?: (client: StacksApiSocketClient | null) => void
 ) {
   const subscription = useRef<Subscription | undefined>(undefined);
   const { stacksApiSocketClientInfo } = useGlobalContext();
-  const { client, connect, disconnect } = stacksApiSocketClientInfo || {};
+  const { connect, disconnect } = stacksApiSocketClientInfo || {};
 
   useEffect(() => {
     const subscribe = async (client: StacksApiSocketClient) => {
@@ -28,15 +29,15 @@ export function useSubscribeBlocks(
       });
     };
 
-    if (liveUpdates && !client?.socket.connected) {
-      connect?.(client => subscribe(client));
+    if (liveUpdates) {
+      connect?.(client => subscribe(client), handleError);
     }
-    if (!liveUpdates && client?.socket.connected) {
+    if (!liveUpdates) {
       disconnect?.();
     }
     return () => {
       disconnect?.();
     };
-  }, [client, handleBlock, connect, liveUpdates, disconnect]);
+  }, [handleBlock, connect, liveUpdates, disconnect, handleError]);
   return subscription;
 }