From ca5981fc337d0497bd734dae13728f88c2ea3dcf Mon Sep 17 00:00:00 2001
From: azaleta <24407500@qq.com>
Date: Mon, 5 Sep 2022 11:36:43 +0800
Subject: [PATCH] feat(useWebSocket): enhence heartbeat

---
 packages/core/useWebSocket/index.md |  1 +
 packages/core/useWebSocket/index.ts | 36 +++++++++++++++++++++++++++--
 2 files changed, 35 insertions(+), 2 deletions(-)

diff --git a/packages/core/useWebSocket/index.md b/packages/core/useWebSocket/index.md
index e38d5239be5..68489eff6ae 100644
--- a/packages/core/useWebSocket/index.md
+++ b/packages/core/useWebSocket/index.md
@@ -71,6 +71,7 @@ const { status, data, close } = useWebSocket('ws://websocketurl', {
   heartbeat: {
     message: 'ping',
     interval: 1000,
+    pongTimeout: 1000,
   },
 })
 ```
diff --git a/packages/core/useWebSocket/index.ts b/packages/core/useWebSocket/index.ts
index 708a6787098..bd689dc561d 100644
--- a/packages/core/useWebSocket/index.ts
+++ b/packages/core/useWebSocket/index.ts
@@ -6,6 +6,8 @@ import { useEventListener } from '../useEventListener'
 
 export type WebSocketStatus = 'OPEN' | 'CONNECTING' | 'CLOSED'
 
+const DEFAULT_PING_MESSAGE = 'ping'
+
 export interface UseWebSocketOptions {
   onConnected?: (ws: WebSocket) => void
   onDisconnected?: (ws: WebSocket, event: CloseEvent) => void
@@ -31,6 +33,13 @@ export interface UseWebSocketOptions {
      * @default 1000
      */
     interval?: number
+
+    /**
+     * Heartbeat response timeout, in milliseconds
+     *
+     * @default 1000
+     */
+    pongTimeout?: number
   }
 
   /**
@@ -159,6 +168,8 @@ export function useWebSocket<Data = any>(
 
   let bufferedData: (string | ArrayBuffer | Blob)[] = []
 
+  let pongTimeoutWait: ReturnType<typeof setTimeout>
+
   // Status code 1000 -> Normal Closure https://developer.mozilla.org/en-US/docs/Web/API/CloseEvent/code
   const close: WebSocket['close'] = (code = 1000, reason) => {
     if (!wsRef.value)
@@ -176,6 +187,10 @@ export function useWebSocket<Data = any>(
     }
   }
 
+  const resetHeartbeat = () => {
+    clearTimeout(pongTimeoutWait)
+  }
+
   const send = (data: string | ArrayBuffer | Blob, useBuffer = true) => {
     if (!wsRef.value || status.value !== 'OPEN') {
       if (useBuffer)
@@ -227,6 +242,16 @@ export function useWebSocket<Data = any>(
     }
 
     ws.onmessage = (e: MessageEvent) => {
+      resetHeartbeat()
+      // Heartbeat response will be skipped
+      if (options.heartbeat) {
+        const {
+          message = DEFAULT_PING_MESSAGE,
+        } = resolveNestedOptions(options.heartbeat)
+        if (e.data === message)
+          return
+      }
+
       data.value = e.data
       onMessage?.(ws!, e)
     }
@@ -234,12 +259,19 @@ export function useWebSocket<Data = any>(
 
   if (options.heartbeat) {
     const {
-      message = 'ping',
+      message = DEFAULT_PING_MESSAGE,
       interval = 1000,
+      pongTimeout = 1000,
     } = resolveNestedOptions(options.heartbeat)
 
     const { pause, resume } = useIntervalFn(
-      () => send(message, false),
+      () => {
+        send(message, false)
+        pongTimeoutWait = setTimeout(() => {
+          // auto-reconnect will be trigger with ws.onclose()
+          close()
+        }, pongTimeout)
+      },
       interval,
       { immediate: false },
     )