From d804263bf07f8f9c0c7d1385f6ee5f011893577f Mon Sep 17 00:00:00 2001 From: OrcaWin <293788423+OrcaWin@users.noreply.github.com> Date: Mon, 27 Jul 2026 19:41:15 -0700 Subject: [PATCH] fix(relay): bound PTY output publication --- src/relay/dispatcher-client-writer.test.ts | 318 +++++++++++ src/relay/dispatcher-client-writer.ts | 319 +++++++++++ src/relay/dispatcher-writer-admission.ts | 155 ++++++ src/relay/dispatcher-writer-lane-scheduler.ts | 74 +++ src/relay/dispatcher-writer-sink.ts | 64 +++ src/relay/dispatcher.test.ts | 58 +- src/relay/dispatcher.ts | 498 ++++++++++++++---- src/relay/fs-handler-file-read.ts | 9 +- src/relay/git-response-stream.ts | 15 +- .../legacy-relay-publication-ledger.test.ts | 48 ++ src/relay/legacy-relay-publication-ledger.ts | 104 ++++ ...-handler-output-drain-differential.test.ts | 4 +- src/relay/pty-handler.test.ts | 99 +++- src/relay/pty-handler.ts | 277 ++++++++-- src/relay/relay-watcher-event-emitter.test.ts | 53 ++ src/relay/relay-watcher-event-emitter.ts | 1 + src/relay/relay.ts | 144 ++++- src/relay/wsl-agent-hook-relay.ts | 26 +- 18 files changed, 2061 insertions(+), 205 deletions(-) create mode 100644 src/relay/dispatcher-client-writer.test.ts create mode 100644 src/relay/dispatcher-client-writer.ts create mode 100644 src/relay/dispatcher-writer-admission.ts create mode 100644 src/relay/dispatcher-writer-lane-scheduler.ts create mode 100644 src/relay/dispatcher-writer-sink.ts create mode 100644 src/relay/legacy-relay-publication-ledger.test.ts create mode 100644 src/relay/legacy-relay-publication-ledger.ts create mode 100644 src/relay/relay-watcher-event-emitter.test.ts diff --git a/src/relay/dispatcher-client-writer.test.ts b/src/relay/dispatcher-client-writer.test.ts new file mode 100644 index 00000000000..4fcf758492c --- /dev/null +++ b/src/relay/dispatcher-client-writer.test.ts @@ -0,0 +1,318 @@ +import { describe, expect, it, vi } from 'vitest' +import { + DispatcherClientWriter, + type DispatcherWriterLane, + type SinkWriteSettlement +} from './dispatcher-client-writer' + +type AcceptedWrite = { + data: string + settle: (result: SinkWriteSettlement) => void +} + +class FakeSink { + readonly accepted: AcceptedWrite[] = [] + readonly drainWaiters: (() => void)[] = [] + writableLength = 0 + highWaterMark = 4096 + saturateNext = false + closed = false + + write = (data: Buffer, settle: (result: SinkWriteSettlement) => void): boolean | void => { + this.accepted.push({ data: data.toString(), settle }) + this.writableLength += data.length + if (this.saturateNext) { + this.saturateNext = false + return false + } + return true + } + + drain(): void { + this.writableLength = 0 + for (const waiter of this.drainWaiters.splice(0)) { + waiter() + } + } +} + +function createWriter( + sink: FakeSink, + onClosed = vi.fn(), + producerQueueMaxBytes = 1024 +): DispatcherClientWriter { + return new DispatcherClientWriter( + sink.write, + { + supportsWriteCallback: true, + writableLength: () => sink.writableLength, + writableHighWaterMark: () => sink.highWaterMark, + waitWriteDrain: (callback) => { + sink.drainWaiters.push(callback) + return () => { + const index = sink.drainWaiters.indexOf(callback) + if (index >= 0) { + sink.drainWaiters.splice(index, 1) + } + } + }, + close: () => { + sink.closed = true + } + }, + onClosed, + producerQueueMaxBytes + ) +} + +function enqueue( + writer: DispatcherClientWriter, + lane: DispatcherWriterLane, + value: string, + settled = vi.fn() +): ReturnType { + expect(writer.enqueue(lane, () => Buffer.from(value), Buffer.byteLength(value), settled)).toBe( + true + ) + return settled +} + +describe('DispatcherClientWriter', () => { + it('accepts write(false) once and stops ordinary writes until drain', () => { + const sink = new FakeSink() + const writer = createWriter(sink) + sink.saturateNext = true + + const firstSettled = enqueue(writer, 'ordinary', 'first') + enqueue(writer, 'ordinary', 'second') + + expect(sink.accepted.map((write) => write.data)).toEqual(['first']) + sink.accepted[0].settle({ ok: true }) + expect(firstSettled).toHaveBeenCalledWith({ ok: true }) + expect(sink.accepted.map((write) => write.data)).toEqual(['first']) + + sink.drain() + expect(sink.accepted.map((write) => write.data)).toEqual(['first', 'second']) + }) + + it('closes on a callback error without settling queued work twice', () => { + const sink = new FakeSink() + const onClosed = vi.fn() + const writer = createWriter(sink, onClosed) + sink.saturateNext = true + + const firstSettled = enqueue(writer, 'ordinary', 'first') + const secondSettled = enqueue(writer, 'ordinary', 'second') + const error = new Error('async write failed') + sink.accepted[0].settle({ ok: false, error }) + sink.accepted[0].settle({ ok: true }) + + expect(firstSettled).toHaveBeenCalledTimes(1) + expect(firstSettled).toHaveBeenCalledWith({ ok: false, error }) + expect(secondSettled).toHaveBeenCalledTimes(1) + expect(onClosed).toHaveBeenCalledWith(error) + expect(sink.closed).toBe(true) + }) + + it('preserves FIFO within lanes and rechecks control before producer work', () => { + const sink = new FakeSink() + const writer = createWriter(sink) + sink.saturateNext = true + + enqueue(writer, 'ordinary', 'blocker') + enqueue(writer, 'bulk', 'bulk-1') + enqueue(writer, 'ordinary', 'pty-1') + enqueue(writer, 'control', 'control-1') + enqueue(writer, 'ordinary', 'pty-2') + enqueue(writer, 'control', 'control-2') + expect(sink.accepted.map((write) => write.data)).toEqual(['blocker']) + + sink.drain() + + expect(sink.accepted.map((write) => write.data)).toEqual([ + 'blocker', + 'control-1', + 'control-2', + 'pty-1', + 'pty-2', + 'bulk-1' + ]) + }) + + it('allows one coalesced liveness bypass during a saturated epoch', () => { + const sink = new FakeSink() + const writer = createWriter(sink) + sink.saturateNext = true + + enqueue(writer, 'ordinary', 'ordinary') + const first = enqueue(writer, 'liveness', 'keepalive-1') + const replaced = enqueue(writer, 'liveness', 'keepalive-2') + const latest = enqueue(writer, 'liveness', 'keepalive-3') + + expect(sink.accepted.map((write) => write.data)).toEqual(['ordinary', 'keepalive-1']) + expect(first).not.toHaveBeenCalled() + expect(replaced).toHaveBeenCalledOnce() + expect(replaced).toHaveBeenCalledWith({ ok: true }) + expect(latest).not.toHaveBeenCalled() + + sink.drain() + expect(sink.accepted.map((write) => write.data)).toEqual([ + 'ordinary', + 'keepalive-1', + 'keepalive-3' + ]) + }) + + it('gives queued bulk a turn after four interactive or ordinary writes', () => { + const sink = new FakeSink() + const writer = createWriter(sink) + sink.saturateNext = true + + enqueue(writer, 'ordinary', 'blocker') + enqueue(writer, 'bulk', 'bulk') + enqueue(writer, 'interactive', 'interactive-1') + enqueue(writer, 'ordinary', 'ordinary-1') + enqueue(writer, 'interactive', 'interactive-2') + enqueue(writer, 'ordinary', 'ordinary-2') + enqueue(writer, 'interactive', 'interactive-3') + + sink.drain() + + expect(sink.accepted.map((write) => write.data)).toEqual([ + 'blocker', + 'interactive-1', + 'interactive-2', + 'interactive-3', + 'bulk', + 'ordinary-1', + 'ordinary-2' + ]) + }) + + it('gives ordinary and bulk lanes bounded turns under sustained interactive traffic', () => { + const sink = new FakeSink() + const writer = createWriter(sink) + sink.saturateNext = true + + enqueue(writer, 'ordinary', 'blocker') + enqueue(writer, 'ordinary', 'ordinary') + enqueue(writer, 'bulk', 'bulk') + for (let index = 1; index <= 12; index++) { + enqueue(writer, 'interactive', `interactive-${index}`) + } + + sink.drain() + + expect(sink.accepted.findIndex((write) => write.data === 'ordinary')).toBeLessThanOrEqual(6) + expect(sink.accepted.findIndex((write) => write.data === 'bulk')).toBeLessThanOrEqual(5) + }) + + it('settles callback-less saturated writes when drain fires during registration', () => { + const settled = vi.fn() + const writer = new DispatcherClientWriter( + () => false, + { + waitWriteDrain: (callback) => callback(), + writableHighWaterMark: () => 4096 + }, + vi.fn(), + 64 + ) + + expect(writer.enqueue('ordinary', () => Buffer.from('accepted'), 8, settled)).toBe(true) + expect(settled).toHaveBeenCalledExactlyOnceWith({ ok: true }) + expect(writer.retainedProducerBytes).toBe(0) + }) + + it('bounds queued producer bytes after write(false) settles before drain', () => { + const sink = new FakeSink() + sink.saturateNext = true + const writer = createWriter(sink, vi.fn(), 64) + + enqueue(writer, 'ordinary', 'accepted-before-drain') + sink.accepted[0].settle({ ok: true }) + + for (let index = 0; index < 4; index++) { + expect(writer.enqueue('ordinary', () => Buffer.alloc(16), 16)).toBe(true) + } + expect(writer.enqueue('ordinary', () => Buffer.alloc(1), 1)).toBe(false) + expect(writer.retainedProducerBytes).toBe(64) + expect(sink.accepted).toHaveLength(1) + expect(sink.writableLength + writer.retainedProducerBytes).toBeLessThanOrEqual( + sink.highWaterMark + 64 + ) + + sink.drain() + expect(sink.accepted).toHaveLength(5) + }) + + it('cleans queued and in-flight callbacks once on close', () => { + const sink = new FakeSink() + const writer = createWriter(sink) + sink.saturateNext = true + const first = enqueue(writer, 'ordinary', 'first') + const second = enqueue(writer, 'bulk', 'second') + + writer.close() + sink.drain() + sink.accepted[0].settle({ ok: true }) + + expect(first).toHaveBeenCalledTimes(1) + expect(second).toHaveBeenCalledTimes(1) + expect(writer.retainedProducerBytes).toBe(0) + expect(sink.drainWaiters).toHaveLength(0) + }) + + it('settles the idle fence after accepted writes finish or close cancels them', async () => { + const sink = new FakeSink() + const writer = createWriter(sink) + enqueue(writer, 'ordinary', 'first') + let idle = false + const waiting = writer.waitForIdle().then(() => { + idle = true + }) + + await Promise.resolve() + expect(idle).toBe(false) + sink.accepted[0].settle({ ok: true }) + await waiting + expect(idle).toBe(true) + + enqueue(writer, 'ordinary', 'second') + const canceled = writer.waitForIdle() + writer.close() + await canceled + }) + + it('bounds producer retention and preserves capacity for control', () => { + const sink = new FakeSink() + sink.highWaterMark = 4096 + sink.saturateNext = true + const writer = createWriter(sink, vi.fn(), 2048) + + expect(writer.producerFrameCapacity).toBe(3072) + expect(writer.enqueue('bulk', () => Buffer.alloc(2048), 2048)).toBe(true) + expect(writer.enqueue('bulk', () => Buffer.alloc(1), 1)).toBe(false) + expect(writer.retainedProducerBytes).toBe(2048) + expect(writer.enqueue('control', () => Buffer.from('cancel'), 6)).toBe(true) + + sink.drain() + expect(sink.accepted[1].data).toBe('cancel') + }) + + it('gives independent writers progress when one client stays saturated', () => { + const slowSink = new FakeSink() + const fastSink = new FakeSink() + const slow = createWriter(slowSink) + const fast = createWriter(fastSink) + slowSink.saturateNext = true + + enqueue(slow, 'ordinary', 'slow-1') + enqueue(slow, 'ordinary', 'slow-2') + enqueue(fast, 'ordinary', 'fast-1') + enqueue(fast, 'ordinary', 'fast-2') + + expect(slowSink.accepted.map((write) => write.data)).toEqual(['slow-1']) + expect(fastSink.accepted.map((write) => write.data)).toEqual(['fast-1', 'fast-2']) + }) +}) diff --git a/src/relay/dispatcher-client-writer.ts b/src/relay/dispatcher-client-writer.ts new file mode 100644 index 00000000000..7a27e670ceb --- /dev/null +++ b/src/relay/dispatcher-client-writer.ts @@ -0,0 +1,319 @@ +import { + DEFAULT_PRODUCER_QUEUE_MAX_BYTES, + DispatcherWriterAdmission, + onceDispatcherWriterSettlement, + type DispatcherWriterEntry, + type DispatcherWriterLane +} from './dispatcher-writer-admission' +import { DispatcherWriterLaneScheduler } from './dispatcher-writer-lane-scheduler' +import { + DispatcherWriterSink, + type RelayClientSinkOptions, + type RelayClientWrite, + type SinkWriteSettlement +} from './dispatcher-writer-sink' + +export { + DEFAULT_PRODUCER_QUEUE_MAX_BYTES, + DISPATCHER_CONTROL_QUEUE_MAX_BYTES, + relayWriterControlReserve +} from './dispatcher-writer-admission' +export type { + RelayClientSinkOptions, + RelayClientWrite, + SinkWriteSettlement +} from './dispatcher-writer-sink' +export type { DispatcherWriterLane } from './dispatcher-writer-admission' + +export class DispatcherClientWriter { + private readonly admission: DispatcherWriterAdmission + private readonly sink: DispatcherWriterSink + private readonly laneScheduler = new DispatcherWriterLaneScheduler() + private readonly inFlight = new Set() + private readonly settleOnDrain = new Set() + private readonly capacityListeners = new Set<() => void>() + private readonly idleWaiters = new Set<() => void>() + private saturated = false + private drainArmed = false + private removeDrainListener: (() => void) | null = null + private livenessBypassOutstanding = false + private pumping = false + private closed = false + private closeNotified = false + + constructor( + write: RelayClientWrite, + sinkOptions: RelayClientSinkOptions = {}, + private readonly onClosed: (error: Error) => void = () => {}, + producerQueueMaxBytes = DEFAULT_PRODUCER_QUEUE_MAX_BYTES + ) { + this.admission = new DispatcherWriterAdmission(producerQueueMaxBytes) + this.sink = new DispatcherWriterSink(write, sinkOptions) + } + + get retainedProducerBytes(): number { + return this.admission.retainedProducerBytes + } + + get producerFrameCapacity(): number { + return this.sink.producerFrameCapacity + } + + canEnqueueProducer(bytes: number): boolean { + return !this.closed && this.admission.canAdmitProducer(bytes, this.producerFrameCapacity) + } + + onCapacity(listener: () => void): () => void { + this.capacityListeners.add(listener) + return () => this.capacityListeners.delete(listener) + } + + waitForIdle(): Promise { + if (this.isIdle()) { + return Promise.resolve() + } + return new Promise((resolve) => this.idleWaiters.add(resolve)) + } + + enqueue( + lane: DispatcherWriterLane, + encode: () => Buffer, + estimatedBytes: number, + onSettled: (result: SinkWriteSettlement) => void = () => {} + ): boolean { + if (this.closed) { + onSettled({ ok: false, error: new Error('Relay writer is closed') }) + return false + } + const entry: DispatcherWriterEntry = { + lane, + encode, + estimatedBytes, + onSettled: onceDispatcherWriterSettlement(onSettled), + settled: false + } + const admission = this.admission.admit(entry, this.producerFrameCapacity) + if (!admission.accepted) { + if (admission.error) { + entry.onSettled({ ok: false, error: admission.error }) + this.close(admission.error) + } else if (lane === 'liveness') { + entry.onSettled({ ok: true }) + } + return false + } + if (admission.replaced) { + admission.replaced.settled = true + admission.replaced.onSettled({ ok: true }) + } + this.pump() + return true + } + + close(error = new Error('Relay writer closed')): void { + if (this.closed) { + return + } + this.closed = true + this.saturated = false + this.drainArmed = false + this.clearDrainListener() + for (const entry of this.admission.takeQueued()) { + this.releaseEntry(entry, { ok: false, error }) + } + for (const entry of Array.from(this.inFlight)) { + this.releaseEntry(entry, { ok: false, error }) + } + this.settleOnDrain.clear() + this.capacityListeners.clear() + this.notifyIdle() + try { + this.sink.close() + } catch { + // The writer state is already closed. + } + if (!this.closeNotified) { + this.closeNotified = true + this.onClosed(error) + } + } + + private pump(): void { + if (this.pumping || this.closed) { + return + } + this.pumping = true + try { + while (!this.closed) { + const entry = this.selectNext() + if (!entry) { + break + } + this.writeEntry(entry) + if (this.saturated && entry.lane !== 'liveness') { + break + } + } + } finally { + this.pumping = false + } + } + + private selectNext(): DispatcherWriterEntry | undefined { + if (this.saturated) { + if (this.livenessBypassOutstanding) { + return undefined + } + const bypass = this.admission.shift('liveness') + if (bypass) { + this.livenessBypassOutstanding = true + } + return bypass + } + return this.laneScheduler.select( + this.admission, + (entry) => this.canWriteControl(entry), + (entry) => this.canWriteProducer(entry) + ) + } + + private canWriteControl(entry: DispatcherWriterEntry): boolean { + const highWaterMark = this.sink.highWaterMark + if (!Number.isFinite(highWaterMark)) { + return true + } + const writableLength = this.sink.writableLength + return writableLength + entry.estimatedBytes <= highWaterMark || writableLength === 0 + } + + private canWriteProducer(entry: DispatcherWriterEntry): boolean { + const highWaterMark = this.sink.highWaterMark + if (!Number.isFinite(highWaterMark)) { + return true + } + return this.sink.writableLength + entry.estimatedBytes <= this.sink.producerFrameCapacity + } + + private writeEntry(entry: DispatcherWriterEntry): void { + this.laneScheduler.recordWrite(entry.lane) + this.inFlight.add(entry) + let callbackResult: SinkWriteSettlement | undefined + let writeReturned = false + const onWriteSettled = (result: SinkWriteSettlement): void => { + if (!writeReturned) { + callbackResult = result + return + } + this.handleWriteSettlement(entry, result) + } + try { + const accepted = this.sink.write(entry.encode(), onWriteSettled) + writeReturned = true + if (accepted === false) { + this.saturated = true + if (!this.sink.supportsWriteCallback) { + this.settleOnDrain.add(entry) + } + this.armDrain() + } else if (!this.sink.supportsWriteCallback) { + this.handleWriteSettlement(entry, { ok: true }) + } + if (callbackResult) { + this.handleWriteSettlement(entry, callbackResult) + } + } catch (error) { + writeReturned = true + this.handleWriteSettlement(entry, { + ok: false, + error: error instanceof Error ? error : new Error(String(error)) + }) + } + } + + private handleWriteSettlement(entry: DispatcherWriterEntry, result: SinkWriteSettlement): void { + if (!result.ok) { + this.releaseEntry(entry, result) + this.close(result.error) + return + } + this.releaseEntry(entry, result) + if (entry.lane === 'liveness') { + this.livenessBypassOutstanding = false + } + this.notifyCapacity() + this.pump() + } + + private armDrain(): void { + if (this.drainArmed) { + return + } + this.drainArmed = true + try { + const registration = this.sink.registerDrain(() => this.handleDrain()) + if (!registration.registered) { + this.drainArmed = false + } else if (this.drainArmed) { + this.removeDrainListener = registration.remove + } else { + registration.remove() + } + } catch (error) { + this.close(error instanceof Error ? error : new Error(String(error))) + } + } + + private handleDrain(): void { + if (this.closed || !this.drainArmed) { + return + } + this.drainArmed = false + this.clearDrainListener() + this.saturated = false + this.livenessBypassOutstanding = false + for (const entry of Array.from(this.settleOnDrain)) { + this.settleOnDrain.delete(entry) + this.releaseEntry(entry, { ok: true }) + } + this.notifyCapacity() + this.pump() + } + + private releaseEntry(entry: DispatcherWriterEntry, result: SinkWriteSettlement): void { + if (entry.settled) { + return + } + entry.settled = true + this.inFlight.delete(entry) + this.settleOnDrain.delete(entry) + this.admission.release(entry) + entry.onSettled(result) + this.notifyIdle() + } + + private notifyCapacity(): void { + for (const listener of this.capacityListeners) { + listener() + } + } + + private clearDrainListener(): void { + const remove = this.removeDrainListener + this.removeDrainListener = null + remove?.() + } + + private isIdle(): boolean { + return this.inFlight.size === 0 && this.admission.queuedEntries === 0 + } + + private notifyIdle(): void { + if (!this.isIdle()) { + return + } + for (const resolve of Array.from(this.idleWaiters)) { + this.idleWaiters.delete(resolve) + resolve() + } + } +} diff --git a/src/relay/dispatcher-writer-admission.ts b/src/relay/dispatcher-writer-admission.ts new file mode 100644 index 00000000000..32a3c9cee75 --- /dev/null +++ b/src/relay/dispatcher-writer-admission.ts @@ -0,0 +1,155 @@ +export type DispatcherWriterLane = + | 'liveness' + | 'control' + | 'legacy-response' + | 'interactive' + | 'ordinary' + | 'bulk' + +export type DispatcherWriterSettlement = { ok: true } | { ok: false; error: Error } + +export type DispatcherWriterEntry = { + lane: DispatcherWriterLane + encode: () => Buffer + estimatedBytes: number + onSettled: (result: DispatcherWriterSettlement) => void + settled: boolean +} + +type AdmissionResult = + | { accepted: true; replaced?: DispatcherWriterEntry } + | { accepted: false; error?: Error } + +const CONTROL_QUEUE_MAX_FRAMES = 256 +export const DISPATCHER_CONTROL_QUEUE_MAX_BYTES = 1024 * 1024 +const LIVENESS_QUEUE_MAX_FRAMES = 2 + +export const DEFAULT_PRODUCER_QUEUE_MAX_BYTES = 2 * 1024 * 1024 + +export function onceDispatcherWriterSettlement( + callback: (result: DispatcherWriterSettlement) => void +): (result: DispatcherWriterSettlement) => void { + let settled = false + return (result) => { + if (settled) { + return + } + settled = true + callback(result) + } +} + +export function relayWriterControlReserve(highWaterMark: number): number { + if (!Number.isFinite(highWaterMark)) { + return 0 + } + return Math.min(64 * 1024, Math.max(1024, Math.floor(highWaterMark / 4))) +} + +export class DispatcherWriterAdmission { + private readonly queues: Record = { + liveness: [], + control: [], + 'legacy-response': [], + interactive: [], + ordinary: [], + bulk: [] + } + private controlBytes = 0 + private controlFrames = 0 + private producerBytes = 0 + private livenessOutstanding = 0 + + constructor(private readonly producerQueueMaxBytes: number) {} + + get retainedProducerBytes(): number { + return this.producerBytes + } + + get queuedEntries(): number { + return Object.values(this.queues).reduce((total, queue) => total + queue.length, 0) + } + + canAdmitProducer(bytes: number, producerFrameCapacity: number): boolean { + return ( + bytes <= producerFrameCapacity && this.producerBytes + bytes <= this.producerQueueMaxBytes + ) + } + + admit(entry: DispatcherWriterEntry, producerFrameCapacity: number): AdmissionResult { + if (entry.lane === 'liveness') { + return this.admitLiveness(entry) + } + if (entry.lane === 'control') { + return this.admitControl(entry) + } + if (entry.lane === 'legacy-response') { + if (this.producerBytes + entry.estimatedBytes > this.producerQueueMaxBytes) { + return { accepted: false } + } + this.producerBytes += entry.estimatedBytes + this.queues[entry.lane].push(entry) + return { accepted: true } + } + if (!this.canAdmitProducer(entry.estimatedBytes, producerFrameCapacity)) { + return { accepted: false } + } + this.producerBytes += entry.estimatedBytes + this.queues[entry.lane].push(entry) + return { accepted: true } + } + + peek(lane: DispatcherWriterLane): DispatcherWriterEntry | undefined { + return this.queues[lane][0] + } + + shift(lane: DispatcherWriterLane): DispatcherWriterEntry | undefined { + return this.queues[lane].shift() + } + + takeQueued(): DispatcherWriterEntry[] { + return Object.values(this.queues).flatMap((queue) => queue.splice(0)) + } + + release(entry: DispatcherWriterEntry): void { + if (entry.lane === 'control') { + this.controlFrames-- + this.controlBytes -= entry.estimatedBytes + } else if (entry.lane === 'liveness') { + this.livenessOutstanding-- + } else { + this.producerBytes -= entry.estimatedBytes + } + } + + private admitLiveness(entry: DispatcherWriterEntry): AdmissionResult { + const queue = this.queues.liveness + if (this.livenessOutstanding >= LIVENESS_QUEUE_MAX_FRAMES) { + const replaced = queue.pop() + if (!replaced) { + return { accepted: false } + } + queue.push(entry) + return { accepted: true, replaced } + } + this.livenessOutstanding++ + queue.push(entry) + return { accepted: true } + } + + private admitControl(entry: DispatcherWriterEntry): AdmissionResult { + if ( + this.controlFrames >= CONTROL_QUEUE_MAX_FRAMES || + this.controlBytes + entry.estimatedBytes > DISPATCHER_CONTROL_QUEUE_MAX_BYTES + ) { + return { + accepted: false, + error: new Error('Relay control queue exceeded its bounded capacity') + } + } + this.controlFrames++ + this.controlBytes += entry.estimatedBytes + this.queues.control.push(entry) + return { accepted: true } + } +} diff --git a/src/relay/dispatcher-writer-lane-scheduler.ts b/src/relay/dispatcher-writer-lane-scheduler.ts new file mode 100644 index 00000000000..966875dd7e9 --- /dev/null +++ b/src/relay/dispatcher-writer-lane-scheduler.ts @@ -0,0 +1,74 @@ +import type { + DispatcherWriterAdmission, + DispatcherWriterEntry, + DispatcherWriterLane +} from './dispatcher-writer-admission' + +const PRODUCER_WRITES_BEFORE_BULK = 4 +const INTERACTIVE_WRITES_BEFORE_ORDINARY = 4 + +export class DispatcherWriterLaneScheduler { + private producerWritesSinceBulk = 0 + private interactiveWritesSinceOrdinary = 0 + + select( + admission: DispatcherWriterAdmission, + canWriteControl: (entry: DispatcherWriterEntry) => boolean, + canWriteProducer: (entry: DispatcherWriterEntry) => boolean + ): DispatcherWriterEntry | undefined { + const liveness = admission.shift('liveness') + if (liveness) { + return liveness + } + const control = admission.peek('control') + if (control && canWriteControl(control)) { + return admission.shift('control') + } + const legacyResponse = admission.peek('legacy-response') + if (legacyResponse && canWriteControl(legacyResponse)) { + return admission.shift('legacy-response') + } + const bulk = admission.peek('bulk') + if ( + bulk && + this.producerWritesSinceBulk >= PRODUCER_WRITES_BEFORE_BULK && + canWriteProducer(bulk) + ) { + return admission.shift('bulk') + } + const ordinary = admission.peek('ordinary') + if ( + ordinary && + this.interactiveWritesSinceOrdinary >= INTERACTIVE_WRITES_BEFORE_ORDINARY && + canWriteProducer(ordinary) + ) { + return admission.shift('ordinary') + } + for (const lane of ['interactive', 'ordinary', 'bulk'] as const) { + const candidate = admission.peek(lane) + if (candidate && canWriteProducer(candidate)) { + return admission.shift(lane) + } + } + return undefined + } + + recordWrite(lane: DispatcherWriterLane): void { + if (lane === 'bulk') { + this.producerWritesSinceBulk = 0 + } else if (lane === 'interactive' || lane === 'ordinary') { + this.producerWritesSinceBulk = Math.min( + PRODUCER_WRITES_BEFORE_BULK, + this.producerWritesSinceBulk + 1 + ) + } + if (lane === 'ordinary') { + this.interactiveWritesSinceOrdinary = 0 + } else if (lane === 'interactive') { + this.interactiveWritesSinceOrdinary = Math.min( + INTERACTIVE_WRITES_BEFORE_ORDINARY, + this.interactiveWritesSinceOrdinary + 1 + ) + } + } +} diff --git a/src/relay/dispatcher-writer-sink.ts b/src/relay/dispatcher-writer-sink.ts new file mode 100644 index 00000000000..7c06813155d --- /dev/null +++ b/src/relay/dispatcher-writer-sink.ts @@ -0,0 +1,64 @@ +import { + relayWriterControlReserve, + type DispatcherWriterSettlement +} from './dispatcher-writer-admission' + +export type SinkWriteSettlement = DispatcherWriterSettlement + +export type RelayClientWrite = ( + data: Buffer, + onSettled: (result: SinkWriteSettlement) => void +) => boolean | void + +export type RelayClientSinkOptions = { + waitWriteDrain?: (callback: () => void) => void | (() => void) + writableLength?: () => number + writableHighWaterMark?: () => number + supportsWriteCallback?: boolean + close?: () => void +} + +export class DispatcherWriterSink { + constructor( + readonly write: RelayClientWrite, + private readonly options: RelayClientSinkOptions + ) {} + + get supportsWriteCallback(): boolean { + return this.options.supportsWriteCallback === true + } + + get writableLength(): number { + const value = this.options.writableLength?.() + return typeof value === 'number' && Number.isFinite(value) && value >= 0 ? value : 0 + } + + get highWaterMark(): number { + const value = this.options.writableHighWaterMark?.() + return typeof value === 'number' && Number.isFinite(value) && value > 0 + ? value + : Number.POSITIVE_INFINITY + } + + get producerFrameCapacity(): number { + if (!Number.isFinite(this.highWaterMark)) { + return Number.MAX_SAFE_INTEGER + } + return Math.max(0, this.highWaterMark - relayWriterControlReserve(this.highWaterMark)) + } + + registerDrain(callback: () => void): { registered: boolean; remove: () => void } { + if (!this.options.waitWriteDrain) { + return { registered: false, remove: () => {} } + } + const remove = this.options.waitWriteDrain(callback) + return { + registered: true, + remove: typeof remove === 'function' ? remove : () => {} + } + } + + close(): void { + this.options.close?.() + } +} diff --git a/src/relay/dispatcher.test.ts b/src/relay/dispatcher.test.ts index 87f5cb69053..f76e4ae809f 100644 --- a/src/relay/dispatcher.test.ts +++ b/src/relay/dispatcher.test.ts @@ -480,6 +480,39 @@ describe('RelayDispatcher', () => { } }) + it('preserves broadcast order when synchronous settlement re-enters production', () => { + const primary: string[] = [] + const secondary: string[] = [] + const readData = (frame: Buffer): string => { + const message = JSON.parse(decodeFirstFrame(frame).payload.toString()) as JsonRpcNotification + return String(message.params?.data) + } + const orderedDispatcher = new RelayDispatcher((frame) => { + primary.push(readData(frame)) + return true + }) + orderedDispatcher.attachClient((frame) => { + secondary.push(readData(frame)) + return true + }) + let reentered = false + orderedDispatcher.onLegacyPtyCapacity(() => { + if (reentered) { + return + } + reentered = true + orderedDispatcher.tryNotifyPtyData({ id: 'pty-2', data: 'second' }) + }) + + try { + expect(orderedDispatcher.tryNotifyPtyData({ id: 'pty-1', data: 'first' })).toBe(true) + expect(primary).toEqual(['first', 'second']) + expect(secondary).toEqual(['first', 'second']) + } finally { + orderedDispatcher.dispose() + } + }) + describe('notifyBulk (bulk lane backpressure)', () => { it('resolves immediately when the sink accepts the frame', async () => { const frames: Buffer[] = [] @@ -506,7 +539,12 @@ describe('RelayDispatcher', () => { frames.push(Buffer.from(data)) return false }, - { waitWriteDrain: (cb) => drainWaiters.add(cb) } + { + waitWriteDrain: (callback) => { + drainWaiters.add(callback) + return () => drainWaiters.delete(callback) + } + } ) try { let firstSettled = false @@ -533,14 +571,20 @@ describe('RelayDispatcher', () => { } }) - it('interactive notify() frames are not gated behind a stalled bulk lane', async () => { + it('does not write an interactive frame around a saturated bulk write', async () => { const frames: Buffer[] = [] + const drainWaiters = new Set<() => void>() const bulkDispatcher = new RelayDispatcher( (data) => { frames.push(Buffer.from(data)) return false }, - { waitWriteDrain: () => {} } + { + waitWriteDrain: (callback) => { + drainWaiters.add(callback) + return () => drainWaiters.delete(callback) + } + } ) try { void bulkDispatcher.notifyBulk('bulk.event', { seq: 0 }) @@ -549,7 +593,13 @@ describe('RelayDispatcher', () => { expect(frames).toHaveLength(1) bulkDispatcher.notify('pty.data', { id: 'pty-1', data: 'x' }) - expect(frames).toHaveLength(2) + expect(frames).toHaveLength(1) + for (const callback of Array.from(drainWaiters)) { + drainWaiters.delete(callback) + callback() + } + await vi.advanceTimersByTimeAsync(0) + expect(frames.length).toBeGreaterThanOrEqual(2) const msg = JSON.parse( decodeFirstFrame(frames[1]).payload.toString() ) as JsonRpcNotification diff --git a/src/relay/dispatcher.ts b/src/relay/dispatcher.ts index 912b666910f..9d1752565f9 100644 --- a/src/relay/dispatcher.ts +++ b/src/relay/dispatcher.ts @@ -13,6 +13,25 @@ import { } from './protocol' import { ClientRequestAborts } from './client-request-aborts' import { MAX_TIMER_DELAY_MS, isSafeTimerDelayMs } from '../shared/timer-delay' +import { + DISPATCHER_CONTROL_QUEUE_MAX_BYTES, + DEFAULT_PRODUCER_QUEUE_MAX_BYTES, + DispatcherClientWriter, + type DispatcherWriterLane, + type RelayClientSinkOptions, + type RelayClientWrite, + type SinkWriteSettlement +} from './dispatcher-client-writer' +import { + LegacyRelayPublicationLedger, + type LegacyPublicationLease +} from './legacy-relay-publication-ledger' + +export type { + RelayClientSinkOptions, + RelayClientWrite, + SinkWriteSettlement +} from './dispatcher-client-writer' export type RequestContext = { clientId: number @@ -27,22 +46,10 @@ export type MethodHandler = ( export type NotificationHandler = (params: Record, context: RequestContext) => void -/** Sink write: `false` signals saturation (Node stream semantics); `void`/`true` mean accepted. */ -export type RelayClientWrite = (data: Buffer) => boolean | void - -export type RelayClientSinkOptions = { - /** One-shot: invoke `cb` when the sink can accept more data (drain) or is permanently dead, so waiters never hang. */ - waitWriteDrain?: (cb: () => void) => void -} - type RelayClient = { id: number decoder: FrameDecoder - write: RelayClientWrite - waitWriteDrain?: (cb: () => void) => void - /** Resolvers for bulk sends stalled on sink saturation; flushed so no pump hangs. */ - drainWaiters: Set<() => void> - /** Serializes bulk-lane sends so only one bulk frame is admitted past the sink high-water mark at a time. */ + writer: DispatcherClientWriter bulkChain: Promise nextOutgoingSeq: number highestReceivedSeq: number @@ -64,8 +71,13 @@ export class RelayDispatcher { private requestHandlers = new Map() private notificationHandlers = new Map() private readonly requestAborts = new ClientRequestAborts() + private readonly publicationLedger = new LegacyRelayPublicationLedger() private pendingRelayRequests = new Map() private clientDetachListeners = new Set<(clientId: number) => void>() + private legacyCapacityListeners = new Set<() => void>() + private publicationTransactionDepth = 0 + private deferredLegacyCapacity = false + private deferredForcedLegacyCapacity = false private keepaliveTimer: ReturnType | null = null private disposed = false private nextClientId = 1 @@ -81,21 +93,15 @@ export class RelayDispatcher { // Why: the new client's multiplexer restarts at seq=1, so reset seq/decoder state or acks stall and fire a false connection-dead signal. 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 old sink is gone; wake stalled bulk senders to re-evaluate against the new one. - this.flushDrainWaiters(this.primaryClient) + this.primaryClient.closed = true + this.primaryClient.writer.close(new Error('Relay primary sink replaced')) this.resetClient(this.primaryClient) + this.primaryClient.writer = this.createWriter(this.primaryClient, write, sinkOptions) } // Why: mark in-flight requests stale on disconnect so a late pty.spawn/fs.watch can't create unowned remote state. invalidateClient(): void { - this.requestAborts.abortClient(this.primaryClient.id) - this.primaryClient.generation++ - this.primaryClient.closed = true - this.flushDrainWaiters(this.primaryClient) - this.notifyClientDetached(this.primaryClient.id) + this.closeClient(this.primaryClient, new Error('Relay primary client invalidated'), false) } // Why: seq numbers and request ids are per SSH channel, so each attached client needs independent protocol state. @@ -110,12 +116,7 @@ export class RelayDispatcher { if (!client || client === this.primaryClient) { return } - this.requestAborts.abortClient(clientId) - client.generation++ - client.closed = true - this.flushDrainWaiters(client) - this.clients.delete(clientId) - this.notifyClientDetached(clientId) + this.closeClient(client, new Error('Relay client detached'), true) } feedClient(clientId: number, data: Buffer): void { @@ -139,6 +140,107 @@ export class RelayDispatcher { return () => this.clientDetachListeners.delete(listener) } + onLegacyPtyCapacity(listener: () => void): () => void { + this.legacyCapacityListeners.add(listener) + return () => this.legacyCapacityListeners.delete(listener) + } + + get legacyRetentionBelowLowWater(): boolean { + return this.publicationLedger.belowLowWater(this.activeClientKeys()) + } + + writePrimaryBytes(data: Buffer, lane: 'control' | 'ordinary' = 'control'): boolean { + if (this.disposed || this.primaryClient.closed) { + return false + } + return this.primaryClient.writer.enqueue(lane, () => data, data.length) + } + + maxLegacyPtyDataChars( + params: Record, + data: string, + limit = data.length + ): number { + const clients = this.activeClients() + if (clients.length === 0) { + return Math.min(data.length, limit) + } + let low = 0 + let high = Math.min(data.length, limit) + while (low < high) { + const mid = Math.ceil((low + high) / 2) + const msg: JsonRpcNotification = { + jsonrpc: '2.0', + method: 'pty.data', + params: { ...params, data: data.slice(0, mid) } + } + const bytes = this.estimateFrameBytes(msg) + if (clients.every((client) => bytes <= client.writer.producerFrameCapacity)) { + low = mid + } else { + high = mid - 1 + } + } + return low + } + + tryNotifyPtyData( + params: Record, + options: { interactive?: boolean } = {} + ): boolean { + if (this.disposed) { + return false + } + const msg: JsonRpcNotification = { + jsonrpc: '2.0', + method: 'pty.data', + params + } + return this.tryPublishToClients( + this.activeClients(), + msg, + options.interactive ? 'interactive' : 'ordinary' + ) + } + + tryNotifyPtyExit(params: Record): boolean { + if (this.disposed) { + return false + } + return this.tryPublishToClients( + this.activeClients(), + { + jsonrpc: '2.0', + method: 'pty.exit', + params + }, + 'ordinary' + ) + } + + producerDataBudget( + method: string, + paramsWithoutData: Record, + clientId?: number + ): number { + const targets = + clientId === undefined + ? this.activeClients() + : [this.clients.get(clientId)].filter((client): client is RelayClient => !!client) + if (targets.length === 0) { + return Number.MAX_SAFE_INTEGER + } + const emptyFrameBytes = this.estimateFrameBytes({ + jsonrpc: '2.0', + method, + params: { ...paramsWithoutData, data: '' } + }) + return Math.max( + 0, + Math.min(...targets.map((client) => client.writer.producerFrameCapacity - emptyFrameBytes)) + ) + } + feed(data: Buffer): void { this.feedForClient(this.primaryClient, data) } @@ -165,9 +267,24 @@ export class RelayDispatcher { method, ...(params !== undefined ? { params } : {}) } - for (const client of this.clients.values()) { - this.sendFrame(client, msg) - } + this.runPublicationTransaction(() => { + for (const client of this.clients.values()) { + if (client.closed) { + continue + } + if (method === 'pty.replay') { + this.enqueueFrame(client, msg, 'control') + continue + } + if (!this.publishToClient(client, msg, 'ordinary')) { + this.closeClient( + client, + new Error('Relay ordinary publication capacity exceeded'), + client !== this.primaryClient + ) + } + } + }) } notifyClient(clientId: number, method: string, params?: Record): void { @@ -178,11 +295,15 @@ export class RelayDispatcher { if (!client || client.closed) { return } - this.sendFrame(client, { - jsonrpc: '2.0', - method, - ...(params !== undefined ? { params } : {}) - }) + this.enqueueFrame( + client, + { + jsonrpc: '2.0', + method, + ...(params !== undefined ? { params } : {}) + }, + 'control' + ) } /** @@ -213,17 +334,7 @@ export class RelayDispatcher { if (client.closed) { continue } - // Why: encode inside the chain step, not at call time, so sequence numbers match 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 - }) + const step = client.bulkChain.then(() => this.publishBulkWhenAvailable(client, msg)) client.bulkChain = step.catch(() => {}) waits.push(step) } @@ -233,35 +344,6 @@ export class RelayDispatcher { 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, @@ -315,7 +397,7 @@ export class RelayDispatcher { reject(new Error(`Request "${method}" timed out after ${timeoutMs}ms`)) }, timeoutMs) this.pendingRelayRequests.set(id, { resolve, reject, timer }) - this.sendFrame(client, msg) + this.enqueueFrame(client, msg, 'control') }) } @@ -336,24 +418,29 @@ export class RelayDispatcher { // Why: can't send responses after dispose; abort in-flight work so SSH-side scans/watchers release. this.requestAborts.abortAll() for (const client of this.clients.values()) { - this.flushDrainWaiters(client) + client.closed = true + client.writer.close(new Error('Relay dispatcher disposed')) } + for (const listener of Array.from(this.legacyCapacityListeners)) { + listener() + } + this.legacyCapacityListeners.clear() } private createClient(write: RelayClientWrite, sinkOptions?: RelayClientSinkOptions): RelayClient { const id = this.nextClientId++ - const client: RelayClient = { + const client = { id, - decoder: new FrameDecoder((frame) => this.handleFrame(client, frame)), - write, - waitWriteDrain: sinkOptions?.waitWriteDrain, - drainWaiters: new Set(), + decoder: undefined as unknown as FrameDecoder, + writer: undefined as unknown as DispatcherClientWriter, bulkChain: Promise.resolve(), nextOutgoingSeq: 1, highestReceivedSeq: 0, generation: 0, closed: false - } + } satisfies RelayClient + client.decoder = new FrameDecoder((frame) => this.handleFrame(client, frame)) + client.writer = this.createWriter(client, write, sinkOptions) return client } @@ -484,19 +571,38 @@ export class RelayDispatcher { id, ...(error ? { error } : { result: result ?? null }) } - this.sendFrame(client, msg) + const estimatedBytes = this.estimateFrameBytes(msg) + const lane = estimatedBytes > DISPATCHER_CONTROL_QUEUE_MAX_BYTES ? 'legacy-response' : 'control' + if (!this.enqueueFrame(client, msg, lane)) { + this.closeClient( + client, + new Error( + `Relay response exceeds the bounded ${DEFAULT_PRODUCER_QUEUE_MAX_BYTES}-byte legacy lane` + ), + client !== this.primaryClient + ) + } } - private sendFrame( + private enqueueFrame( client: RelayClient, - msg: JsonRpcRequest | JsonRpcResponse | JsonRpcNotification - ): boolean | void { + msg: JsonRpcRequest | JsonRpcResponse | JsonRpcNotification, + lane: DispatcherWriterLane, + onSettled: (result: SinkWriteSettlement) => void = () => {} + ): boolean { if (this.disposed || client.closed) { - return + return false } - const seq = client.nextOutgoingSeq++ - const frame = encodeJsonRpcFrame(msg, seq, client.highestReceivedSeq) - return this.writeFrame(client, frame) + const estimatedBytes = this.estimateFrameBytes(msg) + return client.writer.enqueue( + lane, + () => { + const seq = client.nextOutgoingSeq++ + return encodeJsonRpcFrame(msg, seq, client.highestReceivedSeq) + }, + estimatedBytes, + onSettled + ) } private startKeepalive(): void { @@ -508,31 +614,203 @@ export class RelayDispatcher { if (client.closed) { continue } - const seq = client.nextOutgoingSeq++ - const frame = encodeKeepAliveFrame(seq, client.highestReceivedSeq) - this.writeFrame(client, frame) + client.writer.enqueue( + 'liveness', + () => { + const seq = client.nextOutgoingSeq++ + return encodeKeepAliveFrame(seq, client.highestReceivedSeq) + }, + 13 + ) } }, KEEPALIVE_SEND_MS) // Why: unref so the keepalive interval doesn't pin the event loop and block process exit. this.keepaliveTimer.unref() } - private writeFrame(client: RelayClient, frame: Buffer): boolean | void { - try { - return client.write(frame) - } catch (err) { - client.closed = true - client.generation++ - this.requestAborts.abortClient(client.id) - this.flushDrainWaiters(client) - // Why: frames have no retransmit buffer; detach now so reconnect/PTY-reattach runs instead of waiting the ~20s keepalive timeout. - if (client !== this.primaryClient) { - this.clients.delete(client.id) + private activeClients(): RelayClient[] { + return Array.from(this.clients.values()).filter((client) => !client.closed) + } + + private activeClientKeys(): string[] { + return this.activeClients().map((client) => this.clientKey(client)) + } + + private clientKey(client: RelayClient): string { + return `${client.id}:${client.generation}` + } + + private estimateFrameBytes(msg: JsonRpcRequest | JsonRpcResponse | JsonRpcNotification): number { + return encodeJsonRpcFrame(msg, 0, 0).length + } + + private tryPublishToClients( + clients: readonly RelayClient[], + msg: JsonRpcNotification, + lane: 'interactive' | 'ordinary' | 'bulk' + ): boolean { + return this.runPublicationTransaction(() => { + if (clients.length === 0) { + return true } - this.notifyClientDetached(client.id) - process.stderr.write( - `[relay] Client write failed: ${err instanceof Error ? err.message : String(err)}\n` + const bytes = this.estimateFrameBytes(msg) + if (clients.some((client) => !client.writer.canEnqueueProducer(bytes))) { + return false + } + const leases = this.publicationLedger.tryReserve( + clients.map((client) => ({ clientKey: this.clientKey(client), bytes })) ) + if (!leases) { + return false + } + for (let index = 0; index < clients.length; index++) { + if (!this.enqueueLeasedFrame(clients[index], msg, lane, leases[index])) { + for (let remaining = index; remaining < leases.length; remaining++) { + leases[remaining].release() + } + return false + } + } + return true + }) + } + + private publishToClient( + client: RelayClient, + msg: JsonRpcNotification, + lane: 'interactive' | 'ordinary' | 'bulk', + onSettled: (result: SinkWriteSettlement) => void = () => {} + ): boolean { + const bytes = this.estimateFrameBytes(msg) + if (!client.writer.canEnqueueProducer(bytes)) { + return false + } + const leases = this.publicationLedger.tryReserve([{ clientKey: this.clientKey(client), bytes }]) + if (!leases) { + return false + } + return this.enqueueLeasedFrame(client, msg, lane, leases[0], onSettled) + } + + private publishBulkWhenAvailable(client: RelayClient, msg: JsonRpcNotification): Promise { + const bytes = this.estimateFrameBytes(msg) + if (bytes > client.writer.producerFrameCapacity) { + return Promise.reject(new Error('Relay bulk frame exceeds sink producer capacity')) + } + return new Promise((resolve, reject) => { + let removeCapacityListener: (() => void) | null = null + const finish = (): void => { + removeCapacityListener?.() + removeCapacityListener = null + } + const tryPublish = (): void => { + if (this.disposed || client.closed) { + finish() + resolve() + return + } + if ( + this.publishToClient(client, msg, 'bulk', (result) => { + finish() + if (result.ok || this.disposed || client.closed) { + resolve() + } else { + reject(result.error) + } + }) + ) { + return + } + if (!removeCapacityListener) { + removeCapacityListener = this.onLegacyPtyCapacity(tryPublish) + } + } + tryPublish() + }) + } + + private enqueueLeasedFrame( + client: RelayClient, + msg: JsonRpcNotification, + lane: 'interactive' | 'ordinary' | 'bulk', + lease: LegacyPublicationLease, + onSettled: (result: SinkWriteSettlement) => void = () => {} + ): boolean { + const accepted = this.enqueueFrame(client, msg, lane, (result) => { + lease.release() + this.notifyLegacyCapacityIfLow() + onSettled(result) + }) + if (!accepted) { + lease.release() + this.notifyLegacyCapacityIfLow() + } + return accepted + } + + private createWriter( + client: RelayClient, + write: RelayClientWrite, + sinkOptions?: RelayClientSinkOptions + ): DispatcherClientWriter { + const writer = new DispatcherClientWriter(write, sinkOptions, (error) => { + this.closeClient(client, error, client !== this.primaryClient) + }) + writer.onCapacity(() => this.notifyLegacyCapacityIfLow()) + return writer + } + + private closeClient(client: RelayClient, error: Error, remove: boolean): void { + if (client.closed) { + return + } + client.closed = true + this.requestAborts.abortClient(client.id) + client.writer.close(error) + client.generation++ + if (remove) { + this.clients.delete(client.id) + } + this.notifyClientDetached(client.id) + this.notifyLegacyCapacity(true) + if (!/^Relay (?:primary client invalidated|client detached)$/.test(error.message)) { + process.stderr.write(`[relay] Client write closed: ${error.message}\n`) + } + } + + private notifyLegacyCapacityIfLow(): void { + this.notifyLegacyCapacity(false) + } + + private notifyLegacyCapacity(force: boolean): void { + if (this.publicationTransactionDepth > 0) { + this.deferredForcedLegacyCapacity ||= force + this.deferredLegacyCapacity ||= !force + return + } + if (!force && !this.publicationLedger.belowLowWater(this.activeClientKeys())) { + return + } + for (const listener of this.legacyCapacityListeners) { + listener() + } + } + + private runPublicationTransaction(operation: () => T): T { + this.publicationTransactionDepth++ + try { + return operation() + } finally { + this.publicationTransactionDepth-- + if (this.publicationTransactionDepth === 0) { + const force = this.deferredForcedLegacyCapacity + const low = this.deferredLegacyCapacity + this.deferredForcedLegacyCapacity = false + this.deferredLegacyCapacity = false + if (force || low) { + this.notifyLegacyCapacity(force) + } + } } } diff --git a/src/relay/fs-handler-file-read.ts b/src/relay/fs-handler-file-read.ts index 003172c08fd..cb8630331e4 100644 --- a/src/relay/fs-handler-file-read.ts +++ b/src/relay/fs-handler-file-read.ts @@ -183,7 +183,14 @@ async function pumpChunks( break } } - const want = Math.min(STREAM_CHUNK_SIZE, totalSize - offset) + const base64Budget = + dispatcher.producerDataBudget?.( + 'fs.streamChunk', + { streamId, seq }, + pumpOptions.clientId + ) ?? Number.MAX_SAFE_INTEGER + const sinkChunkBytes = Math.max(1, Math.floor(Math.max(0, base64Budget - 32) / 4) * 3) + const want = Math.min(STREAM_CHUNK_SIZE, sinkChunkBytes, totalSize - offset) const bytesRead = await readFullStreamChunk(entry.handle, buffer, want, offset) if (bytesRead !== want) { endReason = 'error' diff --git a/src/relay/git-response-stream.ts b/src/relay/git-response-stream.ts index bced3b47e46..058bb3aebb6 100644 --- a/src/relay/git-response-stream.ts +++ b/src/relay/git-response-stream.ts @@ -22,10 +22,10 @@ type GitResponseStreamEntry = { /** Serialized git responses are chunked as base64 so multi-byte UTF-8 * sequences never split across a chunk boundary (the client concatenates the * decoded bytes and parses once). */ -function encodeChunks(payload: Buffer): string[] { +function encodeChunks(payload: Buffer, chunkBytes = GIT_RESPONSE_CHUNK_SIZE): string[] { const chunks: string[] = [] - for (let offset = 0; offset < payload.length; offset += GIT_RESPONSE_CHUNK_SIZE) { - chunks.push(payload.subarray(offset, offset + GIT_RESPONSE_CHUNK_SIZE).toString('base64')) + for (let offset = 0; offset < payload.length; offset += chunkBytes) { + chunks.push(payload.subarray(offset, offset + chunkBytes).toString('base64')) } return chunks } @@ -116,7 +116,14 @@ export class GitResponseStreamRegistry { context: RequestContext ): GitResponseStreamMarker { const streamId = this.register(context.clientId) - const chunks = encodeChunks(payload) + const base64Budget = + dispatcher.producerDataBudget?.( + 'git.responseChunk', + { streamId, seq: 0 }, + context.clientId + ) ?? Number.MAX_SAFE_INTEGER + const sinkChunkBytes = Math.max(1, Math.floor(Math.max(0, base64Budget - 32) / 4) * 3) + const chunks = encodeChunks(payload, Math.min(GIT_RESPONSE_CHUNK_SIZE, sinkChunkBytes)) // Why: kick the pump off the response task so the client sees the sentinel // (and can subscribe/reassemble) before the first chunk frame arrives. setImmediate(() => { diff --git a/src/relay/legacy-relay-publication-ledger.test.ts b/src/relay/legacy-relay-publication-ledger.test.ts new file mode 100644 index 00000000000..04d64a55407 --- /dev/null +++ b/src/relay/legacy-relay-publication-ledger.test.ts @@ -0,0 +1,48 @@ +import { describe, expect, it } from 'vitest' +import { LegacyRelayPublicationLedger } from './legacy-relay-publication-ledger' + +describe('LegacyRelayPublicationLedger', () => { + it('caps each client and aggregate retained bytes', () => { + const ledger = new LegacyRelayPublicationLedger({ + clientHighBytes: 10, + clientLowBytes: 4, + relayHighBytes: 15, + relayLowBytes: 8 + }) + const first = ledger.tryReserve([ + { clientKey: 'a', bytes: 7 }, + { clientKey: 'b', bytes: 7 } + ]) + expect(first).not.toBeNull() + expect(ledger.retainedBytes).toBe(14) + expect(ledger.tryReserve([{ clientKey: 'a', bytes: 4 }])).toBeNull() + expect(ledger.tryReserve([{ clientKey: 'c', bytes: 2 }])).toBeNull() + + first?.[0].release() + expect(ledger.retainedBytes).toBe(7) + expect(ledger.tryReserve([{ clientKey: 'c', bytes: 2 }])).not.toBeNull() + }) + + it('releases leases exactly once and applies both low waters', () => { + const ledger = new LegacyRelayPublicationLedger({ + clientHighBytes: 10, + clientLowBytes: 3, + relayHighBytes: 20, + relayLowBytes: 5 + }) + const leases = ledger.tryReserve([ + { clientKey: 'a', bytes: 4 }, + { clientKey: 'b', bytes: 4 } + ]) + expect(ledger.belowLowWater()).toBe(false) + + leases?.[0].release() + leases?.[0].release() + expect(ledger.retainedBytes).toBe(4) + expect(ledger.belowLowWater()).toBe(false) + + leases?.[1].release() + expect(ledger.retainedBytes).toBe(0) + expect(ledger.belowLowWater()).toBe(true) + }) +}) diff --git a/src/relay/legacy-relay-publication-ledger.ts b/src/relay/legacy-relay-publication-ledger.ts new file mode 100644 index 00000000000..f7002756bc2 --- /dev/null +++ b/src/relay/legacy-relay-publication-ledger.ts @@ -0,0 +1,104 @@ +export const LEGACY_CLIENT_RETAINED_BYTES_HIGH = 2 * 1024 * 1024 +export const LEGACY_CLIENT_RETAINED_BYTES_LOW = 1024 * 1024 +export const LEGACY_RELAY_RETAINED_BYTES_HIGH = 32 * 1024 * 1024 +export const LEGACY_RELAY_RETAINED_BYTES_LOW = 24 * 1024 * 1024 + +export type LegacyPublicationLease = { + clientKey: string + bytes: number + release: () => void +} + +type LegacyPublicationLimits = { + clientHighBytes?: number + clientLowBytes?: number + relayHighBytes?: number + relayLowBytes?: number +} + +export class LegacyRelayPublicationLedger { + private readonly clientBytes = new Map() + private aggregateBytes = 0 + + readonly clientHighBytes: number + readonly clientLowBytes: number + readonly relayHighBytes: number + readonly relayLowBytes: number + + constructor(limits: LegacyPublicationLimits = {}) { + this.clientHighBytes = limits.clientHighBytes ?? LEGACY_CLIENT_RETAINED_BYTES_HIGH + this.clientLowBytes = limits.clientLowBytes ?? LEGACY_CLIENT_RETAINED_BYTES_LOW + this.relayHighBytes = limits.relayHighBytes ?? LEGACY_RELAY_RETAINED_BYTES_HIGH + this.relayLowBytes = limits.relayLowBytes ?? LEGACY_RELAY_RETAINED_BYTES_LOW + } + + get retainedBytes(): number { + return this.aggregateBytes + } + + retainedBytesFor(clientKey: string): number { + return this.clientBytes.get(clientKey) ?? 0 + } + + tryReserve( + memberships: readonly { clientKey: string; bytes: number }[] + ): LegacyPublicationLease[] | null { + let aggregateAdded = 0 + const additions = new Map() + for (const membership of memberships) { + if (!Number.isSafeInteger(membership.bytes) || membership.bytes < 0) { + return null + } + aggregateAdded += membership.bytes + additions.set( + membership.clientKey, + (additions.get(membership.clientKey) ?? 0) + membership.bytes + ) + } + if (this.aggregateBytes + aggregateAdded > this.relayHighBytes) { + return null + } + for (const [clientKey, bytes] of additions) { + if ((this.clientBytes.get(clientKey) ?? 0) + bytes > this.clientHighBytes) { + return null + } + } + this.aggregateBytes += aggregateAdded + for (const [clientKey, bytes] of additions) { + this.clientBytes.set(clientKey, (this.clientBytes.get(clientKey) ?? 0) + bytes) + } + return memberships.map(({ clientKey, bytes }) => { + let released = false + return { + clientKey, + bytes, + release: () => { + if (released) { + return + } + released = true + this.release(clientKey, bytes) + } + } + }) + } + + belowLowWater(clientKeys?: readonly string[]): boolean { + if (this.aggregateBytes > this.relayLowBytes) { + return false + } + const keys = clientKeys ?? Array.from(this.clientBytes.keys()) + return keys.every((clientKey) => (this.clientBytes.get(clientKey) ?? 0) <= this.clientLowBytes) + } + + private release(clientKey: string, bytes: number): void { + const current = this.clientBytes.get(clientKey) ?? 0 + const next = Math.max(0, current - bytes) + this.aggregateBytes = Math.max(0, this.aggregateBytes - Math.min(bytes, current)) + if (next === 0) { + this.clientBytes.delete(clientKey) + } else { + this.clientBytes.set(clientKey, next) + } + } +} diff --git a/src/relay/pty-handler-output-drain-differential.test.ts b/src/relay/pty-handler-output-drain-differential.test.ts index b2527e92ec4..22ba265cd10 100644 --- a/src/relay/pty-handler-output-drain-differential.test.ts +++ b/src/relay/pty-handler-output-drain-differential.test.ts @@ -9,7 +9,9 @@ const { mockPtySpawn, mockPtyInstance } = vi.hoisted(() => ({ write: vi.fn(), resize: vi.fn(), kill: vi.fn(), - clear: vi.fn() + clear: vi.fn(), + pause: vi.fn(), + resume: vi.fn() } })) diff --git a/src/relay/pty-handler.test.ts b/src/relay/pty-handler.test.ts index 3a2f1ae8ebf..f461baf5a54 100644 --- a/src/relay/pty-handler.test.ts +++ b/src/relay/pty-handler.test.ts @@ -24,7 +24,9 @@ const { mockPtySpawn, mockPtyInstance } = vi.hoisted(() => ({ write: vi.fn(), resize: vi.fn(), kill: vi.fn(), - clear: vi.fn() + clear: vi.fn(), + pause: vi.fn(), + resume: vi.fn() } })) @@ -39,7 +41,8 @@ import { attachIdentityMismatches, formatNodePtyUnavailableMessage } from './pty-handler' -import type { RelayDispatcher } from './dispatcher' +import { RelayDispatcher } from './dispatcher' +import { encodeJsonRpcFrame } from './protocol' type TestRequestContext = { isStale: () => boolean @@ -127,6 +130,8 @@ describe('PtyHandler', () => { mockPtyInstance.resize.mockReset() mockPtyInstance.kill.mockReset() mockPtyInstance.clear.mockReset() + mockPtyInstance.pause.mockReset() + mockPtyInstance.resume.mockReset() mockPtySpawn.mockReturnValue({ ...mockPtyInstance }) @@ -162,6 +167,60 @@ describe('PtyHandler', () => { expect(notifMethods).toContain('pty.ackData') }) + it('pauses native output at the producer hard water and resumes after retained writes settle', async () => { + let onData: ((data: string) => void) | undefined + const pause = vi.fn() + const resume = vi.fn() + mockPtySpawn.mockReturnValueOnce({ + ...mockPtyInstance, + pause, + resume, + onData: vi.fn((callback: (data: string) => void) => { + onData = callback + }) + }) + const writeCallbacks: (() => void)[] = [] + let writableLength = 0 + const boundedDispatcher = new RelayDispatcher( + (data, settle) => { + writableLength += data.length + writeCallbacks.push(() => { + writableLength -= data.length + settle({ ok: true }) + }) + return true + }, + { + supportsWriteCallback: true, + writableLength: () => writableLength, + writableHighWaterMark: () => 4 * 1024 * 1024 + } + ) + const boundedHandler = new PtyHandler(boundedDispatcher) + try { + boundedDispatcher.feed( + encodeJsonRpcFrame({ jsonrpc: '2.0', id: 1, method: 'pty.spawn', params: {} }, 1, 0) + ) + await vi.advanceTimersByTimeAsync(0) + expect(onData).toBeTypeOf('function') + + onData?.('x'.repeat(1536 * 1024)) + expect(pause).toHaveBeenCalledTimes(1) + await vi.advanceTimersByTimeAsync(300) + + expect(writeCallbacks.length).toBeGreaterThan(50) + expect(resume).not.toHaveBeenCalled() + for (const settle of writeCallbacks.splice(0)) { + settle() + } + await vi.advanceTimersByTimeAsync(0) + expect(resume).toHaveBeenCalledTimes(1) + } finally { + await boundedHandler.dispose({ waitForPhysicalExit: false }).catch(() => {}) + boundedDispatcher.dispose() + } + }) + it('rejects strict process inspection for a missing relay PTY', async () => { await expect(dispatcher.callRequest('pty.inspectProcess', { id: 'missing' })).rejects.toThrow( 'terminal_gone' @@ -1549,6 +1608,42 @@ describe('PtyHandler', () => { expect(handler.activePtyCount).toBe(0) }) + it('retains PTY exit until the ordinary writer admits it', async () => { + await handler.dispose({ waitForPhysicalExit: false }) + let capacityListener: (() => void) | undefined + const tryNotifyPtyExit = vi.fn().mockReturnValueOnce(false).mockReturnValueOnce(true) + Object.assign(dispatcher, { + onLegacyPtyCapacity: vi.fn((listener: () => void) => { + capacityListener = listener + return vi.fn() + }), + tryNotifyPtyData: vi.fn(() => true), + tryNotifyPtyExit, + legacyRetentionBelowLowWater: true + }) + handler = new PtyHandler(dispatcher as unknown as RelayDispatcher) + let exitCallback: ((info: { exitCode: number }) => void) | undefined + mockPtySpawn.mockReturnValue({ + ...mockPtyInstance, + onData: vi.fn(), + onExit: vi.fn((callback: (info: { exitCode: number }) => void) => { + exitCallback = callback + }) + }) + + const spawn = await spawnPty() + exitCallback?.({ exitCode: 0 }) + expect(tryNotifyPtyExit).toHaveBeenCalledOnce() + + capacityListener?.() + expect(tryNotifyPtyExit).toHaveBeenLastCalledWith({ + id: 'pty-1', + code: 0, + incarnationId: spawn.incarnationId + }) + expect(tryNotifyPtyExit).toHaveBeenCalledTimes(2) + }) + it('flushes pending PTY output before notifying exit', async () => { let dataCallback: ((data: string) => void) | undefined let exitCallback: ((info: { exitCode: number }) => void) | undefined diff --git a/src/relay/pty-handler.ts b/src/relay/pty-handler.ts index 578858d16b8..c0a07ac22ee 100644 --- a/src/relay/pty-handler.ts +++ b/src/relay/pty-handler.ts @@ -129,6 +129,7 @@ type PendingPtyOutput = { rawLength?: number transformed?: boolean seq?: number + interactive?: boolean } type ManagedStartupCommand = { @@ -191,6 +192,8 @@ const PTY_OUTPUT_BATCH_INTERVAL_MS = 8 const PTY_OUTPUT_DRAIN_CONTINUE_MS = 1 const PTY_OUTPUT_FLUSH_CHUNK_CHARS = 16 * 1024 const PTY_OUTPUT_FLUSH_MAX_WRITES = 2 +const PTY_OUTPUT_PRODUCER_HIGH_BYTES = 128 * 1024 +const PTY_OUTPUT_PRODUCER_LOW_BYTES = 64 * 1024 const INTERACTIVE_OUTPUT_WINDOW_MS = 100 const INTERACTIVE_OUTPUT_MAX_CHARS = 1024 const INTERACTIVE_REDRAW_MAX_CHARS = PTY_OUTPUT_FLUSH_CHUNK_CHARS @@ -309,7 +312,10 @@ export class PtyHandler { private graceTimeMs: number private graceTimer: ReturnType | null = null private outputFlushTimer: ReturnType | null = null - private pendingOutputByPty = new Map() + private pendingOutputByPty = new Map() + private pendingExitByPty = new Map() + private pausedOutputPtys = new Set() + private removeLegacyCapacityListener: (() => void) | null = null private lastInputAtByPty = new Map() private interactiveOutputCharsByPty = new Map() private pendingSpawnCount = 0 @@ -335,6 +341,8 @@ export class PtyHandler { this.dispatcher = dispatcher this.graceTimeMs = graceTimeMs this.registerHandlers() + this.removeLegacyCapacityListener = + this.dispatcher.onLegacyPtyCapacity?.(() => this.handleLegacyCapacity()) ?? null } private async loadPty(): Promise { @@ -578,16 +586,18 @@ export class PtyHandler { } this.clearStartupCommandTimer(managed) this.releaseRelayIngress(managed) + this.pausedOutputPtys.delete(managed.id) this.flushPtyOutput(managed.id) - this.dispatcher.notify('pty.exit', { + this.pendingExitByPty.set(managed.id, { id: managed.id, code: exitCode, incarnationId: managed.incarnationId }) + this.publishPendingExit(managed.id) this.notifyExitListener(managed) this.agentSessionOwners.release(managed.id) this.ptys.delete(managed.id) - this.clearPtyFlowState(managed.id) + this.clearPtyInputState(managed.id) // Why: release the ptmx fd on natural exit, else the master fd leaks until GC (docs/fix-pty-fd-leak.md). disposeManagedPty(managed) }) @@ -697,31 +707,48 @@ export class PtyHandler { data: string, meta: { rawLength?: number; transformed?: boolean; seq?: number } = {} ): void { - const existing = this.pendingOutputByPty.get(id) + const queue = this.pendingOutputByPty.get(id) ?? [] + const existing = queue.at(-1) if (meta.transformed === true) { - // Why: transformed spans lack a raw-to-clean slice mapping, so they can't be folded into the output batch. - if (existing) { - this.flushPtyOutput(id) + if (queue.length === 0) { + const transformed = { data, ...meta } + if (this.publishPtyOutput(id, transformed, false)) { + return + } + queue.push(transformed) + } else if (existing?.transformed) { + existing.data += data + existing.rawLength = (existing.rawLength ?? 0) + (meta.rawLength ?? data.length) + existing.seq = meta.seq + } else { + queue.push({ data, ...meta }) } - this.dispatcher.notify('pty.data', { id, data, ...meta }) + this.pendingOutputByPty.set(id, queue) + this.pausePtyOutput(id) return } - const pending: PendingPtyOutput = { data: (existing?.data ?? '') + data } + const pending: PendingPtyOutput = existing && !existing.transformed ? existing : { data: '' } + const previousLength = pending.data.length + pending.data += data if (existing?.rawLength !== undefined || meta.rawLength !== undefined) { - pending.rawLength = - (existing?.rawLength ?? existing?.data.length ?? 0) + (meta.rawLength ?? data.length) + pending.rawLength = (existing?.rawLength ?? previousLength) + (meta.rawLength ?? data.length) } if (meta.seq !== undefined) { pending.seq = meta.seq } - if (this.shouldSendInteractiveOutputNow(id, pending.data)) { - this.pendingOutputByPty.delete(id) - this.clearOutputFlushTimerIfIdle() - // Why: send interactive echo immediately — batching must not add visible input delay for TUIs. - this.dispatcher.notify('pty.data', { id, ...pending }) - return + if (!existing || existing.transformed) { + queue.push(pending) + } + this.pendingOutputByPty.set(id, queue) + if (queue.length === 1 && this.shouldSendInteractiveOutputNow(id, pending.data)) { + pending.interactive = true + if (this.flushPtyOutput(id)) { + return + } + } + if (this.pendingProducerBytes(id) >= PTY_OUTPUT_PRODUCER_HIGH_BYTES) { + this.pausePtyOutput(id) } - this.pendingOutputByPty.set(id, pending) this.scheduleOutputFlush(PTY_OUTPUT_BATCH_INTERVAL_MS) } @@ -737,56 +764,102 @@ export class PtyHandler { // Why batch before the first send: a re-entrant sink must read the values a whole-map snapshot // would have frozen. Why the raw iterator: `for...of` would consume one entry past the limit. const pendingEntries = this.pendingOutputByPty[Symbol.iterator]() - const batch: [string, PendingPtyOutput][] = [] + const batch: [string, PendingPtyOutput[]][] = [] while (batch.length < PTY_OUTPUT_FLUSH_MAX_WRITES) { const next = pendingEntries.next() if (next.done === true) { break } - batch.push(next.value) + batch.push([next.value[0], next.value[1].map((pending) => ({ ...pending }))]) } - for (const [id, pending] of batch) { + let writes = 0 + for (const [id, queue] of batch) { this.pendingOutputByPty.delete(id) - const chunk = pending.transformed - ? pending.data - : pending.data.slice(0, PTY_OUTPUT_FLUSH_CHUNK_CHARS) - const remaining = pending.transformed ? '' : pending.data.slice(PTY_OUTPUT_FLUSH_CHUNK_CHARS) - if (remaining) { - this.pendingOutputByPty.set(id, { - data: remaining, - ...(pending.rawLength === undefined ? {} : { rawLength: remaining.length }), - seq: pending.seq - }) + if (this.flushPtyOutput(id, queue)) { + writes++ } - const chunkRawLength = pending.transformed - ? pending.rawLength - : pending.rawLength === undefined - ? undefined - : chunk.length - const chunkSeq = - pending.seq === undefined ? undefined : pending.seq - (pending.data.length - chunk.length) - this.dispatcher.notify('pty.data', { - id, - data: chunk, - ...(chunkSeq === undefined ? {} : { seq: chunkSeq }), - ...(chunkRawLength === undefined ? {} : { rawLength: chunkRawLength }), - ...(pending.transformed ? { transformed: true } : {}) - }) } - if (this.pendingOutputByPty.size > 0 && batch.length > 0) { + if (this.pendingOutputByPty.size > 0 && writes > 0) { // Why: yield between slices of a large chunk so client input and control frames can interleave. this.scheduleOutputFlush(PTY_OUTPUT_DRAIN_CONTINUE_MS) } } - private flushPtyOutput(id: string): void { - const pending = this.pendingOutputByPty.get(id) - if (!pending) { - return + private flushPtyOutput(id: string, capturedQueue?: PendingPtyOutput[]): boolean { + const queue = capturedQueue ?? this.pendingOutputByPty.get(id) + const pending = queue?.[0] + if (!queue || !pending) { + this.publishPendingExit(id) + return true } - this.pendingOutputByPty.delete(id) - this.dispatcher.notify('pty.data', { id, ...pending }) + const desiredChars = pending.transformed + ? pending.data.length + : Math.min(pending.data.length, PTY_OUTPUT_FLUSH_CHUNK_CHARS) + const paramsWithoutData = { + id, + ...(pending.seq === undefined ? {} : { seq: pending.seq }), + ...(pending.rawLength === undefined ? {} : { rawLength: pending.rawLength }), + ...(pending.transformed ? { transformed: true } : {}) + } + let chunkChars = pending.transformed + ? desiredChars + : (this.dispatcher.maxLegacyPtyDataChars?.(paramsWithoutData, pending.data, desiredChars) ?? + desiredChars) + if ( + chunkChars > 0 && + chunkChars < pending.data.length && + pending.data.charCodeAt(chunkChars - 1) >= 0xd800 && + pending.data.charCodeAt(chunkChars - 1) <= 0xdbff + ) { + chunkChars-- + } + if (chunkChars <= 0 || (pending.transformed && chunkChars !== pending.data.length)) { + this.pendingOutputByPty.set(id, queue) + this.pausePtyOutput(id) + return false + } + const chunk = pending.data.slice(0, chunkChars) + const remaining = pending.data.slice(chunkChars) + const chunkRawLength = pending.transformed + ? pending.rawLength + : pending.rawLength === undefined + ? undefined + : chunk.length + const chunkSeq = + pending.seq === undefined ? undefined : pending.seq - (pending.data.length - chunk.length) + const published = this.publishPtyOutput( + id, + { + data: chunk, + ...(chunkSeq === undefined ? {} : { seq: chunkSeq }), + ...(chunkRawLength === undefined ? {} : { rawLength: chunkRawLength }), + ...(pending.transformed ? { transformed: true } : {}) + }, + pending.interactive === true + ) + if (!published) { + this.pendingOutputByPty.set(id, queue) + this.pausePtyOutput(id) + return false + } + if (remaining) { + queue[0] = { + data: remaining, + ...(pending.rawLength === undefined ? {} : { rawLength: remaining.length }), + seq: pending.seq + } + } else { + queue.shift() + } + if (queue.length === 0) { + this.pendingOutputByPty.delete(id) + this.publishPendingExit(id) + } else { + this.pendingOutputByPty.set(id, queue) + } + this.maybeResumePtyOutput(id) this.clearOutputFlushTimerIfIdle() + return true } private clearOutputFlushTimerIfIdle(): void { @@ -799,9 +872,102 @@ export class PtyHandler { private clearPtyFlowState(id: string): void { this.pendingOutputByPty.delete(id) + this.pendingExitByPty.delete(id) + this.pausedOutputPtys.delete(id) + this.clearPtyInputState(id) + this.clearOutputFlushTimerIfIdle() + } + + private clearPtyInputState(id: string): void { this.lastInputAtByPty.delete(id) this.interactiveOutputCharsByPty.delete(id) - this.clearOutputFlushTimerIfIdle() + } + + private publishPtyOutput(id: string, output: PendingPtyOutput, interactive: boolean): boolean { + if (this.dispatcher.tryNotifyPtyData) { + return this.dispatcher.tryNotifyPtyData( + { + id, + data: output.data, + ...(output.seq === undefined ? {} : { seq: output.seq }), + ...(output.rawLength === undefined ? {} : { rawLength: output.rawLength }), + ...(output.transformed ? { transformed: true } : {}) + }, + { interactive } + ) + } + this.dispatcher.notify('pty.data', { + id, + data: output.data, + ...(output.seq === undefined ? {} : { seq: output.seq }), + ...(output.rawLength === undefined ? {} : { rawLength: output.rawLength }), + ...(output.transformed ? { transformed: true } : {}) + }) + return true + } + + private publishPendingExit(id: string): void { + if (this.pendingOutputByPty.has(id)) { + return + } + const exit = this.pendingExitByPty.get(id) + if (!exit) { + return + } + const published = this.dispatcher.tryNotifyPtyExit + ? this.dispatcher.tryNotifyPtyExit(exit) + : (this.dispatcher.notify('pty.exit', exit), true) + if (!published) { + return + } + this.pendingExitByPty.delete(id) + } + + private pendingProducerBytes(id: string): number { + return (this.pendingOutputByPty.get(id) ?? []).reduce( + (total, pending) => + total + Math.max(Buffer.byteLength(pending.data, 'utf8'), 2 * pending.data.length) + 128, + 0 + ) + } + + private pausePtyOutput(id: string): void { + if (this.pausedOutputPtys.has(id)) { + return + } + const managed = this.ptys.get(id) + if (!managed || managed.disposed) { + return + } + this.pausedOutputPtys.add(id) + managed.pty.pause() + } + + private maybeResumePtyOutput(id: string): void { + if ( + !this.pausedOutputPtys.has(id) || + this.pendingProducerBytes(id) > PTY_OUTPUT_PRODUCER_LOW_BYTES || + this.dispatcher.legacyRetentionBelowLowWater === false + ) { + return + } + const managed = this.ptys.get(id) + this.pausedOutputPtys.delete(id) + if (managed && !managed.disposed) { + managed.pty.resume() + } + } + + private handleLegacyCapacity(): void { + if (this.pendingOutputByPty.size > 0) { + this.scheduleOutputFlush(0) + } + for (const id of Array.from(this.pendingExitByPty.keys())) { + this.publishPendingExit(id) + } + for (const id of Array.from(this.pausedOutputPtys)) { + this.maybeResumePtyOutput(id) + } } private beginPtyCreation(operationPaths: readonly (string | undefined)[]): () => void { @@ -1205,6 +1371,7 @@ export class PtyHandler { // Why: drop pending batched bytes already in the replay buffer so attach doesn't render them twice. this.pendingOutputByPty.delete(id) this.clearOutputFlushTimerIfIdle() + this.maybeResumePtyOutput(id) if (params.suppressReplayNotification) { return { incarnationId: managed.incarnationId, replay } } @@ -1603,6 +1770,8 @@ export class PtyHandler { if (this.disposePromise) { return this.disposePromise } + this.removeLegacyCapacityListener?.() + this.removeLegacyCapacityListener = null this.agentSessionCreateOperations.clear() const disposePromise = this.disposePtys(options.waitForPhysicalExit !== false) this.disposePromise = disposePromise @@ -1627,6 +1796,8 @@ export class PtyHandler { this.outputFlushTimer = null } this.pendingOutputByPty.clear() + this.pendingExitByPty.clear() + this.pausedOutputPtys.clear() this.lastInputAtByPty.clear() this.interactiveOutputCharsByPty.clear() const results = await Promise.allSettled( diff --git a/src/relay/relay-watcher-event-emitter.test.ts b/src/relay/relay-watcher-event-emitter.test.ts new file mode 100644 index 00000000000..ae5781c5745 --- /dev/null +++ b/src/relay/relay-watcher-event-emitter.test.ts @@ -0,0 +1,53 @@ +import { describe, expect, it } from 'vitest' +import { RelayDispatcher } from './dispatcher' +import { HEADER_LENGTH, parseJsonRpcMessage } from './protocol' +import { emitRelayWatcherEvents, emitRelayWatcherOverflow } from './relay-watcher-event-emitter' + +function frameMethod(frame: Buffer): string { + const payloadLength = frame.readUInt32BE(9) + const message = parseJsonRpcMessage(frame.subarray(HEADER_LENGTH, HEADER_LENGTH + payloadLength)) + return 'method' in message ? message.method : '' +} + +describe('relay watcher writer admission', () => { + it('keeps watcher batches on the bounded ordinary lane under saturation', () => { + const frames: Buffer[] = [] + const drainWaiters: (() => void)[] = [] + let saturate = true + const dispatcher = new RelayDispatcher( + (frame) => { + frames.push(Buffer.from(frame)) + if (saturate) { + saturate = false + return false + } + return true + }, + { + waitWriteDrain: (callback) => { + drainWaiters.push(callback) + return () => { + const index = drainWaiters.indexOf(callback) + if (index >= 0) { + drainWaiters.splice(index, 1) + } + } + } + } + ) + + try { + emitRelayWatcherEvents(dispatcher, false, [ + { type: 'create', path: '/workspace/first', isDirectory: false } + ]) + emitRelayWatcherOverflow(dispatcher, '/workspace', false) + dispatcher.notifyClient(1, 'control.event') + + expect(frames.map(frameMethod)).toEqual(['fs.changed']) + drainWaiters.shift()?.() + expect(frames.map(frameMethod)).toEqual(['fs.changed', 'control.event', 'fs.changed']) + } finally { + dispatcher.dispose() + } + }) +}) diff --git a/src/relay/relay-watcher-event-emitter.ts b/src/relay/relay-watcher-event-emitter.ts index 46515e67959..1a575c360fa 100644 --- a/src/relay/relay-watcher-event-emitter.ts +++ b/src/relay/relay-watcher-event-emitter.ts @@ -9,6 +9,7 @@ export function emitRelayWatcherEvents( if (closed || events.length === 0) { return } + // Watcher floods stay on the bounded producer lane; its cap closes the stale generation. dispatcher.notify('fs.changed', { events: events.map((event) => ({ kind: event.type, diff --git a/src/relay/relay.ts b/src/relay/relay.ts index 82110b58db2..c45345ef71d 100644 --- a/src/relay/relay.ts +++ b/src/relay/relay.ts @@ -53,6 +53,7 @@ import { remoteCliRequestTimeoutMs } from './remote-cli-timeout' import { shouldReadRemoteCliStdin } from './remote-cli-stdin' import { registerManagedHookInstaller } from './managed-hook-installer' import { registerRelayPluginHostCallHandlers } from './plugin-host-call-handler' +import { DispatcherClientWriter } from './dispatcher-client-writer' const DEFAULT_GRACE_MS = DEFAULT_SSH_RELAY_GRACE_PERIOD_SECONDS * 1000 const SOCK_NAME = 'relay.sock' @@ -151,6 +152,25 @@ function parseArgs(argv: string[]): { function runConnectMode(sockPath: string): void { const myVersion = readLaunchVersion() const sock = createConnection({ path: sockPath }) + const stdoutWriter = new DispatcherClientWriter( + (data, onSettled) => + process.stdout.write(data, (error) => { + onSettled(error ? { ok: false, error } : { ok: true }) + }), + { + supportsWriteCallback: true, + writableLength: () => process.stdout.writableLength, + writableHighWaterMark: () => process.stdout.writableHighWaterMark, + waitWriteDrain: (callback) => { + process.stdout.once('drain', callback) + return () => process.stdout.off('drain', callback) + } + }, + () => { + sock.destroy() + process.exit(1) + } + ) const connectTimeout = setTimeout(() => { process.stderr.write(`[relay-connect] Connection timed out after ${CONNECT_TIMEOUT_MS}ms\n`) @@ -162,22 +182,51 @@ function runConnectMode(sockPath: string): void { clearTimeout(connectTimeout) runConnectHandshake(sock, myVersion, { onAccepted: (leftover: Buffer) => { - // Why: write RELAY_SENTINEL only after the handshake passes, so a version mismatch is a clean exit-42 instead of a false sentinel + channel drop. - process.stdout.write(RELAY_SENTINEL) - // Why: forward handshake-buffered leftover bytes before sock.pipe(process.stdout) so the downstream mux sees them in order. + stdoutWriter.enqueue('control', () => Buffer.from(RELAY_SENTINEL), RELAY_SENTINEL.length) if (leftover.length > 0) { - process.stdout.write(leftover) + stdoutWriter.enqueue('control', () => leftover, leftover.length) } process.stdin.pipe(sock) - sock.pipe(process.stdout) + sock.on('data', (data: Buffer) => { + sock.pause() + let offset = 0 + const writeNext = (): void => { + if (offset >= data.length) { + sock.resume() + return + } + const bytes = Math.min(stdoutWriter.producerFrameCapacity, data.length - offset) + if (bytes <= 0) { + stdoutWriter.close(new Error('Relay stdout has no producer capacity')) + return + } + const chunk = data.subarray(offset, offset + bytes) + if ( + !stdoutWriter.enqueue( + 'ordinary', + () => chunk, + chunk.length, + (result) => { + if (!result.ok) { + return + } + offset += bytes + writeNext() + } + ) + ) { + stdoutWriter.close(new Error('Relay stdout bridge capacity exceeded')) + } + } + writeNext() + }) } }) }) // Why: Node swallows EPIPE on stdout, so the bridge would zombie and drop frames; exit on stdout error so the relay enters grace promptly. process.stdout.on('error', () => { - sock.destroy() - process.exit(1) + stdoutWriter.close(new Error('Relay stdout closed')) }) sock.on('error', (err) => { @@ -186,7 +235,8 @@ function runConnectMode(sockPath: string): void { process.exit(1) }) - sock.on('close', () => { + sock.on('close', async () => { + await stdoutWriter.waitForIdle() process.exit(0) }) } @@ -195,6 +245,22 @@ async function runOrcaCliMode(sockPath: string, argv: string[]): Promise { const myVersion = readLaunchVersion() const stdin = shouldReadRemoteCliStdin(argv) ? await readOrcaCliStdin() : undefined const sock = createConnection({ path: sockPath }) + const stdoutWriter = new DispatcherClientWriter( + (data, onSettled) => + process.stdout.write(data, (error) => { + onSettled(error ? { ok: false, error } : { ok: true }) + }), + { + supportsWriteCallback: true, + writableLength: () => process.stdout.writableLength, + writableHighWaterMark: () => process.stdout.writableHighWaterMark, + waitWriteDrain: (callback) => { + process.stdout.once('drain', callback) + return () => process.stdout.off('drain', callback) + } + }, + () => process.exit(1) + ) let nextSeq = 1 let highestReceivedSeq = 0 const requestId = 1 @@ -241,14 +307,24 @@ async function runOrcaCliMode(sockPath: string, argv: string[]): Promise { stderr?: unknown exitCode?: unknown } - if (typeof result.stdout === 'string' && result.stdout.length > 0) { - process.stdout.write(result.stdout) - } if (typeof result.stderr === 'string' && result.stderr.length > 0) { process.stderr.write(result.stderr) } sock.destroy() - process.exit(typeof result.exitCode === 'number' ? result.exitCode : 0) + const exitCode = typeof result.exitCode === 'number' ? result.exitCode : 0 + if (typeof result.stdout === 'string' && result.stdout.length > 0) { + const output = Buffer.from(result.stdout) + stdoutWriter.enqueue( + 'control', + () => output, + output.length, + (settlement) => { + process.exit(settlement.ok ? exitCode : 1) + } + ) + } else { + process.exit(exitCode) + } }) const connectTimeout = setTimeout(() => { @@ -357,26 +433,36 @@ async function main(): Promise { } process.stdout.on('drain', flushStdoutDrainWaiters) const dispatcher = new RelayDispatcher( - (data) => { + (data, onSettled) => { if (!stdoutAlive) { - return + onSettled({ ok: false, error: new Error('Relay stdout is closed') }) + return false } try { - // Why: surface Node's backpressure so bulk frames (fs.streamChunk) wait for drain instead of queueing ahead of interactive pty.data. - return process.stdout.write(data) - } catch { + return process.stdout.write(data, (error) => { + onSettled(error ? { ok: false, error } : { ok: true }) + }) + } catch (error) { stdoutAlive = false flushStdoutDrainWaiters() - return undefined + onSettled({ + ok: false, + error: error instanceof Error ? error : new Error(String(error)) + }) + return false } }, { + supportsWriteCallback: true, + writableLength: () => process.stdout.writableLength, + writableHighWaterMark: () => process.stdout.writableHighWaterMark, waitWriteDrain: (cb) => { if (!stdoutAlive) { cb() return } stdoutDrainWaiters.add(cb) + return () => stdoutDrainWaiters.delete(cb) } } ) @@ -640,19 +726,27 @@ async function main(): Promise { sock.on('close', flushSockDrainWaiters) sock.on('error', flushSockDrainWaiters) const clientId = dispatcher.attachClient( - (data) => { + (data, onSettled) => { if (!sock.destroyed) { - return sock.write(data) + return sock.write(data, (error) => { + onSettled(error ? { ok: false, error } : { ok: true }) + }) } - return undefined + onSettled({ ok: false, error: new Error('Relay socket is closed') }) + return false }, { + supportsWriteCallback: true, + writableLength: () => sock.writableLength, + writableHighWaterMark: () => sock.writableHighWaterMark, + close: () => sock.destroy(), waitWriteDrain: (cb) => { if (sock.destroyed) { cb() return } sockDrainWaiters.add(cb) + return () => sockDrainWaiters.delete(cb) } } ) @@ -865,7 +959,6 @@ async function main(): Promise { if (detached) { // Why: detached stdin is /dev/null, so listening would EOF → grace → shutdown before --connect arrives; use the socket instead. - stdoutAlive = false startGrace('detached startup') } else { process.stdin.on('data', (chunk: Buffer) => { @@ -937,8 +1030,11 @@ async function main(): Promise { relayLogLine(`[relay] Process exiting with code ${code}`) }) - // Why: the client waits for this exact sentinel string before sending framed data. - process.stdout.write(RELAY_SENTINEL) + dispatcher.writePrimaryBytes(Buffer.from(RELAY_SENTINEL)) + if (detached) { + stdoutAlive = false + dispatcher.invalidateClient() + } } function cleanupSocket(sockPath: string): void { diff --git a/src/relay/wsl-agent-hook-relay.ts b/src/relay/wsl-agent-hook-relay.ts index bda345806b9..534c388e07c 100644 --- a/src/relay/wsl-agent-hook-relay.ts +++ b/src/relay/wsl-agent-hook-relay.ts @@ -38,12 +38,26 @@ async function main(): Promise { } let stdoutAlive = true - const dispatcher = new RelayDispatcher((data) => { - if (!stdoutAlive) { - return + const dispatcher = new RelayDispatcher( + (data, onSettled) => { + if (!stdoutAlive) { + onSettled({ ok: false, error: new Error('WSL relay stdout is closed') }) + return false + } + return process.stdout.write(data, (error) => { + onSettled(error ? { ok: false, error } : { ok: true }) + }) + }, + { + supportsWriteCallback: true, + writableLength: () => process.stdout.writableLength, + writableHighWaterMark: () => process.stdout.writableHighWaterMark, + waitWriteDrain: (callback) => { + process.stdout.once('drain', callback) + return () => process.stdout.off('drain', callback) + } } - return process.stdout.write(data) - }) + ) // Why: restart-stable instance key keeps the endpoint file at one path // across app restarts so surviving agents re-coordinate off its rewrite. @@ -120,7 +134,7 @@ async function main(): Promise { // Signal readiness — the host watches for this exact string before // sending framed data (same contract as the SSH relay). - process.stdout.write(RELAY_SENTINEL) + dispatcher.writePrimaryBytes(Buffer.from(RELAY_SENTINEL)) } void main()