在业务中使用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 }
}