Improve heartbeat

 Check interval between ping messages more often (interval / 10).
 This is to prevent some configurations that may cause timeouts
 and connection close. For example, when timeout and interval are
 set to very close values and there was data message recived right
 after connection is established.

 Use one heartbeat interval in shared mode, add test.

 Close #255

Author: w666

src/lib/attach-shared-listeners.ts

@@ -46,7 +46,8 @@ const bindOpenHandler = (
   heartbeatOptions?: boolean | HeartbeatOptions
 ) => {
   webSocketInstance.onopen = (event: WebSocketEventMap['open']) => {
-    getSubscribers(url).forEach(subscriber => {
+    const subscribers = getSubscribers(url);
+    subscribers.forEach(subscriber => {
       subscriber.reconnectCount.current = 0;
       if (subscriber.optionsRef.current.onOpen) {
         subscriber.optionsRef.current.onOpen(event);
@@ -58,9 +59,11 @@ const bindOpenHandler = (
 
       if (heartbeatOptions && webSocketInstance instanceof WebSocket) {
         subscriber.lastMessageTime.current = Date.now();
-        heartbeat(webSocketInstance, subscriber.lastMessageTime, typeof heartbeatOptions === 'boolean' ? undefined : heartbeatOptions,);
       }
     });
+    if (heartbeatOptions && webSocketInstance instanceof WebSocket) {
+      heartbeat(webSocketInstance, subscribers.map(subscriber => subscriber.lastMessageTime), typeof heartbeatOptions === 'boolean' ? undefined : heartbeatOptions,);
+    }
   };
 };
 

src/lib/heartbeat.ts

@@ -2,38 +2,54 @@ import { MutableRefObject } from "react";
 import { DEFAULT_HEARTBEAT } from "./constants";
 import { HeartbeatOptions } from "./types";
 
-export function heartbeat(ws: WebSocket, lastMessageTime: MutableRefObject<number>, options?: HeartbeatOptions): () => void {
+function getLastMessageTime(lastMessageTime: MutableRefObject<number> | MutableRefObject<number>[]): number {
+  if (Array.isArray(lastMessageTime)) {
+    return lastMessageTime.reduce((p, c) => { return (p.current > c.current) ? p : c; }).current;
+  }
+  return lastMessageTime.current
+}
+
+export function heartbeat(ws: WebSocket, lastMessageTime: MutableRefObject<number> | MutableRefObject<number>[], options?: HeartbeatOptions): () => void {
   const {
     interval = DEFAULT_HEARTBEAT.interval,
     timeout = DEFAULT_HEARTBEAT.timeout,
     message = DEFAULT_HEARTBEAT.message,
   } = options || {};
 
+  // how often check interval between ping messages
+  // minimum is 100ms
+  // maximum is ${interval / 10}ms
+  const intervalCheck = Math.max(100, interval / 10);
+
+  let lastPingSentAt = Date.now();
+
   const heartbeatInterval = setInterval(() => {
-    if (lastMessageTime.current + timeout <= Date.now()) {
-      console.warn(`Heartbeat timed out, closing connection, last message was seen ${Date.now() - lastMessageTime.current}ms ago`);
+    const timeNow = Date.now();
+    const lastMessageReceivedAt = getLastMessageTime(lastMessageTime);
+    if (lastMessageReceivedAt + timeout <= timeNow) {
+      console.warn(`Heartbeat timed out, closing connection, last message received ${timeNow - lastMessageReceivedAt}ms ago, last ping sent ${timeNow - lastPingSentAt}ms ago`);
       ws.close();
     } else {
-      if (lastMessageTime.current + interval <= Date.now()) {
+      if (lastMessageReceivedAt + interval <= timeNow && lastPingSentAt + interval <= timeNow) {
         try {
           if (typeof message === 'function') {
             ws.send(message());
           } else {
             ws.send(message);
           }
+          lastPingSentAt = timeNow;
         } catch (err: unknown) {
           console.error(`Heartbeat failed, closing connection`, err instanceof Error ? err.message : err);
           ws.close();
         }
 
       }
     }
-  }, interval);
+  }, intervalCheck);
 
   ws.addEventListener("close", () => {
     clearInterval(heartbeatInterval);
   });
 
-
   return () => { };
 }

src/lib/use-websocket.test.ts

@@ -614,7 +614,7 @@ test.each([false, true])('Options#heartbeat, if provided, close websocket if no
   options.heartbeat = {
     message: 'ping',
     timeout: 1000,
-    interval: 500,
+    interval: 400,
   };
   options.share = shareOption;
 
@@ -627,6 +627,7 @@ test.each([false, true])('Options#heartbeat, if provided, close websocket if no
   }
   await server.connected;
   await sleep(1600);
+  expect(server.messages).toEqual(['ping', 'ping'])
   expect(result.current.readyState).toBe(WebSocket.CLOSED);
 });
 
@@ -682,4 +683,40 @@ test.each([false, true])('Options#heartbeat, if provided, lastMessage is updated
   expect(result.current.lastMessage?.data).toBe('ping');
 });
 
+test.each([false, true])('Options#heartbeat, can handle case when interval is very close to timeout', async (shareOption) => {
+  options.heartbeat = {
+    message: "ping",
+    returnMessage: "pong",
+    timeout: 1000,
+    interval: 800,
+  };
+  options.share = shareOption;
+
+  const { result } = renderHook(() => useWebSocket(URL, options));
+
+  if (shareOption) {
+    renderHook(() => useWebSocket(URL, options));
+  }
+
+  await server.connected;
+  await sleep(50);
+  expect(result.current.readyState).toEqual(ReadyState.OPEN);
+
+  result.current.sendMessage("token");
+  await sleep(50);
+  expect(server).toHaveReceivedMessages(["token"]);
+
+  server.send("authorized");
+  await sleep(50);
+  expect(result.current.lastMessage?.data).toEqual("authorized");
+
+  await sleep(850);
+  // it does not matter is it shared mode or not, in both cases only one timer is used that means only one ping message is sent for all subscribers
+  expect(server.messages).toEqual(["token", "ping"]);
+  server.send("pong");
+
+  expect(result.current.readyState).toBe(WebSocket.OPEN);
+}
+);
+
 // //TODO: Write companion tests for useSocketIO