ENG-7948: Fix hanging frontend and other reconnection woes (#5884)

* ENG-7948: Fix hanging frontend and other reconnection woes

* Disable socketio automatic reconnection logic -- replace with manual
  reconnect logic
* Try to keep the socketio socket allocated, even when it gets disconnected
  for better resource management.

* Introduce new `socket.current.reconnect()` helper
  * Update query param token when reconnecting
  * Only allow one concurrent connection attempt

* ensureSocketConnected only reconnects if the event loop is mounted -- this
  avoids issues during hot reload where some Event-emitting components may
  still have a reference to the _old_ `addEvents` and `ensureSocketConnected`
  was reconnecting the old socket.

* Use `.off` and null the `socket` ref when cleaning up to free memory and
  avoid unwanted reconnections when the socket will be recreated.

Also addresses the following ticket:
ENG-7762: [reconnect] client is given new_token when wifi disconnects and reconnects

This issue was the result of automatic reconnect logic and manual reconnect
logic conflicting with each other and opening two sockets with the same token,
which reflex reassigns the latter to a new token, and this is typically the
socket that gets used in the app.

* useRef for `mounted` to play nice in the react world

Author: masenf

reflex/.templates/web/utils/state.js

@@ -529,6 +529,14 @@ export const connect = async (
   navigate,
   params,
 ) => {
+  // Socket already allocated, just reconnect it if needed.
+  if (socket.current) {
+    if (!socket.current.connected) {
+      socket.current.reconnect();
+    }
+    return;
+  }
+
   // Get backend URL object from the endpoint.
   const endpoint = getBackendURL(EVENTURL);
   const on_hydrated_queue = [];
@@ -540,7 +548,9 @@ export const connect = async (
     protocols: [reflexEnvironment.version],
     autoUnref: false,
     query: { token: getToken() },
+    reconnection: false, // Reconnection will be handled manually.
   });
+  socket.current.wait_connect = !socket.current.connected;
   // Ensure undefined fields in events are sent as null instead of removed
   socket.current.io.encoder.replacer = (k, v) => (v === undefined ? null : v);
   socket.current.io.decoder.tryParse = (str) => {
@@ -550,6 +560,18 @@ export const connect = async (
       return false;
     }
   };
+  // Set up a reconnect helper function
+  socket.current.reconnect = () => {
+    if (
+      socket.current &&
+      !socket.current.connected &&
+      !socket.current.wait_connect
+    ) {
+      socket.current.wait_connect = true;
+      socket.current.io.opts.query = { token: getToken() }; // Update token for reconnect.
+      socket.current.connect();
+    }
+  };
 
   function checkVisibility() {
     if (document.visibilityState === "visible") {
@@ -565,7 +587,7 @@ export const connect = async (
         );
       } else if (!socket.current.connected) {
         console.log("Socket is disconnected, attempting to reconnect ");
-        socket.current.connect();
+        socket.current.reconnect();
       } else {
         console.log("Socket is reconnected ");
       }
@@ -588,6 +610,7 @@ export const connect = async (
 
   // Once the socket is open, hydrate the page.
   socket.current.on("connect", async () => {
+    socket.current.wait_connect = false;
     setConnectErrors([]);
     window.addEventListener("pagehide", pagehideHandler);
     window.addEventListener("beforeunload", disconnectTrigger);
@@ -599,16 +622,33 @@ export const connect = async (
   });
 
   socket.current.on("connect_error", (error) => {
-    setConnectErrors((connectErrors) => [connectErrors.slice(-9), error]);
+    socket.current.wait_connect = false;
+    let n_connect_errors = 0;
+    setConnectErrors((connectErrors) => {
+      const new_errors = [...connectErrors.slice(-9), error];
+      n_connect_errors = new_errors.length;
+      return new_errors;
+    });
+    window.setTimeout(() => {
+      if (socket.current && !socket.current.connected) {
+        socket.current.reconnect();
+      }
+    }, 200 * n_connect_errors); // Incremental backoff
   });
 
   // When the socket disconnects reset the event_processing flag
-  socket.current.on("disconnect", () => {
-    socket.current = null; // allow reconnect to occur automatically
+  socket.current.on("disconnect", (reason, details) => {
+    socket.current.wait_connect = false;
+    const try_reconnect =
+      reason !== "io server disconnect" && reason !== "io client disconnect";
     event_processing = false;
     window.removeEventListener("unload", disconnectTrigger);
     window.removeEventListener("beforeunload", disconnectTrigger);
     window.removeEventListener("pagehide", pagehideHandler);
+    if (try_reconnect) {
+      // Attempt to reconnect transient non-intentional disconnects.
+      socket.current.reconnect();
+    }
   });
 
   // On each received message, queue the updates and events.
@@ -785,6 +825,7 @@ export const useEventLoop = (
   const [searchParams] = useSearchParams();
   const [connectErrors, setConnectErrors] = useState([]);
   const params = useRef(paramsR);
+  const mounted = useRef(false);
 
   useEffect(() => {
     const { "*": splat, ...remainingParams } = paramsR;
@@ -796,11 +837,16 @@ export const useEventLoop = (
   }, [paramsR]);
 
   const ensureSocketConnected = useCallback(async () => {
+    if (!mounted.current) {
+      // During hot reload, some components may still have a reference to
+      // addEvents, so avoid reconnecting the socket of an unmounted event loop.
+      return;
+    }
     // only use websockets if state is present and backend is not disabled (reflex cloud).
     if (
       Object.keys(initialState).length > 1 &&
       !isBackendDisabled() &&
-      !socket.current
+      !socket.current?.connected
     ) {
       // Initialize the websocket connection.
       await connect(
@@ -813,13 +859,23 @@ export const useEventLoop = (
         () => params.current,
       );
     }
-  }, [socket, dispatch, setConnectErrors, client_storage, navigate, params]);
+  }, [
+    socket,
+    dispatch,
+    setConnectErrors,
+    client_storage,
+    navigate,
+    params,
+    mounted,
+  ]);
 
   // Function to add new events to the event queue.
   const addEvents = useCallback((events, args, event_actions) => {
     const _events = events.filter((e) => e !== undefined && e !== null);
-
-    ensureSocketConnected();
+    if (!event_actions?.temporal) {
+      // Reconnect socket if needed for non-temporal events.
+      ensureSocketConnected();
+    }
 
     if (!(args instanceof Array)) {
       args = [args];
@@ -914,12 +970,16 @@ export const useEventLoop = (
   // Handle socket connect/disconnect.
   useEffect(() => {
     // Initialize the websocket connection.
+    mounted.current = true;
     ensureSocketConnected();
 
     // Cleanup function.
     return () => {
+      mounted.current = false;
       if (socket.current) {
         socket.current.disconnect();
+        socket.current.off();
+        socket.current = null;
       }
     };
   }, []);