一个简单的websocket hook

一个简单的 vue3 的 websocket hook. 有以下基础功能：

• 创建链接
• 失败重连
• 发送心跳包

import { ref } from "vue";

export interface WS_CONFIG {
  url: string; // ws链接地址
  sendData?: Record<string, any>; // 请求数据
  reconnectSec?: number; // 请求异常重连间隔
  needHeart?: boolean; // 是否需要心跳
  heartSec?: number; // 心跳间隔
  heartData?: unknown; // 心跳时发送的数据
}

export default function <T>({
  url,
  sendData = {},
  heartSec = 1000 * 30,
  needHeart = false,
  reconnectSec = 1000 * 30,
  heartData = "ping",
}: WS_CONFIG) {
  // 响应数据
  const wsData = ref<T | undefined>();
  // 重试次数
  const retryNumber = ref<number>(0);

  let socket: WebSocket | null = null;
  let heart_Interval: number | null = null;
  let lockConnect = false;
  let number = 0;

  // 创建链接
  function connect() {
    // 重连之前，关闭上次的连接
    if (socket) {
      socket.close();
    }

    socket = new WebSocket(url);
    retryNumber.value = number++;

    // 连接成功回调
    socket.onopen = () => {
      console.log(`${url} open event - ${new Date()}`);
      // 维持心跳
      needHeart && heart();

      // 清空重试次数
      number = 0;
      retryNumber.value = 0;
      // 发送数据
      socket?.send(JSON.stringify(sendData));
    };

    // 连接断开回调
    socket.onclose = (e) => {
      console.log(`
${new Date()}
${url} close event
code: ${e.code};
reason: ${e.reason};
isTrusted: ${e.isTrusted};
        `);

      console.log(`reconnect - ${new Date()}`);
      // 进行重连
      reconnect();
    };

    // 连接失败回调
    socket.onerror = (e) => {
      console.log(`${url} error event`, e);
      // 进行重连
      reconnect();
    };

    // 接收到数据
    socket.onmessage = (e: MessageEvent<string>) => {
      const { data } = e;

      const _data = (typeof data === "string" ? JSON.parse(data) : data) as T;
      wsData.value = _data;
    };
  }

  // 心跳函数，维持链接不断开
  function heart() {
    clearHeart();
    heart_Interval = setInterval(() => {
      socket?.send(
        typeof heartData === "string" ? heartData : JSON.stringify(heartData)
      );
    }, heartSec);
  }

  // 断开心跳
  function clearHeart() {
    heart_Interval && clearInterval(heart_Interval);
  }

  // 重连
  function reconnect() {
    // 控制重连频率
    if (lockConnect) {
      return;
    }
    lockConnect = true;

    // 清除心跳
    clearHeart();

    setTimeout(() => {
      lockConnect = false;
      connect();
    }, reconnectSec);
  }

  connect();

  // 关闭链接
  function clearWS() {
    socket?.close();
    clearHeart();
  }

  return {
    wsData,
    retryNumber,
    clearWS,
  };
}