From 5d3b968283ee7ade9898bad7ea0603fc3fd5f2fd Mon Sep 17 00:00:00 2001 From: Neil <4138956+nwparker@users.noreply.github.com> Date: Wed, 22 Jul 2026 16:21:54 -0700 Subject: [PATCH] perf(runtime): arm the WS-server heartbeat only while clients are connected (#9885) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit * perf(runtime): idle websocket heartbeat without clients Co-authored-by: Orca * 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 --------- Co-authored-by: Orca --- .../remote-runtime-server-heartbeat.test.ts | 32 ++++++- .../rpc/remote-runtime-server-heartbeat.ts | 5 ++ src/main/runtime/rpc/ws-transport.test.ts | 87 +++++++++++++++++++ src/main/runtime/rpc/ws-transport.ts | 17 +++- 4 files changed, 134 insertions(+), 7 deletions(-) diff --git a/src/main/runtime/rpc/remote-runtime-server-heartbeat.test.ts b/src/main/runtime/rpc/remote-runtime-server-heartbeat.test.ts index c9739c344..a34b43c3c 100644 --- a/src/main/runtime/rpc/remote-runtime-server-heartbeat.test.ts +++ b/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 diff --git a/src/main/runtime/rpc/remote-runtime-server-heartbeat.ts b/src/main/runtime/rpc/remote-runtime-server-heartbeat.ts index e8ff47e90..3b5bb9760 100644 --- a/src/main/runtime/rpc/remote-runtime-server-heartbeat.ts +++ b/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 { diff --git a/src/main/runtime/rpc/ws-transport.test.ts b/src/main/runtime/rpc/ws-transport.test.ts index 5c02194a8..d7fef1428 100644 --- a/src/main/runtime/rpc/ws-transport.test.ts +++ b/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 | null + alive: WeakSet + } + heartbeatConnections: Set + wss: { clients: Set } + handleConnection(ws: WebSocket): void + } +} + +async function waitForHeartbeatLifecycle( + transport: WebSocketTransport, + connectionCount: number, + armed: boolean +): Promise { + 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) diff --git a/src/main/runtime/rpc/ws-transport.ts b/src/main/runtime/rpc/ws-transport.ts index 54a3e85f1..6933657b5 100644 --- a/src/main/runtime/rpc/ws-transport.ts +++ b/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() + private heartbeatConnections = new Set() private preAuthTimers = new WeakMap>() 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)