From 9bed9bbd34c039a219097de37eabfc139db28fc4 Mon Sep 17 00:00:00 2001 From: Jinwoo Hong <73622457+Jinwoo-H@users.noreply.github.com> Date: Mon, 6 Jul 2026 17:10:04 -0700 Subject: [PATCH] perf(ssh): bound relay bulk-stream backlog so PTY echo is not head-of-line blocked (#7601) Co-authored-by: Orca --- .../ssh-typing-latency-under-relay-load.md | 115 +++++++ .../ssh-filesystem-provider-stream.test.ts | 3 +- src/main/ssh/relay-protocol.test.ts | 70 ++++- src/main/ssh/relay-protocol.ts | 113 +++++-- src/main/ssh/ssh-filesystem-stream-reader.ts | 8 +- src/relay/dispatcher.test.ts | 115 +++++++ src/relay/dispatcher.ts | 137 +++++++- src/relay/fs-handler-file-read.ts | 50 ++- src/relay/fs-handler-stream.test.ts | 58 ++++ src/relay/fs-handler.ts | 25 +- ...-pty-echo-backpressure.integration.test.ts | 295 ++++++++++++++++++ src/relay/fs-stream-registry.ts | 69 +++- src/relay/protocol.ts | 128 ++++++-- src/relay/relay.ts | 72 ++++- tests/e2e/ssh-docker-relay-perf.spec.ts | 98 ++++++ 15 files changed, 1285 insertions(+), 71 deletions(-) create mode 100644 docs/reference/ssh-typing-latency-under-relay-load.md create mode 100644 src/relay/fs-stream-pty-echo-backpressure.integration.test.ts diff --git a/docs/reference/ssh-typing-latency-under-relay-load.md b/docs/reference/ssh-typing-latency-under-relay-load.md new file mode 100644 index 000000000..f7ae8e0b3 --- /dev/null +++ b/docs/reference/ssh-typing-latency-under-relay-load.md @@ -0,0 +1,115 @@ +# SSH Typing Latency Under Relay Load + +Date: 2026-07-06 + +## Symptom + +Typing in an SSH-host terminal can become extremely slow (hundreds of ms to +seconds per echoed keystroke) while other SSH work is active — file previews, +source-control refresh, search. A quiet SSH shell stays snappy, which made the +reports look unreproducible. + +This is distinct from the local warm-switch lag documented in +`terminal-switch-typing-lag-investigation.md` (daemon `listSessions()` +snapshot cost — already fixed; `TerminalHost.listSessions()` is now +metadata-only via `getAppliedSize()`). + +## Root Cause + +The relay and the Electron client share ONE ordered SSH channel for all +JSON-RPC traffic: PTY input/output, file streams, git responses, search, port +scans. Two mechanisms turned bulk traffic into typing latency: + +1. **Relay-side head-of-line blocking (primary).** The `fs.readFileStream` + pump (`src/relay/fs-handler-file-read.ts`) wrote every 256KB chunk (~340KB + framed) into the relay's stdout as fast as local disk reads completed, + ignoring the `write() === false` backpressure signal. A 10MB preview + enqueued ~13.6MB into the pipe at once; a `pty.data` echo emitted + mid-stream queued behind ALL of it. At 2MB/s WAN that is multiple seconds + of echo delay per open file. The reproduced measurement: 4,195,592 bytes + (the entire remaining 3MB test file, framed) queued ahead of one echo + frame. + +2. **Client-side O(n²) frame buffering (secondary).** `FrameDecoder.feed()` + in both `src/main/ssh/relay-protocol.ts` and `src/relay/protocol.ts` did + `Buffer.concat([buffered, chunk])` per data event, re-copying the whole + backlog for every ~32KB TCP chunk — ~2MB of memcpy per 340KB frame, + ~80MB per 10MB file — on the Electron main thread, delaying `pty:write` + IPC dispatch and echo delivery to the renderer. + +## Fix + +- **Bulk lane with sink backpressure** (`RelayDispatcher.notifyBulk`): + `fs.streamChunk` frames are serialized per client and each send waits for + the sink's `drain` when `write()` returns `false`. Interactive frames + (`pty.data`) still use plain `notify()` and are admitted immediately — + because the bulk lane keeps the outbound buffer at ~1 frame, an echo jumps + ahead of every not-yet-admitted chunk. Sinks (`relay.ts` stdout + Unix + socket clients) surface `write()`'s boolean and a one-shot + `waitWriteDrain`; every death path (EPIPE, stdin end, detach, dispose) + flushes parked waiters so a pump can never hang on a dead pipe. + +- **Credit-based flow control** (`fs.streamAck`): the client requests + streams with `flowControl: 'ack'` and acks each processed chunk; the relay + caps unacked chunks at `STREAM_ACK_WINDOW_CHUNKS = 4` (~1MB raw). This + bounds in-flight bulk bytes even past the relay's own pipe (sshd/TCP + buffers) and paces the relay's base64/JSON encode loop so incoming + keystroke frames get event-loop turns. A parked pump wakes on ack, cancel, + release, client detach, and a 1s staleness recheck. + +- **Cross-version compatibility**: old client + new relay → no + `flowControl` param → legacy unpaced pump (still drain-bounded, which is + transparent). New client + old relay → `fs.streamAck` is an unknown + notification and is silently ignored. `STREAM_CHUNK_SIZE` (256KB) is + intentionally unchanged — both sides bake it into chunk offset math, so + changing it would corrupt cross-version streams. + +- **FrameDecoder rewrite** (both mirrored protocol files): chunk-list + buffering; each frame is assembled exactly once (one copy) instead of + re-concatenating the backlog per feed. + +## What Stays Fast / Unchanged + +- Relay PTY output batching (interactive echo fast path in + `src/relay/pty-handler.ts`) is untouched. +- No `listSessions()` / provider inventory was added anywhere near typing, + focus, switch, resume, resize, or render paths. +- Remote file streaming correctness: per-chunk length checks, chunk-count and + byte-count invariants, cancel paths, and the 12MB round-trip integration + test all pass unchanged. + +## Regression Coverage + +- `src/relay/fs-stream-pty-echo-backpressure.integration.test.ts` — real + mux ↔ dispatcher ↔ FsHandler over a congestible in-memory pipe; asserts + deterministic BYTE bounds (not wall-clock): echo queues behind < 2 framed + chunks when congested (pre-fix: whole file), ack window caps in-flight + chunks, legacy no-ack clients still get complete streams. +- `src/relay/fs-handler-stream.test.ts` — pump parks at the ack window, + resumes per ack, and releases its file handle when cancelled while parked. +- `src/relay/dispatcher.test.ts` — notifyBulk semantics: drain gating, + per-client targeting, dispose releases parked senders, interactive + notify() not gated behind a stalled bulk lane. +- `src/main/ssh/relay-protocol.test.ts` — decoder byte-at-a-time boundary + straddling, oversized-frame resync across odd chunk sizes, and a + no-`Buffer.concat`-during-feed guard that locks out the O(n²) shape. +- `tests/e2e/ssh-docker-relay-perf.spec.ts` — new "busy relay" scenario: + types while two 8MB remote file-read loops and a git.status loop run; + budgets median < 500ms, worst < 2000ms. (On loopback Docker the pre-fix + HOL is only tens of ms — the in-process byte-bound test is the + authoritative red/green; the e2e guards the end-to-end wiring.) + +## Residual Gaps + +- **Large single-frame responses**: `git.*` responses (stdout capped at + `MAX_GIT_BUFFER` = 10MB → ~13MB framed) and other request/response results + are one atomic frame on the wire; a huge diff can still delay an echo by + its own transfer time. Fixing this requires response chunking (a protocol + change); bounded by the 16MB `MAX_MESSAGE_SIZE`. +- **Outbound direction**: a large `fs.writeFile` request frame (client → + relay) can queue ahead of keystroke `pty.data` frames on `channel.stdin`. + Same shape, opposite direction, rarer trigger (saving a large remote file + while typing). +- `pty.ackData` flow control for PTY output remains unenforced + (`src/relay/pty-handler.ts` — "not yet enforced"); PTY output floods are + already paced by the relay's 8ms batch/16KB slice scheduler. diff --git a/src/main/providers/ssh-filesystem-provider-stream.test.ts b/src/main/providers/ssh-filesystem-provider-stream.test.ts index 29daf5af5..b66b56335 100644 --- a/src/main/providers/ssh-filesystem-provider-stream.test.ts +++ b/src/main/providers/ssh-filesystem-provider-stream.test.ts @@ -80,7 +80,8 @@ describe('SshFilesystemProvider readFile streaming', () => { const result = await provider.readFile('/home/user/file.txt') expect(mux.request).toHaveBeenCalledWith('fs.readFileStream', { - filePath: '/home/user/file.txt' + filePath: '/home/user/file.txt', + flowControl: 'ack' }) expect(result).toEqual({ content: text, isBinary: false }) }) diff --git a/src/main/ssh/relay-protocol.test.ts b/src/main/ssh/relay-protocol.test.ts index 8369c3c02..c4e9f0897 100644 --- a/src/main/ssh/relay-protocol.test.ts +++ b/src/main/ssh/relay-protocol.test.ts @@ -1,4 +1,4 @@ -import { describe, expect, it } from 'vitest' +import { describe, expect, it, vi } from 'vitest' import { HEADER_LENGTH, MessageType, @@ -160,6 +160,74 @@ describe('FrameDecoder', () => { expect(frames).toHaveLength(1) expect(frames[0].id).toBe(2) }) + + it('decodes frames fed one byte at a time (worst-case boundary straddling)', () => { + const frames: DecodedFrame[] = [] + const decoder = new FrameDecoder((f) => frames.push(f)) + + const frame1 = encodeFrame(MessageType.Regular, 1, 0, Buffer.from('first payload')) + const frame2 = encodeFrame(MessageType.Regular, 2, 1, Buffer.from('second')) + const combined = Buffer.concat([frame1, frame2]) + + for (let i = 0; i < combined.length; i += 1) { + decoder.feed(combined.subarray(i, i + 1)) + } + + expect(frames).toHaveLength(2) + expect(frames[0].payload.toString()).toBe('first payload') + expect(frames[1].id).toBe(2) + expect(frames[1].payload.toString()).toBe('second') + }) + + it('skips an oversized frame fed in odd-sized chunks and resynchronizes', () => { + const errors: Error[] = [] + const frames: DecodedFrame[] = [] + const decoder = new FrameDecoder( + (f) => frames.push(f), + (err) => errors.push(err) + ) + + const oversizedLength = 17 * 1024 * 1024 + const header = Buffer.alloc(HEADER_LENGTH) + header[0] = MessageType.Regular + header.writeUInt32BE(1, 1) + header.writeUInt32BE(0, 5) + header.writeUInt32BE(oversizedLength, 9) + const oversized = Buffer.concat([header, Buffer.alloc(oversizedLength)]) + const valid = encodeFrame(MessageType.Regular, 2, 0, Buffer.from('after')) + const combined = Buffer.concat([oversized, valid]) + + const chunkSize = 1024 * 1024 - 7 + for (let i = 0; i < combined.length; i += chunkSize) { + decoder.feed(combined.subarray(i, i + chunkSize)) + } + + expect(errors).toHaveLength(1) + expect(frames).toHaveLength(1) + expect(frames[0].payload.toString()).toBe('after') + }) + + it('never rebuilds the buffered stream per feed while assembling a large frame', () => { + // Regression: feed() used Buffer.concat([buffered, chunk]) per data event, + // re-copying the whole backlog for every TCP chunk — O(n²) memcpy on the + // Electron main thread while fs.streamChunk frames arrive (SSH typing lag). + const frames: DecodedFrame[] = [] + const decoder = new FrameDecoder((f) => frames.push(f)) + const frame = encodeFrame(MessageType.Regular, 1, 0, Buffer.alloc(512 * 1024, 0x61)) + + const concatSpy = vi.spyOn(Buffer, 'concat') + try { + for (let i = 0; i < frame.length; i += 32 * 1024) { + decoder.feed(frame.subarray(i, i + 32 * 1024)) + } + } finally { + concatSpy.mockRestore() + } + + expect(frames).toHaveLength(1) + expect(frames[0].payload.length).toBe(512 * 1024) + expect(concatSpy).not.toHaveBeenCalled() + }) }) describe('parseJsonRpcMessage', () => { diff --git a/src/main/ssh/relay-protocol.ts b/src/main/ssh/relay-protocol.ts index 3b8b63af1..db57c4d19 100644 --- a/src/main/ssh/relay-protocol.ts +++ b/src/main/ssh/relay-protocol.ts @@ -127,7 +127,13 @@ export type DecodedFrame = { * Incremental frame parser. Feed it chunks of data; it emits complete frames. */ export class FrameDecoder { - private buffer = Buffer.alloc(0) + // Why: feed() runs on the Electron main thread for every SSH channel data + // event. Rebuilding one contiguous buffer per feed (Buffer.concat) re-copies + // every already-buffered byte for each incoming ~32KB TCP chunk — O(n²) per + // large frame (a 340KB fs.streamChunk frame cost ~2MB of memcpy). A chunk + // list assembles each frame exactly once instead. + private chunks: Buffer[] = [] + private bufferedLength = 0 private onFrame: (frame: DecodedFrame) => void private onError: ((err: Error) => void) | null @@ -137,21 +143,31 @@ export class FrameDecoder { } feed(chunk: Buffer | Uint8Array): void { - this.buffer = Buffer.concat([this.buffer, chunk]) + const buf = Buffer.isBuffer(chunk) + ? chunk + : Buffer.from(chunk.buffer, chunk.byteOffset, chunk.byteLength) + if (buf.length > 0) { + this.chunks.push(buf) + this.bufferedLength += buf.length + } - while (this.buffer.length >= HEADER_LENGTH) { - const length = this.buffer.readUInt32BE(9) + while (this.bufferedLength >= HEADER_LENGTH) { + const header = this.peekBytes(HEADER_LENGTH) + const length = header.readUInt32BE(9) const totalLength = HEADER_LENGTH + length + if (this.bufferedLength < totalLength) { + // Not fully received yet (also holds oversized frames until they can + // be skipped whole, keeping the decoder synchronized). + break + } + // Why: throwing here would leave the buffer in a partially consumed // state — subsequent feed() calls would try to parse leftover payload // bytes as a new header, corrupting every future frame. Instead we // skip the entire oversized frame so the decoder stays synchronized. if (length > MAX_MESSAGE_SIZE) { - if (this.buffer.length < totalLength) { - break - } - this.buffer = this.buffer.subarray(totalLength) + this.discardBytes(totalLength) const err = new Error(`Frame payload too large: ${length} bytes — discarded`) if (this.onError) { this.onError(err) @@ -159,24 +175,83 @@ export class FrameDecoder { continue } - if (this.buffer.length < totalLength) { - break - } - + const framed = this.takeBytes(totalLength) const frame: DecodedFrame = { - type: this.buffer[0], - id: this.buffer.readUInt32BE(1), - ack: this.buffer.readUInt32BE(5), - payload: this.buffer.subarray(HEADER_LENGTH, totalLength) + type: framed[0], + id: framed.readUInt32BE(1), + ack: framed.readUInt32BE(5), + payload: framed.subarray(HEADER_LENGTH, totalLength) } - - this.buffer = this.buffer.subarray(totalLength) this.onFrame(frame) } } reset(): void { - this.buffer = Buffer.alloc(0) + this.chunks = [] + this.bufferedLength = 0 + } + + /** View of the first `count` buffered bytes without consuming them. */ + private peekBytes(count: number): Buffer { + const first = this.chunks[0] + if (first.length >= count) { + return first + } + const out = Buffer.allocUnsafe(count) + let copied = 0 + for (const part of this.chunks) { + copied += part.copy(out, copied, 0, Math.min(part.length, count - copied)) + if (copied >= count) { + break + } + } + return out + } + + /** Consume and return the first `count` buffered bytes (single copy). */ + private takeBytes(count: number): Buffer { + const first = this.chunks[0] + if (first.length === count) { + this.chunks.shift() + this.bufferedLength -= count + return first + } + if (first.length > count) { + this.chunks[0] = first.subarray(count) + this.bufferedLength -= count + return first.subarray(0, count) + } + const out = Buffer.allocUnsafe(count) + let copied = 0 + while (copied < count) { + const part = this.chunks[0] + const take = Math.min(part.length, count - copied) + part.copy(out, copied, 0, take) + copied += take + if (take === part.length) { + this.chunks.shift() + } else { + this.chunks[0] = part.subarray(take) + } + } + this.bufferedLength -= count + return out + } + + /** Consume the first `count` buffered bytes without assembling them. */ + private discardBytes(count: number): void { + let remaining = count + while (remaining > 0) { + const part = this.chunks[0] + if (part.length <= remaining) { + this.chunks.shift() + remaining -= part.length + } else { + this.chunks[0] = part.subarray(remaining) + remaining = 0 + } + } + this.bufferedLength -= count } } diff --git a/src/main/ssh/ssh-filesystem-stream-reader.ts b/src/main/ssh/ssh-filesystem-stream-reader.ts index cda490c85..2147ed0ef 100644 --- a/src/main/ssh/ssh-filesystem-stream-reader.ts +++ b/src/main/ssh/ssh-filesystem-stream-reader.ts @@ -148,6 +148,10 @@ export async function readFileViaStream( expectedSeq += 1 receivedChunks += 1 bytesReceived += decoded.length + // Why: credit-based flow control — the relay caps unacked chunks so bulk + // stream frames cannot queue unbounded ahead of interactive pty.data + // frames on the shared SSH channel. Old relays ignore this notification. + mux.notify('fs.streamAck', { streamId: id, seq }) } const handleEnd = (params: Record): void => { @@ -260,7 +264,9 @@ export async function readFileViaStream( unsubscribers.push(onDispose) void mux - .request('fs.readFileStream', { filePath }) + // Why: flowControl declares this client acks each chunk, letting a new + // relay pace the pump. Old relays ignore the extra param and flood. + .request('fs.readFileStream', { filePath, flowControl: 'ack' }) .then((rawMetadata) => { if (settled) { return diff --git a/src/relay/dispatcher.test.ts b/src/relay/dispatcher.test.ts index a87e37c60..5ab997c9d 100644 --- a/src/relay/dispatcher.test.ts +++ b/src/relay/dispatcher.test.ts @@ -395,4 +395,119 @@ describe('RelayDispatcher', () => { expect(listener).toHaveBeenCalledWith(1) }) + + describe('notifyBulk (bulk lane backpressure)', () => { + it('resolves immediately when the sink accepts the frame', async () => { + const frames: Buffer[] = [] + const bulkDispatcher = new RelayDispatcher((data) => { + frames.push(Buffer.from(data)) + return true + }) + try { + await bulkDispatcher.notifyBulk('bulk.event', { seq: 0 }) + expect(frames).toHaveLength(1) + const frame = decodeFirstFrame(frames[0]) + const msg = JSON.parse(frame.payload.toString()) as JsonRpcNotification + expect(msg.method).toBe('bulk.event') + } finally { + bulkDispatcher.dispose() + } + }) + + it('holds the next bulk frame until the saturated sink drains', async () => { + const frames: Buffer[] = [] + const drainWaiters = new Set<() => void>() + const bulkDispatcher = new RelayDispatcher( + (data) => { + frames.push(Buffer.from(data)) + return false + }, + { waitWriteDrain: (cb) => drainWaiters.add(cb) } + ) + try { + let firstSettled = false + const first = bulkDispatcher.notifyBulk('bulk.event', { seq: 0 }).then(() => { + firstSettled = true + }) + void bulkDispatcher.notifyBulk('bulk.event', { seq: 1 }) + await vi.advanceTimersByTimeAsync(0) + + // First frame written, but its send has not settled and the second + // frame is not admitted while the sink stays saturated. + expect(frames).toHaveLength(1) + expect(firstSettled).toBe(false) + + for (const cb of Array.from(drainWaiters)) { + drainWaiters.delete(cb) + cb() + } + await first + await vi.advanceTimersByTimeAsync(0) + expect(frames).toHaveLength(2) + } finally { + bulkDispatcher.dispose() + } + }) + + it('interactive notify() frames are not gated behind a stalled bulk lane', async () => { + const frames: Buffer[] = [] + const bulkDispatcher = new RelayDispatcher( + (data) => { + frames.push(Buffer.from(data)) + return false + }, + { waitWriteDrain: () => {} } + ) + try { + void bulkDispatcher.notifyBulk('bulk.event', { seq: 0 }) + void bulkDispatcher.notifyBulk('bulk.event', { seq: 1 }) + await vi.advanceTimersByTimeAsync(0) + expect(frames).toHaveLength(1) + + bulkDispatcher.notify('pty.data', { id: 'pty-1', data: 'x' }) + expect(frames).toHaveLength(2) + const msg = JSON.parse( + decodeFirstFrame(frames[1]).payload.toString() + ) as JsonRpcNotification + expect(msg.method).toBe('pty.data') + } finally { + bulkDispatcher.dispose() + } + }) + + it('releases a parked bulk send when the dispatcher is disposed', async () => { + const bulkDispatcher = new RelayDispatcher(() => false, { waitWriteDrain: () => {} }) + const pending = bulkDispatcher.notifyBulk('bulk.event', { seq: 0 }) + await vi.advanceTimersByTimeAsync(0) + bulkDispatcher.dispose() + await expect(pending).resolves.toBeUndefined() + }) + + it('targets only the requested client and resolves for missing clients', async () => { + const primaryFrames: Buffer[] = [] + const secondaryFrames: Buffer[] = [] + const bulkDispatcher = new RelayDispatcher((data) => { + primaryFrames.push(Buffer.from(data)) + return true + }) + try { + const secondaryId = bulkDispatcher.attachClient((data) => { + secondaryFrames.push(Buffer.from(data)) + return true + }) + + await bulkDispatcher.notifyBulk('bulk.event', { seq: 0 }, { clientId: secondaryId }) + expect(primaryFrames).toHaveLength(0) + expect(secondaryFrames).toHaveLength(1) + + await expect( + bulkDispatcher.notifyBulk('bulk.event', { seq: 1 }, { clientId: 999 }) + ).resolves.toBeUndefined() + expect(primaryFrames).toHaveLength(0) + expect(secondaryFrames).toHaveLength(1) + } finally { + bulkDispatcher.dispose() + } + }) + }) }) diff --git a/src/relay/dispatcher.ts b/src/relay/dispatcher.ts index 41809a4d7..427d90e6a 100644 --- a/src/relay/dispatcher.ts +++ b/src/relay/dispatcher.ts @@ -27,10 +27,28 @@ export type MethodHandler = ( export type NotificationHandler = (params: Record, context: RequestContext) => void +/** Sink write. Returning literal `false` signals saturation (Node stream + * semantics); `void`/`true` mean the frame was accepted. */ +export type RelayClientWrite = (data: Buffer) => boolean | void + +export type RelayClientSinkOptions = { + /** One-shot: invoke `cb` once when the sink can accept more data again + * (stream 'drain'), or when the sink is permanently dead (error/close) so + * bulk senders waiting on it never hang. */ + waitWriteDrain?: (cb: () => void) => void +} + type RelayClient = { id: number decoder: FrameDecoder - write: (data: Buffer) => void + write: RelayClientWrite + waitWriteDrain?: (cb: () => void) => void + /** Pending resolvers for bulk sends stalled on sink saturation. Flushed on + * drain, write failure, detach, setWrite, and dispose so no pump hangs. */ + drainWaiters: Set<() => void> + /** Serializes bulk-lane sends per client so at most one bulk frame is + * admitted past the sink's high-water mark at a time. */ + bulkChain: Promise nextOutgoingSeq: number highestReceivedSeq: number generation: number @@ -58,8 +76,8 @@ export class RelayDispatcher { private nextClientId = 1 private nextRequestId = 1 - constructor(write: (data: Buffer) => void) { - this.primaryClient = this.createClient(write) + constructor(write: RelayClientWrite, sinkOptions?: RelayClientSinkOptions) { + this.primaryClient = this.createClient(write, sinkOptions) this.clients.set(this.primaryClient.id, this.primaryClient) this.startKeepalive() } @@ -75,10 +93,14 @@ export class RelayDispatcher { // never acks the new client's frames until the new client's seq catches // up - causing the client's unacked-timeout checker to accumulate stale // timestamps that could eventually fire a false connection-dead signal. - setWrite(write: (data: Buffer) => void): void { + setWrite(write: RelayClientWrite, sinkOptions?: RelayClientSinkOptions): void { this.requestAborts.abortClient(this.primaryClient.id) this.primaryClient.write = write + this.primaryClient.waitWriteDrain = sinkOptions?.waitWriteDrain this.primaryClient.closed = false + // Why: the saturated sink the waiters were parked on no longer exists; + // wake stalled bulk senders so they re-evaluate against the new sink. + this.flushDrainWaiters(this.primaryClient) this.resetClient(this.primaryClient) } @@ -89,14 +111,15 @@ export class RelayDispatcher { this.requestAborts.abortClient(this.primaryClient.id) this.primaryClient.generation++ this.primaryClient.closed = true + this.flushDrainWaiters(this.primaryClient) this.notifyClientDetached(this.primaryClient.id) } // Why: synced remote workspaces can have more than one Orca client attached // to the same relay. Frame sequence numbers and JSON-RPC request ids are per // SSH channel, so each socket client needs independent protocol state. - attachClient(write: (data: Buffer) => void): number { - const client = this.createClient(write) + attachClient(write: RelayClientWrite, sinkOptions?: RelayClientSinkOptions): number { + const client = this.createClient(write, sinkOptions) this.clients.set(client.id, client) return client.id } @@ -109,6 +132,7 @@ export class RelayDispatcher { this.requestAborts.abortClient(clientId) client.generation++ client.closed = true + this.flushDrainWaiters(client) this.clients.delete(clientId) this.notifyClientDetached(clientId) } @@ -165,6 +189,90 @@ export class RelayDispatcher { } } + /** + * Bulk-lane notification. Sends are serialized per client and the returned + * promise resolves only after the sink accepted the frame without reporting + * saturation (or the client went away). Bulk producers (file streams) await + * this between frames so interactive frames (pty.data echo) never queue + * behind an unbounded backlog on the shared SSH channel. + * + * With `clientId`, the frame goes only to that client — stream chunks have + * exactly one consumer, and broadcasting them would let one slow secondary + * client stall everyone. A missing/closed target resolves immediately; the + * caller's staleness check owns aborting the stream. + */ + notifyBulk( + method: string, + params?: Record, + opts?: { clientId?: number } + ): Promise { + 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 + } + // Why: the frame is encoded inside the chain step, not at call time — + // sequence numbers must be assigned in actual write order. + const step = client.bulkChain.then(() => { + if (this.disposed || client.closed) { + return + } + const accepted = this.sendFrame(client, msg) + if (accepted === false) { + return this.waitForClientDrain(client) + } + return undefined + }) + client.bulkChain = step.catch(() => {}) + waits.push(step) + } + if (waits.length === 0) { + return Promise.resolve() + } + return Promise.all(waits).then(() => {}) + } + + private waitForClientDrain(client: RelayClient): Promise { + if (this.disposed || client.closed || !client.waitWriteDrain) { + return Promise.resolve() + } + return new Promise((resolve) => { + let settled = false + const finish = (): void => { + if (settled) { + return + } + settled = true + client.drainWaiters.delete(finish) + resolve() + } + client.drainWaiters.add(finish) + try { + client.waitWriteDrain!(finish) + } catch { + finish() + } + }) + } + + private flushDrainWaiters(client: RelayClient): void { + for (const waiter of Array.from(client.drainWaiters)) { + waiter() + } + } + requestPrimary( method: string, params?: Record, @@ -236,14 +344,20 @@ export class RelayDispatcher { // Why: dispose means this relay instance cannot send responses anymore; // abort in-flight request work so stale SSH-side scans/watchers release. this.requestAborts.abortAll() + for (const client of this.clients.values()) { + this.flushDrainWaiters(client) + } } - private createClient(write: (data: Buffer) => void): RelayClient { + private createClient(write: RelayClientWrite, sinkOptions?: RelayClientSinkOptions): RelayClient { const id = this.nextClientId++ const client: RelayClient = { id, decoder: new FrameDecoder((frame) => this.handleFrame(client, frame)), write, + waitWriteDrain: sinkOptions?.waitWriteDrain, + drainWaiters: new Set(), + bulkChain: Promise.resolve(), nextOutgoingSeq: 1, highestReceivedSeq: 0, generation: 0, @@ -387,13 +501,13 @@ export class RelayDispatcher { private sendFrame( client: RelayClient, msg: JsonRpcRequest | JsonRpcResponse | JsonRpcNotification - ): void { + ): boolean | void { if (this.disposed || client.closed) { return } const seq = client.nextOutgoingSeq++ const frame = encodeJsonRpcFrame(msg, seq, client.highestReceivedSeq) - this.writeFrame(client, frame) + return this.writeFrame(client, frame) } private startKeepalive(): void { @@ -416,12 +530,13 @@ export class RelayDispatcher { this.keepaliveTimer.unref() } - private writeFrame(client: RelayClient, frame: Buffer): void { + private writeFrame(client: RelayClient, frame: Buffer): boolean | void { try { - client.write(frame) + return client.write(frame) } catch (err) { client.closed = true client.generation++ + this.flushDrainWaiters(client) if (client !== this.primaryClient) { this.clients.delete(client.id) this.notifyClientDetached(client.id) diff --git a/src/relay/fs-handler-file-read.ts b/src/relay/fs-handler-file-read.ts index 45a0685b8..003172c08 100644 --- a/src/relay/fs-handler-file-read.ts +++ b/src/relay/fs-handler-file-read.ts @@ -1,7 +1,7 @@ import { open, readFile, stat } from 'node:fs/promises' import { extname } from 'node:path' import type { RelayDispatcher, RequestContext } from './dispatcher' -import { STREAM_CHUNK_SIZE, RelayErrorCode } from './protocol' +import { STREAM_ACK_WINDOW_CHUNKS, STREAM_CHUNK_SIZE, RelayErrorCode } from './protocol' import type { RelayStreamRegistry, TooManyStreamsError } from './fs-stream-registry' import { BINARY_PROBE_BYTES, @@ -61,11 +61,21 @@ type StreamChunkReader = { ): Promise<{ bytesRead: number }> } +export type StreamPumpOptions = { + /** Client that requested the stream. Chunks go only to it — broadcasting + * bulk frames would let one slow secondary client stall the requester. */ + clientId?: number + /** True when the client declared `flowControl: 'ack'` — it sends + * fs.streamAck per processed chunk and the pump caps unacked chunks. */ + paceWithAcks: boolean +} + export async function readRelayFileStreamMetadata( filePath: string, dispatcher: RelayDispatcher, registry: RelayStreamRegistry, - context: RequestContext + context: RequestContext, + pumpOptions?: StreamPumpOptions ): Promise { const stats = await stat(filePath) const mimeType = IMAGE_MIME_TYPES[extname(filePath).toLowerCase()] @@ -106,8 +116,9 @@ export async function readRelayFileStreamMetadata( // Why: pumpChunks owns its own try/finally for handle release; the outer // setImmediate kicks the pump off the metadata-response task so the client // sees the response before the first chunk frame. + const resolvedPumpOptions = pumpOptions ?? { paceWithAcks: false } setImmediate(() => { - void pumpChunks(streamId, stats.size, dispatcher, registry, context) + void pumpChunks(streamId, stats.size, dispatcher, registry, context, resolvedPumpOptions) }) return { @@ -126,7 +137,8 @@ async function pumpChunks( totalSize: number, dispatcher: RelayDispatcher, registry: RelayStreamRegistry, - context: RequestContext + context: RequestContext, + pumpOptions: StreamPumpOptions ): Promise { const entry = registry.get(streamId) if (!entry) { @@ -150,6 +162,27 @@ async function pumpChunks( endReason = 'aborted' break } + // Why: credit window — bulk chunks share one ordered SSH channel with + // interactive pty.data frames. Waiting for client acks bounds how many + // stream bytes a keystroke echo can queue behind, and yields the relay + // event loop so incoming keystrokes are handled between chunks. + if (pumpOptions.paceWithAcks) { + while ( + seq - registry.ackedThroughSeq(streamId) > STREAM_ACK_WINDOW_CHUNKS && + !context.isStale() && + !registry.isAborted(streamId) + ) { + await registry.waitForAck(streamId) + } + if (context.isStale()) { + endReason = 'stale' + break + } + if (registry.isAborted(streamId)) { + endReason = 'aborted' + break + } + } const want = Math.min(STREAM_CHUNK_SIZE, totalSize - offset) const bytesRead = await readFullStreamChunk(entry.handle, buffer, want, offset) if (bytesRead !== want) { @@ -167,7 +200,14 @@ async function pumpChunks( break } const data = buffer.subarray(0, bytesRead).toString('base64') - dispatcher.notify('fs.streamChunk', { streamId, seq, data }) + // Why: the bulk lane waits out sink saturation, so a flood of chunk + // frames cannot pile up in the outbound pipe ahead of interactive + // pty.data frames written via plain notify(). + await dispatcher.notifyBulk( + 'fs.streamChunk', + { streamId, seq, data }, + pumpOptions.clientId !== undefined ? { clientId: pumpOptions.clientId } : undefined + ) offset += bytesRead seq += 1 } diff --git a/src/relay/fs-handler-stream.test.ts b/src/relay/fs-handler-stream.test.ts index d71dbf8c0..c7624afce 100644 --- a/src/relay/fs-handler-stream.test.ts +++ b/src/relay/fs-handler-stream.test.ts @@ -34,6 +34,9 @@ function createMockDispatcher() { notify: vi.fn((method: string, params?: Record) => { notifications.push({ method, params }) }), + notifyBulk: vi.fn(async (method: string, params?: Record): Promise => { + notifications.push({ method, params }) + }), _notifications: notifications, callRequest( method: string, @@ -278,6 +281,61 @@ describe('FsHandler readFileStream', () => { expect(err).toBeNull() }) + it('parks the pump at the ack credit window and resumes on fs.streamAck', async () => { + const filePath = path.join(tmpDir, 'paced.png') + const content = Buffer.alloc(1536 * 1024, 0x42) // 6 chunks + writeFileSync(filePath, content) + + const meta = (await dispatcher.callRequest( + 'fs.readFileStream', + { filePath, flowControl: 'ack' }, + { isStale: () => false } + )) as { streamId: number } + + // Window of 4 admits seqs 0..3; seq 4 must wait for an ack. + await waitFor(() => collectStream(dispatcher).chunks.length === 4) + await flush(10) + expect(collectStream(dispatcher).chunks.length).toBe(4) + expect(collectStream(dispatcher).end).toBeNull() + + // Ack one chunk → exactly one more is admitted. + dispatcher.callNotification('fs.streamAck', { streamId: meta.streamId, seq: 0 }) + await waitFor(() => collectStream(dispatcher).chunks.length === 5) + await flush(10) + expect(collectStream(dispatcher).chunks.length).toBe(5) + + for (let seq = 1; seq < 6; seq += 1) { + dispatcher.callNotification('fs.streamAck', { streamId: meta.streamId, seq }) + } + await waitFor(() => collectStream(dispatcher).end !== null) + + const { chunks, err } = collectStream(dispatcher) + expect(err).toBeNull() + const reassembled = Buffer.concat(chunks.map((c) => Buffer.from(c.data, 'base64'))) + expect(reassembled.equals(content)).toBe(true) + }) + + it('releases a pump parked on the ack window when the stream is cancelled', async () => { + const filePath = path.join(tmpDir, 'parked-cancel.png') + writeFileSync(filePath, Buffer.alloc(4 * 1024 * 1024, 0x42)) // 16 chunks + + const meta = (await dispatcher.callRequest( + 'fs.readFileStream', + { filePath, flowControl: 'ack' }, + { isStale: () => false } + )) as { streamId: number } + + await waitFor(() => collectStream(dispatcher).chunks.length === 4) + dispatcher.callNotification('fs.cancelStream', { streamId: meta.streamId }) + + const registry = (handler as unknown as { streamRegistry: { size(): number } }).streamRegistry + await waitFor(() => registry.size() === 0) + + const { end, err } = collectStream(dispatcher) + expect(end).toBeNull() + expect(err).toBeNull() + }) + it('rejects the 17th concurrent stream with TooManyStreams', async () => { const paths: string[] = [] for (let i = 0; i < 17; i++) { diff --git a/src/relay/fs-handler.ts b/src/relay/fs-handler.ts index 0c4a263f1..c6564c51f 100644 --- a/src/relay/fs-handler.ts +++ b/src/relay/fs-handler.ts @@ -86,7 +86,13 @@ export class FsHandler { constructor(dispatcher: RelayDispatcher, _context: RelayContext) { this.dispatcher = dispatcher this.registerHandlers() - this.dispatcher.onClientDetached?.((clientId) => this.releaseClientWatches(clientId)) + this.dispatcher.onClientDetached?.((clientId) => { + this.releaseClientWatches(clientId) + // Why: a detached client's fs.streamAck frames will never arrive; wake + // any pump parked on the ack window so it re-checks staleness and exits + // instead of stranding its open file handle. + this.streamRegistry.wakeAllAckWaiters() + }) } private registerHandlers(): void { @@ -113,6 +119,7 @@ export class FsHandler { this.dispatcher.onRequest('fs.watch', (p, context) => this.watch(p, context)) this.dispatcher.onNotification('fs.unwatch', (p, context) => this.unwatch(p, context)) this.dispatcher.onNotification('fs.cancelStream', (p) => this.cancelStream(p)) + this.dispatcher.onNotification('fs.streamAck', (p) => this.streamAck(p)) } private async readDir(params: Record) { @@ -148,7 +155,13 @@ export class FsHandler { private async readFileStream(params: Record, context?: RequestContext) { const filePath = expandTilde(params.filePath as string) const ctx = context ?? { clientId: 0, isStale: () => false } - return readRelayFileStreamMetadata(filePath, this.dispatcher, this.streamRegistry, ctx) + return readRelayFileStreamMetadata(filePath, this.dispatcher, this.streamRegistry, ctx, { + // Why: only target the requesting client when the dispatcher actually + // routed this request (context present) — direct-call tests and legacy + // paths keep broadcast semantics. + ...(context ? { clientId: context.clientId } : {}), + paceWithAcks: params.flowControl === 'ack' + }) } private async tempDir(): Promise { @@ -162,6 +175,14 @@ export class FsHandler { } } + private streamAck(params: Record): void { + const streamId = params.streamId as number | undefined + const seq = params.seq as number | undefined + if (typeof streamId === 'number' && typeof seq === 'number') { + this.streamRegistry.recordAck(streamId, seq) + } + } + private async writeFile(params: Record) { const filePath = expandTilde(params.filePath as string) const content = params.content as string diff --git a/src/relay/fs-stream-pty-echo-backpressure.integration.test.ts b/src/relay/fs-stream-pty-echo-backpressure.integration.test.ts new file mode 100644 index 000000000..8cca27878 --- /dev/null +++ b/src/relay/fs-stream-pty-echo-backpressure.integration.test.ts @@ -0,0 +1,295 @@ +/** + * Regression: SSH typing latency under bulk file-stream load. + * + * The relay and the client share ONE ordered SSH channel. If the relay + * enqueues an entire file's fs.streamChunk frames into the outbound pipe + * at once, an interactive pty.data echo emitted mid-stream queues behind + * megabytes of bulk data and typing feels seconds-slow. + * + * These tests model the SSH channel as a congestible in-memory pipe and + * assert deterministic byte bounds instead of wall-clock latency: + * - with a saturated sink, the relay stalls the pump on write backpressure + * so at most ~1 chunk frame sits ahead of a pty echo; + * - with a fast sink but an unresponsive client, the fs.streamAck credit + * window bounds the in-flight backlog; + * - legacy clients that never ack still receive the full stream. + */ +import { describe, expect, it, beforeEach, afterEach } from 'vitest' +import { mkdtempSync, writeFileSync } from 'node:fs' +import { rm } from 'node:fs/promises' +import * as path from 'node:path' +import { tmpdir } from 'node:os' +import { randomBytes } from 'node:crypto' + +import { + SshChannelMultiplexer, + type MultiplexerTransport +} from '../main/ssh/ssh-channel-multiplexer' +import { readFileViaStream } from '../main/ssh/ssh-filesystem-stream-reader' + +import { RelayDispatcher } from './dispatcher' +import { RelayContext } from './context' +import { FsHandler } from './fs-handler' +import { STREAM_CHUNK_SIZE } from './protocol' + +// One framed fs.streamChunk: 256KB raw → base64 (4/3) + JSON envelope + header. +const FRAMED_CHUNK_BYTES = Math.ceil((STREAM_CHUNK_SIZE * 4) / 3) + 512 +// Node pipe/socket sinks report saturation via write() === false past the HWM. +const SINK_HIGH_WATER_MARK = 64 * 1024 + +async function waitUntil( + predicate: () => boolean, + what: string, + timeoutMs = 10_000 +): Promise { + const deadline = Date.now() + timeoutMs + while (!predicate()) { + if (Date.now() > deadline) { + throw new Error(`waitUntil timed out: ${what}`) + } + await new Promise((r) => setImmediate(r)) + } +} + +/** Waits until `read()` stops changing for `stableTurns` macrotask turns. */ +async function waitUntilSettled(read: () => number, stableTurns = 25): Promise { + let last = read() + let stable = 0 + while (stable < stableTurns) { + await new Promise((r) => setImmediate(r)) + const current = read() + if (current === last) { + stable += 1 + } else { + stable = 0 + last = current + } + } +} + +type Harness = { + mux: SshChannelMultiplexer + dispatcher: RelayDispatcher + fsHandler: FsHandler + queuedBytes: () => number + /** Deliver every queued relay→client buffer to the client mux. */ + deliverAll: () => void + /** Start delivering continuously on each macrotask turn. */ + startAutoDeliver: () => void + dispose: () => void +} + +function createHarness(opts: { congested: boolean }): Harness { + let relayFeed: ((data: Buffer) => void) | null = null + const clientDataCallbacks: ((data: Buffer) => void)[] = [] + + const clientTransport: MultiplexerTransport = { + write: (data: Buffer) => { + // Client → relay: keystrokes and acks flow on the opposite direction of + // the duplex channel; they are not blocked by relay→client congestion. + setImmediate(() => relayFeed?.(data)) + }, + onData: (cb) => { + clientDataCallbacks.push(cb) + }, + onClose: () => {} + } + + const outQueue: Buffer[] = [] + let queuedBytes = 0 + const drainWaiters = new Set<() => void>() + const fireDrainIfIdle = (): void => { + if (queuedBytes > 0) { + return + } + for (const cb of Array.from(drainWaiters)) { + drainWaiters.delete(cb) + cb() + } + } + + const dispatcher = new RelayDispatcher( + (data: Buffer) => { + outQueue.push(data) + queuedBytes += data.length + if (!opts.congested) { + return true + } + return queuedBytes < SINK_HIGH_WATER_MARK + }, + { + waitWriteDrain: (cb: () => void) => { + drainWaiters.add(cb) + fireDrainIfIdle() + } + } + ) + relayFeed = (data: Buffer) => dispatcher.feed(data) + + const deliverAll = (): void => { + while (outQueue.length > 0) { + const buf = outQueue.shift()! + queuedBytes -= buf.length + for (const cb of clientDataCallbacks) { + cb(buf) + } + } + fireDrainIfIdle() + } + + let autoDeliverTimer: ReturnType | null = null + const startAutoDeliver = (): void => { + if (autoDeliverTimer) { + return + } + autoDeliverTimer = setInterval(deliverAll, 1) + } + + const context = new RelayContext() + const fsHandler = new FsHandler(dispatcher, context) + const mux = new SshChannelMultiplexer(clientTransport) + + return { + mux, + dispatcher, + fsHandler, + queuedBytes: () => queuedBytes, + deliverAll, + startAutoDeliver, + dispose: () => { + if (autoDeliverTimer) { + clearInterval(autoDeliverTimer) + } + mux.dispose() + dispatcher.dispose() + fsHandler.dispose() + } + } +} + +describe('fs.readFileStream vs pty.data echo head-of-line blocking', () => { + let tmpDir: string + + beforeEach(() => { + tmpDir = mkdtempSync(path.join(tmpdir(), 'relay-stream-hol-')) + }) + + afterEach(async () => { + await rm(tmpDir, { recursive: true, force: true }) + }) + + it('bounds bulk bytes queued ahead of a pty echo when the channel is congested', async () => { + const harness = createHarness({ congested: true }) + try { + const filePath = path.join(tmpDir, 'big.png') + const original = randomBytes(3 * 1024 * 1024) // 12 chunks + writeFileSync(filePath, original) + + // Relay-side fake PTY: echoes input back immediately, mirroring + // PtyHandler's interactive fast path (dispatcher.notify on echo). + let queuedBytesAheadOfEcho = -1 + harness.dispatcher.onNotification('pty.data', (params) => { + queuedBytesAheadOfEcho = harness.queuedBytes() + harness.dispatcher.notify('pty.data', { id: params.id, data: params.data }) + }) + + const readPromise = readFileViaStream(harness.mux, filePath) + // Deliver until the client has the stream metadata, then congest fully. + await waitUntil(() => harness.queuedBytes() > 0, 'metadata response queued') + harness.deliverAll() + // Let the pump run to whatever bound it enforces (pre-fix: whole file). + await waitUntil(() => harness.queuedBytes() > 0, 'first chunk queued') + await waitUntilSettled(() => harness.queuedBytes()) + + // Type one key while the stream is congested. + harness.mux.notify('pty.data', { id: 'pty-1', data: 'x' }) + await waitUntil(() => queuedBytesAheadOfEcho >= 0, 'echo emitted by relay') + + // The echo must not sit behind an unbounded chunk backlog: at most one + // in-flight bulk frame (the write that saturated the sink) plus slack. + expect(queuedBytesAheadOfEcho).toBeLessThan(2 * FRAMED_CHUNK_BYTES) + + // Un-congest: the stream must still complete with intact content. + harness.startAutoDeliver() + const result = await readPromise + expect(Buffer.from(result.content, 'base64').equals(original)).toBe(true) + } finally { + harness.dispose() + } + }, 30_000) + + it('caps in-flight chunks via the fs.streamAck credit window when the client stalls', async () => { + const harness = createHarness({ congested: false }) + try { + const filePath = path.join(tmpDir, 'big.png') + writeFileSync(filePath, randomBytes(3 * 1024 * 1024)) // 12 chunks + + const receivedSeqs: number[] = [] + harness.mux.onNotificationByMethod('fs.streamChunk', (params) => { + receivedSeqs.push(params.seq as number) + }) + let streamEnded = false + harness.mux.onNotificationByMethod('fs.streamEnd', () => { + streamEnded = true + }) + + harness.startAutoDeliver() + // Raw ack-capable request without sending any acks: models a client + // whose main thread is too busy to process chunks. + const metadata = (await harness.mux.request('fs.readFileStream', { + filePath, + flowControl: 'ack' + })) as { streamId: number } + + await waitUntil(() => receivedSeqs.length > 0, 'first chunk received') + await waitUntilSettled(() => receivedSeqs.length) + + // Without acks the relay must stop at the credit window, not flood + // the remaining chunks. + expect(receivedSeqs.length).toBeLessThanOrEqual(5) + expect(streamEnded).toBe(false) + + // Acking releases the window and the stream completes. + const totalChunks = 12 + for (let seq = 0; seq < totalChunks; seq += 1) { + harness.mux.notify('fs.streamAck', { streamId: metadata.streamId, seq }) + } + await waitUntil(() => streamEnded, 'stream completed after acks') + expect(receivedSeqs).toEqual(Array.from({ length: totalChunks }, (_, i) => i)) + } finally { + harness.dispose() + } + }, 30_000) + + it('still delivers the full stream to legacy clients that never ack', async () => { + const harness = createHarness({ congested: false }) + try { + const filePath = path.join(tmpDir, 'legacy.png') + const original = randomBytes(1024 * 1024 + 12345) + writeFileSync(filePath, original) + + const chunks = new Map() + let streamEnded = false + harness.mux.onNotificationByMethod('fs.streamChunk', (params) => { + chunks.set(params.seq as number, Buffer.from(params.data as string, 'base64')) + }) + harness.mux.onNotificationByMethod('fs.streamEnd', () => { + streamEnded = true + }) + + harness.startAutoDeliver() + // Legacy request shape: no flowControl param, and no acks ever sent. + await harness.mux.request('fs.readFileStream', { filePath }) + await waitUntil(() => streamEnded, 'legacy stream completed') + + const reassembled = Buffer.concat( + Array.from(chunks.entries()) + .sort(([a], [b]) => a - b) + .map(([, buf]) => buf) + ) + expect(reassembled.equals(original)).toBe(true) + } finally { + harness.dispose() + } + }, 30_000) +}) diff --git a/src/relay/fs-stream-registry.ts b/src/relay/fs-stream-registry.ts index e0bf2a5ed..2bca76c94 100644 --- a/src/relay/fs-stream-registry.ts +++ b/src/relay/fs-stream-registry.ts @@ -1,9 +1,14 @@ import type { FileHandle } from 'node:fs/promises' -import { MAX_CONCURRENT_STREAMS, RelayErrorCode } from './protocol' +import { MAX_CONCURRENT_STREAMS, RelayErrorCode, STREAM_ACK_STALL_RECHECK_MS } from './protocol' type StreamEntry = { handle: FileHandle aborted: boolean + /** Highest chunk seq the client acknowledged (in-order; -1 = none yet). */ + ackedThroughSeq: number + /** Pumps parked on the ack credit window. Woken by acks, abort, release, + * and a periodic stall recheck so a vanished client cannot strand a pump. */ + ackWaiters: Set<() => void> } export class TooManyStreamsError extends Error { @@ -22,7 +27,12 @@ export class RelayStreamRegistry { throw new TooManyStreamsError() } const streamId = this.nextId++ - this.streams.set(streamId, { handle, aborted: false }) + this.streams.set(streamId, { + handle, + aborted: false, + ackedThroughSeq: -1, + ackWaiters: new Set() + }) return streamId } @@ -30,6 +40,7 @@ export class RelayStreamRegistry { const entry = this.streams.get(streamId) if (entry) { entry.aborted = true + this.wakeAckWaiters(entry) } } @@ -41,11 +52,65 @@ export class RelayStreamRegistry { return this.streams.get(streamId) } + recordAck(streamId: number, seq: number): void { + const entry = this.streams.get(streamId) + if (!entry || typeof seq !== 'number' || !Number.isFinite(seq)) { + return + } + if (seq > entry.ackedThroughSeq) { + entry.ackedThroughSeq = seq + } + this.wakeAckWaiters(entry) + } + + ackedThroughSeq(streamId: number): number { + return this.streams.get(streamId)?.ackedThroughSeq ?? Number.MAX_SAFE_INTEGER + } + + /** Resolves on the next ack/abort/release for this stream, or after the + * stall-recheck interval so callers can re-evaluate staleness. */ + waitForAck(streamId: number): Promise { + const entry = this.streams.get(streamId) + if (!entry || entry.aborted) { + return Promise.resolve() + } + return new Promise((resolve) => { + let settled = false + const timer = setTimeout(() => finish(), STREAM_ACK_STALL_RECHECK_MS) + timer.unref?.() + const finish = (): void => { + if (settled) { + return + } + settled = true + clearTimeout(timer) + entry.ackWaiters.delete(finish) + resolve() + } + entry.ackWaiters.add(finish) + }) + } + + /** Wake every parked pump (all streams) so it re-checks staleness — used + * when a client detaches and its acks will never arrive. */ + wakeAllAckWaiters(): void { + for (const entry of this.streams.values()) { + this.wakeAckWaiters(entry) + } + } + + private wakeAckWaiters(entry: StreamEntry): void { + for (const waiter of Array.from(entry.ackWaiters)) { + waiter() + } + } + async release(streamId: number): Promise { const entry = this.streams.get(streamId) if (!entry) { return } + this.wakeAckWaiters(entry) this.streams.delete(streamId) try { await entry.handle.close() diff --git a/src/relay/protocol.ts b/src/relay/protocol.ts index 66818efaf..267f6ef48 100644 --- a/src/relay/protocol.ts +++ b/src/relay/protocol.ts @@ -48,6 +48,17 @@ export const TIMEOUT_MS = 20_000 export const STREAM_CHUNK_SIZE = 256 * 1024 export const MAX_CONCURRENT_STREAMS = 16 +/** Max unacked fs.streamChunk frames in flight per stream when the client + * requested `flowControl: 'ack'`. Bounds how many bulk bytes an interactive + * pty.data frame can queue behind on the shared SSH channel (~1MB raw) while + * keeping the pipe full across one ack round-trip on fast links. */ +export const STREAM_ACK_WINDOW_CHUNKS = 4 + +/** Safety-valve poll interval for a pump stalled on acks: re-checks stream + * abort/staleness so a client that vanished mid-stream cannot park the pump + * (and its open file handle) forever. */ +export const STREAM_ACK_STALL_RECHECK_MS = 1_000 + export const RelayErrorCode = { TooManyStreams: -33006, StreamProtocolError: -33007 @@ -109,7 +120,12 @@ export function encodeKeepAliveFrame(id: number, ack: number): Buffer { } export class FrameDecoder { - private buffer = Buffer.alloc(0) + // Why: feed() sits on the hot receive path. Rebuilding one contiguous + // buffer per feed (Buffer.concat) re-copies every already-buffered byte for + // each incoming chunk — O(n²) per large frame. A chunk list assembles each + // frame exactly once instead. + private chunks: Buffer[] = [] + private bufferedLength = 0 private onFrame: (frame: DecodedFrame) => void private onError: ((err: Error) => void) | null @@ -119,23 +135,32 @@ export class FrameDecoder { } feed(chunk: Buffer | Uint8Array): void { - this.buffer = Buffer.concat([this.buffer, chunk]) + const buf = Buffer.isBuffer(chunk) + ? chunk + : Buffer.from(chunk.buffer, chunk.byteOffset, chunk.byteLength) + if (buf.length > 0) { + this.chunks.push(buf) + this.bufferedLength += buf.length + } - while (this.buffer.length >= HEADER_LENGTH) { - const length = this.buffer.readUInt32BE(9) + while (this.bufferedLength >= HEADER_LENGTH) { + const header = this.peekBytes(HEADER_LENGTH) + const length = header.readUInt32BE(9) const totalLength = HEADER_LENGTH + length + if (this.bufferedLength < totalLength) { + // Not fully received yet (also holds oversized frames until they can + // be skipped whole, keeping the decoder synchronized). + break + } + if (length > MAX_MESSAGE_SIZE) { // Why: Throwing here would leave the buffer in a partially consumed // state — subsequent feed() calls would try to parse the leftover // payload bytes as a new header, corrupting every future frame. // Instead we skip the entire oversized frame so the decoder stays // synchronized with the stream. - if (this.buffer.length < totalLength) { - // Haven't received the full oversized payload yet; wait for more data. - break - } - this.buffer = this.buffer.subarray(totalLength) + this.discardBytes(totalLength) const err = new Error(`Frame payload too large: ${length} bytes — discarded`) if (this.onError) { this.onError(err) @@ -145,23 +170,20 @@ export class FrameDecoder { continue } - if (this.buffer.length < totalLength) { - break - } - + const framed = this.takeBytes(totalLength) const frame: DecodedFrame = { - type: this.buffer[0], - id: this.buffer.readUInt32BE(1), - ack: this.buffer.readUInt32BE(5), - payload: this.buffer.subarray(HEADER_LENGTH, totalLength) + type: framed[0], + id: framed.readUInt32BE(1), + ack: framed.readUInt32BE(5), + payload: framed.subarray(HEADER_LENGTH, totalLength) } - this.buffer = this.buffer.subarray(totalLength) this.onFrame(frame) } } reset(): void { - this.buffer = Buffer.alloc(0) + this.chunks = [] + this.bufferedLength = 0 } // Why: at the handshake → dispatcher transition, the next consumer must @@ -169,10 +191,74 @@ export class FrameDecoder { // frame. This returns and clears the decoder's internal residue so the // caller can hand it to the dispatcher (or stdout pipe) without loss. drain(): Buffer { - const out = this.buffer - this.buffer = Buffer.alloc(0) + const out = + this.chunks.length === 1 ? this.chunks[0] : Buffer.concat(this.chunks, this.bufferedLength) + this.reset() return out } + + /** View of the first `count` buffered bytes without consuming them. */ + private peekBytes(count: number): Buffer { + const first = this.chunks[0] + if (first.length >= count) { + return first + } + const out = Buffer.allocUnsafe(count) + let copied = 0 + for (const part of this.chunks) { + copied += part.copy(out, copied, 0, Math.min(part.length, count - copied)) + if (copied >= count) { + break + } + } + return out + } + + /** Consume and return the first `count` buffered bytes (single copy). */ + private takeBytes(count: number): Buffer { + const first = this.chunks[0] + if (first.length === count) { + this.chunks.shift() + this.bufferedLength -= count + return first + } + if (first.length > count) { + this.chunks[0] = first.subarray(count) + this.bufferedLength -= count + return first.subarray(0, count) + } + const out = Buffer.allocUnsafe(count) + let copied = 0 + while (copied < count) { + const part = this.chunks[0] + const take = Math.min(part.length, count - copied) + part.copy(out, copied, 0, take) + copied += take + if (take === part.length) { + this.chunks.shift() + } else { + this.chunks[0] = part.subarray(take) + } + } + this.bufferedLength -= count + return out + } + + /** Consume the first `count` buffered bytes without assembling them. */ + private discardBytes(count: number): void { + let remaining = count + while (remaining > 0) { + const part = this.chunks[0] + if (part.length <= remaining) { + this.chunks.shift() + remaining -= part.length + } else { + this.chunks[0] = part.subarray(remaining) + remaining = 0 + } + } + this.bufferedLength -= count + } } export function parseJsonRpcMessage(payload: Buffer): JsonRpcMessage { diff --git a/src/relay/relay.ts b/src/relay/relay.ts index c930a244b..992252279 100644 --- a/src/relay/relay.ts +++ b/src/relay/relay.ts @@ -373,15 +373,44 @@ async function main(): Promise { // write to a dead pipe, silently failing or throwing EPIPE. When a // socket client reconnects, setWrite swaps the callback to the socket. let stdoutAlive = true - const dispatcher = new RelayDispatcher((data) => { - if (stdoutAlive) { + // Why: one-shot waiters parked by the dispatcher's bulk lane when stdout + // reports saturation (write() === false). Flushed on 'drain' and on every + // stdout-death path so a stalled file-stream pump never outlives the pipe + // it was waiting on. + const stdoutDrainWaiters = new Set<() => void>() + const flushStdoutDrainWaiters = (): void => { + for (const cb of Array.from(stdoutDrainWaiters)) { + stdoutDrainWaiters.delete(cb) + cb() + } + } + process.stdout.on('drain', flushStdoutDrainWaiters) + const dispatcher = new RelayDispatcher( + (data) => { + if (!stdoutAlive) { + return + } try { - process.stdout.write(data) + // Why: surface Node's backpressure signal to the dispatcher so bulk + // frames (fs.streamChunk) wait for drain instead of queueing megabytes + // ahead of interactive pty.data frames on the SSH channel. + return process.stdout.write(data) } catch { stdoutAlive = false + flushStdoutDrainWaiters() + return undefined + } + }, + { + waitWriteDrain: (cb) => { + if (!stdoutAlive) { + cb() + return + } + stdoutDrainWaiters.add(cb) } } - }) + ) const context = new RelayContext() @@ -680,11 +709,35 @@ async function main(): Promise { ) cancelGrace('socket client accepted') - const clientId = dispatcher.attachClient((data) => { - if (!sock.destroyed) { - sock.write(data) + // Why: same backpressure surface as the stdout sink — bulk frames wait + // for the socket to drain so they cannot bury interactive PTY frames. + const sockDrainWaiters = new Set<() => void>() + const flushSockDrainWaiters = (): void => { + for (const cb of Array.from(sockDrainWaiters)) { + sockDrainWaiters.delete(cb) + cb() } - }) + } + sock.on('drain', flushSockDrainWaiters) + sock.on('close', flushSockDrainWaiters) + sock.on('error', flushSockDrainWaiters) + const clientId = dispatcher.attachClient( + (data) => { + if (!sock.destroyed) { + return sock.write(data) + } + return undefined + }, + { + waitWriteDrain: (cb) => { + if (sock.destroyed) { + cb() + return + } + sockDrainWaiters.add(cb) + } + } + ) socketClients.set(sock, clientId) // Why: bytes that arrived in the same TCP send as the handshake frame @@ -889,6 +942,7 @@ async function main(): Promise { // before the grace period starts. process.stdout.on('error', () => { stdoutAlive = false + flushStdoutDrainWaiters() dispatcher.invalidateClient() }) @@ -934,6 +988,7 @@ async function main(): Promise { // dead so the primary client write callback becomes a no-op while // socket clients, if any, keep their own live transports. stdoutAlive = false + flushStdoutDrainWaiters() dispatcher.invalidateClient() if (socketClients.size === 0) { startGrace('stdin ended') @@ -942,6 +997,7 @@ async function main(): Promise { process.stdin.on('error', () => { stdoutAlive = false + flushStdoutDrainWaiters() dispatcher.invalidateClient() if (socketClients.size === 0) { startGrace('stdin error') diff --git a/tests/e2e/ssh-docker-relay-perf.spec.ts b/tests/e2e/ssh-docker-relay-perf.spec.ts index ddcbee655..c8e3beab6 100644 --- a/tests/e2e/ssh-docker-relay-perf.spec.ts +++ b/tests/e2e/ssh-docker-relay-perf.spec.ts @@ -207,6 +207,104 @@ test.describe('Docker SSH relay perf', () => { } }) + test('keeps remote typing responsive while relay file streams and git churn are active', async ({ + orcaPage + }, testInfo) => { + test.slow() + let target: DockerSshRelayTarget | null = null + try { + target = startDockerSshRelayTarget(testInfo) + await waitForSessionReady(orcaPage) + await waitForActiveWorktree(orcaPage) + const remote = await connectDockerRemote(orcaPage, target) + await ensureTerminalVisible(orcaPage, 45_000) + await waitForActiveTerminalManager(orcaPage, 60_000) + const ptyId = await waitForActivePanePtyId(orcaPage, 60_000) + + const runId = String(Date.now()) + // Large remote binaries: each read streams ~8MB of fs.streamChunk frames + // over the same SSH channel that carries the pty echo. + const loadFiles = [ + `${DOCKER_SSH_RELAY_REMOTE_REPO_PATH}/stream-load-a.png`, + `${DOCKER_SSH_RELAY_REMOTE_REPO_PATH}/stream-load-b.png` + ] + await execInTerminal( + orcaPage, + ptyId, + `dd if=/dev/urandom of=${shellQuote(loadFiles[0])} bs=1M count=8 status=none && ` + + `dd if=/dev/urandom of=${shellQuote(loadFiles[1])} bs=1M count=8 status=none && ` + + `echo LOAD_FILES_READY_${runId}` + ) + await waitForTerminalOutput(orcaPage, `LOAD_FILES_READY_${runId}`, 60_000, 80_000) + + await execInTerminal(orcaPage, ptyId, `node -e ${shellQuote(remoteTypingLoadScript(runId))}`) + await waitForTerminalOutput(orcaPage, `REMOTE_TUI_READY_${runId}`, 30_000, 80_000) + + // Background relay pressure: continuous large file reads plus git status + // refreshes, mirroring file preview + source-control churn while typing. + await orcaPage.evaluate( + ({ targetId, files, repoPath }) => { + const state = { stopped: false, reads: 0, errors: [] as string[] } + ;(window as unknown as { __sshRelayLoad: typeof state }).__sshRelayLoad = state + const loop = async (run: () => Promise): Promise => { + while (!state.stopped) { + try { + await run() + state.reads += 1 + } catch (err) { + state.errors.push(String(err)) + await new Promise((r) => setTimeout(r, 100)) + } + } + } + for (const filePath of files) { + void loop(() => window.api.fs.readFile({ filePath, connectionId: targetId })) + } + void loop(() => window.api.git.status({ worktreePath: repoPath, connectionId: targetId })) + }, + { + targetId: remote.targetId, + files: loadFiles, + repoPath: DOCKER_SSH_RELAY_REMOTE_REPO_PATH + } + ) + // Let the bulk load ramp before measuring. + await orcaPage.waitForTimeout(1_000) + + const measurement = await measureRemoteTyping(orcaPage, ptyId, runId) + const load = await orcaPage.evaluate(() => { + const state = ( + window as unknown as { + __sshRelayLoad: { stopped: boolean; reads: number; errors: string[] } + } + ).__sshRelayLoad + state.stopped = true + return { reads: state.reads, errors: state.errors.slice(0, 3) } + }) + + const summary = + `median=${measurement.medianLatencyMs.toFixed(1)}ms ` + + `worst=${measurement.worstLatencyMs.toFixed(1)}ms ` + + `bulkReads=${load.reads} ` + + `samples=${measurement.latencies.map((value) => value.toFixed(1)).join(',')}` + console.log(`[docker-ssh-relay-perf:busy] ${summary}`) + testInfo.annotations.push({ + type: 'docker-ssh-relay-typing-busy', + description: summary + }) + + // The load must actually have been streaming and error-free, otherwise + // the latency numbers prove nothing. + expect(load.errors).toEqual([]) + expect(load.reads).toBeGreaterThan(0) + expect(measurement.medianLatencyMs).toBeLessThan(MAX_MEDIAN_KEY_LATENCY_MS) + expect(measurement.worstLatencyMs).toBeLessThan(MAX_WORST_KEY_LATENCY_MS) + await stopRemoteLoad(orcaPage, ptyId) + } finally { + cleanupDockerSshRelayTarget(target) + } + }) + test('keeps an SSH workspace terminal usable after disconnect and reconnect', async ({ orcaPage }, testInfo) => {