Replace per-message heartbeat timeout with a watchdog

package.json

@@ -1,7 +1,7 @@
 {
   "name": "@replit/river",
   "description": "It's like tRPC but... with JSON Schema Support, duplex streaming and support for service multiplexing. Transport agnostic!",
-  "version": "0.220.0",
+  "version": "0.220.1",
   "repository": {
     "type": "git",
     "url": "git+https://github.com/replit/river.git"

transport/sessionStateMachine/SessionConnected.ts

@@ -55,24 +55,20 @@ export class SessionConnected<
   listeners: SessionConnectedListeners;
 
   private heartbeatHandle?: ReturnType<typeof setInterval> | undefined;
-  private heartbeatMissTimeout?: ReturnType<typeof setTimeout> | undefined;
+  private heartbeatWatchdog?: ReturnType<typeof setInterval> | undefined;
+  private lastInboundAt = Date.now();
   private isActivelyHeartbeating = false;
   private rehandshakeTimer?: ReturnType<typeof setTimeout> | undefined;
   private credentialExpiry?: number | undefined;
 
   updateBookkeeping(ack: number, seq: number) {
     this.sendBuffer = this.sendBuffer.filter((unacked) => unacked.seq >= ack);
     this.ack = seq + 1;
+    this.lastInboundAt = Date.now();
 
     if (this.sendBuffer.length < this.options.sendBufferHighWaterMark) {
       this.notifySendBufferDrain();
     }
-
-    if (this.heartbeatMissTimeout) {
-      clearTimeout(this.heartbeatMissTimeout);
-    }
-
-    this.startMissingHeartbeatTimeout();
   }
 
   private assertSendOrdering(encodedMsg: EncodedTransportMessage) {
@@ -163,10 +159,26 @@ export class SessionConnected<
     };
   }
 
-  startMissingHeartbeatTimeout() {
+  /**
+   * Arms the watchdog that closes the connection once the peer stops sending.
+   *
+   * A single interval for the lifetime of the session compares wall time against
+   * the last inbound message, so a busy session does not allocate a timer per
+   * frame. Elapsed time comes from {@link Date.now}, thus a throttled or suspended
+   * timer can only delay detection of a dead connection, never report a heartbeat
+   * as missed while messages keep arriving.
+   */
+  startHeartbeatWatchdog() {
     const maxMisses = this.options.heartbeatsUntilDead;
     const missDuration = maxMisses * this.options.heartbeatIntervalMs;
-    this.heartbeatMissTimeout = setTimeout(() => {
+    this.heartbeatWatchdog = setInterval(() => {
+      if (Date.now() - this.lastInboundAt < missDuration) {
+        return;
+      }
+
+      // the peer is gone, so nothing will refresh the deadline: stop checking
+      this.clearHeartbeatWatchdog();
+
       this.log?.info(
         `closing connection to ${this.to} due to inactivity (missed ${maxMisses} heartbeats which is ${missDuration}ms)`,
         this.loggingMetadata,
@@ -176,7 +188,14 @@ export class SessionConnected<
       );
 
       this.conn.close();
-    }, missDuration);
+    }, this.options.heartbeatIntervalMs);
+  }
+
+  private clearHeartbeatWatchdog() {
+    if (this.heartbeatWatchdog) {
+      clearInterval(this.heartbeatWatchdog);
+      this.heartbeatWatchdog = undefined;
+    }
   }
 
   startActiveHeartbeat() {
@@ -369,11 +388,7 @@ export class SessionConnected<
       this.heartbeatHandle = undefined;
     }
 
-    if (this.heartbeatMissTimeout) {
-      clearTimeout(this.heartbeatMissTimeout);
-      this.heartbeatMissTimeout = undefined;
-    }
-
+    this.clearHeartbeatWatchdog();
     this.clearRehandshakeTimer();
   }
 

transport/sessionStateMachine/stateMachine.test.ts

@@ -3,7 +3,11 @@ import {
   payloadToTransportMessage,
   testingSessionOptions,
 } from '../../testUtil';
-import { waitFor } from '../../testUtil/fixtures/cleanup';
+import {
+  advanceFakeTimersByDisconnectGrace,
+  advanceFakeTimersByHeartbeat,
+  waitFor,
+} from '../../testUtil/fixtures/cleanup';
 import {
   ControlFlags,
   ControlMessageAckSchema,
@@ -2046,6 +2050,48 @@ describe('session state machine', () => {
       });
     });
 
+    test('closes the connection once nothing arrives before the deadline', async () => {
+      const sessionHandle = await createSessionConnected();
+
+      // a heartbeat of silence is still well inside the deadline
+      await advanceFakeTimersByHeartbeat();
+      expect(sessionHandle.onConnectionClosed).not.toHaveBeenCalled();
+
+      await advanceFakeTimersByDisconnectGrace();
+
+      // the watchdog stops itself once it fires, so the peer is only declared
+      // dead once rather than on every subsequent check
+      expect(sessionHandle.onConnectionClosed).toHaveBeenCalledTimes(1);
+    });
+
+    test('inbound messages keep the connection past the deadline', async () => {
+      const sessionHandle = await createSessionConnected();
+      const session = sessionHandle.session;
+
+      // a message every heartbeat for well past the deadline: the session only
+      // survives if each one pushes the deadline out
+      for (let seq = 0; seq < 5; seq++) {
+        session.conn.onData(
+          session.options.codec.toBuffer({
+            id: `msgid-${seq}`,
+            to: session.from,
+            from: session.to,
+            seq,
+            ack: 0,
+            streamId: 'heartbeat',
+            controlFlags: ControlFlags.AckBit,
+            payload: {
+              type: 'ACK',
+            } satisfies Static<typeof ControlMessageAckSchema>,
+          }),
+        );
+
+        await advanceFakeTimersByHeartbeat();
+      }
+
+      expect(sessionHandle.onConnectionClosed).not.toHaveBeenCalled();
+    });
+
     test('does not dispatch acks', async () => {
       const sessionHandle = await createSessionConnected();
       const session = sessionHandle.session;

transport/sessionStateMachine/transitions.ts

@@ -234,7 +234,7 @@ export const SessionStateGraph = {
         ...carriedState,
       });
 
-      session.startMissingHeartbeatTimeout();
+      session.startHeartbeatWatchdog();
 
       session.log?.info(
         `session ${session.id} transition from Handshaking to Connected`,
@@ -293,7 +293,7 @@ export const SessionStateGraph = {
         ...carriedState,
       });
 
-      session.startMissingHeartbeatTimeout();
+      session.startHeartbeatWatchdog();
 
       conn.telemetry = createConnectionTelemetryInfo(
         session.tracer,