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

import React from "react";

export type UseWebsocketHookParams = {
    debounce?: number;
    url: string;
    disableReconnect?: boolean;
    keepAliveDuration?: number;
    refreshConnection?: number;
};

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

let tries = 0;

/**
 * # 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.addEventLiatener("wsDataEvent", (e)=>{
 *    console.log(e.detail.data) // type object
 * })
 * @example window.addEventLiatener("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 || 200;
    const KEEP_ALIVE_DURATION = keepAliveDuration || 1000 * 30;
    const KEEP_ALIVE_TIMEOUT = 1000 * 60 * 3;

    const KEEP_ALIVE_MESSAGE = "twui::ping";

    let uptime = 0;

    let reconnectInterval: any;
    let msgInterval: any;
    let sendInterval: any;
    let keepAliveInterval: any;

    const [socket, setSocket] = React.useState<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 wsURL = url;
        if (!wsURL) return;

        let ws = new WebSocket(wsURL);

        ws.onopen = (ev) => {
            window.clearInterval(reconnectInterval);
            window.clearInterval(keepAliveInterval);
            keepAliveInterval = setInterval(() => {
                if (ws.readyState === WebSocket.OPEN) {
                    ws.send(KEEP_ALIVE_MESSAGE);
                    uptime += KEEP_ALIVE_DURATION;
                    if (uptime >= KEEP_ALIVE_TIMEOUT) {
                        console.log("Websocket connection timed out ...");
                        window.clearInterval(keepAliveInterval);
                        ws.close();
                    }
                }
            }, KEEP_ALIVE_DURATION);
            setSocket(ws);
            tries = 0;
            console.log(`Websocket connected to ${wsURL}`);
            uptime = 0;
        };

        ws.onmessage = (ev) => {
            window.clearInterval(msgInterval);
            messageQueueRef.current.push(ev.data);
            msgInterval = setInterval(handleReceivedMessageQueue, DEBOUNCE);
            if (ev.data !== KEEP_ALIVE_MESSAGE) {
                uptime = 0;
            }
        };

        ws.onclose = (ev) => {
            console.log("Websocket closed!");

            if (disableReconnect) return;

            console.log("Attempting to reconnect ...");
            console.log("URL:", url);
            window.clearInterval(keepAliveInterval);

            reconnectInterval = setInterval(() => {
                if (tries >= 3) {
                    return window.clearInterval(reconnectInterval);
                }

                console.log("Attempting to reconnect ...");
                tries++;

                connect();
            }, 1000);
        };
    }, []);

    /**
     * # Window Close Handler
     */
    const handleWindowClose = React.useCallback(() => {
        console.log("Window Unloaded ...");
    }, [socket]);

    /**
     * # Window Focus Handler
     */
    const handleWindowFocus = React.useCallback(() => {
        if (socket?.readyState === WebSocket.CLOSED) {
            console.log("Websocket closed ... Attempting to reconnect ...");
            connect();
        }
        if (socket?.readyState === WebSocket.OPEN) {
            console.log("Websocket connection alive ...");
            socket.send(KEEP_ALIVE_MESSAGE);
            uptime = 0;
        }
    }, [socket]);

    /**
     * # Initial Connection
     */
    React.useEffect(() => {
        connect();

        return function () {
            window.clearInterval(reconnectInterval);
        };
    }, []);

    /**
     * # Window Close and Focus Handlers
     */
    React.useEffect(() => {
        if (!socket) return;

        window.addEventListener("beforeunload", handleWindowClose, {
            once: true,
        });
        window.addEventListener("focus", handleWindowFocus);

        return function () {
            window.removeEventListener("focus", handleWindowFocus);
            window.removeEventListener("beforeunload", handleWindowClose);
        };
    }, [socket]);

    /**
     * # Refresh Connection
     */
    React.useEffect(() => {
        console.log("Refreshing connection ...");

        if (!socket) return;
        if (socket.readyState !== WebSocket.CLOSED) {
            socket?.close();
        }
        connect();
    }, [refreshConnection]);

    /**
     * Received Message Queue Handler
     */
    const handleReceivedMessageQueue = React.useCallback(() => {
        if (messageQueueRef.current.length > 0) {
            const newMessage = messageQueueRef.current.shift();
            if (!newMessage) return;
            try {
                const jsonData = JSON.parse(newMessage);
                dispatchCustomEvent("wsMessageEvent", newMessage);
                dispatchCustomEvent("wsDataEvent", jsonData);
            } catch (error) {
                console.log("Unable to parse string. Returning string.");
            }
        } else {
            window.clearInterval(msgInterval);
            uptime = 0;
        }
    }, []);

    /**
     * Send Message Queue Handler
     */
    const handleSendMessageQueue = React.useCallback(() => {
        if (sendMessageQueueRef.current.length > 0) {
            const newMessage = sendMessageQueueRef.current.shift();
            if (!newMessage) return;
            socket?.send(newMessage);
        } else {
            window.clearInterval(sendInterval);
        }
    }, [socket]);

    /**
     * # Send Data Function
     */
    const sendData = React.useCallback(
        (data: T) => {
            try {
                window.clearInterval(sendInterval);
                sendMessageQueueRef.current.push(JSON.stringify(data));
                sendInterval = setInterval(handleSendMessageQueue, DEBOUNCE);
            } catch (error: any) {
                console.log("Error Sending socket message", error.message);
            }
        },
        [socket]
    );

    return { socket, sendData };
}