From 3383688edd1a5ac3e41b9558066efbeba50611fe Mon Sep 17 00:00:00 2001
From: goodjun <goodjun@163.com>
Date: Thu, 25 Apr 2024 17:47:11 +0800
Subject: [PATCH 01/11] feat(useWebSocket): add heartbeat feature

---
 packages/hooks/src/useWebSocket/index.ts | 24 +++++++++++++++++++++++-
 1 file changed, 23 insertions(+), 1 deletion(-)

diff --git a/packages/hooks/src/useWebSocket/index.ts b/packages/hooks/src/useWebSocket/index.ts
index 8260195133..afc881e665 100644
--- a/packages/hooks/src/useWebSocket/index.ts
+++ b/packages/hooks/src/useWebSocket/index.ts
@@ -10,6 +10,12 @@ export enum ReadyState {
   Closed = 3,
 }
 
+export interface HeartbeatOptions {
+  message?: string;
+  returnMessage?: string;
+  interval?: number;
+}
+
 export interface Options {
   reconnectLimit?: number;
   reconnectInterval?: number;
@@ -18,8 +24,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 +47,12 @@ export default function useWebSocket(socketUrl: string, options: Options = {}):
     onMessage,
     onError,
     protocols,
+    heartbeat = false,
   } = options;
 
+  const { message: heartbeatMessage = 'ping', interval: heartbeatInterval = 60 * 1000 } =
+    typeof heartbeat === 'object' ? heartbeat : {};
+
   const onOpenRef = useLatest(onOpen);
   const onCloseRef = useLatest(onClose);
   const onMessageRef = useLatest(onMessage);
@@ -51,6 +61,7 @@ export default function useWebSocket(socketUrl: string, options: Options = {}):
   const reconnectTimesRef = useRef(0);
   const reconnectTimerRef = useRef<ReturnType<typeof setTimeout>>();
   const websocketRef = useRef<WebSocket>();
+  const heartbeatTimerRef = useRef<ReturnType<typeof setTimeout>>();
 
   const [latestMessage, setLatestMessage] = useState<WebSocketEventMap['message']>();
   const [readyState, setReadyState] = useState<ReadyState>(ReadyState.Closed);
@@ -98,12 +109,20 @@ export default function useWebSocket(socketUrl: string, options: Options = {}):
       }
       onOpenRef.current?.(event, ws);
       reconnectTimesRef.current = 0;
+      if (heartbeat) {
+        heartbeatTimerRef.current = setInterval(() => {
+          ws.send(heartbeatMessage);
+        }, heartbeatInterval);
+      }
       setReadyState(ws.readyState || ReadyState.Open);
     };
     ws.onmessage = (message: WebSocketEventMap['message']) => {
       if (websocketRef.current !== ws) {
         return;
       }
+      if (heartbeat && typeof heartbeat !== 'boolean' && heartbeat.returnMessage === message.data) {
+        return;
+      }
       onMessageRef.current?.(message, ws);
       setLatestMessage(message);
     };
@@ -116,6 +135,9 @@ export default function useWebSocket(socketUrl: string, options: Options = {}):
       // closed by disconnect or closed by server
       if (!websocketRef.current || websocketRef.current === ws) {
         setReadyState(ws.readyState || ReadyState.Closed);
+        if (heartbeat) {
+          clearInterval(heartbeatTimerRef.current);
+        }
       }
     };
 

From 732fc96ffcb26255c59e3f5a80ce6fdbe1430abe Mon Sep 17 00:00:00 2001
From: goodjun <goodjun@163.com>
Date: Thu, 25 Apr 2024 17:47:33 +0800
Subject: [PATCH 02/11] feat(useWebSocket): add test case

---
 .../src/useWebSocket/__tests__/index.test.ts  | 53 +++++++++++++++++++
 1 file changed, 53 insertions(+)

