Fixed set state race conditions on useWebsockets

dashboard/src/shared/hooks/useWebsockets.ts

@@ -1,4 +1,4 @@
-import { useCallback, useEffect, useState } from "react"
+import { useCallback, useEffect, useRef, useState } from "react"
 
 interface NewWebsocketOptions {
   onopen?: () => void;
@@ -20,9 +20,8 @@ type WebsocketMap = {
 }
 
 export const useWebsockets = () => {
-  const [websocketMap, setWebsocketMap] = useState<WebsocketMap>({});
-  const [websocketConfigMap, setWebsocketConfigMap] = useState<WebsocketConfigMap>({});
-
+  const websocketMap = useRef<WebsocketMap>({});
+  const websocketConfigMap = useRef<WebsocketConfigMap>({})
   
   const newWebsocket = (id: string, apiEndpoint: string, options: NewWebsocketOptions): WebsocketConfig => {
     
@@ -30,68 +29,58 @@ export const useWebsockets = () => {
 
     const url = `${protocol}://${window.location.host}${apiEndpoint}`
 
-    const mockFunction = (method: string) => () => {
-      console.log(`${method} not implemented`);
-    }
+    const mockFunction = () => {}
     
     const wsConfig: WebsocketConfig = {
       url,
-      onopen: options.onopen || mockFunction("onopen"),
-      onmessage: options.onmessage || mockFunction("onmessage"),
-      onerror: options.onerror || mockFunction("onerror"),
-      onclose: options.onclose || mockFunction("onclose"),
+      onopen: options.onopen || mockFunction,
+      onmessage: options.onmessage || mockFunction,
+      onerror: options.onerror || mockFunction,
+      onclose: options.onclose || mockFunction,
     }
     
-    setWebsocketConfigMap((oldWebsocketConfigMap) => {
-      console.log(`Save config for ${id}`)
-      return ({
-      ...oldWebsocketConfigMap,
+    websocketConfigMap.current = {
+      ...websocketConfigMap.current,
       [id]: wsConfig,
-    })});
-
+    }
     return wsConfig;
   }
 
-  const openWebsocket = 
-    (id: string) => {
-      const wsConfig = websocketConfigMap[id];
-      // debugger;
-      if (!wsConfig) {
-        console.log("Couldn't find ws config")
-        return;
-      }
-      const ws = new WebSocket(wsConfig.url);
-      ws.onopen = wsConfig.onopen;
-      ws.onmessage = wsConfig.onmessage;
-      ws.onerror = wsConfig.onerror;
-      ws.onclose = wsConfig.onclose;
-  
-      setWebsocketMap((oldWebsocketMap) => ({
-        ...oldWebsocketMap,
-        [id]: ws,
-      }))
+  const openWebsocket = (id: string) => {
+    const wsConfig = websocketConfigMap.current[id];
+
+    if (!wsConfig) {
+      console.log("Couldn't find ws config")
+      return;
+    }
+    const ws = new WebSocket(wsConfig.url);
+    ws.onopen = wsConfig.onopen;
+    ws.onmessage = wsConfig.onmessage;
+    ws.onerror = wsConfig.onerror;
+    ws.onclose = wsConfig.onclose;
+    
+    websocketMap.current = {
+      ...websocketMap.current,
+      [id]: ws,
     }
+  }
 
   const closeWebsocket = (id: string, code?: number, reason?: string) => {
-    const ws = websocketMap[id];
+    const ws = websocketMap.current[id];
 
     ws.close(code, reason);
   }
 
   const closeAllWebsockets = () => {
-    Object.keys(websocketMap).forEach(key => {
+    Object.keys(websocketMap.current).forEach(key => {
       closeWebsocket(key);
     })
   }
 
   const getWebsocket = (id: string) => {
-    return websocketMap[id];
+    return websocketMap.current[id];
   }
 
-  useEffect(() => {
-    console.log(websocketConfigMap);
-  }, [websocketConfigMap])
-
   return {
     newWebsocket,
     openWebsocket,