ahooks/es/useWebSocket/index.js

react hooks library

import { useEffect, useRef, useState } from 'react';
import useLatest from '../useLatest';
import useMemoizedFn from '../useMemoizedFn';
import useUnmount from '../useUnmount';
export var ReadyState;
(function (ReadyState) {
    ReadyState[ReadyState["Connecting"] = 0] = "Connecting";
    ReadyState[ReadyState["Open"] = 1] = "Open";
    ReadyState[ReadyState["Closing"] = 2] = "Closing";
    ReadyState[ReadyState["Closed"] = 3] = "Closed";
})(ReadyState || (ReadyState = {}));
function useWebSocket(socketUrl, options = {}) {
    const { reconnectLimit = 3, reconnectInterval = 3 * 1000, manual = false, onOpen, onClose, onMessage, onError, protocols, } = options;
    const [latestMessage, setLatestMessage] = useState();
    const [readyState, setReadyState] = useState(ReadyState.Closed);
    const onOpenRef = useLatest(onOpen);
    const onCloseRef = useLatest(onClose);
    const onMessageRef = useLatest(onMessage);
    const onErrorRef = useLatest(onError);
    const readyStateRef = useLatest(readyState);
    const reconnectTimesRef = useRef(0);
    const reconnectTimerRef = useRef(undefined);
    const websocketRef = useRef(undefined);
    const reconnect = () => {
        var _a;
        if (reconnectTimesRef.current < reconnectLimit &&
            ((_a = websocketRef.current) === null || _a === void 0 ? void 0 : _a.readyState) !== ReadyState.Open) {
            if (reconnectTimerRef.current) {
                clearTimeout(reconnectTimerRef.current);
            }
            reconnectTimerRef.current = setTimeout(() => {
                // eslint-disable-next-line @typescript-eslint/no-use-before-define
                connectWs();
                reconnectTimesRef.current++;
            }, reconnectInterval);
        }
    };
    const connectWs = () => {
        if (reconnectTimerRef.current) {
            clearTimeout(reconnectTimerRef.current);
        }
        if (websocketRef.current) {
            websocketRef.current.close();
        }
        const ws = new WebSocket(socketUrl, protocols);
        setReadyState(ReadyState.Connecting);
        ws.onerror = (event) => {
            var _a;
            if (websocketRef.current !== ws) {
                return;
            }
            reconnect();
            (_a = onErrorRef.current) === null || _a === void 0 ? void 0 : _a.call(onErrorRef, event, ws);
            setReadyState(ws.readyState || ReadyState.Closed);
        };
        ws.onopen = (event) => {
            var _a;
            if (websocketRef.current !== ws) {
                return;
            }
            (_a = onOpenRef.current) === null || _a === void 0 ? void 0 : _a.call(onOpenRef, event, ws);
            reconnectTimesRef.current = 0;
            setReadyState(ws.readyState || ReadyState.Open);
        };
        ws.onmessage = (message) => {
            var _a;
            if (websocketRef.current !== ws) {
                return;
            }
            (_a = onMessageRef.current) === null || _a === void 0 ? void 0 : _a.call(onMessageRef, message, ws);
            setLatestMessage(message);
        };
        ws.onclose = (event) => {
            var _a;
            (_a = onCloseRef.current) === null || _a === void 0 ? void 0 : _a.call(onCloseRef, event, ws);
            // closed by server
            if (websocketRef.current === ws) {
                reconnect();
            }
            // closed by disconnect or closed by server
            if (!websocketRef.current || websocketRef.current === ws) {
                setReadyState(ws.readyState || ReadyState.Closed);
            }
        };
        websocketRef.current = ws;
    };
    const sendMessage = (message) => {
        var _a;
        if (readyStateRef.current === ReadyState.Open) {
            (_a = websocketRef.current) === null || _a === void 0 ? void 0 : _a.send(message);
        }
        else {
            throw new Error('WebSocket disconnected');
        }
    };
    const connect = () => {
        reconnectTimesRef.current = 0;
        connectWs();
    };
    const disconnect = () => {
        var _a;
        if (reconnectTimerRef.current) {
            clearTimeout(reconnectTimerRef.current);
        }
        reconnectTimesRef.current = reconnectLimit;
        (_a = websocketRef.current) === null || _a === void 0 ? void 0 : _a.close();
        websocketRef.current = undefined;
    };
    useEffect(() => {
        if (!manual && socketUrl) {
            connect();
        }
    }, [socketUrl, manual]);
    useUnmount(() => {
        disconnect();
    });
    return {
        latestMessage,
        sendMessage: useMemoizedFn(sendMessage),
        connect: useMemoizedFn(connect),
        disconnect: useMemoizedFn(disconnect),
        readyState,
        webSocketIns: websocketRef.current,
    };
}
export default useWebSocket;