diff --git a/packages/hooks/src/useWebSocket/__tests__/index.test.ts b/packages/hooks/src/useWebSocket/__tests__/index.test.ts
index ad6949577b..be1a61d070 100644
--- a/packages/hooks/src/useWebSocket/__tests__/index.test.ts
+++ b/packages/hooks/src/useWebSocket/__tests__/index.test.ts
@@ -144,4 +144,57 @@ describe('useWebSocket', () => {
     act(() => wsServer1.close());
     act(() => wsServer2.close());
   });
+
+  it('should send heartbeat message periodically', async () => {
+    jest.spyOn(global, 'clearInterval');
+
+    const wsServer = new WS(wsUrl);
+    renderHook(() => useWebSocket(wsUrl, { heartbeat: { interval: 100 } }));
+    await act(async () => {
+      await wsServer.connected;
+      return promise;
+    });
+
+    jest.advanceTimersByTime(100);
+    await expect(wsServer).toReceiveMessage('ping');
+
+    jest.advanceTimersByTime(100);
+    await expect(wsServer).toReceiveMessage('ping');
+
+    expect(wsServer).toHaveReceivedMessages(['ping', 'ping']);
+
+    act(() => wsServer.close());
+    await act(async () => {
+      await wsServer.closed;
+      return promise;
+    });
+    expect(clearInterval).toHaveBeenCalledTimes(1);
+  });
+
+  it('should ignore heartbeat response message', async () => {
+    const wsServer = new WS(wsUrl);
+    const hooks = renderHook(() =>
+      useWebSocket(wsUrl, { heartbeat: { interval: 100, returnMessage: 'pong' } }),
+    );
+    await act(async () => {
+      await wsServer.connected;
+      return promise;
+    });
+
+    jest.advanceTimersByTime(100);
+    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());
+  });
 });

From 9890c1d6add3f142e7e075e305dce355da98d964 Mon Sep 17 00:00:00 2001
From: goodjun <goodjun@163.com>
Date: Thu, 25 Apr 2024 17:48:54 +0800
Subject: [PATCH 03/11] docs(useWebSocket): update doc

---
 packages/hooks/src/useWebSocket/index.en-US.md |  9 +++++++++
 packages/hooks/src/useWebSocket/index.zh-CN.md | 16 ++++++++++++++++
 2 files changed, 25 insertions(+)

