perf(runtime): arm the WS-server heartbeat only while clients are connected (#9885)

* perf(runtime): idle websocket heartbeat without clients

Co-authored-by: Orca <help@stably.ai>

* fix(runtime): probe immediately when the WS heartbeat arms

Arming the heartbeat on the first accepted connection started a fresh interval,
so the first liveness ping was a full interval (~15s) out — a socket that died
right after connecting went unprobed for that window. Run one sweep synchronously
in start() so the first ping goes out at arm time; the seeded socket is pinged
(never reaped on the arm sweep) and reaped on the next tick only if it never pongs.
Tests updated for the earlier first probe.

Co-authored-by: Orca <help@stably.ai>

---------

Co-authored-by: Orca <help@stably.ai>

src/main/runtime/rpc/remote-runtime-server-heartbeat.test.ts

@@ -7,20 +7,46 @@ afterEach(() => {
 })
 
 describe('RemoteRuntimeServerHeartbeat', () => {
+  it('still reaps a client that misses a probe while another remains alive', async () => {
+    vi.useFakeTimers()
+    let now = 1_000
+    const responsiveSocket = { ping: vi.fn(), terminate: vi.fn() } as unknown as WebSocket
+    const deadSocket = { ping: vi.fn(), terminate: vi.fn() } as unknown as WebSocket
+    const heartbeat = new RemoteRuntimeServerHeartbeat(100, () => now)
+    heartbeat.noteAlive(responsiveSocket)
+    heartbeat.noteAlive(deadSocket)
+    // start() probes immediately: both are pinged now (probe #1) and cleared to await a pong.
+    heartbeat.start(() => [responsiveSocket, deadSocket])
+    // Only the responsive socket pongs the immediate probe.
+    heartbeat.noteAlive(responsiveSocket)
+
+    now += 100
+    await vi.advanceTimersByTimeAsync(100)
+
+    expect(responsiveSocket.ping).toHaveBeenCalledTimes(2)
+    expect(responsiveSocket.terminate).not.toHaveBeenCalled()
+    expect(deadSocket.ping).toHaveBeenCalledTimes(1)
+    expect(deadSocket.terminate).toHaveBeenCalledTimes(1)
+    heartbeat.stop()
+  })
+
   it('grants clients a fresh probe after the server event loop resumes', async () => {
     vi.useFakeTimers()
     let now = 1_000
     const socket = { ping: vi.fn(), terminate: vi.fn() } as unknown as WebSocket
     const heartbeat = new RemoteRuntimeServerHeartbeat(100, () => now)
     heartbeat.noteAlive(socket)
+    // start() probes immediately (ping #1); the socket pongs it.
     heartbeat.start(() => [socket])
+    heartbeat.noteAlive(socket)
 
     now += 100
-    await vi.advanceTimersByTimeAsync(100)
+    await vi.advanceTimersByTimeAsync(100) // ping #2, socket pongs
+    heartbeat.noteAlive(socket)
     now += 3_600_000
-    await vi.advanceTimersByTimeAsync(100)
+    await vi.advanceTimersByTimeAsync(100) // resumed-from-pause: re-grants a probe (ping #3), no reap
 
-    expect(socket.ping).toHaveBeenCalledTimes(2)
+    expect(socket.ping).toHaveBeenCalledTimes(3)
     expect(socket.terminate).not.toHaveBeenCalled()
 
     now += 100

src/main/runtime/rpc/remote-runtime-server-heartbeat.ts

@@ -22,6 +22,11 @@ export class RemoteRuntimeServerHeartbeat {
     this.lastTickAt = this.now()
     this.timer = setInterval(() => this.sweep(getClients()), this.intervalMs)
     this.timer.unref?.()
+    // Why: the interval's first tick is a full intervalMs (~15s) out, so arming on the first accepted
+    // connection would leave that socket unprobed for the whole window. Sweep once now so the first
+    // liveness ping goes out immediately; seeded-alive sockets are pinged (not reaped) and have until
+    // the next tick to pong. WS pong is answered at the protocol level, so a live socket always survives.
+    this.sweep(getClients())
   }
 
   stop(): void {

src/main/runtime/rpc/ws-transport.test.ts

@@ -1,3 +1,4 @@
+import { EventEmitter } from 'node:events'
 import { mkdtempSync } from 'node:fs'
 import { tmpdir } from 'node:os'
 import { join } from 'node:path'
@@ -14,6 +15,30 @@ function makeTls() {
   return loadOrCreateTlsCertificate(userDataPath)
 }
 
+function heartbeatLifecycle(transport: WebSocketTransport) {
+  return transport as unknown as {
+    heartbeat: {
+      timer: ReturnType<typeof setInterval> | null
+      alive: WeakSet<WebSocket>
+    }
+    heartbeatConnections: Set<WebSocket>
+    wss: { clients: Set<WebSocket> }
+    handleConnection(ws: WebSocket): void
+  }
+}
+
+async function waitForHeartbeatLifecycle(
+  transport: WebSocketTransport,
+  connectionCount: number,
+  armed: boolean
+): Promise<void> {
+  await vi.waitFor(() => {
+    const lifecycle = heartbeatLifecycle(transport)
+    expect(lifecycle.heartbeatConnections.size).toBe(connectionCount)
+    expect(lifecycle.heartbeat.timer !== null).toBe(armed)
+  })
+}
+
 describe('WebSocketTransport', () => {
   const transports: WebSocketTransport[] = []
 
@@ -70,6 +95,68 @@ describe('WebSocketTransport', () => {
     await transport.stop()
   })
 
+  it('arms heartbeat only while accepted connections exist', async () => {
+    const { transport } = await createTransport()
+    await transport.start()
+
+    const lifecycle = heartbeatLifecycle(transport)
+    expect(lifecycle.heartbeat.timer).toBeNull()
+    expect(lifecycle.heartbeatConnections.size).toBe(0)
+
+    const firstClient = await connectWs(transport)
+    await waitForHeartbeatLifecycle(transport, 1, true)
+    const firstServerSocket = Array.from(lifecycle.wss.clients)[0]
+    expect(firstServerSocket).toBeDefined()
+    // Note: arming probes immediately, so `alive` membership is racy here (the client's protocol-level
+    // pong re-adds the socket right after the arm sweep clears it). Assert the arm/disarm lifecycle only.
+    const firstTimer = lifecycle.heartbeat.timer
+
+    const secondClient = await connectWs(transport)
+    await waitForHeartbeatLifecycle(transport, 2, true)
+    expect(lifecycle.heartbeat.timer).toBe(firstTimer)
+
+    firstClient.close()
+    await waitForHeartbeatLifecycle(transport, 1, true)
+    expect(lifecycle.heartbeat.timer).toBe(firstTimer)
+
+    secondClient.close()
+    await waitForHeartbeatLifecycle(transport, 0, false)
+
+    const thirdClient = await connectWs(transport)
+    await waitForHeartbeatLifecycle(transport, 1, true)
+    expect(lifecycle.heartbeat.timer).not.toBe(firstTimer)
+
+    thirdClient.close()
+    await waitForHeartbeatLifecycle(transport, 0, false)
+  })
+
+  it('finalizes heartbeat membership once when error and close race', () => {
+    const transport = new WebSocketTransport({ host: '127.0.0.1', port: 0 })
+    transports.push(transport)
+    const lifecycle = heartbeatLifecycle(transport)
+    const socket = Object.assign(new EventEmitter(), {
+      OPEN: WebSocket.OPEN,
+      readyState: WebSocket.OPEN,
+      close: vi.fn(),
+      ping: vi.fn(),
+      terminate: vi.fn()
+    }) as unknown as WebSocket
+    const closeHandler = vi.fn()
+    transport.onConnectionClose(closeHandler)
+
+    lifecycle.handleConnection(socket)
+    expect(lifecycle.heartbeatConnections.size).toBe(1)
+    expect(lifecycle.heartbeat.timer).not.toBeNull()
+
+    socket.emit('error', new Error('connection reset'))
+    socket.emit('close')
+
+    expect(closeHandler).toHaveBeenCalledTimes(1)
+    expect(socket.close).toHaveBeenCalledTimes(1)
+    expect(lifecycle.heartbeatConnections.size).toBe(0)
+    expect(lifecycle.heartbeat.timer).toBeNull()
+  })
+
   it('handles request/response round-trip', async () => {
     const { transport } = await createTransport((msg, reply) => {
       const request = JSON.parse(msg)

src/main/runtime/rpc/ws-transport.ts

@@ -61,6 +61,7 @@ export class WebSocketTransport implements RpcTransport {
     | null = null
   // Why: maps each socket to its authenticated clientId so close can report which device disconnected.
   private wsClientIds = new Map<WebSocket, string>()
+  private heartbeatConnections = new Set<WebSocket>()
   private preAuthTimers = new WeakMap<WebSocket, ReturnType<typeof setTimeout>>()
 
   constructor({
@@ -199,7 +200,6 @@ export class WebSocketTransport implements RpcTransport {
 
     this.httpServer = httpServer
     this.wss = wss
-    this.heartbeat.start(() => this.wss?.clients ?? [])
   }
 
   // Why: force-terminate soon after the 1013 close since a half-open phone may never ack and would hold the descriptor past the WS cap; the 'error' listener absorbs a reset while closing.
@@ -217,6 +217,7 @@ export class WebSocketTransport implements RpcTransport {
     this.wss = null
     this.httpServer = null
     this.heartbeat.stop()
+    this.heartbeatConnections.clear()
 
     if (wss) {
       for (const client of wss.clients) {
@@ -280,6 +281,10 @@ export class WebSocketTransport implements RpcTransport {
       ws.off('close', finalizeConnection)
       ws.off('error', onError)
       this.clearPreAuthTimer(ws)
+      this.heartbeatConnections.delete(ws)
+      if (this.heartbeatConnections.size === 0) {
+        this.heartbeat.stop()
+      }
       const clientId = this.wsClientIds.get(ws) ?? null
       this.wsClientIds.delete(ws)
       const hasOtherConnections =
@@ -287,6 +292,13 @@ export class WebSocketTransport implements RpcTransport {
       this.connectionCloseHandler?.(clientId, ws, hasOtherConnections)
     }
 
+    // Why: seed before arming so a fresh first socket survives its initial sweep.
+    this.heartbeatConnections.add(ws)
+    this.heartbeat.noteAlive(ws)
+    if (this.heartbeatConnections.size === 1) {
+      this.heartbeat.start(() => this.wss?.clients ?? [])
+    }
+
     const preAuthTimer = setTimeout(() => {
       if (!this.wsClientIds.has(ws)) {
         // Why: a silent auto-ponging client would otherwise hold a finite mobile slot forever without starting the E2EE handshake.
@@ -298,9 +310,6 @@ export class WebSocketTransport implements RpcTransport {
     }
     this.preAuthTimers.set(ws, preAuthTimer)
 
-    // Why: seed alive so the first heartbeat tick doesn't reap a fresh socket before its first pong.
-    this.heartbeat.noteAlive(ws)
-
     ws.on('pong', onPong)
     ws.on('message', onMessage)