mirror of
https://github.com/stablyai/orca.git
synced 2026-10-20 00:03:45 +00:00
fix(relay): bound PTY output publication
This commit is contained in:
OrcaWin
@@ -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<typeof vi.fn> {
|
||||
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'])
|
||||
})
|
||||
})
|
||||
@@ -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<DispatcherWriterEntry>()
|
||||
private readonly settleOnDrain = new Set<DispatcherWriterEntry>()
|
||||
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<void> {
|
||||
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()
|
||||
}
|
||||
}
|
||||
}
|
||||
@@ -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<DispatcherWriterLane, DispatcherWriterEntry[]> = {
|
||||
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 }
|
||||
}
|
||||
}
|
||||
@@ -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
|
||||
)
|
||||
}
|
||||
}
|
||||
}
|
||||
@@ -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?.()
|
||||
}
|
||||
}
|
||||
@@ -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
|
||||
|
||||
+388
-110
@@ -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<string, unknown>, 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<void>
|
||||
nextOutgoingSeq: number
|
||||
highestReceivedSeq: number
|
||||
@@ -64,8 +71,13 @@ export class RelayDispatcher {
|
||||
private requestHandlers = new Map<string, MethodHandler>()
|
||||
private notificationHandlers = new Map<string, NotificationHandler>()
|
||||
private readonly requestAborts = new ClientRequestAborts()
|
||||
private readonly publicationLedger = new LegacyRelayPublicationLedger()
|
||||
private pendingRelayRequests = new Map<number, PendingRelayRequest>()
|
||||
private clientDetachListeners = new Set<(clientId: number) => void>()
|
||||
private legacyCapacityListeners = new Set<() => void>()
|
||||
private publicationTransactionDepth = 0
|
||||
private deferredLegacyCapacity = false
|
||||
private deferredForcedLegacyCapacity = false
|
||||
private keepaliveTimer: ReturnType<typeof setInterval> | 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<string, unknown>,
|
||||
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<string, unknown>,
|
||||
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<string, unknown>): boolean {
|
||||
if (this.disposed) {
|
||||
return false
|
||||
}
|
||||
return this.tryPublishToClients(
|
||||
this.activeClients(),
|
||||
{
|
||||
jsonrpc: '2.0',
|
||||
method: 'pty.exit',
|
||||
params
|
||||
},
|
||||
'ordinary'
|
||||
)
|
||||
}
|
||||
|
||||
producerDataBudget(
|
||||
method: string,
|
||||
paramsWithoutData: Record<string, unknown>,
|
||||
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<string, unknown>): 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<void> {
|
||||
if (this.disposed || client.closed || !client.waitWriteDrain) {
|
||||
return Promise.resolve()
|
||||
}
|
||||
return new Promise<void>((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<string, unknown>,
|
||||
@@ -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<void> {
|
||||
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<void>((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<T>(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)
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
|
||||
@@ -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'
|
||||
|
||||
@@ -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(() => {
|
||||
|
||||
@@ -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)
|
||||
})
|
||||
})
|
||||
@@ -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<string, number>()
|
||||
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<string, number>()
|
||||
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)
|
||||
}
|
||||
}
|
||||
}
|
||||
@@ -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()
|
||||
}
|
||||
}))
|
||||
|
||||
|
||||
@@ -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
|
||||
|
||||
+224
-53
@@ -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<typeof setTimeout> | null = null
|
||||
private outputFlushTimer: ReturnType<typeof setTimeout> | null = null
|
||||
private pendingOutputByPty = new Map<string, PendingPtyOutput>()
|
||||
private pendingOutputByPty = new Map<string, PendingPtyOutput[]>()
|
||||
private pendingExitByPty = new Map<string, { id: string; code: number; incarnationId: string }>()
|
||||
private pausedOutputPtys = new Set<string>()
|
||||
private removeLegacyCapacityListener: (() => void) | null = null
|
||||
private lastInputAtByPty = new Map<string, number>()
|
||||
private interactiveOutputCharsByPty = new Map<string, number>()
|
||||
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<typeof NodePty | null> {
|
||||
@@ -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(
|
||||
|
||||
@@ -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()
|
||||
}
|
||||
})
|
||||
})
|
||||
@@ -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,
|
||||
|
||||
+120
-24
@@ -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<void> {
|
||||
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<void> {
|
||||
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<void> {
|
||||
}
|
||||
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<void> {
|
||||
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<void> {
|
||||
|
||||
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<void> {
|
||||
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 {
|
||||
|
||||
@@ -38,12 +38,26 @@ async function main(): Promise<void> {
|
||||
}
|
||||
|
||||
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<void> {
|
||||
|
||||
// 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()
|
||||
|
||||
Reference in New Issue
Block a user