diff --git a/packages/hooks/src/useWebSocket/index.en-US.md b/packages/hooks/src/useWebSocket/index.en-US.md
index 763a18681e..62f3d2eaf7 100644
--- a/packages/hooks/src/useWebSocket/index.en-US.md
+++ b/packages/hooks/src/useWebSocket/index.en-US.md
@@ -64,6 +64,15 @@ 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` |
+| returnMessage | Heartbeat response message, `latestMessage` will ignore this message | `string` | -      |
+| interval      | Heartbeat Interval(ms)                                               | `number` | `6000` |
 
 ### Result
 
diff --git a/packages/hooks/src/useWebSocket/index.zh-CN.md b/packages/hooks/src/useWebSocket/index.zh-CN.md
index 4d7f52a84d..6d5297a803 100644
--- a/packages/hooks/src/useWebSocket/index.zh-CN.md
+++ b/packages/hooks/src/useWebSocket/index.zh-CN.md
@@ -23,6 +23,12 @@ enum ReadyState {
   Closed = 3,
 }
 
+interface HeartbeatOptions{
+  message?: string;
+  returnMessage?: string;
+  interval?: number;
+}
+
 interface Options {
   reconnectLimit?: number;
   reconnectInterval?: number;
@@ -31,6 +37,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 {
@@ -64,6 +71,15 @@ useWebSocket(socketUrl: string, options?: Options): Result;
 | reconnectInterval | 重试时间间隔（ms）     | `number`                                                               | `3000`  |
 | manual            | 手动启动连接           | `boolean`                                                              | `false` |
 | protocols         | 子协议                 | `string` \| `string[]`                                                 | -       |
+| heartbeat         | 心跳                   | `boolean` \| `HeartbeatOptions`                                        | `false` |
+
+### HeartbeatOptions
+
+| 参数          | 说明                                       | 类型     | 默认值 |
+| ------------- | ------------------------------------------ | -------- | ------ |
+| message       | 心跳消息                                   | `string` | `ping` |
+| returnMessage | 心跳回复消息，`latestMessage` 会忽略该消息 | `string` | -      |
+| interval      | 心跳时间间隔（ms）                         | `number` | `6000` |
 
 ### Result
 

From c86813b68efde9b4ef50697e26fed791499c08af Mon Sep 17 00:00:00 2001
From: goodjun <goodjun@163.com>
Date: Sat, 27 Apr 2024 08:04:49 +0800
Subject: [PATCH 04/11] fix: heartbeat timer

---
 packages/hooks/src/useWebSocket/index.ts | 6 +++++-
 1 file changed, 5 insertions(+), 1 deletion(-)

diff --git a/packages/hooks/src/useWebSocket/index.ts b/packages/hooks/src/useWebSocket/index.ts
index afc881e665..e082ee3d99 100644
--- a/packages/hooks/src/useWebSocket/index.ts
+++ b/packages/hooks/src/useWebSocket/index.ts
@@ -61,7 +61,7 @@ export default function useWebSocket(socketUrl: string, options: Options = {}):
   const reconnectTimesRef = useRef(0);
   const reconnectTimerRef = useRef<ReturnType<typeof setTimeout>>();
   const websocketRef = useRef<WebSocket>();
-  const heartbeatTimerRef = useRef<ReturnType<typeof setTimeout>>();
+  const heartbeatTimerRef = useRef<ReturnType<typeof setInterval>>();
 
   const [latestMessage, setLatestMessage] = useState<WebSocketEventMap['message']>();
   const [readyState, setReadyState] = useState<ReadyState>(ReadyState.Closed);
@@ -162,6 +162,10 @@ export default function useWebSocket(socketUrl: string, options: Options = {}):
       clearTimeout(reconnectTimerRef.current);
     }
 
+    if (heartbeatTimerRef.current) {
+      clearInterval(heartbeatTimerRef.current);
+    }
+
     reconnectTimesRef.current = reconnectLimit;
     websocketRef.current?.close();
     websocketRef.current = undefined;

From 497a58003c87a44b9035c1b85076037a212b411a Mon Sep 17 00:00:00 2001
From: goodjun <goodjun@163.com>
Date: Sat, 27 Apr 2024 08:27:39 +0800
Subject: [PATCH 05/11] fix(useWebSocket): fix test case

---
 packages/hooks/src/useWebSocket/__tests__/index.test.ts | 3 ---
 1 file changed, 3 deletions(-)

diff --git a/packages/hooks/src/useWebSocket/__tests__/index.test.ts b/packages/hooks/src/useWebSocket/__tests__/index.test.ts
index be1a61d070..d8846b5b0e 100644
--- a/packages/hooks/src/useWebSocket/__tests__/index.test.ts
+++ b/packages/hooks/src/useWebSocket/__tests__/index.test.ts
@@ -155,10 +155,8 @@ describe('useWebSocket', () => {
       return promise;
     });
 
-    jest.advanceTimersByTime(100);
     await expect(wsServer).toReceiveMessage('ping');
 
-    jest.advanceTimersByTime(100);
     await expect(wsServer).toReceiveMessage('ping');
 
     expect(wsServer).toHaveReceivedMessages(['ping', 'ping']);
@@ -181,7 +179,6 @@ describe('useWebSocket', () => {
       return promise;
     });
 
-    jest.advanceTimersByTime(100);
     await expect(wsServer).toReceiveMessage('ping');
 
     act(() => {

From 532286c3e0a3b899f4a62d156986fed958fffcb2 Mon Sep 17 00:00:00 2001
From: liuyib <1656081615@qq.com>
Date: Thu, 25 Apr 2024 19:13:06 +0800
Subject: [PATCH 06/11] docs: update title

---
 packages/hooks/src/useWebSocket/index.en-US.md | 2 +-
 packages/hooks/src/useWebSocket/index.zh-CN.md | 2 +-
 2 files changed, 2 insertions(+), 2 deletions(-)

diff --git a/packages/hooks/src/useWebSocket/index.en-US.md b/packages/hooks/src/useWebSocket/index.en-US.md
index 62f3d2eaf7..1c71b9110e 100644
--- a/packages/hooks/src/useWebSocket/index.en-US.md
+++ b/packages/hooks/src/useWebSocket/index.en-US.md
@@ -52,7 +52,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 |
 | ----------------- | ---------------------------------- | ---------------------------------------------------------------------- | ------- |
diff --git a/packages/hooks/src/useWebSocket/index.zh-CN.md b/packages/hooks/src/useWebSocket/index.zh-CN.md
index 6d5297a803..40b22802f6 100644
--- a/packages/hooks/src/useWebSocket/index.zh-CN.md
+++ b/packages/hooks/src/useWebSocket/index.zh-CN.md
@@ -59,7 +59,7 @@ useWebSocket(socketUrl: string, options?: Options): Result;
 | socketUrl | 必填，webSocket 地址 | `string`  | -      |
 | options   | 可选，连接配置项     | `Options` | -      |
 
-#### Options
+### Options
 
 | 参数              | 说明                   | 类型                                                                   | 默认值  |
 | ----------------- | ---------------------- | ---------------------------------------------------------------------- | ------- |

From 9e32b815bfba84661489fec242710bfe06df9f8b Mon Sep 17 00:00:00 2001
From: liuyib <1656081615@qq.com>
Date: Thu, 12 Sep 2024 09:38:47 +0800
Subject: [PATCH 07/11] refactor: add heartbeat params `responseTimeout`

---
 .../src/useWebSocket/__tests__/index.test.ts  |  37 +++++--
 .../hooks/src/useWebSocket/demo/demo2.tsx     |  61 ++++++++++
 .../hooks/src/useWebSocket/index.en-US.md     |  10 +-
 packages/hooks/src/useWebSocket/index.ts      | 104 ++++++++++--------
 .../hooks/src/useWebSocket/index.zh-CN.md     |  16 ++-
 5 files changed, 163 insertions(+), 65 deletions(-)
 create mode 100644 packages/hooks/src/useWebSocket/demo/demo2.tsx

diff --git a/packages/hooks/src/useWebSocket/__tests__/index.test.ts b/packages/hooks/src/useWebSocket/__tests__/index.test.ts
index d8846b5b0e..c4c08be719 100644
--- a/packages/hooks/src/useWebSocket/__tests__/index.test.ts
+++ b/packages/hooks/src/useWebSocket/__tests__/index.test.ts
@@ -149,36 +149,53 @@ describe('useWebSocket', () => {
     jest.spyOn(global, 'clearInterval');
 
     const wsServer = new WS(wsUrl);
-    renderHook(() => useWebSocket(wsUrl, { heartbeat: { interval: 100 } }));
+    renderHook(() =>
+      useWebSocket(wsUrl, {
+        heartbeat: {
+          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(['ping']);
 
-    await expect(wsServer).toReceiveMessage('ping');
-
-    await expect(wsServer).toReceiveMessage('ping');
-
-    expect(wsServer).toHaveReceivedMessages(['ping', 'ping']);
+    await act(async () => {
+      await sleep(110);
+    });
+    expect(wsServer.messages).toStrictEqual(['ping', 'ping']);
 
-    act(() => wsServer.close());
+    expect(clearInterval).toHaveBeenCalledTimes(1);
     await act(async () => {
+      wsServer.close();
       await wsServer.closed;
+      await sleep(110);
       return promise;
     });
-    expect(clearInterval).toHaveBeenCalledTimes(1);
+    expect(clearInterval).toHaveBeenCalledTimes(2);
   });
 
+  // TODO: 更详细的测试心跳相关的所有参数
+  // TODO: 心跳逻辑 demo 完善、demo 测试
+
   it('should ignore heartbeat response message', async () => {
     const wsServer = new WS(wsUrl);
     const hooks = renderHook(() =>
-      useWebSocket(wsUrl, { heartbeat: { interval: 100, returnMessage: 'pong' } }),
+      useWebSocket(wsUrl, { heartbeat: { interval: 100, responseMessage: 'pong' } }),
     );
+
     await act(async () => {
       await wsServer.connected;
       return promise;
     });
-
     await expect(wsServer).toReceiveMessage('ping');
 
     act(() => {
diff --git a/packages/hooks/src/useWebSocket/demo/demo2.tsx b/packages/hooks/src/useWebSocket/demo/demo2.tsx
new file mode 100644
index 0000000000..4df5dd7b40
--- /dev/null
+++ b/packages/hooks/src/useWebSocket/demo/demo2.tsx
@@ -0,0 +1,61 @@
+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: {
+        interval: 1000 * 3,
+      },
+    },
+  );
+
+  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>
+  );
+};
diff --git a/packages/hooks/src/useWebSocket/index.en-US.md b/packages/hooks/src/useWebSocket/index.en-US.md
index 1c71b9110e..bbb4aeb90b 100644
--- a/packages/hooks/src/useWebSocket/index.en-US.md
+++ b/packages/hooks/src/useWebSocket/index.en-US.md
@@ -68,11 +68,11 @@ useWebSocket(socketUrl: string, options?: Options): Result;
 
 ### HeartbeatOptions
 
-| 参数          | 说明                                                                 | 类型     | 默认值 |
-| ------------- | -------------------------------------------------------------------- | -------- | ------ |
-| message       | Heartbeat message                                                    | `string` | `ping` |
-| returnMessage | Heartbeat response message, `latestMessage` will ignore this message | `string` | -      |
-| interval      | Heartbeat Interval(ms)                                               | `number` | `6000` |
+| 参数            | 说明                                                                 | 类型     | 默认值 |
+| --------------- | -------------------------------------------------------------------- | -------- | ------ |
+| message         | Heartbeat message                                                    | `string` | `ping` |
+| responseMessage | Heartbeat response message, `latestMessage` will ignore this message | `string` | -      |
+| interval        | Heartbeat Interval(ms)                                               | `number` | `6000` |
 
 ### Result
 
diff --git a/packages/hooks/src/useWebSocket/index.ts b/packages/hooks/src/useWebSocket/index.ts
index e082ee3d99..77131b561b 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,10 +17,13 @@ export enum ReadyState {
   Closed = 3,
 }
 
+export type HeartbeatMessage = Parameters<WebSocket['send']>[0];
+
 export interface HeartbeatOptions {
-  message?: string;
-  returnMessage?: string;
+  message?: HeartbeatMessage;
+  responseMessage?: HeartbeatMessage;
   interval?: number;
+  responseTimeout?: number;
 }
 
 export interface Options {
@@ -50,8 +60,12 @@ export default function useWebSocket(socketUrl: string, options: Options = {}):
     heartbeat = false,
   } = options;
 
-  const { message: heartbeatMessage = 'ping', interval: heartbeatInterval = 60 * 1000 } =
-    typeof heartbeat === 'object' ? heartbeat : {};
+  const {
+    message: heartbeatMessage = DEFAULT_MESSAGE.PING,
+    responseMessage = DEFAULT_MESSAGE.PONG,
+    interval = 1000,
+    responseTimeout = 1000,
+  } = isObject(heartbeat) ? heartbeat : {};
 
   const onOpenRef = useLatest(onOpen);
   const onCloseRef = useLatest(onClose);
@@ -62,6 +76,7 @@ export default function useWebSocket(socketUrl: string, options: Options = {}):
   const reconnectTimerRef = useRef<ReturnType<typeof setTimeout>>();
   const websocketRef = useRef<WebSocket>();
   const heartbeatTimerRef = useRef<ReturnType<typeof setInterval>>();
+  const heartbeatTimeoutRef = useRef<ReturnType<typeof setTimeout>>();
 
   const [latestMessage, setLatestMessage] = useState<WebSocketEventMap['message']>();
   const [readyState, setReadyState] = useState<ReadyState>(ReadyState.Closed);
@@ -71,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();
@@ -83,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);
 
@@ -109,68 +132,61 @@ 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(() => {
-          ws.send(heartbeatMessage);
-        }, heartbeatInterval);
+          if (ws.readyState === ReadyState.Open) {
+            ws.send(heartbeatMessage);
+          }
+          if (!isNil(heartbeatTimeoutRef.current)) {
+            return;
+          }
+
+          heartbeatTimeoutRef.current = setTimeout(() => {
+            disconnect();
+          }, responseTimeout);
+        }, interval);
       }
-      setReadyState(ws.readyState || ReadyState.Open);
     };
     ws.onmessage = (message: WebSocketEventMap['message']) => {
       if (websocketRef.current !== ws) {
         return;
       }
-      if (heartbeat && typeof heartbeat !== 'boolean' && heartbeat.returnMessage === message.data) {
-        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
       if (!websocketRef.current || websocketRef.current === ws) {
         setReadyState(ws.readyState || ReadyState.Closed);
-        if (heartbeat) {
-          clearInterval(heartbeatTimerRef.current);
-        }
       }
     };
 
     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);
-    }
-
-    if (heartbeatTimerRef.current) {
-      clearInterval(heartbeatTimerRef.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 40b22802f6..51e9ddf701 100644
--- a/packages/hooks/src/useWebSocket/index.zh-CN.md
+++ b/packages/hooks/src/useWebSocket/index.zh-CN.md
@@ -13,6 +13,10 @@ nav:
 
 <code src="./demo/demo1.tsx" />
 
+### 测试示例
+
+<code src="./demo/demo2.tsx" />
+
 ## API
 
 ```typescript
