feat(useWebSocket): add heartbeat feature

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());
+  });
 });

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<any[]>([]);
+
+  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 (
+    <div>
+      {/* send message */}
+      <button
+        onClick={() => sendMessage && sendMessage(`${Date.now()}`)}
+        disabled={readyState !== ReadyState.Open}
+        style={{ marginRight: 8 }}
+      >
+        ✉️ send
+      </button>
+      {/* disconnect */}
+      <button
+        onClick={() => disconnect && disconnect()}
+        disabled={readyState !== ReadyState.Open}
+        style={{ marginRight: 8 }}
+      >
+        ❌ disconnect
+      </button>
+      {/* connect */}
+      <button onClick={() => connect && connect()} disabled={readyState === ReadyState.Open}>
+        {readyState === ReadyState.Connecting ? 'connecting' : '📞 connect'}
+      </button>
+      <div style={{ marginTop: 8 }}>readyState: {readyState}</div>
+      <div style={{ marginTop: 8 }}>
+        <p>received message: </p>
+        {messageHistory.current.map((message, index) => (
+          <p key={index} style={{ wordWrap: 'break-word' }}>
+            {message?.data}
+          </p>
+        ))}
+      </div>
+    </div>
+  );
+};

packages/hooks/src/useWebSocket/index.en-US.md

@@ -13,6 +13,28 @@ A hook for WebSocket.
 
 <code src="./demo/demo1.tsx" />
 
+### 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,
+    }
+  }
+);
+```
+
+<code src="./demo/demo2.tsx" />
+
 ## 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