diff --git a/config/scripts/relay-json-payload-benchmark.ts b/config/scripts/relay-json-payload-benchmark.ts new file mode 100644 index 00000000000..8c05fea0227 --- /dev/null +++ b/config/scripts/relay-json-payload-benchmark.ts @@ -0,0 +1,166 @@ +import { performance } from 'node:perf_hooks' +import { platform, release } from 'node:os' +import { + HEADER_LENGTH, + MAX_MESSAGE_SIZE, + MessageType, + encodeFrame, + encodePreparedJsonRpcFrame, + prepareJsonRpcPayload, + type JsonRpcNotification +} from '../../src/relay/protocol' + +type BenchmarkCase = { + name: string + message: JsonRpcNotification + targets: number + iterations: number +} + +type BenchmarkResult = { + name: string + targets: number + payloadBytes: number + legacyMicros: number + preparedMicros: number + reductionPercent: number + speedup: number +} + +const SAMPLES = 17 +let checksum = 0 + +function encodeLegacyJsonRpcFrame(message: JsonRpcNotification, id: number, ack: number): Buffer { + const payload = Buffer.from(JSON.stringify(message), 'utf8') + if (payload.length > MAX_MESSAGE_SIZE) { + throw new Error(`Message too large: ${payload.length} bytes`) + } + return encodeFrame(MessageType.Regular, id, ack, payload) +} + +function runLegacy(message: JsonRpcNotification, targets: number, iteration: number): void { + const estimated = encodeLegacyJsonRpcFrame(message, 0, 0).length + checksum = (checksum + estimated) >>> 0 + for (let target = 0; target < targets; target++) { + consume(encodeLegacyJsonRpcFrame(message, iteration + target + 1, iteration)) + } +} + +function runPrepared(message: JsonRpcNotification, targets: number, iteration: number): void { + const payload = prepareJsonRpcPayload(message) + checksum = (checksum + HEADER_LENGTH + payload.byteLength) >>> 0 + for (let target = 0; target < targets; target++) { + consume(encodePreparedJsonRpcFrame(payload, iteration + target + 1, iteration)) + } +} + +function consume(frame: Buffer): void { + checksum = (checksum + frame.length + frame[0] + frame.at(-1)!) >>> 0 +} + +function measure(operation: (iteration: number) => void, iterations: number): number { + const startedAt = performance.now() + for (let iteration = 0; iteration < iterations; iteration++) { + operation(iteration) + } + return ((performance.now() - startedAt) * 1000) / iterations +} + +function median(values: readonly number[]): number { + const sorted = [...values].sort((left, right) => left - right) + return sorted[Math.floor(sorted.length / 2)] +} + +function benchmarkCase(entry: BenchmarkCase): BenchmarkResult { + const legacy = (iteration: number): void => runLegacy(entry.message, entry.targets, iteration) + const prepared = (iteration: number): void => runPrepared(entry.message, entry.targets, iteration) + for (let warmup = 0; warmup < 3; warmup++) { + measure(legacy, Math.min(20, entry.iterations)) + measure(prepared, Math.min(20, entry.iterations)) + } + const legacySamples: number[] = [] + const preparedSamples: number[] = [] + for (let sample = 0; sample < SAMPLES; sample++) { + if (sample % 2 === 0) { + legacySamples.push(measure(legacy, entry.iterations)) + preparedSamples.push(measure(prepared, entry.iterations)) + } else { + preparedSamples.push(measure(prepared, entry.iterations)) + legacySamples.push(measure(legacy, entry.iterations)) + } + } + const legacyMicros = median(legacySamples) + const preparedMicros = median(preparedSamples) + return { + name: entry.name, + targets: entry.targets, + payloadBytes: Buffer.byteLength(JSON.stringify(entry.message)), + legacyMicros, + preparedMicros, + reductionPercent: ((legacyMicros - preparedMicros) / legacyMicros) * 100, + speedup: legacyMicros / preparedMicros + } +} + +function terminalData(bytes: number): string { + const row = '\u001b[38;5;45mcompile src/renderer/pane.ts\u001b[0m\r\n' + return row.repeat(Math.ceil(bytes / row.length)).slice(0, bytes) +} + +function watcherEvents(count: number): { + kind: string + absolutePath: string + isDirectory: boolean +}[] { + return Array.from({ length: count }, (_, index) => ({ + kind: index % 7 === 0 ? 'create' : 'update', + absolutePath: `/repo/src/features/terminal/generated-${String(index).padStart(4, '0')}.ts`, + isDirectory: false + })) +} + +const ptyMessage: JsonRpcNotification = { + jsonrpc: '2.0', + method: 'pty.data', + params: { id: 'pty-42', data: terminalData(16 * 1024), seq: 18_432, rawLength: 16 * 1024 } +} +const watcherMessage: JsonRpcNotification = { + jsonrpc: '2.0', + method: 'fs.changed', + params: { events: watcherEvents(256) } +} +const fileMessage: JsonRpcNotification = { + jsonrpc: '2.0', + method: 'fs.streamChunk', + params: { + streamId: 37, + seq: 11, + data: Buffer.alloc(256 * 1024, 0x61).toString('base64') + } +} + +const cases: BenchmarkCase[] = [ + { name: 'PTY 16KiB', message: ptyMessage, targets: 1, iterations: 800 }, + { name: 'PTY 16KiB', message: ptyMessage, targets: 2, iterations: 600 }, + { name: 'watcher 256 events', message: watcherMessage, targets: 1, iterations: 400 }, + { name: 'watcher 256 events', message: watcherMessage, targets: 2, iterations: 300 }, + { name: 'file 256KiB raw', message: fileMessage, targets: 1, iterations: 50 }, + { name: 'file 256KiB raw', message: fileMessage, targets: 2, iterations: 35 } +] + +const results = cases.map(benchmarkCase) +const lines = [ + `relay JSON payload benchmark: Node ${process.version}, ${platform()} ${release()}, ${SAMPLES} interleaved samples`, + 'case | targets | JSON bytes | legacy median us/op | prepared median us/op | reduction | speedup', + ...results.map( + (result) => + `${result.name} | ${result.targets} | ${result.payloadBytes} | ` + + `${result.legacyMicros.toFixed(2)} | ${result.preparedMicros.toFixed(2)} | ` + + `${result.reductionPercent.toFixed(1)}% | ${result.speedup.toFixed(2)}x` + ), + `checksum=${checksum}`, + 'Limitations: in-process framing only; excludes dispatcher scheduling, sink/network I/O, and tail-GC.', + 'Legacy is the exact pre-change framing composition retained in this script, not a separate checkout.', + 'The two-target rows model dispatcher fan-out; watcher batches and file streams are commonly single-target.' +] +process.stdout.write(`${lines.join('\n')}\n`) diff --git a/src/main/ssh/relay-protocol.ts b/src/main/ssh/relay-protocol.ts index 8c9f35b96d1..f675ac6ed9d 100644 --- a/src/main/ssh/relay-protocol.ts +++ b/src/main/ssh/relay-protocol.ts @@ -115,6 +115,13 @@ export type JsonRpcNotification = { export type JsonRpcMessage = JsonRpcRequest | JsonRpcResponse | JsonRpcNotification +const JSON_RPC_PAYLOAD_BYTES = Symbol('jsonRpcPayloadBytes') + +export type PreparedJsonRpcPayload = Readonly<{ + byteLength: number + [JSON_RPC_PAYLOAD_BYTES]: Buffer +}> + // ── Framing: encode / decode ──────────────────────────────────────── /** @@ -141,11 +148,23 @@ export function encodeFrame( } export function encodeJsonRpcFrame(msg: JsonRpcMessage, id: number, ack: number): Buffer { + return encodePreparedJsonRpcFrame(prepareJsonRpcPayload(msg), id, ack) +} + +export function prepareJsonRpcPayload(msg: JsonRpcMessage): PreparedJsonRpcPayload { const payload = Buffer.from(JSON.stringify(msg), 'utf-8') if (payload.length > MAX_MESSAGE_SIZE) { throw new Error(`Message too large: ${payload.length} bytes (max ${MAX_MESSAGE_SIZE})`) } - return encodeFrame(MessageType.Regular, id, ack, payload) + return Object.freeze({ byteLength: payload.length, [JSON_RPC_PAYLOAD_BYTES]: payload }) +} + +export function encodePreparedJsonRpcFrame( + payload: PreparedJsonRpcPayload, + id: number, + ack: number +): Buffer { + return encodeFrame(MessageType.Regular, id, ack, payload[JSON_RPC_PAYLOAD_BYTES]) } export function encodeKeepAliveFrame(id: number, ack: number): Buffer { diff --git a/src/relay/dispatcher-capacity-degradation.test.ts b/src/relay/dispatcher-capacity-degradation.test.ts index a546989d5df..7bcf81674dc 100644 --- a/src/relay/dispatcher-capacity-degradation.test.ts +++ b/src/relay/dispatcher-capacity-degradation.test.ts @@ -9,16 +9,15 @@ import { LEGACY_CLIENT_RETAINED_BYTES_LOW } from './legacy-relay-publication-led import type * as ProtocolModule from './protocol' import { encodeJsonRpcFrame, RelayErrorCode } from './protocol' -// Counts every frame encode (including the estimate-only ones) so a redundant re-encode is observable. -const encodeCalls = vi.hoisted(() => ({ count: 0 })) +const preparationCalls = vi.hoisted(() => ({ count: 0 })) vi.mock('./protocol', async (importOriginal) => { const actual = await importOriginal() return { ...actual, - encodeJsonRpcFrame: (...args: Parameters) => { - encodeCalls.count++ - return actual.encodeJsonRpcFrame(...args) + prepareJsonRpcPayload: (...args: Parameters) => { + preparationCalls.count++ + return actual.prepareJsonRpcPayload(...args) } } }) @@ -344,7 +343,7 @@ describe('RelayDispatcher bounded-capacity degradation', () => { } }) - it('does not re-encode a dropped frame when its log line is suppressed', () => { + it('prepares a dropped frame once when its log line is suppressed', () => { const primary = makeBoundedClient(16384) const stderr = vi.spyOn(process.stderr, 'write').mockReturnValue(true) const bounded = new RelayDispatcher(primary.write, primary.options) @@ -353,11 +352,10 @@ describe('RelayDispatcher bounded-capacity degradation', () => { bounded.notify('fs.changed', flood) expect(stderr).toHaveBeenCalledTimes(1) - // The suppressed drop must size the frame once, not once to publish and again to log. - encodeCalls.count = 0 + preparationCalls.count = 0 bounded.notify('fs.changed', flood) expect(stderr).toHaveBeenCalledTimes(1) - expect(encodeCalls.count).toBe(1) + expect(preparationCalls.count).toBe(1) } finally { stderr.mockRestore() bounded.dispose() diff --git a/src/relay/dispatcher-json-payload.test.ts b/src/relay/dispatcher-json-payload.test.ts new file mode 100644 index 00000000000..4965fe6898a --- /dev/null +++ b/src/relay/dispatcher-json-payload.test.ts @@ -0,0 +1,328 @@ +import { afterEach, describe, expect, it, vi } from 'vitest' +import type * as ProtocolModule from './protocol' +import { + RelayDispatcher, + type RelayClientSinkOptions, + type SinkWriteSettlement +} from './dispatcher' +import { encodeKeepAliveFrame, type JsonRpcNotification } from './protocol' + +const protocolCalls = vi.hoisted(() => ({ preparations: 0, encodes: 0 })) + +vi.mock('./protocol', async (importOriginal) => { + const actual = await importOriginal() + return { + ...actual, + prepareJsonRpcPayload: (...args: Parameters) => { + protocolCalls.preparations++ + return actual.prepareJsonRpcPayload(...args) + }, + encodePreparedJsonRpcFrame: (...args: Parameters) => { + protocolCalls.encodes++ + return actual.encodePreparedJsonRpcFrame(...args) + } + } +}) + +type DecodedFrame = { + id: number + ack: number + message: JsonRpcNotification + payload: Buffer +} + +function decodeFrame(frame: Buffer): DecodedFrame { + const length = frame.readUInt32BE(9) + const payload = frame.subarray(13, 13 + length) + return { + id: frame.readUInt32BE(1), + ack: frame.readUInt32BE(5), + message: JSON.parse(payload.toString('utf8')), + payload + } +} + +class DrainSink { + readonly frames: Buffer[] = [] + private readonly drainWaiters = new Set<() => void>() + private writableBytes = 0 + private blocked = true + + constructor( + highWaterMark = 4 * 1024 * 1024, + private readonly mutateFrames = false + ) { + this.options = { + writableLength: () => this.writableBytes, + writableHighWaterMark: () => highWaterMark, + waitWriteDrain: (callback) => { + this.drainWaiters.add(callback) + return () => this.drainWaiters.delete(callback) + } + } + } + + readonly options: RelayClientSinkOptions + + write = (data: Buffer): boolean => { + this.frames.push(Buffer.from(data)) + if (this.mutateFrames) { + data.fill(0x78, 13) + } + if (!this.blocked) { + return true + } + this.writableBytes = data.length + return false + } + + drain(): void { + this.blocked = false + this.writableBytes = 0 + for (const callback of Array.from(this.drainWaiters)) { + callback() + } + } +} + +describe('RelayDispatcher prepared JSON payloads', () => { + afterEach(() => { + protocolCalls.preparations = 0 + protocolCalls.encodes = 0 + }) + + it('shares one payload while allocating independent drain-time sequence and ACK headers', () => { + const primary = new DrainSink(undefined, true) + const secondary = new DrainSink() + const dispatcher = new RelayDispatcher(primary.write, primary.options) + const secondaryId = dispatcher.attachClient(secondary.write, secondary.options) + try { + dispatcher.notify('test.blocker') + dispatcher.notifyClient(1, 'test.primary-only') + protocolCalls.preparations = 0 + protocolCalls.encodes = 0 + + dispatcher.notify('workspace.changed', { revision: 7 }) + expect(protocolCalls).toEqual({ preparations: 1, encodes: 0 }) + dispatcher.feed(encodeKeepAliveFrame(41, 0)) + dispatcher.feedClient(secondaryId, encodeKeepAliveFrame(73, 0)) + + primary.drain() + secondary.drain() + + const primaryShared = primary.frames + .map(decodeFrame) + .find((frame) => frame.message.method === 'workspace.changed') + const secondaryShared = secondary.frames + .map(decodeFrame) + .find((frame) => frame.message.method === 'workspace.changed') + expect(primaryShared).toMatchObject({ id: 3, ack: 41 }) + expect(secondaryShared).toMatchObject({ id: 2, ack: 73 }) + expect(primaryShared?.payload.equals(secondaryShared!.payload)).toBe(true) + expect(protocolCalls).toEqual({ preparations: 1, encodes: 3 }) + } finally { + dispatcher.dispose() + } + }) + + it('prepares a rejected producer frame once without allocating a header', () => { + const frames: Buffer[] = [] + const dispatcher = new RelayDispatcher( + (frame) => { + frames.push(Buffer.from(frame)) + return true + }, + { writableLength: () => 0, writableHighWaterMark: () => 1024 } + ) + try { + expect( + dispatcher.publishProducerNotification( + 1, + 'fs.changed', + { events: [{ path: '/repo/a' }] }, + { + logDrop: false + } + ) + ).toBe(false) + expect(frames).toEqual([]) + expect(protocolCalls).toEqual({ preparations: 1, encodes: 0 }) + } finally { + dispatcher.dispose() + } + }) + + it('snapshots wire data and retains only PTY admission identity while queued', () => { + const sink = new DrainSink() + const dispatcher = new RelayDispatcher(sink.write, sink.options) + const params: Record = { + id: 'pty-1', + data: 'before'.repeat(4096), + deliveryToken: 'token-before', + clientGeneration: 3, + ownerGeneration: 5, + ptyIncarnation: 'incarnation-before' + } + const internals = dispatcher as unknown as { + prepareFrame: (message: JsonRpcNotification) => { + ptyDataAdmissionParams: Readonly> | null + } + } + const prepared = internals.prepareFrame({ jsonrpc: '2.0', method: 'pty.data', params }) + expect(prepared.ptyDataAdmissionParams).toEqual({ + id: 'pty-1', + deliveryToken: 'token-before', + clientGeneration: 3, + ownerGeneration: 5, + ptyIncarnation: 'incarnation-before' + }) + expect(prepared.ptyDataAdmissionParams).not.toHaveProperty('data') + + const admissions: Readonly>[] = [] + dispatcher.registerPtyDataPublicationAdmission((_clientId, admissionParams) => { + admissions.push(admissionParams) + return admissionParams.deliveryToken === 'token-before' + }) + dispatcher.notify('test.blocker') + const settled = vi.fn<(result: SinkWriteSettlement) => void>() + expect(dispatcher.tryNotifyPtyDataToClient(1, params, settled)).toBe(true) + admissions.length = 0 + params.data = 'after' + params.deliveryToken = 'token-after' + params.clientGeneration = 99 + params.ownerGeneration = 100 + params.ptyIncarnation = 'incarnation-after' + + sink.drain() + + const publication = sink.frames + .map(decodeFrame) + .find((frame) => frame.message.method === 'pty.data') + expect(publication?.message.params).toMatchObject({ + data: 'before'.repeat(4096), + deliveryToken: 'token-before', + clientGeneration: 3, + ownerGeneration: 5, + ptyIncarnation: 'incarnation-before' + }) + expect(admissions).toHaveLength(1) + expect(admissions[0]).not.toHaveProperty('data') + expect(admissions[0]).toMatchObject({ deliveryToken: 'token-before', clientGeneration: 3 }) + expect(settled).toHaveBeenCalledExactlyOnceWith({ ok: true }) + dispatcher.dispose() + }) + + it('retires queued PTY data without consuming sequence or losing settlement', () => { + const sink = new DrainSink() + const dispatcher = new RelayDispatcher(sink.write, sink.options) + let admitted = true + dispatcher.registerPtyDataPublicationAdmission(() => admitted) + try { + dispatcher.notify('test.blocker') + const retired = vi.fn<(result: SinkWriteSettlement) => void>() + expect(dispatcher.tryNotifyPtyDataToClient(1, { id: 'pty-1', data: 'retire' }, retired)).toBe( + true + ) + admitted = false + sink.drain() + expect(retired).toHaveBeenCalledExactlyOnceWith({ + ok: false, + error: expect.objectContaining({ message: 'PTY publication retired' }) + }) + admitted = true + expect(dispatcher.tryNotifyPtyDataToClient(1, { id: 'pty-1', data: 'next' }, vi.fn())).toBe( + true + ) + const published = sink.frames + .map(decodeFrame) + .find((frame) => frame.message.method === 'pty.data') + expect(published).toMatchObject({ id: 2 }) + expect(published?.message.params?.data).toBe('next') + } finally { + dispatcher.dispose() + } + }) + + it('reuses one snapshot when a 256 KiB file frame retries after producer drain', async () => { + const sink = new DrainSink() + const dispatcher = new RelayDispatcher(sink.write, sink.options) + try { + dispatcher.notify('test.blocker', { data: 'x'.repeat(16 * 1024) }) + protocolCalls.preparations = 0 + protocolCalls.encodes = 0 + const originalData = Buffer.alloc(256 * 1024, 0x61).toString('base64') + const params = { streamId: 7, seq: 2, data: originalData } + const pending = dispatcher.notifyBulk('fs.streamChunk', params) + params.data = 'mutated' + await Promise.resolve() + + expect(protocolCalls).toEqual({ preparations: 1, encodes: 0 }) + sink.drain() + await pending + + const chunk = sink.frames + .map(decodeFrame) + .find((frame) => frame.message.method === 'fs.streamChunk') + expect(chunk?.message.params).toEqual({ streamId: 7, seq: 2, data: originalData }) + expect(protocolCalls).toEqual({ preparations: 1, encodes: 1 }) + } finally { + dispatcher.dispose() + } + }) + + it('prepares queued 256 KiB bulk payloads only when their chain step becomes active', async () => { + const sink = new DrainSink() + const dispatcher = new RelayDispatcher(sink.write, sink.options) + try { + dispatcher.notify('test.blocker', { data: 'x'.repeat(16 * 1024) }) + protocolCalls.preparations = 0 + protocolCalls.encodes = 0 + const originalData = Buffer.alloc(256 * 1024, 0x61).toString('base64') + const pending = Array.from({ length: 16 }, (_, index) => { + const params = { streamId: index + 1, seq: 1, data: originalData } + const publication = dispatcher.notifyBulk('fs.streamChunk', params) + params.data = 'mutated' + return publication + }) + + expect(protocolCalls).toEqual({ preparations: 0, encodes: 0 }) + await Promise.resolve() + expect(protocolCalls).toEqual({ preparations: 1, encodes: 0 }) + + sink.drain() + await Promise.all(pending) + + const chunks = sink.frames + .map(decodeFrame) + .filter((frame) => frame.message.method === 'fs.streamChunk') + expect(chunks).toHaveLength(16) + expect(chunks.every((frame) => frame.message.params?.data === originalData)).toBe(true) + expect(protocolCalls).toEqual({ preparations: 16, encodes: 16 }) + } finally { + dispatcher.dispose() + } + }) + + it('lazily shares one bulk payload across broadcast clients', async () => { + const primary = new DrainSink() + const secondary = new DrainSink() + primary.drain() + secondary.drain() + const dispatcher = new RelayDispatcher(primary.write, primary.options) + dispatcher.attachClient(secondary.write, secondary.options) + try { + const params = { streamId: 4, seq: 8, data: 'before' } + const pending = dispatcher.notifyBulk('git.responseChunk', params) + params.data = 'after' + + expect(protocolCalls).toEqual({ preparations: 0, encodes: 0 }) + await pending + + expect(decodeFrame(primary.frames[0]).message.params?.data).toBe('before') + expect(decodeFrame(secondary.frames[0]).message.params?.data).toBe('before') + expect(protocolCalls).toEqual({ preparations: 1, encodes: 2 }) + } finally { + dispatcher.dispose() + } + }) +}) diff --git a/src/relay/dispatcher.test.ts b/src/relay/dispatcher.test.ts index f38d3833dca..280f9351200 100644 --- a/src/relay/dispatcher.test.ts +++ b/src/relay/dispatcher.test.ts @@ -725,12 +725,18 @@ describe('RelayDispatcher', () => { type DispatcherInternals = { primaryClient: object estimateFrameBytes: (msg: JsonRpcNotification) => number + prepareFrame: (msg: JsonRpcNotification) => object enqueueFrame: ( client: object, msg: JsonRpcNotification, lane: string, - onSettled?: (result: SinkWriteSettlement) => void, - estimatedBytes?: number + onSettled?: (result: SinkWriteSettlement) => void + ) => boolean + enqueuePreparedFrame: ( + client: object, + frame: object, + lane: string, + onSettled?: (result: SinkWriteSettlement) => void ) => boolean } @@ -840,14 +846,14 @@ describe('RelayDispatcher', () => { } }) - it('publishes PTY data with a single frame estimate', () => { + it('publishes PTY data with a single frame preparation', () => { const frames: Buffer[] = [] const publisher = new RelayDispatcher((data) => { frames.push(Buffer.from(data)) return true }) try { - const spy = vi.spyOn(publisher as unknown as DispatcherInternals, 'estimateFrameBytes') + const spy = vi.spyOn(publisher as unknown as DispatcherInternals, 'prepareFrame') expect(publisher.tryNotifyPtyData({ id: 'pty-1', data: 'hello' })).toBe(true) expect(frames).toHaveLength(1) expect(spy).toHaveBeenCalledTimes(1) @@ -856,7 +862,7 @@ describe('RelayDispatcher', () => { } }) - it('enqueueFrame with a caller-supplied estimate matches the computed default', () => { + it('a prepared enqueue matches the composition wrapper', () => { const frames: Buffer[] = [] const publisher = new RelayDispatcher((data) => { frames.push(Buffer.from(data)) @@ -871,12 +877,10 @@ describe('RelayDispatcher', () => { } expect(internals.enqueueFrame(internals.primaryClient, msg, 'ordinary')).toBe(true) expect( - internals.enqueueFrame( + internals.enqueuePreparedFrame( internals.primaryClient, - msg, - 'ordinary', - undefined, - internals.estimateFrameBytes(msg) + internals.prepareFrame(msg), + 'ordinary' ) ).toBe(true) expect(frames).toHaveLength(2) @@ -888,7 +892,7 @@ describe('RelayDispatcher', () => { } }) - it('enqueueFrame rejects identically with and without a caller-supplied estimate', () => { + it('a prepared enqueue rejects identically to the composition wrapper', () => { const { sized } = makeDispatcher([1030]) try { const internals = sized as unknown as DispatcherInternals @@ -899,12 +903,10 @@ describe('RelayDispatcher', () => { } expect(internals.enqueueFrame(internals.primaryClient, msg, 'ordinary')).toBe(false) expect( - internals.enqueueFrame( + internals.enqueuePreparedFrame( internals.primaryClient, - msg, - 'ordinary', - undefined, - internals.estimateFrameBytes(msg) + internals.prepareFrame(msg), + 'ordinary' ) ).toBe(false) } finally { diff --git a/src/relay/dispatcher.ts b/src/relay/dispatcher.ts index 5620b44f0d7..29d1297d283 100644 --- a/src/relay/dispatcher.ts +++ b/src/relay/dispatcher.ts @@ -1,16 +1,19 @@ /* eslint-disable max-lines -- dispatcher keeps client routing, cancellation, and framing state together */ import { FrameDecoder, + HEADER_LENGTH, MessageType, - encodeJsonRpcFrame, + encodePreparedJsonRpcFrame, encodeKeepAliveFrame, parseJsonRpcMessage, + prepareJsonRpcPayload, KEEPALIVE_SEND_MS, RelayErrorCode, type DecodedFrame, type JsonRpcRequest, type JsonRpcNotification, - type JsonRpcResponse + type JsonRpcResponse, + type PreparedJsonRpcPayload } from './protocol' import { ClientRequestAborts } from './client-request-aborts' import { MAX_TIMER_DELAY_MS, isSafeTimerDelayMs } from '../shared/timer-delay' @@ -80,6 +83,14 @@ type RelayClient = { sessionIdentity: RelayClientSessionIdentity } +type OutgoingJsonRpcMessage = JsonRpcRequest | JsonRpcResponse | JsonRpcNotification + +type PreparedRelayFrame = Readonly<{ + payload: PreparedJsonRpcPayload + frameBytes: number + ptyDataAdmissionParams: Readonly> | null +}> + // Why: the log key set is rebuilt per generation, but a producer minting synthetic method names would still // grow it inside one generation — cap it well above the fixed relay method vocabulary. const DROPPED_NOTIFICATION_LOG_KEY_LIMIT = 64 @@ -571,7 +582,8 @@ export class RelayDispatcher { method, ...(params !== undefined ? { params } : {}) } - const frameBytes = this.estimateFrameBytes(msg) + const frame = this.prepareFrame(msg) + const frameBytes = frame.frameBytes this.runPublicationTransaction(() => { for (const client of this.clients.values()) { if (client.closed) { @@ -583,12 +595,12 @@ export class RelayDispatcher { if (method === 'pty.replay') { // Why: replay is never re-sent, so it takes the control lane where overflow is fatal — the // writer closes the client and reconnect reloads history rather than stranding a short buffer. - this.enqueueFrame(client, msg, 'control', undefined, frameBytes) + this.enqueuePreparedFrame(client, frame, 'control') continue } // Why: closing can never make an oversized frame sendable — the producer regenerates it after // reattach and re-kills the link, turning a recoverable drop into an endless reconnect loop. - if (!this.publishToClient(client, msg, 'ordinary', undefined, frameBytes)) { + if (!this.publishPreparedToClient(client, frame, 'ordinary')) { this.logDroppedProducerNotification(client, method, frameBytes) } } @@ -620,13 +632,13 @@ export class RelayDispatcher { if (method === 'pty.data' && !this.admitsPtyDataPublication(client.id, params ?? {})) { return false } - const frameBytes = this.estimateFrameBytes(msg) - if (this.publishToClient(client, msg, 'ordinary', undefined, frameBytes)) { + const frame = this.prepareFrame(msg) + if (this.publishPreparedToClient(client, frame, 'ordinary')) { return true } // Why: same diagnostics as notify() — a producer that drops here must not do so silently. if (options?.logDrop !== false) { - this.logDroppedProducerNotification(client, method, frameBytes) + this.logDroppedProducerNotification(client, method, frame.frameBytes) } return false } @@ -664,7 +676,6 @@ export class RelayDispatcher { onSettled: (result: SinkWriteSettlement) => void = () => {}, options: { controlOverflow?: 'close-client' | 'reject' - estimatedBytes?: number } = {} ): boolean { if (this.disposed) { @@ -685,7 +696,6 @@ export class RelayDispatcher { }, 'control', onSettled, - options.estimatedBytes, options.controlOverflow ) } @@ -699,8 +709,13 @@ export class RelayDispatcher { method, ...(params !== undefined ? { params } : {}) } - for (const client of this.activeClients()) { - if (!this.enqueueFrame(client, msg, 'control')) { + const clients = this.activeClients() + if (clients.length === 0) { + return + } + const frame = this.prepareFrame(msg) + for (const client of clients) { + if (!this.enqueuePreparedFrame(client, frame, 'control')) { this.closeClient( client, new Error('Relay control publication capacity exceeded'), @@ -724,27 +739,48 @@ export class RelayDispatcher { if (this.disposed) { return Promise.resolve() } - const msg: JsonRpcNotification = { - jsonrpc: '2.0', - method, - ...(params !== undefined ? { params } : {}) - } const targets = opts?.clientId !== undefined ? [this.clients.get(opts.clientId)].filter((c): c is RelayClient => c !== undefined) : Array.from(this.clients.values()) - const waits: Promise[] = [] - for (const client of targets) { - if (client.closed) { - continue + const activeTargets = targets.filter((client) => !client.closed) + if (activeTargets.length === 0) { + return Promise.resolve() + } + const msg: JsonRpcNotification = { + jsonrpc: '2.0', + method, + ...(params !== undefined ? { params: { ...params } } : {}) + } + let prepared: + | { ok: true; frame: PreparedRelayFrame } + | { ok: false; error: unknown } + | undefined + const prepareOnce = (): PreparedRelayFrame => { + if (!prepared) { + try { + prepared = { ok: true, frame: this.prepareFrame(msg) } + } catch (error) { + prepared = { ok: false, error } + } } - const step = client.bulkChain.then(() => this.publishBulkWhenAvailable(client, msg)) + if (!prepared.ok) { + throw prepared.error + } + return prepared.frame + } + const lane = method === 'fs.streamChunk' ? 'fixed-bulk' : 'bulk' + const waits: Promise[] = [] + for (const client of activeTargets) { + const step = client.bulkChain.then(() => { + if (this.disposed || client.closed) { + return + } + return this.publishBulkWhenAvailable(client, prepareOnce(), lane) + }) client.bulkChain = step.catch(() => {}) waits.push(step) } - if (waits.length === 0) { - return Promise.resolve() - } return Promise.all(waits).then(() => {}) } @@ -801,7 +837,7 @@ export class RelayDispatcher { reject(new Error(`Request "${method}" timed out after ${timeoutMs}ms`)) }, timeoutMs) this.pendingRelayRequests.set(id, { resolve, reject, timer }) - if (!this.enqueueFrame(client, msg, 'control', () => {}, undefined, 'reject')) { + if (!this.enqueueFrame(client, msg, 'control', () => {}, 'reject')) { clearTimeout(timer) this.pendingRelayRequests.delete(id) reject(new Error(`Request "${method}" exceeded the relay control transport capacity`)) @@ -1048,24 +1084,25 @@ export class RelayDispatcher { id, ...(error ? { error } : { result: result ?? null }) } - const estimatedBytes = this.estimateFrameBytes(msg) - const lane = estimatedBytes > DISPATCHER_CONTROL_QUEUE_MAX_BYTES ? 'legacy-response' : 'control' - const accepted = this.enqueueFrame(client, msg, lane, onSettled) + const frame = this.prepareFrame(msg) + const lane = + frame.frameBytes > DISPATCHER_CONTROL_QUEUE_MAX_BYTES ? 'legacy-response' : 'control' + const accepted = this.enqueuePreparedFrame(client, frame, lane, onSettled) if (accepted) { return true } // Why: an oversized response must fail its own request; closing would kill every pane on the host. // A rejected first enqueue either left onSettled untouched or closed the client, so exactly one settlement happens. - return this.enqueueFrame( + return this.enqueuePreparedFrame( client, - { + this.prepareFrame({ jsonrpc: '2.0', id, error: { code: RelayErrorCode.ResponseOverCapacity, message: RESPONSE_OVER_CAPACITY_MESSAGE } - }, + }), 'control', // Why: writing the substitute is not delivering the result — a settlement fence must never read // the capacity error's successful write as "the peer received your result". @@ -1080,29 +1117,45 @@ export class RelayDispatcher { private enqueueFrame( client: RelayClient, - msg: JsonRpcRequest | JsonRpcResponse | JsonRpcNotification, + msg: OutgoingJsonRpcMessage, + lane: DispatcherWriterLane, + onSettled: (result: SinkWriteSettlement) => void = () => {}, + controlOverflow: 'close-client' | 'reject' = 'close-client' + ): boolean { + if (this.disposed || client.closed) { + return false + } + return this.enqueuePreparedFrame( + client, + this.prepareFrame(msg), + lane, + onSettled, + controlOverflow + ) + } + + private enqueuePreparedFrame( + client: RelayClient, + frame: PreparedRelayFrame, lane: DispatcherWriterLane, onSettled: (result: SinkWriteSettlement) => void = () => {}, - // Why: publish paths already sized the frame; avoid a redundant encode. - estimatedBytes?: number, controlOverflow: 'close-client' | 'reject' = 'close-client' ): boolean { if (this.disposed || client.closed) { return false } - const frameBytes = estimatedBytes ?? this.estimateFrameBytes(msg) const encode = (): Buffer => { const seq = client.nextOutgoingSeq++ - return encodeJsonRpcFrame(msg, seq, client.highestReceivedSeq) + return encodePreparedJsonRpcFrame(frame.payload, seq, client.highestReceivedSeq) } - const isStillAdmitted = - 'method' in msg && msg.method === 'pty.data' - ? () => this.admitsPtyDataPublication(client.id, msg.params ?? {}) - : undefined + const admissionParams = frame.ptyDataAdmissionParams + const isStillAdmitted = admissionParams + ? () => this.admitsPtyDataPublication(client.id, admissionParams) + : undefined return client.writer.enqueue( lane, encode, - frameBytes, + frame.frameBytes, onSettled, lane === 'control' && controlOverflow === 'reject', isStillAdmitted @@ -1151,8 +1204,27 @@ export class RelayDispatcher { return `${client.id}:${client.generation}` } - private estimateFrameBytes(msg: JsonRpcRequest | JsonRpcResponse | JsonRpcNotification): number { - return encodeJsonRpcFrame(msg, 0, 0).length + private prepareFrame(msg: OutgoingJsonRpcMessage): PreparedRelayFrame { + const payload = prepareJsonRpcPayload(msg) + const params = 'method' in msg && msg.method === 'pty.data' ? (msg.params ?? {}) : null + return Object.freeze({ + payload, + frameBytes: HEADER_LENGTH + payload.byteLength, + ptyDataAdmissionParams: + params === null + ? null + : Object.freeze({ + id: params.id, + deliveryToken: params.deliveryToken, + clientGeneration: params.clientGeneration, + ownerGeneration: params.ownerGeneration, + ptyIncarnation: params.ptyIncarnation + }) + }) + } + + private estimateFrameBytes(msg: OutgoingJsonRpcMessage): number { + return HEADER_LENGTH + prepareJsonRpcPayload(msg).byteLength } private tryPublishToClients( @@ -1164,7 +1236,8 @@ export class RelayDispatcher { if (clients.length === 0) { return true } - const bytes = this.estimateFrameBytes(msg) + const frame = this.prepareFrame(msg) + const bytes = frame.frameBytes if (clients.some((client) => !client.writer.canEnqueueProducer(bytes))) { return false } @@ -1175,7 +1248,7 @@ export class RelayDispatcher { return false } for (let index = 0; index < clients.length; index++) { - if (!this.enqueueLeasedFrame(clients[index], msg, lane, leases[index], bytes)) { + if (!this.enqueueLeasedFrame(clients[index], frame, lane, leases[index])) { if (this.disposed || clients[index].closed) { continue } @@ -1195,8 +1268,12 @@ export class RelayDispatcher { lane: 'interactive' | 'ordinary' ): boolean { return this.runPublicationTransaction(() => { + if (clients.length === 0) { + return true + } + const frame = this.prepareFrame(msg) for (const client of clients) { - if (client.closed || this.publishToClient(client, msg, lane)) { + if (client.closed || this.publishPreparedToClient(client, frame, lane)) { continue } this.closeClient( @@ -1213,11 +1290,21 @@ export class RelayDispatcher { client: RelayClient, msg: JsonRpcNotification, lane: 'interactive' | 'ordinary' | 'fixed-bulk' | 'bulk', - onSettled: (result: SinkWriteSettlement) => void = () => {}, - // Why: broadcast callers size the frame once for every client; avoid a redundant encode. - estimatedBytes?: number + onSettled: (result: SinkWriteSettlement) => void = () => {} ): boolean { - const bytes = estimatedBytes ?? this.estimateFrameBytes(msg) + if (this.disposed || client.closed) { + return false + } + return this.publishPreparedToClient(client, this.prepareFrame(msg), lane, onSettled) + } + + private publishPreparedToClient( + client: RelayClient, + frame: PreparedRelayFrame, + lane: 'interactive' | 'ordinary' | 'fixed-bulk' | 'bulk', + onSettled: (result: SinkWriteSettlement) => void = () => {} + ): boolean { + const bytes = frame.frameBytes const fixedBlocked = lane === 'fixed-bulk' && (client.writer.retainedProducerBytes > 0 || bytes > client.writer.fixedFrameCapacity) @@ -1228,12 +1315,15 @@ export class RelayDispatcher { if (!leases) { return false } - return this.enqueueLeasedFrame(client, msg, lane, leases[0], bytes, onSettled) + return this.enqueueLeasedFrame(client, frame, lane, leases[0], onSettled) } - private publishBulkWhenAvailable(client: RelayClient, msg: JsonRpcNotification): Promise { - const bytes = this.estimateFrameBytes(msg) - const lane = msg.method === 'fs.streamChunk' ? 'fixed-bulk' : 'bulk' + private publishBulkWhenAvailable( + client: RelayClient, + frame: PreparedRelayFrame, + lane: 'fixed-bulk' | 'bulk' + ): Promise { + const bytes = frame.frameBytes if (bytes > DEFAULT_PRODUCER_QUEUE_MAX_BYTES) { return Promise.reject(new Error('Relay bulk frame exceeds sink producer capacity')) } @@ -1253,7 +1343,7 @@ export class RelayDispatcher { return } if ( - this.publishToClient(client, msg, lane, (result) => { + this.publishPreparedToClient(client, frame, lane, (result) => { finish() if (result.ok || this.disposed || client.closed) { resolve() @@ -1274,23 +1364,16 @@ export class RelayDispatcher { private enqueueLeasedFrame( client: RelayClient, - msg: JsonRpcNotification, + frame: PreparedRelayFrame, lane: 'interactive' | 'ordinary' | 'fixed-bulk' | 'bulk', lease: LegacyPublicationLease, - estimatedBytes: number, onSettled: (result: SinkWriteSettlement) => void = () => {} ): boolean { - const accepted = this.enqueueFrame( - client, - msg, - lane, - (result) => { - lease.release() - onSettled(result) - this.notifyLegacyCapacityIfLow() - }, - estimatedBytes - ) + const accepted = this.enqueuePreparedFrame(client, frame, lane, (result) => { + lease.release() + onSettled(result) + this.notifyLegacyCapacityIfLow() + }) if (!accepted) { lease.release() this.notifyLegacyCapacityIfLow() diff --git a/src/relay/protocol-json-payload.test.ts b/src/relay/protocol-json-payload.test.ts new file mode 100644 index 00000000000..3abdb47058e --- /dev/null +++ b/src/relay/protocol-json-payload.test.ts @@ -0,0 +1,62 @@ +import { describe, expect, it } from 'vitest' +import * as mainProtocol from '../main/ssh/relay-protocol' +import { + HEADER_LENGTH, + MAX_MESSAGE_SIZE, + MessageType, + encodeFrame, + encodeJsonRpcFrame, + encodePreparedJsonRpcFrame, + prepareJsonRpcPayload, + type JsonRpcNotification +} from './protocol' + +describe('prepared relay JSON payload framing', () => { + it('is byte-equivalent to direct composition for every header field', () => { + const message: JsonRpcNotification = { + jsonrpc: '2.0', + method: 'pty.data', + params: { id: 'pty-1', data: 'héllo "𝄞"\\\n\uD800', seq: 42 } + } + const payload = Buffer.from(JSON.stringify(message), 'utf8') + const relayPrepared = prepareJsonRpcPayload(message) + const mainPrepared = mainProtocol.prepareJsonRpcPayload(message) + + for (const [id, ack] of [ + [0, 0], + [19, 7], + [0xffffffff, 0xfffffffe] + ]) { + const reference = encodeFrame(MessageType.Regular, id, ack, payload) + expect(encodePreparedJsonRpcFrame(relayPrepared, id, ack).equals(reference)).toBe(true) + expect(encodeJsonRpcFrame(message, id, ack).equals(reference)).toBe(true) + expect(mainProtocol.encodePreparedJsonRpcFrame(mainPrepared, id, ack).equals(reference)).toBe( + true + ) + expect(mainProtocol.encodeJsonRpcFrame(message, id, ack).equals(reference)).toBe(true) + } + }) + + it('accepts the exact payload maximum and rejects one byte more', () => { + const base: JsonRpcNotification = { + jsonrpc: '2.0', + method: 'x', + params: { data: '' } + } + const overhead = Buffer.byteLength(JSON.stringify(base)) + const exact = { ...base, params: { data: 'a'.repeat(MAX_MESSAGE_SIZE - overhead) } } + const oversized = { ...base, params: { data: 'a'.repeat(MAX_MESSAGE_SIZE - overhead + 1) } } + + const relayPrepared = prepareJsonRpcPayload(exact) + const mainPrepared = mainProtocol.prepareJsonRpcPayload(exact) + expect(relayPrepared.byteLength).toBe(MAX_MESSAGE_SIZE) + expect(mainPrepared.byteLength).toBe(MAX_MESSAGE_SIZE) + expect(encodePreparedJsonRpcFrame(relayPrepared, 1, 0)).toHaveLength( + HEADER_LENGTH + MAX_MESSAGE_SIZE + ) + expect(() => prepareJsonRpcPayload(oversized)).toThrow('Message too large') + expect(() => mainProtocol.prepareJsonRpcPayload(oversized)).toThrow('Message too large') + expect(() => encodeJsonRpcFrame(oversized, 1, 0)).toThrow('Message too large') + expect(() => mainProtocol.encodeJsonRpcFrame(oversized, 1, 0)).toThrow('Message too large') + }) +}) diff --git a/src/relay/protocol.ts b/src/relay/protocol.ts index 8cb4ab57a75..25e158dc1e2 100644 --- a/src/relay/protocol.ts +++ b/src/relay/protocol.ts @@ -130,6 +130,13 @@ export type JsonRpcNotification = { export type JsonRpcMessage = JsonRpcRequest | JsonRpcResponse | JsonRpcNotification +const JSON_RPC_PAYLOAD_BYTES = Symbol('jsonRpcPayloadBytes') + +export type PreparedJsonRpcPayload = Readonly<{ + byteLength: number + [JSON_RPC_PAYLOAD_BYTES]: Buffer +}> + export function encodeFrame( type: number, id: number, @@ -145,11 +152,23 @@ export function encodeFrame( } export function encodeJsonRpcFrame(msg: JsonRpcMessage, id: number, ack: number): Buffer { + return encodePreparedJsonRpcFrame(prepareJsonRpcPayload(msg), id, ack) +} + +export function prepareJsonRpcPayload(msg: JsonRpcMessage): PreparedJsonRpcPayload { const payload = Buffer.from(JSON.stringify(msg), 'utf-8') if (payload.length > MAX_MESSAGE_SIZE) { throw new Error(`Message too large: ${payload.length} bytes`) } - return encodeFrame(MessageType.Regular, id, ack, payload) + return Object.freeze({ byteLength: payload.length, [JSON_RPC_PAYLOAD_BYTES]: payload }) +} + +export function encodePreparedJsonRpcFrame( + payload: PreparedJsonRpcPayload, + id: number, + ack: number +): Buffer { + return encodeFrame(MessageType.Regular, id, ack, payload[JSON_RPC_PAYLOAD_BYTES]) } export function encodeKeepAliveFrame(id: number, ack: number): Buffer { diff --git a/src/relay/relay-watcher-event-emitter.ts b/src/relay/relay-watcher-event-emitter.ts index 6bfef6a97f0..739ae89b2f8 100644 --- a/src/relay/relay-watcher-event-emitter.ts +++ b/src/relay/relay-watcher-event-emitter.ts @@ -104,7 +104,7 @@ function publishWatcherBatchToClient( const publish = (events: readonly MappedWatcherEvent[]): boolean => dispatcher.publishProducerNotification(clientId, 'fs.changed', { events }) - // Fast path: publish the whole batch first — two encodes, the same cost as an unchunked emit. + // Fast path: publish the whole batch before paying to group or size individual events. // logDrop:false because rejection here is a measurement, not an outcome: the batch is re-sent in // chunks below, so logging it would report a drop for events that all arrive. if ( @@ -244,7 +244,7 @@ function publishOverflowMarker( // it through onClientDetached, which fires after this settlement. retainOverflowMarker(dispatcher, state, clientId, rootPath, frameBytes) }, - { controlOverflow: 'reject', estimatedBytes: frameBytes } + { controlOverflow: 'reject' } ) if (accepted || settled) { return