@@ -25,7 +29,7 @@ enum ReadyState {
 
 interface HeartbeatOptions{
   message?: string;
-  returnMessage?: string;
+  responseMessage?: string;
   interval?: number;
 }
 
@@ -75,11 +79,11 @@ useWebSocket(socketUrl: string, options?: Options): Result;
 
 ### HeartbeatOptions
 
-| 参数          | 说明                                       | 类型     | 默认值 |
-| ------------- | ------------------------------------------ | -------- | ------ |
-| message       | 心跳消息                                   | `string` | `ping` |
-| returnMessage | 心跳回复消息，`latestMessage` 会忽略该消息 | `string` | -      |
-| interval      | 心跳时间间隔（ms）                         | `number` | `6000` |
+| 参数            | 说明                                       | 类型     | 默认值 |
+| --------------- | ------------------------------------------ | -------- | ------ |
+| message         | 心跳消息                                   | `string` | `ping` |
+| responseMessage | 心跳回复消息，`latestMessage` 会忽略该消息 | `string` | `pong` |
+| interval        | 心跳时间间隔（ms）                         | `number` | `6000` |
 
 ### Result
 

From cca313b132a089938bee8fcdf56dccd286756e15 Mon Sep 17 00:00:00 2001
From: goodjun <goodjun@163.com>
Date: Mon, 23 Sep 2024 22:07:29 +0800
Subject: [PATCH 08/11] test: update heartbeat test case

---
 .../src/useWebSocket/__tests__/index.test.ts  | 33 ++++++++++++++++---
 1 file changed, 29 insertions(+), 4 deletions(-)

diff --git a/packages/hooks/src/useWebSocket/__tests__/index.test.ts b/packages/hooks/src/useWebSocket/__tests__/index.test.ts
index c4c08be719..6d599ed968 100644
--- a/packages/hooks/src/useWebSocket/__tests__/index.test.ts
+++ b/packages/hooks/src/useWebSocket/__tests__/index.test.ts
@@ -148,10 +148,13 @@ describe('useWebSocket', () => {
   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,
         },
@@ -166,12 +169,12 @@ describe('useWebSocket', () => {
       await sleep(110);
       return promise;
     });
-    expect(wsServer.messages).toStrictEqual(['ping']);
+    expect(wsServer.messages).toStrictEqual([pingMessage]);
 
     await act(async () => {
       await sleep(110);
     });
-    expect(wsServer.messages).toStrictEqual(['ping', 'ping']);
+    expect(wsServer.messages).toStrictEqual([pingMessage, pingMessage]);
 
     expect(clearInterval).toHaveBeenCalledTimes(1);
     await act(async () => {
@@ -183,8 +186,30 @@ describe('useWebSocket', () => {
     expect(clearInterval).toHaveBeenCalledTimes(2);
   });
 
-  // TODO: 更详细的测试心跳相关的所有参数
-  // TODO: 心跳逻辑 demo 完善、demo 测试
+  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(350);
+      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);

From 36deda1a37fac9c54f0e77850773dd259fd95297 Mon Sep 17 00:00:00 2001
From: goodjun <goodjun@163.com>
Date: Tue, 24 Sep 2024 08:11:13 +0800
Subject: [PATCH 09/11] test: update heartbeat test case

---
 packages/hooks/src/useWebSocket/__tests__/index.test.ts | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/packages/hooks/src/useWebSocket/__tests__/index.test.ts b/packages/hooks/src/useWebSocket/__tests__/index.test.ts
index 6d599ed968..a2d088717f 100644
--- a/packages/hooks/src/useWebSocket/__tests__/index.test.ts
+++ b/packages/hooks/src/useWebSocket/__tests__/index.test.ts
@@ -199,7 +199,7 @@ describe('useWebSocket', () => {
 
     await act(async () => {
       await wsServer.connected;
-      await sleep(350);
+      await sleep(310);
       return promise;
     });
 

From 03d87d31276f95567c08974329fc7c200b78eea9 Mon Sep 17 00:00:00 2001
From: goodjun <goodjun@163.com>
Date: Tue, 24 Sep 2024 13:55:33 +0800
Subject: [PATCH 10/11] feat(useWebSocket): update default values for
 `interval` and `responseTimeout`

---
 packages/hooks/src/useWebSocket/index.ts | 4 ++--
 1 file changed, 2 insertions(+), 2 deletions(-)

diff --git a/packages/hooks/src/useWebSocket/index.ts b/packages/hooks/src/useWebSocket/index.ts
index 77131b561b..08c1552ca8 100644
--- a/packages/hooks/src/useWebSocket/index.ts
+++ b/packages/hooks/src/useWebSocket/index.ts
@@ -63,8 +63,8 @@ export default function useWebSocket(socketUrl: string, options: Options = {}):
   const {
     message: heartbeatMessage = DEFAULT_MESSAGE.PING,
     responseMessage = DEFAULT_MESSAGE.PONG,
-    interval = 1000,
-    responseTimeout = 1000,
+    interval = 5 * 1000,
+    responseTimeout = 10 * 1000,
   } = isObject(heartbeat) ? heartbeat : {};
 
   const onOpenRef = useLatest(onOpen);

From 4fd86453f177cfe194209e33ddd5acd4767dba95 Mon Sep 17 00:00:00 2001
From: goodjun <goodjun@163.com>
Date: Tue, 24 Sep 2024 13:56:51 +0800
Subject: [PATCH 11/11] docs(useWebSocket): update heartbeat doc

---
 .../hooks/src/useWebSocket/demo/demo2.tsx     |  5 ++-
 .../hooks/src/useWebSocket/index.en-US.md     | 41 ++++++++++++++++---
 .../hooks/src/useWebSocket/index.zh-CN.md     | 32 ++++++++++++---
 3 files changed, 66 insertions(+), 12 deletions(-)

diff --git a/packages/hooks/src/useWebSocket/demo/demo2.tsx b/packages/hooks/src/useWebSocket/demo/demo2.tsx
index 4df5dd7b40..152cdfa041 100644
--- a/packages/hooks/src/useWebSocket/demo/demo2.tsx
+++ b/packages/hooks/src/useWebSocket/demo/demo2.tsx
@@ -15,7 +15,10 @@ export default () => {
     'wss://ws.postman-echo.com/raw',
     {
       heartbeat: {
-        interval: 1000 * 3,
+        message: 'ping',
+        responseMessage: 'pong',
+        interval: 3000,
+        responseTimeout: 10000,
       },
     },
   );
diff --git a/packages/hooks/src/useWebSocket/index.en-US.md b/packages/hooks/src/useWebSocket/index.en-US.md
index bbb4aeb90b..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.
 
 <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 {
@@ -68,11 +98,12 @@ useWebSocket(socketUrl: string, options?: Options): Result;
 
 ### HeartbeatOptions
 
-| 参数            | 说明                                                                 | 类型     | 默认值 |
-| --------------- | -------------------------------------------------------------------- | -------- | ------ |
-| message         | Heartbeat message                                                    | `string` | `ping` |
-| responseMessage | Heartbeat response message, `latestMessage` will ignore this message | `string` | -      |
-| interval        | Heartbeat Interval(ms)                                               | `number` | `6000` |
+| 参数            | 说明                                                                                                                                                                                     | 类型     | 默认值  |
+| --------------- | ---------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------- | -------- | ------- |
+| 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.zh-CN.md b/packages/hooks/src/useWebSocket/index.zh-CN.md
index 51e9ddf701..94738693e2 100644
--- a/packages/hooks/src/useWebSocket/index.zh-CN.md
+++ b/packages/hooks/src/useWebSocket/index.zh-CN.md
@@ -13,7 +13,25 @@ nav:
 
 <code src="./demo/demo1.tsx" />
 
-### 测试示例
+### 心跳示例
+
+通过设置 `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,
+    }
+  }
+);
+```
 
 <code src="./demo/demo2.tsx" />
 
@@ -31,6 +49,7 @@ interface HeartbeatOptions{
   message?: string;
   responseMessage?: string;
   interval?: number;
+  responseTimeout? :number;
 }
 
 interface Options {
@@ -79,11 +98,12 @@ useWebSocket(socketUrl: string, options?: Options): Result;
 
 ### HeartbeatOptions
 
-| 参数            | 说明                                       | 类型     | 默认值 |
-| --------------- | ------------------------------------------ | -------- | ------ |
-| message         | 心跳消息                                   | `string` | `ping` |
-| responseMessage | 心跳回复消息，`latestMessage` 会忽略该消息 | `string` | `pong` |
-| interval        | 心跳时间间隔（ms）                         | `number` | `6000` |
+| 参数            | 说明                                                                         | 类型     | 默认值  |
+| --------------- | ---------------------------------------------------------------------------- | -------- | ------- |
+| message         | 心跳消息                                                                     | `string` | `ping`  |
+| responseMessage | 心跳回复消息，`onMessage`、`latestMessage` 会忽略该消息                      | `string` | `pong`  |
+| interval        | 心跳时间间隔（ms）                                                           | `number` | `5000`  |
+| responseTimeout | 心跳超时时间(ms)，超过此时间未收到心跳消息或其他消息将视为连接异常并断开连接 | `number` | `10000` |
 
 ### Result