ahooks-v2/es/useWebSocket/index.js

react hooks library

var __read = this && this.__read || function (o, n) {
  var m = typeof Symbol === "function" && o[Symbol.iterator];
  if (!m) return o;
  var i = m.call(o),
      r,
      ar = [],
      e;

  try {
    while ((n === void 0 || n-- > 0) && !(r = i.next()).done) {
      ar.push(r.value);
    }
  } catch (error) {
    e = {
      error: error
    };
  } finally {
    try {
      if (r && !r.done && (m = i["return"])) m.call(i);
    } finally {
      if (e) throw e.error;
    }
  }

  return ar;
};

import useUnmount from '../useUnmount';
import usePersistFn from '../usePersistFn';
import { useEffect, useRef, useState } from 'react';
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 = {}));

export default function useWebSocket(socketUrl, options) {
  if (options === void 0) {
    options = {};
  }

  var _a = options.reconnectLimit,
      reconnectLimit = _a === void 0 ? 3 : _a,
      _b = options.reconnectInterval,
      reconnectInterval = _b === void 0 ? 3 * 1000 : _b,
      _c = options.manual,
      manual = _c === void 0 ? false : _c,
      onOpen = options.onOpen,
      onClose = options.onClose,
      onMessage = options.onMessage,
      onError = options.onError;
  var reconnectTimesRef = useRef(0);
  var reconnectTimerRef = useRef();
  var websocketRef = useRef();
  var unmountedRef = useRef(false);

  var _d = __read(useState(), 2),
      latestMessage = _d[0],
      setLatestMessage = _d[1];

  var _e = __read(useState(ReadyState.Closed), 2),
      readyState = _e[0],
      setReadyState = _e[1];
  /**
   * 重连
   */


  var reconnect = usePersistFn(function () {
    var _a;

    if (reconnectTimesRef.current < reconnectLimit && ((_a = websocketRef.current) === null || _a === void 0 ? void 0 : _a.readyState) !== ReadyState.Open) {
      reconnectTimerRef.current && clearTimeout(reconnectTimerRef.current);
      reconnectTimerRef.current = setTimeout(function () {
        connectWs();
        reconnectTimesRef.current++;
      }, reconnectInterval);
    }
  });
  var connectWs = usePersistFn(function () {
    reconnectTimerRef.current && clearTimeout(reconnectTimerRef.current);

    if (websocketRef.current) {
      websocketRef.current.close();
    }

    try {
      websocketRef.current = new WebSocket(socketUrl);

      websocketRef.current.onerror = function (event) {
        var _a;

        reconnect();
        onError && onError(event);
        setReadyState(((_a = websocketRef.current) === null || _a === void 0 ? void 0 : _a.readyState) || ReadyState.Closed);
      };

      websocketRef.current.onopen = function (event) {
        var _a;

        onOpen && onOpen(event);
        reconnectTimesRef.current = 0;
        setReadyState(((_a = websocketRef.current) === null || _a === void 0 ? void 0 : _a.readyState) || ReadyState.Closed);
      };

      websocketRef.current.onmessage = function (message) {
        onMessage && onMessage(message);
        setLatestMessage(message);
      };

      websocketRef.current.onclose = function (event) {
        var _a;

        if (unmountedRef.current) {
          return;
        }

        reconnect();
        onClose && onClose(event);
        setReadyState(((_a = websocketRef.current) === null || _a === void 0 ? void 0 : _a.readyState) || ReadyState.Closed);
      };
    } catch (error) {
      throw error;
    }
  });
  /**
   * 发送消息
   * @param message
   */

  var sendMessage = usePersistFn(function (message) {
    var _a;

    if (readyState === ReadyState.Open) {
      (_a = websocketRef.current) === null || _a === void 0 ? void 0 : _a.send(message);
    } else {
      throw new Error('WebSocket disconnected');
    }
  });
  /**
   * 手动 connect
   */

  var connect = usePersistFn(function () {
    reconnectTimesRef.current = 0;
    connectWs();
  });
  /**
   * disconnect websocket
   */

  var disconnect = usePersistFn(function () {
    var _a;

    reconnectTimerRef.current && clearTimeout(reconnectTimerRef.current);
    reconnectTimesRef.current = reconnectLimit;
    (_a = websocketRef.current) === null || _a === void 0 ? void 0 : _a.close();
  });
  useEffect(function () {
    // 初始连接
    if (!manual) {
      connect();
    }
  }, [socketUrl, manual]);
  useUnmount(function () {
    unmountedRef.current = true;
    disconnect();
  });
  return {
    latestMessage: latestMessage,
    sendMessage: sendMessage,
    connect: connect,
    disconnect: disconnect,
    readyState: readyState,
    webSocketIns: websocketRef.current
  };
}