Fix useWs opening multiple connections at once

todoblue/src/app/(api)/useWs.ts

@@ -1,10 +1,9 @@
 'use client';
 
-import {useState, useCallback, useEffect} from "react"
+import {useState, useCallback, useEffect, useReducer, Reducer} from "react"
 
 export interface WebSocketHandlerParams<E extends Event> {
 	event: E,
-	openWebSocketAfterBackoff: () => void,
 	closeWebSocket: () => void,
 }
 
@@ -15,11 +14,45 @@ export interface WebSocketHandlers {
 	onopen?: (params: WebSocketHandlerParams<Event>) => void,
 }
 
+interface WebSocketState {
+	webSocket: WebSocket | undefined,
+	webSocketState: number | undefined,
+	webSocketBackoffMs: number | undefined,
+	nextBackOffMs: number
+}
+
+type WebSocketAction = { event: "connect", webSocket: WebSocket } | { event: "disconnect" } | { event: "backoffExpire" } | { event: "stateChange" };
+
+function wsReducer(prevState: WebSocketState, action: WebSocketAction) {
+	switch(action.event) {
+		case "connect":
+			return {
+				webSocket: action.webSocket,
+				webSocketState: action.webSocket.readyState,
+                webSocketBackoffMs: prevState.nextBackOffMs,
+                nextBackOffMs: prevState.nextBackOffMs * 2,
+			}
+		case "disconnect":
+			return {
+                ...prevState,
+				webSocket: undefined,
+				webSocketState: prevState.webSocket?.readyState,
+			}
+		case "backoffExpire":
+			return {
+				...prevState,
+				webSocketBackoffMs: undefined,
+			}
+		case "stateChange":
+			return {
+				...prevState,
+				webSocketState: prevState.webSocket?.readyState
+			}
+	}
+}
+
 export function useWs(url: string | undefined, {onclose, onerror, onmessage, onopen}: WebSocketHandlers) {
-	const [webSocket, setWebSocket] = useState<WebSocket | undefined>(undefined)
+	const [{webSocket, webSocketState, webSocketBackoffMs, nextBackOffMs}, dispatch] = useReducer<Reducer<WebSocketState, WebSocketAction>>(wsReducer, {webSocket: undefined, webSocketState: undefined, webSocketBackoffMs: undefined, nextBackOffMs: 1000})
-	const [webSocketState, setWebSocketState] = useState<number | undefined>(undefined);
-	const [isBackingOff, setBackingOff] = useState<boolean>(false);
-	const [webSocketBackoffMs, setWebSocketBackoffMs] = useState<number>(1);
 
 	const closeWebSocket = useCallback(() => {
 		console.debug("[useWebSocket] Closing WebSocket:", webSocket);
@@ -33,61 +66,52 @@ export function useWs(url: string | undefined, {onclose, onerror, onmessage, ono
 		catch(closeErr) {
 			console.debug("[useWebSocket] Failed to close the websocket (it might be already closed):", closeErr)
 		}
-		setWebSocket(undefined);
+		dispatch({event: "disconnect"})
-		setWebSocketState(WebSocket.CLOSED);
 	}, [webSocket])
 
-	const backOff = useCallback((func: () => void) => async () => {
-		// This WILL cause no-ops, but they're going to be pretty infrequent, so idc
-		console.debug("[useWebSocket] Backing off for:", webSocketBackoffMs, "ms")
-		setBackingOff(true)
-		await new Promise(resolve => setTimeout(resolve, webSocketBackoffMs))
-		setBackingOff(false)
-		func()
-	}, [webSocketBackoffMs])
-
 	const openWebSocket = useCallback(() => {
 		console.debug("[useWebSocket] Opening WebSocket:", url);
 		if(url === undefined) {
 			console.warn("[useWebSocket] Trying to open WebSocket, but no URL has been given; ignoring request...")
 			return;
 		}
-		setWebSocketBackoffMs(prev => prev * 2);  // Workaround for connections that get closed immediately by the server
-		setWebSocketState(WebSocket.CONNECTING);  // Workaround for constructor blocking and giving no feedback to the user
 		const sock = new WebSocket(url);
-		const openWebSocketAfterBackoff = backOff(openWebSocket);
+		dispatch({event: "connect", webSocket: sock})
 		sock.onopen = (event) => {
 			console.debug("[useWebSocket] Opened connection:", event)
-			setWebSocket(sock)
+			dispatch({event: "stateChange"})
-			setWebSocketState(sock.readyState)
+			onopen?.({event, closeWebSocket});
-			onopen?.({event, openWebSocketAfterBackoff, closeWebSocket});
 		}
 		sock.onclose = (event) => {
 			console.debug("[useWebSocket] Closed connection:", event)
-			setWebSocket(undefined)
+            dispatch({event: "stateChange"})
-			setWebSocketState(sock.readyState);
+			onclose?.({event, closeWebSocket});
-			onclose?.({event, openWebSocketAfterBackoff, closeWebSocket});
 		}
 		sock.onerror = (event) => {
 			console.error("[useWebSocket] Error in connection:", event)
-			setWebSocketState(sock.readyState)
+            dispatch({event: "stateChange"})
-			onerror?.({event, openWebSocketAfterBackoff, closeWebSocket});
+			onerror?.({event, closeWebSocket});
 		}
 		sock.onmessage = (event) => {
 			console.debug("[useWebSocket] Received message:", event)
-			onmessage?.({event, openWebSocketAfterBackoff, closeWebSocket});
+			onmessage?.({event, closeWebSocket});
 		}
-	}, [url, onopen, onclose, onerror, onmessage, backOff, closeWebSocket])
+	}, [url, onopen, onclose, onerror, onmessage, closeWebSocket])
+
+	useEffect(() => {
+		if(webSocketBackoffMs === undefined) return;
+		console.debug("[useWebSocket] Backing off for:", webSocketBackoffMs, "ms")
+		new Promise(resolve => setTimeout(resolve, webSocketBackoffMs)).then(() => {
+			dispatch({event: "backoffExpire"})
+		})
+	}, [webSocketBackoffMs])
 
 	useEffect(() => {
-		if(!url) return;
 		if(webSocket !== undefined) return;
-		if(isBackingOff) return;
+		if(webSocketBackoffMs !== undefined) return;
-		console.debug("[useWebSocket] Hook mounted, opening connection as soon as possible...")
+		console.debug("[useWebSocket] Back off expired, opening websocket...")
-		const openWebSocketAfterBackoff = backOff(openWebSocket);
+		openWebSocket()
-		// noinspection JSIgnoredPromiseFromCall
+	}, [webSocket, webSocketBackoffMs, openWebSocket])
-		openWebSocketAfterBackoff()
-	}, [url, webSocket, isBackingOff, webSocketState])
 
 	return {webSocket, webSocketState, webSocketBackoffMs}
 }