React 18使用websocket，在React封装websocket，在React中使用WebSocket

在业务中使用WebSocket是很常见的需求，但是原生WebSocket使用太过复杂，本文将介绍如何在ahooks3.0库 useWebSocket() hook的基础上再次封装，达到app只使用一个WebSocket链接进行通信的目的。

1. src/contexts/WebSocketContext.tsx

import type { ReadyState } from "ahooks/es/useWebSocket"
import { createContext, useContext } from "react"

export interface WebSocketContext {
  sendMessage: (message: string | object) => void
  readyState: ReadyState
  lastMessage: MessageEvent<any> | null
  connect: () => void
  disconnect: () => void
  webSocketIns: WebSocket | null
}

export const WebSocketContext = createContext<WebSocketContext>({} as WebSocketContext)

export const useWebSocketContext = () => useContext(WebSocketContext)

2. src/providers/WebSocketProvider.tsx

import { WebSocketContext } from "@/contexts/WebSocketContext"
import { useWebSocketHeartbeat } from "@/hooks/useWebSocketHeartbeat"
import { useWebSocket } from "ahooks"
import type React from "react"
import { useEffect, useMemo } from "react"

const WS_URL = `${import.meta.env.VITE_WEB_SOCKET_BASE}`

export const WebSocketProvider = ({ children }: { children: React.ReactNode }) => {
  const {
    readyState,
    sendMessage: originalSend,
    connect,
    disconnect,
    latestMessage,
    webSocketIns
  } = useWebSocket(WS_URL, {
    reconnectLimit: 3, // 最大重试次数
    reconnectInterval: 3000, // 重试间隔
    manual: true // 手动启动连接
  })

  // 统一封装发送方法，支持对象自动序列化
  const sendMessage = useMemo(() => {
    return (message: string | object) => {
      let _message = message
      if (typeof _message !== "string") {
        _message = JSON.stringify(message)
      }
      originalSend(_message)
    }
  }, [originalSend])

  // 初始化连接逻辑
  useEffect(() => {
    connect() // 启动连接
    return () => disconnect() // 组件卸载时断开
  }, [connect, disconnect])

  // 可选：使用心跳检测 hook
  const { handlePong } = useWebSocketHeartbeat({
    sendMessage,
    readyState,
    interval: 3000
  })

  // 处理消息
  useEffect(() => {
    if (latestMessage?.data === 'pong') {
      handlePong()
    }
  }, [latestMessage])

  // 暴露给子组件的上下文值
  const contextValue = useMemo(
    () => ({
      sendMessage,
      readyState,
      lastMessage: latestMessage || null,
      connect,
      disconnect,
      webSocketIns: webSocketIns || null
    }),
    [sendMessage, readyState, latestMessage, connect, disconnect]
  )

  return <WebSocketContext.Provider value={contextValue}>{children}</WebSocketContext.Provider>
}

3. 在组件中使用：

import { useWebSocketContext } from "@/contexts/WebSocketContext.tsx"
import { useEffect } from "react"

const Send = () => {
  const { sendMessage, readyState, lastMessage } = useWebSocketContext()

  useEffect(() => {
    if (readyState === 1) {
      // 聊天
      sendMessage("阿巴阿巴")
    }
  }, [readyState])

  useEffect(() => {
    console.log("推送过来的聊天记录：", lastMessage)
  }, [lastMessage])

  return (
    <div>测试WebSocket</div>
  )
}

export default Send

注意⚠️:

可以观察到useWebSocketContext导出了webSocketIns，这个是全局的socket实例，导出是为了直接获取到socket实例做一些操作，操作不能使用直接调用的形式

例如: webSocketIns.onmessage，这会覆盖hook的onmessage从而导致lastMessage不更新等问题❌

使用webSocketIns.addEventListener("message",()=>{}）的方式代替✅

---

可选： 因为业务上有显示服务器延迟的需求，所以自己定义了另一个useWebSocketHeartbeat hook 用来计算服务器延迟

import { useEffect, useRef } from 'react'
import { useSetAtom } from 'jotai'
import { serverPing } from '@/store/statusBar'

interface UseWebSocketHeartbeatOptions {
  sendMessage: (message: string | object) => void
  readyState: number
  interval?: number
}

export const useWebSocketHeartbeat = ({
  sendMessage,
  readyState,
  interval = 3000
}: UseWebSocketHeartbeatOptions) => {
  const pingIntervalRef = useRef<NodeJS.Timeout | null>(null)
  const pingSentTimeRef = useRef<number | null>(null)
  const setServerPing = useSetAtom(serverPing)

  useEffect(() => {
    const startPing = () => {
      if (readyState === WebSocket.OPEN) {
        pingIntervalRef.current = setInterval(() => {
          try {
            pingSentTimeRef.current = Date.now()
            sendMessage('ping')
          } catch (error) {
            console.error('Failed to send ping:', error)
          }
        }, interval)
      }
    }

    const stopPing = () => {
      if (pingIntervalRef.current) {
        clearInterval(pingIntervalRef.current)
        pingIntervalRef.current = null
      }
    }

    // 启动心跳
    startPing()

    // 清理
    return () => stopPing()
  }, [readyState, sendMessage, interval])

  // 处理 pong 响应的函数
  const handlePong = () => {
    if (pingSentTimeRef.current) {
      const latency = Date.now() - pingSentTimeRef.current
      setServerPing(`${latency}`)
    }
  }

  return { handlePong }
}