Updates

components/hooks/useWebSocket.tsx

@@ -6,6 +6,7 @@ export type UseWebsocketHookParams = {
     disableReconnect?: boolean;
     /** Interval to ping the websocket. So that the connection doesn't go down. Default 30000ms (30 seconds)  */
     keepAliveDuration?: number;
+    /** Interval in ms to force-refresh the connection */
     refreshConnection?: number;
 };
 
@@ -16,10 +17,10 @@ export const WebSocketEventNames = ["wsDataEvent", "wsMessageEvent"] as const;
  * @event wsDataEvent Listen for event named `wsDataEvent` on `window` to receive Data events
  * @event wsMessageEvent Listen for event named `wsMessageEvent` on `window` to receive Message events
  *
- * @example window.addEventLiatener("wsDataEvent", (e)=>{
+ * @example window.addEventListener("wsDataEvent", (e)=>{
  *    console.log(e.detail.data) // type object
  * })
- * @example window.addEventLiatener("wsMessageEvent", (e)=>{
+ * @example window.addEventListener("wsMessageEvent", (e)=>{
  *      console.log(e.detail.message) // type string
  * })
  */
@@ -34,22 +35,29 @@ export default function useWebSocket<
 }: UseWebsocketHookParams) {
     const DEBOUNCE = debounce || 500;
     const KEEP_ALIVE_DURATION = keepAliveDuration || 1000 * 30;
-    const KEEP_ALIVE_TIMEOUT = 1000 * 60 * 3;
-
     const KEEP_ALIVE_MESSAGE = "twui::ping";
 
-    let uptime = 0;
+    const tries = useRef(0);
-    let tries = useRef(0);
 
-    // const queue: string[] = [];
+    // Refs to avoid stale closures in callbacks
+    const urlRef = useRef(url);
+    const disableReconnectRef = useRef(disableReconnect);
+    const keepAliveDurationRef = useRef(KEEP_ALIVE_DURATION);
+
+    React.useEffect(() => {
+        urlRef.current = url;
+        disableReconnectRef.current = disableReconnect;
+        keepAliveDurationRef.current = KEEP_ALIVE_DURATION;
+    });
 
     const msgInterval = useRef<any>(null);
     const sendInterval = useRef<any>(null);
     const keepAliveInterval = useRef<any>(null);
+    const refreshInterval = useRef<any>(null);
+    const reconnectTimeout = useRef<any>(null);
 
-    const [socket, setSocket] = React.useState<WebSocket | undefined>(
+    const [socket, setSocket] = React.useState<WebSocket | undefined>(undefined);
-        undefined
+    const socketRef = useRef<WebSocket | undefined>(undefined);
-    );
 
     const messageQueueRef = React.useRef<string[]>([]);
     const sendMessageQueueRef = React.useRef<string[]>([]);
@@ -74,16 +82,17 @@ export default function useWebSocket<
      * # Connect to Websocket
      */
     const connect = React.useCallback(() => {
+        const currentUrl = urlRef.current;
         const domain = window.location.origin;
-        const wsURL = url.startsWith(`ws`)
+        const wsURL = currentUrl.startsWith("ws")
-            ? url
+            ? currentUrl
-            : domain.replace(/^http/, "ws") + ("/" + url).replace(/\/\//g, "/");
+            : domain.replace(/^http/, "ws") + ("/" + currentUrl).replace(/\/\//g, "/");
 
         if (!wsURL) return;
 
-        let ws = new WebSocket(wsURL);
+        const ws = new WebSocket(wsURL);
 
-        ws.onerror = (ev) => {
+        ws.onerror = () => {
             console.log(`Websocket ERROR:`);
         };
 
@@ -91,15 +100,17 @@ export default function useWebSocket<
             messageQueueRef.current.push(ev.data);
         };
 
-        ws.onopen = (ev) => {
+        ws.onopen = () => {
             window.clearInterval(keepAliveInterval.current);
 
             keepAliveInterval.current = window.setInterval(() => {
                 if (ws.readyState === WebSocket.OPEN) {
                     ws.send(KEEP_ALIVE_MESSAGE);
                 }
-            }, KEEP_ALIVE_DURATION);
+            }, keepAliveDurationRef.current);
 
+            tries.current = 0;
+            socketRef.current = ws;
             setSocket(ws);
             console.log(`Websocket connected to ${wsURL}`);
         };
@@ -111,23 +122,21 @@ export default function useWebSocket<
                 wasClean: ev.wasClean,
             });
 
-            if (disableReconnect) return;
-
-            console.log("Attempting to reconnect ...");
-            console.log("URL:", url);
             window.clearInterval(keepAliveInterval.current);
+            socketRef.current = undefined;
+            setSocket(undefined);
 
-            console.log("tries", tries);
+            if (disableReconnectRef.current) return;
 
             if (tries.current >= 3) {
+                console.log("Max reconnect attempts reached.");
                 return;
             }
 
-            console.log("Attempting to reconnect ...");
-
             tries.current += 1;
-
+            const backoff = Math.min(1000 * 2 ** tries.current, 30000);
-            connect();
+            console.log(`Attempting to reconnect in ${backoff}ms... (attempt ${tries.current})`);
+            reconnectTimeout.current = window.setTimeout(connect, backoff);
         };
     }, []);
 
@@ -135,18 +144,40 @@ export default function useWebSocket<
      * # Initial Connection
      */
     React.useEffect(() => {
-        if (socket) return;
-
         connect();
+
+        return () => {
+            window.clearTimeout(reconnectTimeout.current);
+            window.clearInterval(keepAliveInterval.current);
+            window.clearInterval(refreshInterval.current);
+            socketRef.current?.close();
+        };
     }, []);
 
+    /**
+     * # Refresh Connection Interval
+     */
+    React.useEffect(() => {
+        if (!refreshConnection) return;
+
+        refreshInterval.current = window.setInterval(() => {
+            console.log("Refreshing WebSocket connection...");
+            window.clearTimeout(reconnectTimeout.current);
+            socketRef.current?.close();
+            tries.current = 0;
+            connect();
+        }, refreshConnection);
+
+        return () => window.clearInterval(refreshInterval.current);
+    }, [refreshConnection]);
+
     React.useEffect(() => {
         if (!socket) return;
 
         sendInterval.current = setInterval(handleSendMessageQueue, DEBOUNCE);
         msgInterval.current = setInterval(handleReceivedMessageQueue, DEBOUNCE);
 
-        return function () {
+        return () => {
             window.clearInterval(sendInterval.current);
             window.clearInterval(msgInterval.current);
         };
@@ -173,7 +204,8 @@ export default function useWebSocket<
      * Send Message Queue Handler
      */
     const handleSendMessageQueue = React.useCallback(() => {
-        if (!socket || socket.readyState !== WebSocket.OPEN) {
+        const ws = socketRef.current;
+        if (!ws || ws.readyState !== WebSocket.OPEN) {
             window.clearInterval(sendInterval.current);
             return;
         }
@@ -181,19 +213,18 @@ export default function useWebSocket<
         const newMessage = sendMessageQueueRef.current.shift();
         if (!newMessage) return;
 
-        socket.send(newMessage);
+        ws.send(newMessage);
-    }, [socket]);
+    }, []);
 
     /**
      * # Send Data Function
      */
-    const sendData = React.useCallback(
+    const sendData = React.useCallback((data: T) => {
-        (data: T) => {
         try {
             const queueItemJSON = JSON.stringify(data);
 
             const existingQueue = sendMessageQueueRef.current.find(
-                    (q) => q == queueItemJSON
+                (q) => q === queueItemJSON
             );
             if (!existingQueue) {
                 sendMessageQueueRef.current.push(queueItemJSON);
@@ -201,9 +232,7 @@ export default function useWebSocket<
         } catch (error: any) {
             console.log("Error Sending socket message", error.message);
         }
-        },
+    }, []);
-        [socket]
-    );
 
     return { socket, sendData };
 }