【useWebsocket】vue3+vite封装websocket hooks

不废话了，直接上干货。可直接食用
ts刚开始用，不合理的地方欢迎指导，虚心请教

import { onUnmounted, ref } from 'vue'

import { useSystemStore } from '@/store/modules/systemStore/systemStore'

interface QueryType {
  url: string //url
  options?: OptionsType //配置
  params?: object //url参数
  onmessage?: Function //回调函数
}
interface OptionsType {
  maxReconnect?: number //最大重连次数， -1 标识无限重连
  interval?: number //心跳间隔
  timeout?: number //  响应超时时间
  pingMessage?: string // 心跳请求信息
}
const defaultOptions = {
  maxReconnect: 6,
  interval: 30 * 1000,
  timeout: 10 * 1000,
  pingMessage: 'ping'
}

/**
 * websocket封装
 * @param {object} param0
 * @returns  serverMessage, sendMessage, destroySocket
 */
export default function useWebSocket({ url, options = {}, params, onmessage }: QueryType) {
  const OPTION = { ...defaultOptions, ...options }
  let socket: WebSocket | null = null //websocket实例
  let pingTimeoutObj: NodeJS.Timeout | null = null //延时发送心跳的定时器
  let pongTimeoutObj: NodeJS.Timeout | null = null //接收心跳响应的定时器
  let reconnectCount: number = 0 // 重连尝试次数
  let lockReconnect: Boolean = false // 重连锁，避免多次重连

  const token = useSystemStore().getToken

  // 初始化连接socket
  function initWebSocket() {
    const VITE_APP_SOCKET_API = import.meta.env.VITE_APP_SOCKET_API
    const _params = objectToEncodeQuery(params)
    const wsUri = `${VITE_APP_SOCKET_API}${url}?token=${token}` + _params

    socket = new WebSocket(wsUri)

    // 连接成功
    socket.onopen = onOpen
    // 连接错误
    socket.onerror = onError
    // 接收信息
    socket.onmessage = onMessage
    // 连接关闭
    socket.onclose = onClose
  }

  /**
   * 连接成功事件
   */
  function onOpen() {
    // 开启心跳
    startHeartbeat()
    reconnectCount = 0
  }

  /**
   * 连接失败事件
   * @param e
   */
  function onError(e) {
    console.error(`WebSocket connection error：${e.code} ${e.reason} ${e.wasClean}`)
    reconnect()
  }

  /**
   * 连接关闭事件
   * @param e
   */
  function onClose() {
    reconnect()
  }

  /**
   * 重新连接
   */
  function reconnect() {
    if (!token) {
      return
    }

    if (lockReconnect || (OPTION.maxReconnect !== -1 && reconnectCount > OPTION.maxReconnect)) {
      return
    }

    lockReconnect = true
    setTimeout(() => {
      reconnectCount++
      // 建立新连接
      initWebSocket()
      lockReconnect = false
    }, 5000)
  }

  /**
   * 清空定时器
   */
  function clearTimeoutObj() {
    pingTimeoutObj && clearTimeout(pingTimeoutObj)
    pongTimeoutObj && clearTimeout(pongTimeoutObj)
  }

  /**
   * 开启心跳
   */
  function startHeartbeat() {
    console.log('startHeartbeat')
    // 清空定时器
    pingTimeoutObj && clearTimeout(pingTimeoutObj)
    // 延时发送下一次心跳
    pingTimeoutObj = setTimeout(() => {
      sendMessage(OPTION.pingMessage)
    }, OPTION.interval)
  }

  const serverMessage = ref()
  /**
   * 接收服务器推送的信息
   * @param msgEvent
   */
  function onMessage(msgEvent) {
    console.log('msgEvent')
    // 收到服务器信息，心跳重置并发送
    startHeartbeat()
    const text = msgEvent.data

    if (text === 'pong') {
      return
    }
    const data = JSON.parse(msgEvent.data)
    serverMessage.value = data.data
    onmessage && onmessage(data.data)
  }

  function sendMessage(message: string) {
    // 如果连接正常
    if (socket && socket.readyState === 1) {
      // 这里发送一个心跳，后端收到后，返回一个心跳消息，
      socket.send(message)
      // 心跳发送后，如果服务器超时未响应则断开，如果响应了会被重置心跳定时器
      pongTimeoutObj && clearTimeout(pongTimeoutObj)
      pongTimeoutObj = setTimeout(() => {
        socket && socket.close()
      }, OPTION.timeout)
    } else {
      // 否则重连
      reconnect()
    }
  }
  /**
   * 销毁socket
   */
  function destroySocket() {
    if (socket) {
      lockReconnect = true
      socket.close()
      clearTimeoutObj()
    }
  }
  //  断开连接
  onUnmounted(() => {
    destroySocket()
  })

  /**
   * 对象转url参数
   */
  function objectToEncodeQuery(query: object | undefined): string {
    if (!query) return ''

    const params = Object.keys(query)
      .filter(key => query[key])
      .map(key => {
        return encodeURIComponent(key) + '=' + encodeURIComponent(query[key])
      })
      .join('&')

    return params
  }

  initWebSocket()

  return { serverMessage, sendMessage, destroySocket }
}