new-personal-site/components/lib/hooks/useWebSocket.tsx

import React, { useRef } from "react";

export type UseWebsocketHookParams = {
    debounce?: number;
    url: string;
    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;
};

export const WebSocketEventNames = ["wsDataEvent", "wsMessageEvent"] as const;

/**
 * # Use Websocket Hook
 * @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.addEventListener("wsDataEvent", (e)=>{
 *    console.log(e.detail.data) // type object
 * })
 * @example window.addEventListener("wsMessageEvent", (e)=>{
 *      console.log(e.detail.message) // type string
 * })
 */
export default function useWebSocket<
    T extends { [key: string]: any } = { [key: string]: any }
>({
    url,
    debounce,
    disableReconnect,
    keepAliveDuration,
    refreshConnection,
}: UseWebsocketHookParams) {
    const DEBOUNCE = debounce || 500;
    const KEEP_ALIVE_DURATION = keepAliveDuration || 1000 * 30;
    const KEEP_ALIVE_MESSAGE = "twui::ping";

    const tries = useRef(0);

    // 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>(undefined);
    const socketRef = useRef<WebSocket | undefined>(undefined);

    const messageQueueRef = React.useRef<string[]>([]);
    const sendMessageQueueRef = React.useRef<string[]>([]);

    /**
     * # Dispatch Custom Event
     */
    const dispatchCustomEvent = React.useCallback(
        (evtName: (typeof WebSocketEventNames)[number], value: string | T) => {
            const event = new CustomEvent(evtName, {
                detail: {
                    data: value,
                    message: value,
                },
            });
            window.dispatchEvent(event);
        },
        []
    );

    /**
     * # Connect to Websocket
     */
    const connect = React.useCallback(() => {
        const currentUrl = urlRef.current;
        const domain = window.location.origin;
        const wsURL = currentUrl.startsWith("ws")
            ? currentUrl
            : domain.replace(/^http/, "ws") + ("/" + currentUrl).replace(/\/\//g, "/");

        if (!wsURL) return;

        const ws = new WebSocket(wsURL);

        ws.onerror = () => {
            console.log(`Websocket ERROR:`);
        };

        ws.onmessage = (ev) => {
            messageQueueRef.current.push(ev.data);
        };

        ws.onopen = () => {
            window.clearInterval(keepAliveInterval.current);

            keepAliveInterval.current = window.setInterval(() => {
                if (ws.readyState === WebSocket.OPEN) {
                    ws.send(KEEP_ALIVE_MESSAGE);
                }
            }, keepAliveDurationRef.current);

            tries.current = 0;
            socketRef.current = ws;
            setSocket(ws);
            console.log(`Websocket connected to ${wsURL}`);
        };

        ws.onclose = (ev) => {
            console.log("Websocket closed!", {
                code: ev.code,
                reason: ev.reason,
                wasClean: ev.wasClean,
            });

            window.clearInterval(keepAliveInterval.current);
            socketRef.current = undefined;
            setSocket(undefined);

            if (disableReconnectRef.current) return;

            if (tries.current >= 3) {
                console.log("Max reconnect attempts reached.");
                return;
            }

            tries.current += 1;
            const backoff = Math.min(1000 * 2 ** tries.current, 30000);
            console.log(`Attempting to reconnect in ${backoff}ms... (attempt ${tries.current})`);
            reconnectTimeout.current = window.setTimeout(connect, backoff);
        };
    }, []);

    /**
     * # Initial Connection
     */
    React.useEffect(() => {
        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 () => {
            window.clearInterval(sendInterval.current);
            window.clearInterval(msgInterval.current);
        };
    }, [socket]);

    /**
     * Received Message Queue Handler
     */
    const handleReceivedMessageQueue = React.useCallback(() => {
        try {
            const msg = messageQueueRef.current.shift();

            if (!msg) return;

            const jsonData = JSON.parse(msg);
            dispatchCustomEvent("wsMessageEvent", msg);
            dispatchCustomEvent("wsDataEvent", jsonData);
        } catch (error) {
            console.log("Unable to parse string. Returning string.");
        }
    }, []);

    /**
     * Send Message Queue Handler
     */
    const handleSendMessageQueue = React.useCallback(() => {
        const ws = socketRef.current;
        if (!ws || ws.readyState !== WebSocket.OPEN) {
            window.clearInterval(sendInterval.current);
            return;
        }

        const newMessage = sendMessageQueueRef.current.shift();
        if (!newMessage) return;

        ws.send(newMessage);
    }, []);

    /**
     * # Send Data Function
     */
    const sendData = React.useCallback((data: T) => {
        try {
            const queueItemJSON = JSON.stringify(data);

            const existingQueue = sendMessageQueueRef.current.find(
                (q) => q === queueItemJSON
            );
            if (!existingQueue) {
                sendMessageQueueRef.current.push(queueItemJSON);
            }
        } catch (error: any) {
            console.log("Error Sending socket message", error.message);
        }
    }, []);

    return { socket, sendData };
}