diff --git a/packages/hooks/src/useWebSocket/__tests__/index.test.ts b/packages/hooks/src/useWebSocket/__tests__/index.test.ts index ad6949577b..a2d088717f 100644 --- a/packages/hooks/src/useWebSocket/__tests__/index.test.ts +++ b/packages/hooks/src/useWebSocket/__tests__/index.test.ts @@ -144,4 +144,96 @@ describe('useWebSocket', () => { act(() => wsServer1.close()); act(() => wsServer2.close()); }); + + it('should send heartbeat message periodically', async () => { + jest.spyOn(global, 'clearInterval'); + + const pingMessage = Date.now().toString(); + + const wsServer = new WS(wsUrl); + renderHook(() => + useWebSocket(wsUrl, { + heartbeat: { + message: pingMessage, + interval: 100, + responseTimeout: 200, + }, + }), + ); + + // Called on mount + expect(clearInterval).toHaveBeenCalledTimes(1); + + await act(async () => { + await wsServer.connected; + await sleep(110); + return promise; + }); + expect(wsServer.messages).toStrictEqual([pingMessage]); + + await act(async () => { + await sleep(110); + }); + expect(wsServer.messages).toStrictEqual([pingMessage, pingMessage]); + + expect(clearInterval).toHaveBeenCalledTimes(1); + await act(async () => { + wsServer.close(); + await wsServer.closed; + await sleep(110); + return promise; + }); + expect(clearInterval).toHaveBeenCalledTimes(2); + }); + + it('disconnect if no heartbeat message received', async () => { + const wsServer = new WS(wsUrl); + const hooks = renderHook(() => + useWebSocket(wsUrl, { + heartbeat: { + interval: 100, + responseTimeout: 200, + }, + }), + ); + + await act(async () => { + await wsServer.connected; + await sleep(310); + return promise; + }); + + expect(hooks.result.current.readyState).toBe(ReadyState.Closed); + + await act(async () => { + await wsServer.closed; + return promise; + }); + }); + + it('should ignore heartbeat response message', async () => { + const wsServer = new WS(wsUrl); + const hooks = renderHook(() => + useWebSocket(wsUrl, { heartbeat: { interval: 100, responseMessage: 'pong' } }), + ); + + await act(async () => { + await wsServer.connected; + return promise; + }); + await expect(wsServer).toReceiveMessage('ping'); + + act(() => { + wsServer.send('pong'); + }); + expect(hooks.result.current.latestMessage?.data).toBeUndefined(); + + const nowTime = `${Date.now()}`; + act(() => { + wsServer.send(nowTime); + }); + expect(hooks.result.current.latestMessage?.data).toBe(nowTime); + + act(() => wsServer.close()); + }); }); diff --git a/packages/hooks/src/useWebSocket/demo/demo2.tsx b/packages/hooks/src/useWebSocket/demo/demo2.tsx new file mode 100644 index 0000000000..152cdfa041 --- /dev/null +++ b/packages/hooks/src/useWebSocket/demo/demo2.tsx @@ -0,0 +1,64 @@ +import React, { useRef, useMemo } from 'react'; +import { useWebSocket } from 'ahooks'; + +enum ReadyState { + Connecting = 0, + Open = 1, + Closing = 2, + Closed = 3, +} + +export default () => { + const messageHistory = useRef([]); + + const { readyState, sendMessage, latestMessage, disconnect, connect } = useWebSocket( + 'wss://ws.postman-echo.com/raw', + { + heartbeat: { + message: 'ping', + responseMessage: 'pong', + interval: 3000, + responseTimeout: 10000, + }, + }, + ); + + messageHistory.current = useMemo( + () => messageHistory.current.concat(latestMessage), + [latestMessage], + ); + + return ( +
+ {/* send message */} + + {/* disconnect */} + + {/* connect */} + +
readyState: {readyState}
+
+

received message:

+ {messageHistory.current.map((message, index) => ( +

+ {message?.data} +

+ ))} +
+
+ ); +}; diff --git a/packages/hooks/src/useWebSocket/index.en-US.md b/packages/hooks/src/useWebSocket/index.en-US.md index 763a18681e..6c100ce0d4 100644 --- a/packages/hooks/src/useWebSocket/index.en-US.md +++ b/packages/hooks/src/useWebSocket/index.en-US.md @@ -13,6 +13,28 @@ A hook for WebSocket. +### Heartbeat example + +By setting the `heartbeat`, you can enable the heartbeat mechanism. After a successful connection, `useWebSocket` will send a `message` every `interval` milliseconds. If no messages are received within the `responseTimeout`, it may indicate that there is an issue with the connection, and the connection will be closed. + +It is important to note that if a `responseMessage` is defined, it will be ignored, and it will not trigger the `onMessage` event or update the `latestMessage`. + +```tsx | pure +useWebSocket( + 'wss://ws.postman-echo.com/raw', + { + heartbeat: { + message: 'ping', + responseMessage: 'pong', + interval: 3000, + responseTimeout: 10000, + } + } +); +``` + + + ## API ```typescript @@ -23,6 +45,13 @@ enum ReadyState { Closed = 3, } +interface HeartbeatOptions{ + message?: string; + responseMessage?: string; + interval?: number; + responseTimeout? :number; +} + interface Options { reconnectLimit?: number; reconnectInterval?: number; @@ -31,6 +60,7 @@ interface Options { onMessage?: (message: WebSocketEventMap['message'], instance: WebSocket) => void; onError?: (event: WebSocketEventMap['error'], instance: WebSocket) => void; protocols?: string | string[]; + heartbeat?: boolean | HeartbeatOptions; } interface Result { @@ -52,7 +82,7 @@ useWebSocket(socketUrl: string, options?: Options): Result; | socketUrl | Required, webSocket url | `string` | - | | options | connect the configuration item | `Options` | - | -#### Options +### Options | Options Property | Description | Type | Default | | ----------------- | ---------------------------------- | ---------------------------------------------------------------------- | ------- | @@ -64,6 +94,16 @@ useWebSocket(socketUrl: string, options?: Options): Result; | reconnectInterval | Retry interval(ms) | `number` | `3000` | | manual | Manually starts connection | `boolean` | `false` | | protocols | Sub protocols | `string` \| `string[]` | - | +| heartbeat | Heartbeat options | `boolean` \| `HeartbeatOptions` | `false` | + +### HeartbeatOptions + +| 参数 | 说明 | 类型 | 默认值 | +| --------------- | ---------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------- | -------- | ------- | +| message | Heartbeat message | `string` | `ping` | +| responseMessage | Heartbeat response message; the `onMessage` and `latestMessage` will ignore this message. | `string` | `pong` | +| interval | Heartbeat Interval(ms) | `number` | `5000` | +| responseTimeout | The heartbeat timeout (ms) indicates that if no heartbeat messages or other messages are received within this time, the connection will be considered abnormal and will be disconnected. | `number` | `10000` | ### Result diff --git a/packages/hooks/src/useWebSocket/index.ts b/packages/hooks/src/useWebSocket/index.ts index 8260195133..08c1552ca8 100644 --- a/packages/hooks/src/useWebSocket/index.ts +++ b/packages/hooks/src/useWebSocket/index.ts @@ -2,6 +2,13 @@ import { useEffect, useRef, useState } from 'react'; import useLatest from '../useLatest'; import useMemoizedFn from '../useMemoizedFn'; import useUnmount from '../useUnmount'; +import isObject from 'lodash/isObject'; +import isNil from 'lodash/isNil'; + +const DEFAULT_MESSAGE = { + PING: 'ping', + PONG: 'pong', +}; export enum ReadyState { Connecting = 0, @@ -10,6 +17,15 @@ export enum ReadyState { Closed = 3, } +export type HeartbeatMessage = Parameters[0]; + +export interface HeartbeatOptions { + message?: HeartbeatMessage; + responseMessage?: HeartbeatMessage; + interval?: number; + responseTimeout?: number; +} + export interface Options { reconnectLimit?: number; reconnectInterval?: number; @@ -18,8 +34,8 @@ export interface Options { onClose?: (event: WebSocketEventMap['close'], instance: WebSocket) => void; onMessage?: (message: WebSocketEventMap['message'], instance: WebSocket) => void; onError?: (event: WebSocketEventMap['error'], instance: WebSocket) => void; - protocols?: string | string[]; + heartbeat?: boolean | HeartbeatOptions; } export interface Result { @@ -41,8 +57,16 @@ export default function useWebSocket(socketUrl: string, options: Options = {}): onMessage, onError, protocols, + heartbeat = false, } = options; + const { + message: heartbeatMessage = DEFAULT_MESSAGE.PING, + responseMessage = DEFAULT_MESSAGE.PONG, + interval = 5 * 1000, + responseTimeout = 10 * 1000, + } = isObject(heartbeat) ? heartbeat : {}; + const onOpenRef = useLatest(onOpen); const onCloseRef = useLatest(onClose); const onMessageRef = useLatest(onMessage); @@ -51,6 +75,8 @@ export default function useWebSocket(socketUrl: string, options: Options = {}): const reconnectTimesRef = useRef(0); const reconnectTimerRef = useRef>(); const websocketRef = useRef(); + const heartbeatTimerRef = useRef>(); + const heartbeatTimeoutRef = useRef>(); const [latestMessage, setLatestMessage] = useState(); const [readyState, setReadyState] = useState(ReadyState.Closed); @@ -60,10 +86,7 @@ export default function useWebSocket(socketUrl: string, options: Options = {}): reconnectTimesRef.current < reconnectLimit && websocketRef.current?.readyState !== ReadyState.Open ) { - if (reconnectTimerRef.current) { - clearTimeout(reconnectTimerRef.current); - } - + clearTimeout(reconnectTimerRef.current); reconnectTimerRef.current = setTimeout(() => { // eslint-disable-next-line @typescript-eslint/no-use-before-define connectWs(); @@ -72,15 +95,26 @@ export default function useWebSocket(socketUrl: string, options: Options = {}): } }; - const connectWs = () => { - if (reconnectTimerRef.current) { - clearTimeout(reconnectTimerRef.current); - } + // Status code 1000 -> Normal Closure: https://developer.mozilla.org/en-US/docs/Web/API/CloseEvent/code + const disconnect: WebSocket['close'] = (code = 1000, reason) => { + clearTimeout(reconnectTimerRef.current); + clearInterval(heartbeatTimerRef.current); + clearTimeout(heartbeatTimeoutRef.current); + + reconnectTimesRef.current = reconnectLimit; + websocketRef.current?.close(code, reason); + websocketRef.current = undefined; + }; - if (websocketRef.current) { - websocketRef.current.close(); + const sendMessage: WebSocket['send'] = (message) => { + if (readyState === ReadyState.Open) { + websocketRef.current?.send(message); + } else { + throw new Error('WebSocket disconnected'); } + }; + const connectWs = () => { const ws = new WebSocket(socketUrl, protocols); setReadyState(ReadyState.Connecting); @@ -99,18 +133,43 @@ export default function useWebSocket(socketUrl: string, options: Options = {}): onOpenRef.current?.(event, ws); reconnectTimesRef.current = 0; setReadyState(ws.readyState || ReadyState.Open); + + if (heartbeat) { + heartbeatTimerRef.current = setInterval(() => { + if (ws.readyState === ReadyState.Open) { + ws.send(heartbeatMessage); + } + if (!isNil(heartbeatTimeoutRef.current)) { + return; + } + + heartbeatTimeoutRef.current = setTimeout(() => { + disconnect(); + }, responseTimeout); + }, interval); + } }; ws.onmessage = (message: WebSocketEventMap['message']) => { if (websocketRef.current !== ws) { return; } + if (heartbeat) { + clearTimeout(heartbeatTimeoutRef.current); + + if (responseMessage === message.data) { + return; + } + } + onMessageRef.current?.(message, ws); setLatestMessage(message); }; + ws.onclose = (event) => { onCloseRef.current?.(event, ws); // closed by server if (websocketRef.current === ws) { + // ws 关闭后，如果设置了超时重试的参数，则等待重试间隔时间后重试 reconnect(); } // closed by disconnect or closed by server @@ -122,29 +181,12 @@ export default function useWebSocket(socketUrl: string, options: Options = {}): websocketRef.current = ws; }; - const sendMessage: WebSocket['send'] = (message) => { - if (readyState === ReadyState.Open) { - websocketRef.current?.send(message); - } else { - throw new Error('WebSocket disconnected'); - } - }; - const connect = () => { + disconnect(); reconnectTimesRef.current = 0; connectWs(); }; - const disconnect = () => { - if (reconnectTimerRef.current) { - clearTimeout(reconnectTimerRef.current); - } - - reconnectTimesRef.current = reconnectLimit; - websocketRef.current?.close(); - websocketRef.current = undefined; - }; - useEffect(() => { if (!manual && socketUrl) { connect(); diff --git a/packages/hooks/src/useWebSocket/index.zh-CN.md b/packages/hooks/src/useWebSocket/index.zh-CN.md index 4d7f52a84d..94738693e2 100644 --- a/packages/hooks/src/useWebSocket/index.zh-CN.md +++ b/packages/hooks/src/useWebSocket/index.zh-CN.md @@ -13,6 +13,28 @@ nav: +### 心跳示例 + +通过设置 `heartbeat`，可以启用心跳机制，`useWebSocket` 在连接成功后，每隔 `interval` 毫秒发送一个 `message`，如果超过 `responseTimeout` 时间未收到任何消息，可能表示连接出问题，将关闭连接。 + +需要注意的是，如果定义了 `responseMessage`，该消息将被忽略，不会触发 `onMessage` 事件，也不会更新 `latestMessage`。 + +```tsx | pure +useWebSocket( + 'wss://ws.postman-echo.com/raw', + { + heartbeat: { + message: 'ping', + responseMessage: 'pong', + interval: 3000, + responseTimeout: 10000, + } + } +); +``` + + + ## API ```typescript @@ -23,6 +45,13 @@ enum ReadyState { Closed = 3, } +interface HeartbeatOptions{ + message?: string; + responseMessage?: string; + interval?: number; + responseTimeout? :number; +} + interface Options { reconnectLimit?: number; reconnectInterval?: number; @@ -31,6 +60,7 @@ interface Options { onMessage?: (message: WebSocketEventMap['message'], instance: WebSocket) => void; onError?: (event: WebSocketEventMap['error'], instance: WebSocket) => void; protocols?: string | string[]; + heartbeat?: boolean | HeartbeatOptions; } interface Result { @@ -52,7 +82,7 @@ useWebSocket(socketUrl: string, options?: Options): Result; | socketUrl | 必填，webSocket 地址 | `string` | - | | options | 可选，连接配置项 | `Options` | - | -#### Options +### Options | 参数 | 说明 | 类型 | 默认值 | | ----------------- | ---------------------- | ---------------------------------------------------------------------- | ------- | @@ -64,6 +94,16 @@ useWebSocket(socketUrl: string, options?: Options): Result; | reconnectInterval | 重试时间间隔（ms） | `number` | `3000` | | manual | 手动启动连接 | `boolean` | `false` | | protocols | 子协议 | `string` \| `string[]` | - | +| heartbeat | 心跳 | `boolean` \| `HeartbeatOptions` | `false` | + +### HeartbeatOptions + +| 参数 | 说明 | 类型 | 默认值 | +| --------------- | ---------------------------------------------------------------------------- | -------- | ------- | +| message | 心跳消息 | `string` | `ping` | +| responseMessage | 心跳回复消息，`onMessage`、`latestMessage` 会忽略该消息 | `string` | `pong` | +| interval | 心跳时间间隔（ms） | `number` | `5000` | +| responseTimeout | 心跳超时时间(ms)，超过此时间未收到心跳消息或其他消息将视为连接异常并断开连接 | `number` | `10000` | ### Result