useWebSocket自定义hook实现

借鉴ahooks库

源代码

import { useEffect, useRef, useState, useMemo } from 'react';

// 获取最新值hook
const useLatest = value => {
  const ref = useRef(value);
  ref.current = value;
  return ref;
}

// 函数持久化，避免重复声明函数，提升性能
const useMemoizedFn = fn => {
  const fnRef = useRef(fn);
  fnRef.current = useMemo(() => fn, [fn]);
  const memoizedFn = useRef();
  if (!memoizedFn.current) {
    memoizedFn.current = function (...args) {
      return fnRef.current.apply(this, args);
    }
  };
  return memoizedFn.current;
}

// 模拟componentWillUnmount
const useUnmount = fn => {
  const fnRef = useLatest(fn);
  useEffect(() => () => {
    fnRef.current();
  }, []);
}

// websocket连接状态
const READYSTATE = {
  Connecting: 0,
  Open: 1,
  Closing: 2,
  Closed: 3,
}

export default function useWebSocket(socketUrl, options = {}) {
  const {
    reconnectLimit = 3,
    reconnectInterval = 3000,
    manual = false,
    onOpen,
    onClose,
    onMessage,
    onError,
    protocols
  } = options;

  const onOpenRef = useLatest(onOpen);
  const onCloseRef = useLatest(onClose);
  const onMessageRef = useLatest(onMessage);
  const onErrorRef = useLatest(onError);

  const reconnectTimesRef = useRef(0);
  const reconnectTimerRef = useRef();
  const websocketRef = useRef();

  const unmountedRef = useRef(false);

  const [latestMessage, setLatestMessage] = useState();
  const [readyState, setReadyState] = useState(READYSTATE.Closed);

  const createWS = () => {
    if (reconnectTimerRef.current) { clearTimeout(reconnectTimerRef.current); }
    if (websocketRef, current) { websocketRef.current.close(); }

    const ws = new WebSocket(socketUrl, protocols);
    setReadyState(READYSTATE.Connecting);

    ws.onerror = (event) => {
      if (unmountedRef.current) {
        return;
      }
      reconnect();
      onErrorRef.current?.(event, ws);
      setReadyState(ws.readyState || READYSTATE.Closed);
    };
    ws.onopen = (event) => {
      if (unmountedRef.current) {
        return;
      }
      onOpenRef.current?.(event, ws);
      reconnectTimesRef.current = 0;
      setReadyState(ws.readyState || READYSTATE.Open);
    };
    ws.onmessage = (message) => {
      if (unmountedRef.current) {
        return;
      }
      onMessageRef.current?.(message, ws);
      setLatestMessage(message);
    };
    ws.onclose = (event) => {
      if (unmountedRef.current) {
        return;
      }
      reconnect();
      onCloseRef.current?.(event, ws);
      setReadyState(ws.readyState || READYSTATE.Closed);
    }

    websocketRef.current = ws;
  }

  const connect = () => {
    reconnectTimesRef.current = 0;
    createWS();
  };

  const reconnect = () => {
    if (reconnectTimesRef.current < reconnectLimit && 
    websocketRef.current?.readyState !== READYSTATE.Open
    ) {
      if (reconnectTimerRef.current) { clearTimeout(reconnectTimerRef.current); }
      reconnectTimerRef.current = setTimeout(() => {
        createWS();
        reconnectTimesRef.current++;
      }, reconnectInterval);
    }
  };

  const sendMessage = message => {
    if (readyState === READYSTATE.Open) {
      websocketRef.current?.send(message);
    } else {
      throw new Error('websocket disconnected');
    }
  };

  const disconnect = () => {
    if (reconnectTimerRef.current) {
      clearTimeout(reconnectTimerRef.current);
    }

    reconnectTimesRef.current = reconnectLimit;
    websocketRef.current?.close();
  };

  useEffect(() => {
    if (!manual) {
      connect();
    }
  }, [socketUrl, manual]);

  useUnmount(() => {
    unmountedRef.current = true;
    disconnect();
  });

  return {
    latestMessage,
    sendMessage: useMemoizedFn(sendMessage),
    connect: useMemoizedFn(connect),
    disconnect: useMemoizedFn(disconnect),
    readyState,
    webSocketIns: websocketRef.current,
  }
}