[P2] perf(relay): stop re-encoding frames to size PTY chunks (#11620)
* perf(relay): stop re-encoding frames to size PTY chunks maxLegacyPtyDataChars now sizes the full candidate once (the common case) and falls back to a binary search over an exact cheap byte formula instead of fully encoding the frame at every probe. Publish paths thread their already-computed frame estimate into enqueueFrame so each PTY publish encodes the message once for admission instead of twice. Co-authored-by: Orca <help@stably.ai> * fix(relay): preserve dispatcher frame guards --------- Co-authored-by: Orca <help@stably.ai>
This commit is contained in:
Neil
GitHub
parent
854f76cb51
commit
83af12dc3e
|
|
@ -0,0 +1,47 @@
|
|||
import { describe, expect, it, vi } from 'vitest'
|
||||
import { RelayDispatcher } from './dispatcher'
|
||||
import type { JsonRpcNotification } from './protocol'
|
||||
|
||||
type DispatcherInternals = {
|
||||
primaryClient: object
|
||||
estimateFrameBytes: (msg: JsonRpcNotification) => number
|
||||
enqueueFrame: (client: object, msg: JsonRpcNotification, lane: string) => boolean
|
||||
}
|
||||
|
||||
describe('RelayDispatcher frame guards', () => {
|
||||
it('returns zero for invalid active-client limits without encoding', () => {
|
||||
const dispatcher = new RelayDispatcher(() => true, {
|
||||
writableHighWaterMark: () => 1024 * 1024,
|
||||
writableLength: () => 0
|
||||
})
|
||||
try {
|
||||
const spy = vi.spyOn(dispatcher as unknown as DispatcherInternals, 'estimateFrameBytes')
|
||||
for (const limit of [0, -1, Number.NaN]) {
|
||||
expect(dispatcher.maxLegacyPtyDataChars({ id: 'pty-1' }, 'hello', limit)).toBe(0)
|
||||
}
|
||||
expect(spy).not.toHaveBeenCalled()
|
||||
} finally {
|
||||
dispatcher.dispose()
|
||||
}
|
||||
})
|
||||
|
||||
it('does not estimate frames after disposal', () => {
|
||||
const dispatcher = new RelayDispatcher(() => true)
|
||||
const internals = dispatcher as unknown as DispatcherInternals
|
||||
const spy = vi.spyOn(internals, 'estimateFrameBytes')
|
||||
const msg: JsonRpcNotification = {
|
||||
jsonrpc: '2.0',
|
||||
method: 'pty.data',
|
||||
params: {
|
||||
data: {
|
||||
toJSON: () => {
|
||||
throw new Error('must not serialize')
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
dispatcher.dispose()
|
||||
expect(internals.enqueueFrame(internals.primaryClient, msg, 'ordinary')).toBe(false)
|
||||
expect(spy).not.toHaveBeenCalled()
|
||||
})
|
||||
})
|
||||
|
|
@ -1,5 +1,6 @@
|
|||
import { describe, expect, it, vi, beforeEach, afterEach } from 'vitest'
|
||||
import { RelayDispatcher, type SinkWriteSettlement } from './dispatcher'
|
||||
import { relayWriterControlReserve } from './dispatcher-writer-admission'
|
||||
import {
|
||||
encodeJsonRpcFrame,
|
||||
encodeKeepAliveFrame,
|
||||
|
|
@ -716,4 +717,196 @@ describe('RelayDispatcher', () => {
|
|||
}
|
||||
})
|
||||
})
|
||||
|
||||
describe('legacy PTY chunk sizing', () => {
|
||||
type DispatcherInternals = {
|
||||
primaryClient: object
|
||||
estimateFrameBytes: (msg: JsonRpcNotification) => number
|
||||
enqueueFrame: (
|
||||
client: object,
|
||||
msg: JsonRpcNotification,
|
||||
lane: string,
|
||||
onSettled?: (result: SinkWriteSettlement) => void,
|
||||
estimatedBytes?: number
|
||||
) => boolean
|
||||
}
|
||||
|
||||
// Pre-optimization sizing loop, kept verbatim for differential verification.
|
||||
function referenceMaxChars(
|
||||
capacities: number[],
|
||||
params: Record<string, unknown>,
|
||||
data: string,
|
||||
limit: number
|
||||
): number {
|
||||
if (capacities.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 = encodeJsonRpcFrame(msg, 0, 0).length
|
||||
if (capacities.every((capacity) => bytes <= capacity)) {
|
||||
low = mid
|
||||
} else {
|
||||
high = mid - 1
|
||||
}
|
||||
}
|
||||
return low
|
||||
}
|
||||
|
||||
function makeDispatcher(highWaterMarks: number[]): {
|
||||
sized: RelayDispatcher
|
||||
capacities: number[]
|
||||
} {
|
||||
const [primaryHwm, ...rest] = highWaterMarks
|
||||
const sized = new RelayDispatcher(() => true, {
|
||||
writableHighWaterMark: () => primaryHwm,
|
||||
writableLength: () => 0
|
||||
})
|
||||
for (const hwm of rest) {
|
||||
sized.attachClient(() => true, {
|
||||
writableHighWaterMark: () => hwm,
|
||||
writableLength: () => 0
|
||||
})
|
||||
}
|
||||
return {
|
||||
sized,
|
||||
capacities: highWaterMarks.map((hwm) => Math.max(0, hwm - relayWriterControlReserve(hwm)))
|
||||
}
|
||||
}
|
||||
|
||||
function mulberry32(seed: number): () => number {
|
||||
let a = seed
|
||||
return () => {
|
||||
a = (a + 0x6d2b79f5) | 0
|
||||
let t = Math.imul(a ^ (a >>> 15), 1 | a)
|
||||
t = (t + Math.imul(t ^ (t >>> 7), 61 | t)) ^ t
|
||||
return ((t ^ (t >>> 14)) >>> 0) / 4294967296
|
||||
}
|
||||
}
|
||||
|
||||
// Multi-byte UTF-8, astral pairs, lone surrogates, controls, quotes, and backslashes.
|
||||
const alphabet = 'aZ9 "\\\n\r\t é߀中𝄞😀𐀀<F09F9880>'
|
||||
|
||||
it('matches the pre-optimization sizing loop across randomized inputs', () => {
|
||||
const random = mulberry32(0xc0ffee)
|
||||
const hwmChoices = [1030, 1100, 1250, 1400, 2048, 1024 * 1024]
|
||||
for (let trial = 0; trial < 300; trial++) {
|
||||
const length = Math.floor(random() * 240)
|
||||
let data = ''
|
||||
for (let i = 0; i < length; i++) {
|
||||
data += alphabet[Math.floor(random() * alphabet.length)]
|
||||
}
|
||||
const clientCount = 1 + Math.floor(random() * 2)
|
||||
const hwms = Array.from(
|
||||
{ length: clientCount },
|
||||
() => hwmChoices[Math.floor(random() * hwmChoices.length)]
|
||||
)
|
||||
const params = { id: `pty-${trial}`, seq: trial }
|
||||
const { sized, capacities } = makeDispatcher(hwms)
|
||||
try {
|
||||
for (const limit of [0, 1, Math.floor(length / 2), length, length + 17]) {
|
||||
expect(sized.maxLegacyPtyDataChars(params, data, limit)).toBe(
|
||||
referenceMaxChars(capacities, params, data, limit)
|
||||
)
|
||||
}
|
||||
expect(sized.maxLegacyPtyDataChars(params, data)).toBe(
|
||||
referenceMaxChars(capacities, params, data, data.length)
|
||||
)
|
||||
} finally {
|
||||
sized.dispose()
|
||||
}
|
||||
}
|
||||
})
|
||||
|
||||
it('sizes a chunk that fits every client with a single frame encode', () => {
|
||||
const { sized } = makeDispatcher([1024 * 1024])
|
||||
try {
|
||||
const spy = vi.spyOn(sized as unknown as DispatcherInternals, 'estimateFrameBytes')
|
||||
const data = 'x'.repeat(16 * 1024)
|
||||
expect(sized.maxLegacyPtyDataChars({ id: 'pty-1' }, data)).toBe(data.length)
|
||||
expect(spy).toHaveBeenCalledTimes(1)
|
||||
} finally {
|
||||
sized.dispose()
|
||||
}
|
||||
})
|
||||
|
||||
it('publishes PTY data with a single frame estimate', () => {
|
||||
const frames: Buffer[] = []
|
||||
const publisher = new RelayDispatcher((data) => {
|
||||
frames.push(Buffer.from(data))
|
||||
return true
|
||||
})
|
||||
try {
|
||||
const spy = vi.spyOn(publisher as unknown as DispatcherInternals, 'estimateFrameBytes')
|
||||
expect(publisher.tryNotifyPtyData({ id: 'pty-1', data: 'hello' })).toBe(true)
|
||||
expect(frames).toHaveLength(1)
|
||||
expect(spy).toHaveBeenCalledTimes(1)
|
||||
} finally {
|
||||
publisher.dispose()
|
||||
}
|
||||
})
|
||||
|
||||
it('enqueueFrame with a caller-supplied estimate matches the computed default', () => {
|
||||
const frames: Buffer[] = []
|
||||
const publisher = new RelayDispatcher((data) => {
|
||||
frames.push(Buffer.from(data))
|
||||
return true
|
||||
})
|
||||
try {
|
||||
const internals = publisher as unknown as DispatcherInternals
|
||||
const msg: JsonRpcNotification = {
|
||||
jsonrpc: '2.0',
|
||||
method: 'pty.data',
|
||||
params: { id: 'pty-1', data: 'héllo "𝄞"\\\n\uD800' }
|
||||
}
|
||||
expect(internals.enqueueFrame(internals.primaryClient, msg, 'ordinary')).toBe(true)
|
||||
expect(
|
||||
internals.enqueueFrame(
|
||||
internals.primaryClient,
|
||||
msg,
|
||||
'ordinary',
|
||||
undefined,
|
||||
internals.estimateFrameBytes(msg)
|
||||
)
|
||||
).toBe(true)
|
||||
expect(frames).toHaveLength(2)
|
||||
expect(
|
||||
decodeFirstFrame(frames[1]).payload.equals(decodeFirstFrame(frames[0]).payload)
|
||||
).toBe(true)
|
||||
} finally {
|
||||
publisher.dispose()
|
||||
}
|
||||
})
|
||||
|
||||
it('enqueueFrame rejects identically with and without a caller-supplied estimate', () => {
|
||||
const { sized } = makeDispatcher([1030])
|
||||
try {
|
||||
const internals = sized as unknown as DispatcherInternals
|
||||
const msg: JsonRpcNotification = {
|
||||
jsonrpc: '2.0',
|
||||
method: 'pty.data',
|
||||
params: { id: 'pty-1', data: 'x'.repeat(512) }
|
||||
}
|
||||
expect(internals.enqueueFrame(internals.primaryClient, msg, 'ordinary')).toBe(false)
|
||||
expect(
|
||||
internals.enqueueFrame(
|
||||
internals.primaryClient,
|
||||
msg,
|
||||
'ordinary',
|
||||
undefined,
|
||||
internals.estimateFrameBytes(msg)
|
||||
)
|
||||
).toBe(false)
|
||||
} finally {
|
||||
sized.dispose()
|
||||
}
|
||||
})
|
||||
})
|
||||
})
|
||||
|
|
|
|||
|
|
@ -193,20 +193,34 @@ export class RelayDispatcher {
|
|||
limit = data.length
|
||||
): number {
|
||||
const clients = this.activeClients()
|
||||
const max = Math.min(data.length, limit)
|
||||
if (clients.length === 0) {
|
||||
return Math.min(data.length, limit)
|
||||
return max
|
||||
}
|
||||
let low = 0
|
||||
let high = Math.min(data.length, limit)
|
||||
while (low < high) {
|
||||
const mid = Math.ceil((low + high) / 2)
|
||||
const msg: JsonRpcNotification = {
|
||||
if (!(max > 0)) {
|
||||
return 0
|
||||
}
|
||||
const fitsAll = (bytes: number): boolean =>
|
||||
clients.every((client) => bytes <= client.writer.producerFrameCapacity)
|
||||
const sizeFrame = (chunk: string): number =>
|
||||
this.estimateFrameBytes({
|
||||
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)) {
|
||||
params: { ...params, data: chunk }
|
||||
})
|
||||
// Fast path: the whole chunk usually fits — one encode instead of log2(n).
|
||||
if (fitsAll(sizeFrame(data.slice(0, max)))) {
|
||||
return max
|
||||
}
|
||||
// Exact per-step size: only the escaped data string varies; its quotes are in baseBytes.
|
||||
const baseBytes = sizeFrame('')
|
||||
const bytesFor = (chars: number): number =>
|
||||
baseBytes + Buffer.byteLength(JSON.stringify(data.slice(0, chars))) - 2
|
||||
let low = 0
|
||||
let high = max
|
||||
while (low < high) {
|
||||
const mid = Math.ceil((low + high) / 2)
|
||||
if (fitsAll(bytesFor(mid))) {
|
||||
low = mid
|
||||
} else {
|
||||
high = mid - 1
|
||||
|
|
@ -822,19 +836,21 @@ export class RelayDispatcher {
|
|||
client: RelayClient,
|
||||
msg: JsonRpcRequest | JsonRpcResponse | JsonRpcNotification,
|
||||
lane: DispatcherWriterLane,
|
||||
onSettled: (result: SinkWriteSettlement) => void = () => {}
|
||||
onSettled: (result: SinkWriteSettlement) => void = () => {},
|
||||
// Why: publish paths already sized the frame; avoid a redundant encode.
|
||||
estimatedBytes?: number
|
||||
): boolean {
|
||||
if (this.disposed || client.closed) {
|
||||
return false
|
||||
}
|
||||
const estimatedBytes = this.estimateFrameBytes(msg)
|
||||
const frameBytes = estimatedBytes ?? this.estimateFrameBytes(msg)
|
||||
return client.writer.enqueue(
|
||||
lane,
|
||||
() => {
|
||||
const seq = client.nextOutgoingSeq++
|
||||
return encodeJsonRpcFrame(msg, seq, client.highestReceivedSeq)
|
||||
},
|
||||
estimatedBytes,
|
||||
frameBytes,
|
||||
onSettled
|
||||
)
|
||||
}
|
||||
|
|
@ -898,7 +914,7 @@ export class RelayDispatcher {
|
|||
return false
|
||||
}
|
||||
for (let index = 0; index < clients.length; index++) {
|
||||
if (!this.enqueueLeasedFrame(clients[index], msg, lane, leases[index])) {
|
||||
if (!this.enqueueLeasedFrame(clients[index], msg, lane, leases[index], bytes)) {
|
||||
if (this.disposed || clients[index].closed) {
|
||||
continue
|
||||
}
|
||||
|
|
@ -949,7 +965,7 @@ export class RelayDispatcher {
|
|||
if (!leases) {
|
||||
return false
|
||||
}
|
||||
return this.enqueueLeasedFrame(client, msg, lane, leases[0], onSettled)
|
||||
return this.enqueueLeasedFrame(client, msg, lane, leases[0], bytes, onSettled)
|
||||
}
|
||||
|
||||
private publishBulkWhenAvailable(client: RelayClient, msg: JsonRpcNotification): Promise<void> {
|
||||
|
|
@ -998,13 +1014,20 @@ export class RelayDispatcher {
|
|||
msg: JsonRpcNotification,
|
||||
lane: 'interactive' | 'ordinary' | 'fixed-bulk' | 'bulk',
|
||||
lease: LegacyPublicationLease,
|
||||
estimatedBytes: number,
|
||||
onSettled: (result: SinkWriteSettlement) => void = () => {}
|
||||
): boolean {
|
||||
const accepted = this.enqueueFrame(client, msg, lane, (result) => {
|
||||
lease.release()
|
||||
onSettled(result)
|
||||
this.notifyLegacyCapacityIfLow()
|
||||
})
|
||||
const accepted = this.enqueueFrame(
|
||||
client,
|
||||
msg,
|
||||
lane,
|
||||
(result) => {
|
||||
lease.release()
|
||||
onSettled(result)
|
||||
this.notifyLegacyCapacityIfLow()
|
||||
},
|
||||
estimatedBytes
|
||||
)
|
||||
if (!accepted) {
|
||||
lease.release()
|
||||
this.notifyLegacyCapacityIfLow()
|
||||
|
|
|
|||
Loading…
Reference in New Issue