From 9aeb5848911f0fbb65b5530c095be4082ec77d5f Mon Sep 17 00:00:00 2001 From: Armaan Aggarwal <49625265+Armster15@users.noreply.github.com> Date: Wed, 14 Oct 2026 22:38:36 +0000 Subject: [PATCH 1/3] transport: add opt-in receiver ACKs and per-stream backpressure Assisted-by: Replit --- PROTOCOL.md | 36 ++++ __tests__/backpressure.test.ts | 47 ++++- .../properties/backpressure.property.test.ts | 74 ++++++++ __tests__/receiverAck.test.ts | 64 +++++++ router/client.ts | 6 +- router/server.ts | 1 + router/streams.ts | 4 +- .../sessionStateMachine/SessionConnected.ts | 38 +++- transport/sessionStateMachine/common.ts | 99 ++++++++++- .../sessionStateMachine/flowControl.test.ts | 163 ++++++++++++++++++ transport/sessionStateMachine/transitions.ts | 1 + transport/transport.ts | 26 ++- 12 files changed, 537 insertions(+), 22 deletions(-) create mode 100644 __tests__/properties/backpressure.property.test.ts create mode 100644 __tests__/receiverAck.test.ts create mode 100644 transport/sessionStateMachine/flowControl.test.ts diff --git a/PROTOCOL.md b/PROTOCOL.md index 453b50e2..fc27d332 100644 --- a/PROTOCOL.md +++ b/PROTOCOL.md @@ -669,6 +669,42 @@ When sending messages to the other side, the session associated with the connect When receiving messages, the transport MUST ensure that the only messages it validates are those that have a `seq` that is exactly equal to the `ack` of the session. This ensures exactly-once delivery semantics. +#### Optional receiver-driven acknowledgements + +A receiver MAY send an explicit ACK after a configured number of inbound +non-ACK messages or a maximum delay, whichever occurs first. The TypeScript +implementation enables this with `receiverAck: { maxMessages, maxDelayMs }`; +it is disabled by default. An outgoing message carries the current cumulative +ACK and cancels the pending batch. Replayed messages retain their original +ACK field and do not cancel a pending batch. + +These ACKs use the existing `heartbeat` stream, `AckBit`, and `{ type: 'ACK' }` +payload. They participate in sequencing and replay exactly like heartbeat +ACKs. There is no new wire field, capability negotiation, or protocol version. +An older active-heartbeat peer consumes them without replying; an older +passive-heartbeat peer responds with its normal heartbeat ACK. ACK messages +MUST NOT count toward the receiver-driven threshold or arm its timer. The +existing active/passive heartbeat roles remain unchanged, avoiding ACK loops. +Either endpoint may enable receiver-driven ACKs independently. Timers are +cancelled when a connected state exits; heartbeat ACKs still guarantee progress +after reconnect and when receiver-driven ACKs are disabled. + +#### Optional per-stream advisory backpressure + +`streamSendBufferHighWaterMark` opts the TypeScript router's writers into a +per-stream message-count watermark instead of the shared session watermark. +Each writer counts all unacknowledged messages on its stream, including open +and close frames. A busy stream does not consume another stream's allowance. +Cumulative session ACKs free the corresponding stream slots; reconnect carries +both buffered messages and pending drain waiters. Session destruction releases +all waiters. + +This is a local producer policy, not a wire change or a hard memory limit. +Writes remain accepted even when they return `false`. Many streams may +collectively exceed the session watermark in this mode; callers must bound +stream concurrency and retained bytes separately. Unscoped transport accessors +and transports without the option retain session-level backpressure. + #### On disconnect Handling close events is detailed in the 'Creating Connections and Sessions' section above. diff --git a/__tests__/backpressure.test.ts b/__tests__/backpressure.test.ts index d48c69ed..58660158 100644 --- a/__tests__/backpressure.test.ts +++ b/__tests__/backpressure.test.ts @@ -19,12 +19,19 @@ import { TestSetupHelpers } from '../testUtil/fixtures/transports'; const HIGH_WATER_MARK = 4; -describe.each(testMatrix())( - 'writable backpressure ($transport.name transport, $codec.name codec)', - async ({ transport, codec }) => { +describe.each( + testMatrix().flatMap((entry) => + ['session', 'stream'].map((mode) => ({ ...entry, mode })), + ), +)( + '$mode writable backpressure ($transport.name transport, $codec.name codec)', + async ({ transport, codec, mode }) => { const opts = { codec: codec.codec, sendBufferHighWaterMark: HIGH_WATER_MARK, + ...(mode === 'stream' + ? { streamSendBufferHighWaterMark: HIGH_WATER_MARK } + : {}), }; const { addPostTestCleanup, postTestCleanup } = createPostTestCleanups(); @@ -195,5 +202,39 @@ describe.each(testMatrix())( payload: expect.objectContaining({ code: UNEXPECTED_DISCONNECT_CODE }), }); }); + + test('a heavy stream only pressures a quiet stream in session mode', async () => { + const clientTransport = getClientTransport('client'); + const serverTransport = getServerTransport(); + const services = { test: TestServiceSchema }; + createServer(serverTransport, services); + const client = createClient( + clientTransport, + serverTransport.clientId, + ); + addPostTestCleanup(async () => { + await cleanupTransports([clientTransport, serverTransport]); + }); + const heavy = client.test.echo.stream({}); + const quiet = client.test.echo.stream({}); + for (let i = 0; i < HIGH_WATER_MARK + 1; i++) { + heavy.reqWritable.write({ msg: 'heavy', ignore: true }); + } + expect(heavy.reqWritable.write({ msg: 'heavy', ignore: true })).toBe( + false, + ); + expect(quiet.reqWritable.write({ msg: 'quiet', ignore: false })).toBe( + mode === 'stream', + ); + if (mode === 'stream') { + await quiet.reqWritable.waitForWriteReady(); + } + expect(await readNextResult(quiet.resReadable)).toStrictEqual({ + ok: true, + payload: { response: 'quiet' }, + }); + heavy.reqWritable.close(); + quiet.reqWritable.close(); + }); }, ); diff --git a/__tests__/properties/backpressure.property.test.ts b/__tests__/properties/backpressure.property.test.ts new file mode 100644 index 00000000..23d013d4 --- /dev/null +++ b/__tests__/properties/backpressure.property.test.ts @@ -0,0 +1,74 @@ +import { expect, test } from 'vitest'; +import * as hegel from '@hegeldev/hegel'; +import * as gs from '@hegeldev/hegel/generators'; +import { createClient } from '../../router/client'; +import { createServer } from '../../router/server'; +import { closeAllConnections, readNextResult } from '../../testUtil'; +import { cleanupTransports } from '../../testUtil/fixtures/cleanup'; +import { createMockTransportNetwork } from '../../testUtil/fixtures/mockTransport'; +import { TestServiceSchema } from '../../testUtil/fixtures/services'; + +test( + 'generated heavy feeds cannot park an unrelated stream across reconnect', + () => + hegel.testAsync( + async (tc) => { + const mark = tc.draw(gs.integers({ minValue: 3, maxValue: 12 })); + const count = tc.draw( + gs.integers({ minValue: mark, maxValue: mark * 3 }), + ); + const reconnect = tc.draw(gs.booleans()); + const options = { + sendBufferHighWaterMark: mark, + streamSendBufferHighWaterMark: mark, + receiverAck: { maxMessages: 3, maxDelayMs: 20 }, + maxJitterMs: 0, + }; + const network = createMockTransportNetwork({ + client: options, + server: options, + }); + const clientTransport = network.getClientTransport('client'); + const serverTransport = network.getServerTransport(); + const services = { test: TestServiceSchema }; + createServer(serverTransport, services); + const client = createClient( + clientTransport, + serverTransport.clientId, + ); + try { + const heavy = client.test.echo.stream({}); + heavy.reqWritable.write({ msg: 'establish', ignore: false }); + await readNextResult(heavy.resReadable); + if (reconnect) closeAllConnections(clientTransport); + for (let i = 0; i < count; i++) { + heavy.reqWritable.write({ msg: String(i), ignore: false }); + } + const heavyReady = heavy.reqWritable.waitForWriteReady(); + const quiet = client.test.echo.stream({}); + expect(quiet.reqWritable.write({ msg: 'quiet', ignore: false })).toBe( + true, + ); + await quiet.reqWritable.waitForWriteReady(); + expect(await readNextResult(quiet.resReadable)).toStrictEqual({ + ok: true, + payload: { response: 'quiet' }, + }); + await heavyReady; + for (let i = 0; i < count; i++) { + expect(await readNextResult(heavy.resReadable)).toStrictEqual({ + ok: true, + payload: { response: String(i) }, + }); + } + heavy.reqWritable.close(); + quiet.reqWritable.close(); + } finally { + await cleanupTransports([clientTransport, serverTransport]); + await network.cleanup(); + } + }, + { testCases: 20 }, + ), + 30_000, +); diff --git a/__tests__/receiverAck.test.ts b/__tests__/receiverAck.test.ts new file mode 100644 index 00000000..0b071e7d --- /dev/null +++ b/__tests__/receiverAck.test.ts @@ -0,0 +1,64 @@ +import { describe, expect, test } from 'vitest'; +import { createClient } from '../router/client'; +import { createServer } from '../router/server'; +import { readNextResult } from '../testUtil'; +import { cleanupTransports, waitFor } from '../testUtil/fixtures/cleanup'; +import { testMatrix } from '../testUtil/fixtures/matrix'; +import { TestServiceSchema } from '../testUtil/fixtures/services'; + +describe.each( + testMatrix().flatMap((entry) => + ['client', 'server', 'both'].map((enabled) => ({ ...entry, enabled })), + ), +)( + 'receiver ACKs: $enabled ($transport.name, $codec.name)', + ({ transport, codec, enabled }) => { + test('an opt-in receiver drains one-way traffic with a legacy-mode peer', async () => { + const receiverAck = { maxMessages: 3, maxDelayMs: 20 }; + const common = { + codec: codec.codec, + heartbeatIntervalMs: 30_000, + sessionDisconnectGraceMs: 90_000, + }; + const setup = await transport.setup({ + client: { ...common, ...(enabled !== 'server' ? { receiverAck } : {}) }, + server: { ...common, ...(enabled !== 'client' ? { receiverAck } : {}) }, + }); + const clientTransport = setup.getClientTransport('client'); + const serverTransport = setup.getServerTransport(); + const services = { test: TestServiceSchema }; + createServer(serverTransport, services); + const client = createClient( + clientTransport, + serverTransport.clientId, + ); + try { + const call = client.test.echo.stream({}); + for (let i = 0; i < 7; i++) { + call.reqWritable.write({ + msg: `message-${i}`, + ignore: enabled === 'server', + }); + } + if (enabled !== 'server') { + for (let i = 0; i < 7; i++) { + expect(await readNextResult(call.resReadable)).toStrictEqual({ + ok: true, + payload: { response: `message-${i}` }, + }); + } + } + const sender = enabled === 'server' ? clientTransport : serverTransport; + await waitFor(() => { + const session = [...sender.sessions.values()][0]; + expect(session).toBeDefined(); + expect(session.sendBuffer.length).toBeLessThanOrEqual(1); + }); + call.reqWritable.close(); + } finally { + await cleanupTransports([clientTransport, serverTransport]); + await setup.cleanup(); + } + }); + }, +); diff --git a/router/client.ts b/router/client.ts index a14d4454..e64966cc 100644 --- a/router/client.ts +++ b/router/client.ts @@ -342,13 +342,13 @@ function handleProc( serverId, session.id, ); + const procClosesWithInit = procType === 'rpc' || procType === 'subscription'; + const streamId = generateId(); const sessionBackpressure = transport.getSessionBackpressure( serverId, session.id, + streamId, ); - - const procClosesWithInit = procType === 'rpc' || procType === 'subscription'; - const streamId = generateId(); const { span, ctx } = createProcTelemetryInfo( transport.tracer, session, diff --git a/router/server.ts b/router/server.ts index b8ba4616..f29f9552 100644 --- a/router/server.ts +++ b/router/server.ts @@ -464,6 +464,7 @@ class RiverServer< const sessionBackpressure = this.transport.getSessionBackpressure( from, sessionId, + streamId, ); const cancelStream = ( diff --git a/router/streams.ts b/router/streams.ts index f81ed27f..93cd6bd5 100644 --- a/router/streams.ts +++ b/router/streams.ts @@ -92,7 +92,7 @@ export interface Writable { /** * {@link write} writes a value to the pipe. An error is thrown if writing to a closed {@link Writable}. * - * Returns `false` if the underlying session's send buffer is at or above + * Returns `false` if the underlying session or stream send buffer is at or above * its high-water mark, signalling that the producer should stop writing * until {@link waitForWriteReady} resolves. This is purely advisory — * the value is still buffered and will be delivered, exactly like @@ -100,7 +100,7 @@ export interface Writable { */ write(value: T): boolean; /** - * {@link waitForWriteReady} resolves once the underlying session's send + * {@link waitForWriteReady} resolves once the underlying session or stream send * buffer has drained back below its high-water mark and it is productive * to call {@link write} again. * diff --git a/transport/sessionStateMachine/SessionConnected.ts b/transport/sessionStateMachine/SessionConnected.ts index f0ce6a8b..e62b2be6 100644 --- a/transport/sessionStateMachine/SessionConnected.ts +++ b/transport/sessionStateMachine/SessionConnected.ts @@ -60,15 +60,13 @@ export class SessionConnected< private isActivelyHeartbeating = false; private rehandshakeTimer?: ReturnType | undefined; private credentialExpiry?: number | undefined; + private receiverAckTimer?: ReturnType; + private pendingReceiverAckMessages = 0; updateBookkeeping(ack: number, seq: number) { - this.sendBuffer = this.sendBuffer.filter((unacked) => unacked.seq >= ack); + this.acknowledgeSendBuffer(ack); this.ack = seq + 1; this.lastInboundAt = Date.now(); - - if (this.sendBuffer.length < this.options.sendBufferHighWaterMark) { - this.notifySendBufferDrain(); - } } private assertSendOrdering(encodedMsg: EncodedTransportMessage) { @@ -91,7 +89,7 @@ export class SessionConnected< const encodedMsg = encodeResult.value; this.assertSendOrdering(encodedMsg); - this.sendBuffer.push(encodedMsg); + this.bufferMessage(encodedMsg); const sent = this.conn.send(encodedMsg.data); if (!sent) { @@ -105,6 +103,7 @@ export class SessionConnected< } this.seqSent = encodedMsg.seq; + this.clearReceiverAck(); return { ok: true, value: encodedMsg.id }; } @@ -218,6 +217,31 @@ export class SessionConnected< this.send(heartbeat); } + private clearReceiverAck(): void { + if (this.receiverAckTimer !== undefined) { + clearTimeout(this.receiverAckTimer); + this.receiverAckTimer = undefined; + } + this.pendingReceiverAckMessages = 0; + } + + private scheduleReceiverAck(): void { + const options = this.options.receiverAck; + if (!options) { + return; + } + + this.pendingReceiverAckMessages++; + if (this.pendingReceiverAckMessages >= options.maxMessages) { + this.sendHeartbeat(); + return; + } + this.receiverAckTimer ??= setTimeout(() => { + this.receiverAckTimer = undefined; + this.sendHeartbeat(); + }, options.maxDelayMs); + } + /** * Schedules the next proactive re-handshake from the credential's expiry. The * server calls this after each (re)validation, mirroring {@link startActiveHeartbeat}: @@ -351,6 +375,7 @@ export class SessionConnected< // dispatch directly if its not an explicit ack if (!isAck(parsedMsg.controlFlags)) { + this.scheduleReceiverAck(); // re-handshake frames ride a reserved stream and are consumed by the // transport, never surfaced to the router (same as the acks handled below) if (parsedMsg.streamId === RehandshakeStreamId) { @@ -390,6 +415,7 @@ export class SessionConnected< this.clearHeartbeatWatchdog(); this.clearRehandshakeTimer(); + this.clearReceiverAck(); } _handleClose(): void { diff --git a/transport/sessionStateMachine/common.ts b/transport/sessionStateMachine/common.ts index 96a06607..50fcad93 100644 --- a/transport/sessionStateMachine/common.ts +++ b/transport/sessionStateMachine/common.ts @@ -146,6 +146,16 @@ export interface SessionOptions { * session close. */ sendBufferHighWaterMark: number; + /** + * Optional per-stream advisory watermark. When set, stream writers use + * this limit instead of the session watermark. This is not a memory cap. + */ + streamSendBufferHighWaterMark?: number; + /** + * Opt-in ACK batching for inbound non-ACK messages. Outbound messages + * already carry an ACK and cancel any pending batch. + */ + receiverAck?: { maxMessages: number; maxDelayMs: number }; /** * The codec to use for encoding/decoding messages over the wire */ @@ -212,6 +222,7 @@ export interface IdentifiedSessionProps extends CommonSessionProps { seqSent: number; sendBuffer: Array; sendBufferDrainWaiter: PromiseWithResolvers | undefined; + streamSendBufferDrainWaiters?: Map>; telemetry: TelemetryInfo; protocolVersion: ProtocolVersion; listeners: IdentifiedSessionListeners; @@ -246,6 +257,11 @@ export abstract class IdentifiedSession extends CommonSession { * Carried across session state transitions alongside {@link sendBuffer}. */ sendBufferDrainWaiter: PromiseWithResolvers | undefined; + readonly streamSendBufferDrainWaiters: Map< + string, + PromiseWithResolvers + >; + private readonly streamSendBufferSizes = new Map(); constructor(props: IdentifiedSessionProps) { const { @@ -268,6 +284,11 @@ export abstract class IdentifiedSession extends CommonSession { this.ack = ack; this.sendBuffer = sendBuffer; this.sendBufferDrainWaiter = sendBufferDrainWaiter; + this.streamSendBufferDrainWaiters = + props.streamSendBufferDrainWaiters ?? new Map(); + for (const msg of sendBuffer) { + this.addStreamBufferEntry(msg); + } this.telemetry = telemetry; this.log = log; this.protocolVersion = protocolVersion; @@ -332,7 +353,16 @@ export abstract class IdentifiedSession extends CommonSession { return this.sendBuffer.length > 0 ? this.sendBuffer[0].seq : this.seq; } - isSendBufferFull(): boolean { + isSendBufferFull(streamId?: string): boolean { + if ( + streamId !== undefined && + this.options.streamSendBufferHighWaterMark !== undefined + ) { + return ( + (this.streamSendBufferSizes.get(streamId) ?? 0) >= + this.options.streamSendBufferHighWaterMark + ); + } return this.sendBuffer.length >= this.options.sendBufferHighWaterMark; } @@ -342,19 +372,75 @@ export abstract class IdentifiedSession extends CommonSession { * already is. Also resolves (never rejects) when the session closes so * waiting producers don't hang forever. */ - waitForSendBufferDrain(): Promise { - if (!this.isSendBufferFull()) { + waitForSendBufferDrain(streamId?: string): Promise { + if (!this.isSendBufferFull(streamId)) { return Promise.resolve(); } + if ( + streamId !== undefined && + this.options.streamSendBufferHighWaterMark !== undefined + ) { + let waiter = this.streamSendBufferDrainWaiters.get(streamId); + if (!waiter) { + waiter = createPromiseWithResolvers(); + this.streamSendBufferDrainWaiters.set(streamId, waiter); + } + return waiter.promise; + } + this.sendBufferDrainWaiter ??= createPromiseWithResolvers(); return this.sendBufferDrainWaiter.promise; } protected notifySendBufferDrain(): void { - this.sendBufferDrainWaiter?.resolve(); - this.sendBufferDrainWaiter = undefined; + if (!this.isSendBufferFull()) { + this.sendBufferDrainWaiter?.resolve(); + this.sendBufferDrainWaiter = undefined; + } + for (const [streamId, waiter] of this.streamSendBufferDrainWaiters) { + if (!this.isSendBufferFull(streamId)) { + waiter.resolve(); + this.streamSendBufferDrainWaiters.delete(streamId); + } + } + } + + private addStreamBufferEntry(msg: EncodedTransportMessage): void { + if (this.options.streamSendBufferHighWaterMark === undefined) { + return; + } + const streamId = msg.msg.streamId; + this.streamSendBufferSizes.set( + streamId, + (this.streamSendBufferSizes.get(streamId) ?? 0) + 1, + ); + } + + protected bufferMessage(msg: EncodedTransportMessage): void { + this.sendBuffer.push(msg); + this.addStreamBufferEntry(msg); + } + + protected acknowledgeSendBuffer(ack: number): void { + this.sendBuffer = this.sendBuffer.filter((msg) => { + if (msg.seq >= ack) { + return true; + } + if (this.options.streamSendBufferHighWaterMark === undefined) { + return false; + } + const streamId = msg.msg.streamId; + const remaining = (this.streamSendBufferSizes.get(streamId) ?? 1) - 1; + if (remaining === 0) { + this.streamSendBufferSizes.delete(streamId); + } else { + this.streamSendBufferSizes.set(streamId, remaining); + } + return false; + }); + this.notifySendBufferDrain(); } send(msg: PartialTransportMessage): SendResult { @@ -363,7 +449,7 @@ export abstract class IdentifiedSession extends CommonSession { return encodeResult; } - this.sendBuffer.push(encodeResult.value); + this.bufferMessage(encodeResult.value); return { ok: true, @@ -378,6 +464,7 @@ export abstract class IdentifiedSession extends CommonSession { _handleClose(): void { // zero out the buffer this.sendBuffer.length = 0; + this.streamSendBufferSizes.clear(); // wake any producers waiting for drain so they don't hang forever, // they should check isWritable/isSendBufferFull before writing again this.notifySendBufferDrain(); diff --git a/transport/sessionStateMachine/flowControl.test.ts b/transport/sessionStateMachine/flowControl.test.ts new file mode 100644 index 00000000..e5d4e22e --- /dev/null +++ b/transport/sessionStateMachine/flowControl.test.ts @@ -0,0 +1,163 @@ +import { afterEach, describe, expect, test, vi } from 'vitest'; +import { CodecMessageAdapter } from '../../codec'; +import { createSessionTelemetryInfo, getTracer } from '../../tracing'; +import { Connection } from '../connection'; +import { ControlFlags, currentProtocolVersion } from '../message'; +import { defaultTransportOptions, ProvidedTransportOptions } from '../options'; +import { SessionConnected } from './SessionConnected'; +import { SessionStateGraph } from './transitions'; + +class RecordingConnection extends Connection { + send = vi.fn(() => true); + close = vi.fn(); +} + +function connected(options: ProvidedTransportOptions = {}) { + const conn = new RecordingConnection(); + const tracer = getTracer(); + const session = new SessionConnected({ + from: 'receiver', + to: 'sender', + id: 'session', + seq: 0, + seqSent: 0, + ack: 0, + sendBuffer: [], + sendBufferDrainWaiter: undefined, + options: { ...defaultTransportOptions, ...options }, + codec: new CodecMessageAdapter(defaultTransportOptions.codec), + protocolVersion: currentProtocolVersion, + tracer, + log: undefined, + telemetry: createSessionTelemetryInfo( + tracer, + 'session', + 'sender', + 'receiver', + ), + conn, + listeners: { + onMessage: vi.fn(), + onRehandshake: vi.fn(), + onConnectionErrored: vi.fn(), + onConnectionClosed: vi.fn(), + onInvalidMessage: vi.fn(), + onMessageSendFailure: vi.fn(), + }, + }); + const receive = (controlFlags = 0, ack = 0) => { + const data = session.codec.toBuffer({ + id: `inbound-${session.ack}`, + from: 'sender', + to: 'receiver', + streamId: controlFlags ? 'heartbeat' : 'data', + controlFlags, + payload: controlFlags ? { type: 'ACK' } : {}, + seq: session.ack, + ack, + }); + if (!data.ok) throw new Error(data.reason); + session.onMessageData(data.value); + }; + return { session, conn, receive }; +} + +afterEach(() => vi.clearAllTimers()); + +describe('receiver-driven ACKs', () => { + test('count threshold sends a cumulative ACK without waiting for a heartbeat', () => { + const { session, conn, receive } = connected({ + receiverAck: { maxMessages: 3, maxDelayMs: 50 }, + }); + receive(); + receive(); + expect(conn.send).not.toHaveBeenCalled(); + receive(); + expect(conn.send).toHaveBeenCalledTimes(1); + expect(session.sendBuffer[0].msg).toMatchObject({ + streamId: 'heartbeat', + controlFlags: ControlFlags.AckBit, + payload: { type: 'ACK' }, + }); + expect(session.codec.fromBuffer(session.sendBuffer[0].data)).toMatchObject({ + ok: true, + value: { ack: 3 }, + }); + vi.advanceTimersByTime(50); + expect(conn.send).toHaveBeenCalledTimes(1); + session._handleClose(); + }); + + test('delay is bounded from the first message, not reset by later arrivals', () => { + const { session, conn, receive } = connected({ + receiverAck: { maxMessages: 10, maxDelayMs: 50 }, + }); + receive(); + vi.advanceTimersByTime(40); + receive(); + vi.advanceTimersByTime(10); + expect(conn.send).toHaveBeenCalledTimes(1); + session._handleClose(); + }); + + test('piggybacking cancels a batch and connected-state exit cancels its timer', () => { + const { session, conn, receive } = connected({ + receiverAck: { maxMessages: 10, maxDelayMs: 50 }, + }); + receive(); + session.send({ streamId: 'reply', controlFlags: 0, payload: {} }); + vi.advanceTimersByTime(50); + expect(conn.send).toHaveBeenCalledTimes(1); + receive(); + session._handleStateExit(); + vi.advanceTimersByTime(50); + expect(conn.send).toHaveBeenCalledTimes(1); + session._handleClose(); + }); + + test('ACKs do not arm batches; default mode does not send unsolicited ACKs', () => { + const enabled = connected({ + receiverAck: { maxMessages: 1, maxDelayMs: 50 }, + }); + enabled.session.startActiveHeartbeat(); + enabled.receive(ControlFlags.AckBit); + vi.advanceTimersByTime(50); + expect(enabled.conn.send).not.toHaveBeenCalled(); + enabled.session._handleClose(); + + const disabled = connected(); + disabled.receive(); + vi.advanceTimersByTime(50); + expect(disabled.conn.send).not.toHaveBeenCalled(); + disabled.session._handleClose(); + }); +}); + +test('stream waiters drain independently, survive state transitions, and settle on close', async () => { + const { session, receive } = connected({ + sendBufferHighWaterMark: 2, + streamSendBufferHighWaterMark: 2, + }); + for (const streamId of ['light', 'light', 'heavy', 'heavy', 'heavy']) { + session.send({ streamId, controlFlags: 0, payload: {} }); + } + const lightReady = session.waitForSendBufferDrain('light'); + const heavyReady = session.waitForSendBufferDrain('heavy'); + let heavyDrained = false; + void heavyReady.then(() => { + heavyDrained = true; + }); + receive(ControlFlags.AckBit, 2); + await lightReady; + expect(heavyDrained).toBe(false); + expect(session.isSendBufferFull('light')).toBe(false); + expect(session.isSendBufferFull('heavy')).toBe(true); + const disconnected = SessionStateGraph.transition.ConnectedToNoConnection( + session, + { onSessionGracePeriodElapsed: vi.fn(), onMessageSendFailure: vi.fn() }, + ); + expect(disconnected.isSendBufferFull('heavy')).toBe(true); + expect(disconnected.waitForSendBufferDrain('heavy')).toBe(heavyReady); + disconnected._handleClose(); + await heavyReady; +}); diff --git a/transport/sessionStateMachine/transitions.ts b/transport/sessionStateMachine/transitions.ts index 41df4313..fb23c44f 100644 --- a/transport/sessionStateMachine/transitions.ts +++ b/transport/sessionStateMachine/transitions.ts @@ -60,6 +60,7 @@ function inheritSharedSession( seqSent: session.seqSent, sendBuffer: session.sendBuffer, sendBufferDrainWaiter: session.sendBufferDrainWaiter, + streamSendBufferDrainWaiters: session.streamSendBufferDrainWaiters, telemetry: session.telemetry, options: session.options, log: session.log, diff --git a/transport/transport.ts b/transport/transport.ts index 390f5d47..122b89cd 100644 --- a/transport/transport.ts +++ b/transport/transport.ts @@ -118,6 +118,26 @@ export abstract class Transport< providedOptions?: ProvidedTransportOptions, ) { this.options = { ...defaultTransportOptions, ...providedOptions }; + const streamMark = this.options.streamSendBufferHighWaterMark; + if ( + streamMark !== undefined && + (!Number.isSafeInteger(streamMark) || streamMark < 1) + ) { + throw new Error( + 'streamSendBufferHighWaterMark must be a positive integer', + ); + } + const receiverAck = this.options.receiverAck; + if ( + receiverAck && + (!Number.isSafeInteger(receiverAck.maxMessages) || + receiverAck.maxMessages < 1 || + !Number.isFinite(receiverAck.maxDelayMs) || + receiverAck.maxDelayMs <= 0 || + receiverAck.maxDelayMs > 2_147_483_647) + ) { + throw new Error('receiverAck requires positive message and delay limits'); + } this.eventDispatcher = new EventDispatcher< EventTypes, HandshakeFailureCode @@ -414,6 +434,7 @@ export abstract class Transport< getSessionBackpressure( to: TransportClientId, sessionId: SessionId, + streamId?: string, ): SessionBackpressure { const currentSession = () => { const session = this.sessions.get(to); @@ -425,9 +446,10 @@ export abstract class Transport< }; return { - isSendBufferFull: () => currentSession()?.isSendBufferFull() ?? false, + isSendBufferFull: () => + currentSession()?.isSendBufferFull(streamId) ?? false, waitForSendBufferDrain: () => - currentSession()?.waitForSendBufferDrain() ?? Promise.resolve(), + currentSession()?.waitForSendBufferDrain(streamId) ?? Promise.resolve(), }; } } From e54a55534e3b817ecabde3d8f6c3c9667ac12d61 Mon Sep 17 00:00:00 2001 From: Armaan Aggarwal <49625265+Armster15@users.noreply.github.com> Date: Wed, 14 Oct 2026 22:48:16 +0000 Subject: [PATCH 2/3] transport: cover protobuf writers and synchronous ACK batching Assisted-by: Replit --- .../properties/backpressure.property.test.ts | 13 ++- __tests__/protobufBackpressure.test.ts | 109 ++++++++++++++++++ __tests__/receiverAck.test.ts | 4 +- protobuf/client.ts | 4 +- protobuf/server.ts | 1 + .../sessionStateMachine/SessionConnected.ts | 5 +- transport/sessionStateMachine/common.ts | 7 +- .../sessionStateMachine/flowControl.test.ts | 21 ++++ 8 files changed, 158 insertions(+), 6 deletions(-) create mode 100644 __tests__/protobufBackpressure.test.ts diff --git a/__tests__/properties/backpressure.property.test.ts b/__tests__/properties/backpressure.property.test.ts index 23d013d4..e5dba6a8 100644 --- a/__tests__/properties/backpressure.property.test.ts +++ b/__tests__/properties/backpressure.property.test.ts @@ -1,4 +1,4 @@ -import { expect, test } from 'vitest'; +import { expect, test, vi } from 'vitest'; import * as hegel from '@hegeldev/hegel'; import * as gs from '@hegeldev/hegel/generators'; import { createClient } from '../../router/client'; @@ -13,6 +13,7 @@ test( () => hegel.testAsync( async (tc) => { + vi.useRealTimers(); const mark = tc.draw(gs.integers({ minValue: 3, maxValue: 12 })); const count = tc.draw( gs.integers({ minValue: mark, maxValue: mark * 3 }), @@ -30,6 +31,14 @@ test( }); const clientTransport = network.getClientTransport('client'); const serverTransport = network.getServerTransport(); + const violations: Array = []; + for (const transport of [clientTransport, serverTransport]) { + transport.bindLogger((text, meta) => { + if (meta?.tags?.includes('invariant-violation')) { + violations.push(text); + } + }, 'debug'); + } const services = { test: TestServiceSchema }; createServer(serverTransport, services); const client = createClient( @@ -63,9 +72,11 @@ test( } heavy.reqWritable.close(); quiet.reqWritable.close(); + expect(violations).toEqual([]); } finally { await cleanupTransports([clientTransport, serverTransport]); await network.cleanup(); + vi.useFakeTimers({ shouldAdvanceTime: true }); } }, { testCases: 20 }, diff --git a/__tests__/protobufBackpressure.test.ts b/__tests__/protobufBackpressure.test.ts new file mode 100644 index 00000000..0d09006a --- /dev/null +++ b/__tests__/protobufBackpressure.test.ts @@ -0,0 +1,109 @@ +import { describe, expect, test, vi } from 'vitest'; +import { BinaryCodec, NaiveJsonCodec } from '../codec'; +import { + createClient, + createProtoService, + createServer, + Ok, + ProtoCodec, +} from '../protobuf'; +import { cleanupTransports, waitFor } from '../testUtil/fixtures/cleanup'; +import { TestService } from '../testUtil/fixtures/protobuf'; +import { transports } from '../testUtil/fixtures/transports'; +import { isAck } from '../transport/message'; + +const matrix = transports.flatMap((transport) => + [ + { name: 'naive', codec: NaiveJsonCodec }, + { name: 'binary', codec: BinaryCodec }, + { name: 'proto', codec: ProtoCodec }, + ].map((codec) => ({ transport, codec })), +); + +describe.each(matrix)( + 'protobuf stream isolation ($transport.name, $codec.name)', + ({ transport, codec }) => { + test('request and response writers use their own stream watermarks', async () => { + vi.useRealTimers(); + const options = { + codec: codec.codec, + sendBufferHighWaterMark: 4, + streamSendBufferHighWaterMark: 4, + receiverAck: { maxMessages: 3, maxDelayMs: 20 }, + heartbeatIntervalMs: 30_000, + sessionDisconnectGraceMs: 90_000, + }; + const setup = await transport.setup({ client: options, server: options }); + const clientTransport = setup.getClientTransport('client'); + const serverTransport = setup.getServerTransport(); + const violations: Array = []; + for (const transport of [clientTransport, serverTransport]) { + transport.bindLogger((text, meta) => { + if (meta?.tags?.includes('invariant-violation')) { + violations.push(text); + } + }, 'debug'); + } + const responsePressure: Array> = []; + const service = createProtoService().define(TestService, { + countUp: ({ request, resWritable }) => { + const results: Array = []; + for (let value = 1; value <= request.limit; value++) { + results.push(resWritable.write(Ok({ value }))); + } + responsePressure.push(results); + resWritable.close(); + }, + sum: async ({ reqReadable }) => { + let total = 0; + for await (const result of reqReadable) { + if (result.ok) total += result.payload.value; + } + + return Ok({ total }); + }, + }); + createServer(serverTransport, [service]); + const client = createClient( + TestService, + clientTransport, + serverTransport.clientId, + ); + try { + const heavyRequest = client.sum(); + for (let i = 0; i < 6; i++) + heavyRequest.reqWritable.write({ value: i }); + const quietRequest = client.sum(); + expect(quietRequest.reqWritable.write({ value: 1 })).toBe(true); + heavyRequest.reqWritable.close(); + quietRequest.reqWritable.close(); + expect(await heavyRequest.finalize()).toMatchObject({ + ok: true, + payload: { total: 15 }, + }); + expect(await quietRequest.finalize()).toMatchObject({ + ok: true, + payload: { total: 1 }, + }); + + const heavyResponse = client.countUp({ limit: 6 }); + const quietResponse = client.countUp({ limit: 1 }); + expect(await heavyResponse.collect()).toHaveLength(6); + expect(await quietResponse.collect()).toHaveLength(1); + expect(responsePressure[0][5]).toBe(false); + expect(responsePressure[1][0]).toBe(true); + await waitFor(() => { + const session = [...serverTransport.sessions.values()][0]; + expect( + session.sendBuffer.filter((msg) => !isAck(msg.msg.controlFlags)), + ).toHaveLength(0); + }); + expect(violations).toEqual([]); + } finally { + await cleanupTransports([clientTransport, serverTransport]); + await setup.cleanup(); + vi.useFakeTimers({ shouldAdvanceTime: true }); + } + }); + }, +); diff --git a/__tests__/receiverAck.test.ts b/__tests__/receiverAck.test.ts index 0b071e7d..46c1a7dd 100644 --- a/__tests__/receiverAck.test.ts +++ b/__tests__/receiverAck.test.ts @@ -1,4 +1,4 @@ -import { describe, expect, test } from 'vitest'; +import { describe, expect, test, vi } from 'vitest'; import { createClient } from '../router/client'; import { createServer } from '../router/server'; import { readNextResult } from '../testUtil'; @@ -14,6 +14,7 @@ describe.each( 'receiver ACKs: $enabled ($transport.name, $codec.name)', ({ transport, codec, enabled }) => { test('an opt-in receiver drains one-way traffic with a legacy-mode peer', async () => { + vi.useRealTimers(); const receiverAck = { maxMessages: 3, maxDelayMs: 20 }; const common = { codec: codec.codec, @@ -58,6 +59,7 @@ describe.each( } finally { await cleanupTransports([clientTransport, serverTransport]); await setup.cleanup(); + vi.useFakeTimers({ shouldAdvanceTime: true }); } }); }, diff --git a/protobuf/client.ts b/protobuf/client.ts index 00947531..91fcc74c 100644 --- a/protobuf/client.ts +++ b/protobuf/client.ts @@ -274,12 +274,12 @@ function startMethodCall< serverId, session.id, ); + const streamId = generateId(); const sessionBackpressure = transport.getSessionBackpressure( serverId, session.id, + streamId, ); - - const streamId = generateId(); const { span, ctx } = createProcTelemetryInfo( transport.tracer, session, diff --git a/protobuf/server.ts b/protobuf/server.ts index e21af017..898f50c5 100644 --- a/protobuf/server.ts +++ b/protobuf/server.ts @@ -371,6 +371,7 @@ class ProtobufServer< const sessionBackpressure = this.transport.getSessionBackpressure( from, sessionId, + streamId, ); const deferredCleanups: Array<() => void | Promise> = []; diff --git a/transport/sessionStateMachine/SessionConnected.ts b/transport/sessionStateMachine/SessionConnected.ts index e62b2be6..91c1d277 100644 --- a/transport/sessionStateMachine/SessionConnected.ts +++ b/transport/sessionStateMachine/SessionConnected.ts @@ -91,6 +91,9 @@ export class SessionConnected< this.assertSendOrdering(encodedMsg); this.bufferMessage(encodedMsg); + // A synchronous connection can deliver more inbound messages during send. + // Only clear the batch covered by the ACK we just encoded. + this.clearReceiverAck(); const sent = this.conn.send(encodedMsg.data); if (!sent) { const reason = 'failed to send message'; @@ -103,7 +106,6 @@ export class SessionConnected< } this.seqSent = encodedMsg.seq; - this.clearReceiverAck(); return { ok: true, value: encodedMsg.id }; } @@ -234,6 +236,7 @@ export class SessionConnected< this.pendingReceiverAckMessages++; if (this.pendingReceiverAckMessages >= options.maxMessages) { this.sendHeartbeat(); + return; } this.receiverAckTimer ??= setTimeout(() => { diff --git a/transport/sessionStateMachine/common.ts b/transport/sessionStateMachine/common.ts index 50fcad93..3da66aa4 100644 --- a/transport/sessionStateMachine/common.ts +++ b/transport/sessionStateMachine/common.ts @@ -261,6 +261,7 @@ export abstract class IdentifiedSession extends CommonSession { string, PromiseWithResolvers >; + private readonly streamSendBufferSizes = new Map(); constructor(props: IdentifiedSessionProps) { @@ -285,7 +286,8 @@ export abstract class IdentifiedSession extends CommonSession { this.sendBuffer = sendBuffer; this.sendBufferDrainWaiter = sendBufferDrainWaiter; this.streamSendBufferDrainWaiters = - props.streamSendBufferDrainWaiters ?? new Map(); + props.streamSendBufferDrainWaiters ?? + new Map>(); for (const msg of sendBuffer) { this.addStreamBufferEntry(msg); } @@ -363,6 +365,7 @@ export abstract class IdentifiedSession extends CommonSession { this.options.streamSendBufferHighWaterMark ); } + return this.sendBuffer.length >= this.options.sendBufferHighWaterMark; } @@ -386,6 +389,7 @@ export abstract class IdentifiedSession extends CommonSession { waiter = createPromiseWithResolvers(); this.streamSendBufferDrainWaiters.set(streamId, waiter); } + return waiter.promise; } @@ -438,6 +442,7 @@ export abstract class IdentifiedSession extends CommonSession { } else { this.streamSendBufferSizes.set(streamId, remaining); } + return false; }); this.notifySendBufferDrain(); diff --git a/transport/sessionStateMachine/flowControl.test.ts b/transport/sessionStateMachine/flowControl.test.ts index e5d4e22e..d0fccf89 100644 --- a/transport/sessionStateMachine/flowControl.test.ts +++ b/transport/sessionStateMachine/flowControl.test.ts @@ -59,6 +59,7 @@ function connected(options: ProvidedTransportOptions = {}) { if (!data.ok) throw new Error(data.reason); session.onMessageData(data.value); }; + return { session, conn, receive }; } @@ -131,6 +132,26 @@ describe('receiver-driven ACKs', () => { expect(disabled.conn.send).not.toHaveBeenCalled(); disabled.session._handleClose(); }); + + test('inbound messages arriving during a synchronous send keep their ACK deadline', () => { + const { session, conn, receive } = connected({ + receiverAck: { maxMessages: 10, maxDelayMs: 50 }, + }); + receive(); + conn.send.mockImplementationOnce(() => { + receive(); + + return true; + }); + session.send({ streamId: 'reply', controlFlags: 0, payload: {} }); + vi.advanceTimersByTime(50); + expect(conn.send).toHaveBeenCalledTimes(2); + expect(session.codec.fromBuffer(session.sendBuffer[1].data)).toMatchObject({ + ok: true, + value: { ack: 2 }, + }); + session._handleClose(); + }); }); test('stream waiters drain independently, survive state transitions, and settle on close', async () => { From 290f5f8bc15404ba2b9a86a140afff1648835280 Mon Sep 17 00:00:00 2001 From: Armaan Aggarwal <49625265+Armster15@users.noreply.github.com> Date: Thu, 15 Oct 2026 00:05:10 +0000 Subject: [PATCH 3/3] transport: simplify stream accounting and focus flow-control tests Assisted-by: Replit --- PROTOCOL.md | 39 ++-- __tests__/backpressure.test.ts | 47 +---- .../properties/backpressure.property.test.ts | 85 --------- __tests__/protobufBackpressure.test.ts | 167 ++++++++---------- __tests__/receiverAck.test.ts | 110 ++++++------ .../sessionStateMachine/SessionConnected.ts | 16 +- transport/sessionStateMachine/common.ts | 113 ++++-------- .../sessionStateMachine/flowControl.test.ts | 51 ++---- transport/sessionStateMachine/transitions.ts | 2 +- 9 files changed, 204 insertions(+), 426 deletions(-) delete mode 100644 __tests__/properties/backpressure.property.test.ts diff --git a/PROTOCOL.md b/PROTOCOL.md index fc27d332..fc23b77d 100644 --- a/PROTOCOL.md +++ b/PROTOCOL.md @@ -672,38 +672,19 @@ This ensures exactly-once delivery semantics. #### Optional receiver-driven acknowledgements A receiver MAY send an explicit ACK after a configured number of inbound -non-ACK messages or a maximum delay, whichever occurs first. The TypeScript -implementation enables this with `receiverAck: { maxMessages, maxDelayMs }`; -it is disabled by default. An outgoing message carries the current cumulative -ACK and cancels the pending batch. Replayed messages retain their original -ACK field and do not cancel a pending batch. - -These ACKs use the existing `heartbeat` stream, `AckBit`, and `{ type: 'ACK' }` -payload. They participate in sequencing and replay exactly like heartbeat -ACKs. There is no new wire field, capability negotiation, or protocol version. -An older active-heartbeat peer consumes them without replying; an older -passive-heartbeat peer responds with its normal heartbeat ACK. ACK messages -MUST NOT count toward the receiver-driven threshold or arm its timer. The -existing active/passive heartbeat roles remain unchanged, avoiding ACK loops. -Either endpoint may enable receiver-driven ACKs independently. Timers are -cancelled when a connected state exits; heartbeat ACKs still guarantee progress -after reconnect and when receiver-driven ACKs are disabled. +non-ACK messages or a delay: `receiverAck: { maxMessages, maxDelayMs }`. +This is disabled by default. Piggybacked ACKs cancel pending batches; ACK +frames do not schedule batches. These are existing heartbeat ACK frames, so +either peer can opt in without negotiation or a protocol version change. +Existing active/passive heartbeat behavior and reconnect guarantees remain. #### Optional per-stream advisory backpressure -`streamSendBufferHighWaterMark` opts the TypeScript router's writers into a -per-stream message-count watermark instead of the shared session watermark. -Each writer counts all unacknowledged messages on its stream, including open -and close frames. A busy stream does not consume another stream's allowance. -Cumulative session ACKs free the corresponding stream slots; reconnect carries -both buffered messages and pending drain waiters. Session destruction releases -all waiters. - -This is a local producer policy, not a wire change or a hard memory limit. -Writes remain accepted even when they return `false`. Many streams may -collectively exceed the session watermark in this mode; callers must bound -stream concurrency and retained bytes separately. Unscoped transport accessors -and transports without the option retain session-level backpressure. +`streamSendBufferHighWaterMark` makes router writers use their own unacknowledged +message count instead of the shared session watermark. It is disabled by default. +ACKs release slots; reconnect preserves pending waiters and close releases them. +Writes remain advisory and accepted. This is not a memory cap: aggregate occupancy +can exceed the session watermark. Unscoped accessors retain session behavior. #### On disconnect diff --git a/__tests__/backpressure.test.ts b/__tests__/backpressure.test.ts index 58660158..d48c69ed 100644 --- a/__tests__/backpressure.test.ts +++ b/__tests__/backpressure.test.ts @@ -19,19 +19,12 @@ import { TestSetupHelpers } from '../testUtil/fixtures/transports'; const HIGH_WATER_MARK = 4; -describe.each( - testMatrix().flatMap((entry) => - ['session', 'stream'].map((mode) => ({ ...entry, mode })), - ), -)( - '$mode writable backpressure ($transport.name transport, $codec.name codec)', - async ({ transport, codec, mode }) => { +describe.each(testMatrix())( + 'writable backpressure ($transport.name transport, $codec.name codec)', + async ({ transport, codec }) => { const opts = { codec: codec.codec, sendBufferHighWaterMark: HIGH_WATER_MARK, - ...(mode === 'stream' - ? { streamSendBufferHighWaterMark: HIGH_WATER_MARK } - : {}), }; const { addPostTestCleanup, postTestCleanup } = createPostTestCleanups(); @@ -202,39 +195,5 @@ describe.each( payload: expect.objectContaining({ code: UNEXPECTED_DISCONNECT_CODE }), }); }); - - test('a heavy stream only pressures a quiet stream in session mode', async () => { - const clientTransport = getClientTransport('client'); - const serverTransport = getServerTransport(); - const services = { test: TestServiceSchema }; - createServer(serverTransport, services); - const client = createClient( - clientTransport, - serverTransport.clientId, - ); - addPostTestCleanup(async () => { - await cleanupTransports([clientTransport, serverTransport]); - }); - const heavy = client.test.echo.stream({}); - const quiet = client.test.echo.stream({}); - for (let i = 0; i < HIGH_WATER_MARK + 1; i++) { - heavy.reqWritable.write({ msg: 'heavy', ignore: true }); - } - expect(heavy.reqWritable.write({ msg: 'heavy', ignore: true })).toBe( - false, - ); - expect(quiet.reqWritable.write({ msg: 'quiet', ignore: false })).toBe( - mode === 'stream', - ); - if (mode === 'stream') { - await quiet.reqWritable.waitForWriteReady(); - } - expect(await readNextResult(quiet.resReadable)).toStrictEqual({ - ok: true, - payload: { response: 'quiet' }, - }); - heavy.reqWritable.close(); - quiet.reqWritable.close(); - }); }, ); diff --git a/__tests__/properties/backpressure.property.test.ts b/__tests__/properties/backpressure.property.test.ts deleted file mode 100644 index e5dba6a8..00000000 --- a/__tests__/properties/backpressure.property.test.ts +++ /dev/null @@ -1,85 +0,0 @@ -import { expect, test, vi } from 'vitest'; -import * as hegel from '@hegeldev/hegel'; -import * as gs from '@hegeldev/hegel/generators'; -import { createClient } from '../../router/client'; -import { createServer } from '../../router/server'; -import { closeAllConnections, readNextResult } from '../../testUtil'; -import { cleanupTransports } from '../../testUtil/fixtures/cleanup'; -import { createMockTransportNetwork } from '../../testUtil/fixtures/mockTransport'; -import { TestServiceSchema } from '../../testUtil/fixtures/services'; - -test( - 'generated heavy feeds cannot park an unrelated stream across reconnect', - () => - hegel.testAsync( - async (tc) => { - vi.useRealTimers(); - const mark = tc.draw(gs.integers({ minValue: 3, maxValue: 12 })); - const count = tc.draw( - gs.integers({ minValue: mark, maxValue: mark * 3 }), - ); - const reconnect = tc.draw(gs.booleans()); - const options = { - sendBufferHighWaterMark: mark, - streamSendBufferHighWaterMark: mark, - receiverAck: { maxMessages: 3, maxDelayMs: 20 }, - maxJitterMs: 0, - }; - const network = createMockTransportNetwork({ - client: options, - server: options, - }); - const clientTransport = network.getClientTransport('client'); - const serverTransport = network.getServerTransport(); - const violations: Array = []; - for (const transport of [clientTransport, serverTransport]) { - transport.bindLogger((text, meta) => { - if (meta?.tags?.includes('invariant-violation')) { - violations.push(text); - } - }, 'debug'); - } - const services = { test: TestServiceSchema }; - createServer(serverTransport, services); - const client = createClient( - clientTransport, - serverTransport.clientId, - ); - try { - const heavy = client.test.echo.stream({}); - heavy.reqWritable.write({ msg: 'establish', ignore: false }); - await readNextResult(heavy.resReadable); - if (reconnect) closeAllConnections(clientTransport); - for (let i = 0; i < count; i++) { - heavy.reqWritable.write({ msg: String(i), ignore: false }); - } - const heavyReady = heavy.reqWritable.waitForWriteReady(); - const quiet = client.test.echo.stream({}); - expect(quiet.reqWritable.write({ msg: 'quiet', ignore: false })).toBe( - true, - ); - await quiet.reqWritable.waitForWriteReady(); - expect(await readNextResult(quiet.resReadable)).toStrictEqual({ - ok: true, - payload: { response: 'quiet' }, - }); - await heavyReady; - for (let i = 0; i < count; i++) { - expect(await readNextResult(heavy.resReadable)).toStrictEqual({ - ok: true, - payload: { response: String(i) }, - }); - } - heavy.reqWritable.close(); - quiet.reqWritable.close(); - expect(violations).toEqual([]); - } finally { - await cleanupTransports([clientTransport, serverTransport]); - await network.cleanup(); - vi.useFakeTimers({ shouldAdvanceTime: true }); - } - }, - { testCases: 20 }, - ), - 30_000, -); diff --git a/__tests__/protobufBackpressure.test.ts b/__tests__/protobufBackpressure.test.ts index 0d09006a..f6664721 100644 --- a/__tests__/protobufBackpressure.test.ts +++ b/__tests__/protobufBackpressure.test.ts @@ -1,5 +1,4 @@ import { describe, expect, test, vi } from 'vitest'; -import { BinaryCodec, NaiveJsonCodec } from '../codec'; import { createClient, createProtoService, @@ -12,98 +11,80 @@ import { TestService } from '../testUtil/fixtures/protobuf'; import { transports } from '../testUtil/fixtures/transports'; import { isAck } from '../transport/message'; -const matrix = transports.flatMap((transport) => - [ - { name: 'naive', codec: NaiveJsonCodec }, - { name: 'binary', codec: BinaryCodec }, - { name: 'proto', codec: ProtoCodec }, - ].map((codec) => ({ transport, codec })), -); - -describe.each(matrix)( - 'protobuf stream isolation ($transport.name, $codec.name)', - ({ transport, codec }) => { - test('request and response writers use their own stream watermarks', async () => { - vi.useRealTimers(); - const options = { - codec: codec.codec, - sendBufferHighWaterMark: 4, - streamSendBufferHighWaterMark: 4, - receiverAck: { maxMessages: 3, maxDelayMs: 20 }, - heartbeatIntervalMs: 30_000, - sessionDisconnectGraceMs: 90_000, - }; - const setup = await transport.setup({ client: options, server: options }); - const clientTransport = setup.getClientTransport('client'); - const serverTransport = setup.getServerTransport(); - const violations: Array = []; - for (const transport of [clientTransport, serverTransport]) { - transport.bindLogger((text, meta) => { - if (meta?.tags?.includes('invariant-violation')) { - violations.push(text); - } - }, 'debug'); - } - const responsePressure: Array> = []; - const service = createProtoService().define(TestService, { - countUp: ({ request, resWritable }) => { - const results: Array = []; - for (let value = 1; value <= request.limit; value++) { - results.push(resWritable.write(Ok({ value }))); - } - responsePressure.push(results); - resWritable.close(); - }, - sum: async ({ reqReadable }) => { - let total = 0; - for await (const result of reqReadable) { - if (result.ok) total += result.payload.value; - } +describe('protobuf stream isolation', () => { + test('request and response writers use their own stream watermarks', async () => { + vi.useRealTimers(); + const options = { + codec: ProtoCodec, + sendBufferHighWaterMark: 4, + streamSendBufferHighWaterMark: 4, + receiverAck: { maxMessages: 3, maxDelayMs: 20 }, + heartbeatIntervalMs: 30_000, + sessionDisconnectGraceMs: 90_000, + }; + const setup = await transports[0].setup({ + client: options, + server: options, + }); + const clientTransport = setup.getClientTransport('client'); + const serverTransport = setup.getServerTransport(); + const responsePressure: Array> = []; + const service = createProtoService().define(TestService, { + countUp: ({ request, resWritable }) => { + const results: Array = []; + for (let value = 1; value <= request.limit; value++) { + results.push(resWritable.write(Ok({ value }))); + } + responsePressure.push(results); + resWritable.close(); + }, + sum: async ({ reqReadable }) => { + let total = 0; + for await (const result of reqReadable) { + if (result.ok) total += result.payload.value; + } - return Ok({ total }); - }, + return Ok({ total }); + }, + }); + createServer(serverTransport, [service]); + const client = createClient( + TestService, + clientTransport, + serverTransport.clientId, + ); + try { + const heavyRequest = client.sum(); + for (let i = 0; i < 6; i++) heavyRequest.reqWritable.write({ value: i }); + const quietRequest = client.sum(); + expect(quietRequest.reqWritable.write({ value: 1 })).toBe(true); + heavyRequest.reqWritable.close(); + quietRequest.reqWritable.close(); + expect(await heavyRequest.finalize()).toMatchObject({ + ok: true, + payload: { total: 15 }, + }); + expect(await quietRequest.finalize()).toMatchObject({ + ok: true, + payload: { total: 1 }, }); - createServer(serverTransport, [service]); - const client = createClient( - TestService, - clientTransport, - serverTransport.clientId, - ); - try { - const heavyRequest = client.sum(); - for (let i = 0; i < 6; i++) - heavyRequest.reqWritable.write({ value: i }); - const quietRequest = client.sum(); - expect(quietRequest.reqWritable.write({ value: 1 })).toBe(true); - heavyRequest.reqWritable.close(); - quietRequest.reqWritable.close(); - expect(await heavyRequest.finalize()).toMatchObject({ - ok: true, - payload: { total: 15 }, - }); - expect(await quietRequest.finalize()).toMatchObject({ - ok: true, - payload: { total: 1 }, - }); - const heavyResponse = client.countUp({ limit: 6 }); - const quietResponse = client.countUp({ limit: 1 }); - expect(await heavyResponse.collect()).toHaveLength(6); - expect(await quietResponse.collect()).toHaveLength(1); - expect(responsePressure[0][5]).toBe(false); - expect(responsePressure[1][0]).toBe(true); - await waitFor(() => { - const session = [...serverTransport.sessions.values()][0]; - expect( - session.sendBuffer.filter((msg) => !isAck(msg.msg.controlFlags)), - ).toHaveLength(0); - }); - expect(violations).toEqual([]); - } finally { - await cleanupTransports([clientTransport, serverTransport]); - await setup.cleanup(); - vi.useFakeTimers({ shouldAdvanceTime: true }); - } - }); - }, -); + const heavyResponse = client.countUp({ limit: 6 }); + const quietResponse = client.countUp({ limit: 1 }); + expect(await heavyResponse.collect()).toHaveLength(6); + expect(await quietResponse.collect()).toHaveLength(1); + expect(responsePressure[0][5]).toBe(false); + expect(responsePressure[1][0]).toBe(true); + await waitFor(() => { + const session = [...serverTransport.sessions.values()][0]; + expect( + session.sendBuffer.filter((msg) => !isAck(msg.msg.controlFlags)), + ).toHaveLength(0); + }); + } finally { + await cleanupTransports([clientTransport, serverTransport]); + await setup.cleanup(); + vi.useFakeTimers({ shouldAdvanceTime: true }); + } + }); +}); diff --git a/__tests__/receiverAck.test.ts b/__tests__/receiverAck.test.ts index 46c1a7dd..bc71d68c 100644 --- a/__tests__/receiverAck.test.ts +++ b/__tests__/receiverAck.test.ts @@ -1,66 +1,62 @@ -import { describe, expect, test, vi } from 'vitest'; +import { expect, test, vi } from 'vitest'; import { createClient } from '../router/client'; import { createServer } from '../router/server'; -import { readNextResult } from '../testUtil'; -import { cleanupTransports, waitFor } from '../testUtil/fixtures/cleanup'; -import { testMatrix } from '../testUtil/fixtures/matrix'; +import { closeAllConnections, readNextResult } from '../testUtil'; +import { cleanupTransports } from '../testUtil/fixtures/cleanup'; import { TestServiceSchema } from '../testUtil/fixtures/services'; +import { transports } from '../testUtil/fixtures/transports'; -describe.each( - testMatrix().flatMap((entry) => - ['client', 'server', 'both'].map((enabled) => ({ ...entry, enabled })), - ), -)( - 'receiver ACKs: $enabled ($transport.name, $codec.name)', - ({ transport, codec, enabled }) => { - test('an opt-in receiver drains one-way traffic with a legacy-mode peer', async () => { - vi.useRealTimers(); - const receiverAck = { maxMessages: 3, maxDelayMs: 20 }; - const common = { - codec: codec.codec, - heartbeatIntervalMs: 30_000, - sessionDisconnectGraceMs: 90_000, - }; - const setup = await transport.setup({ - client: { ...common, ...(enabled !== 'server' ? { receiverAck } : {}) }, - server: { ...common, ...(enabled !== 'client' ? { receiverAck } : {}) }, - }); - const clientTransport = setup.getClientTransport('client'); - const serverTransport = setup.getServerTransport(); - const services = { test: TestServiceSchema }; - createServer(serverTransport, services); - const client = createClient( - clientTransport, - serverTransport.clientId, +test.each(['client', 'server'])( + 'stream isolation and reconnect work with only the %s opting in', + async (enabled) => { + vi.useRealTimers(); + const options = { + streamSendBufferHighWaterMark: 4, + receiverAck: { maxMessages: 3, maxDelayMs: 20 }, + }; + const setup = await transports[0].setup({ + client: enabled === 'client' ? options : {}, + server: enabled === 'server' ? options : {}, + }); + const clientTransport = setup.getClientTransport('client'); + const serverTransport = setup.getServerTransport(); + const services = { test: TestServiceSchema }; + createServer(serverTransport, services); + const client = createClient( + clientTransport, + serverTransport.clientId, + ); + try { + const heavy = client.test.echo.stream({}); + heavy.reqWritable.write({ msg: 'establish', ignore: false }); + await readNextResult(heavy.resReadable); + closeAllConnections(clientTransport); + const writes = Array.from({ length: 6 }, (_, i) => + heavy.reqWritable.write({ msg: String(i), ignore: false }), + ); + if (enabled === 'client') expect(writes[5]).toBe(false); + const ready = heavy.reqWritable.waitForWriteReady(); + const light = client.test.echo.stream({}); + expect(light.reqWritable.write({ msg: 'light', ignore: false })).toBe( + true, ); - try { - const call = client.test.echo.stream({}); - for (let i = 0; i < 7; i++) { - call.reqWritable.write({ - msg: `message-${i}`, - ignore: enabled === 'server', - }); - } - if (enabled !== 'server') { - for (let i = 0; i < 7; i++) { - expect(await readNextResult(call.resReadable)).toStrictEqual({ - ok: true, - payload: { response: `message-${i}` }, - }); - } - } - const sender = enabled === 'server' ? clientTransport : serverTransport; - await waitFor(() => { - const session = [...sender.sessions.values()][0]; - expect(session).toBeDefined(); - expect(session.sendBuffer.length).toBeLessThanOrEqual(1); + expect(await readNextResult(light.resReadable)).toMatchObject({ + ok: true, + payload: { response: 'light' }, + }); + await ready; + for (let i = 0; i < 6; i++) { + expect(await readNextResult(heavy.resReadable)).toMatchObject({ + ok: true, + payload: { response: String(i) }, }); - call.reqWritable.close(); - } finally { - await cleanupTransports([clientTransport, serverTransport]); - await setup.cleanup(); - vi.useFakeTimers({ shouldAdvanceTime: true }); } - }); + heavy.reqWritable.close(); + light.reqWritable.close(); + } finally { + await cleanupTransports([clientTransport, serverTransport]); + await setup.cleanup(); + vi.useFakeTimers({ shouldAdvanceTime: true }); + } }, ); diff --git a/transport/sessionStateMachine/SessionConnected.ts b/transport/sessionStateMachine/SessionConnected.ts index 91c1d277..7c495106 100644 --- a/transport/sessionStateMachine/SessionConnected.ts +++ b/transport/sessionStateMachine/SessionConnected.ts @@ -64,9 +64,23 @@ export class SessionConnected< private pendingReceiverAckMessages = 0; updateBookkeeping(ack: number, seq: number) { - this.acknowledgeSendBuffer(ack); + this.sendBuffer = this.sendBuffer.filter((msg) => { + if (msg.seq >= ack) return true; + const stream = this.streamSendBuffers.get(msg.msg.streamId); + if (stream) { + stream.size--; + if (!this.isSendBufferFull(msg.msg.streamId)) { + stream.waiter?.resolve(); + stream.waiter = undefined; + } + if (stream.size === 0) this.streamSendBuffers.delete(msg.msg.streamId); + } + + return false; + }); this.ack = seq + 1; this.lastInboundAt = Date.now(); + if (!this.isSendBufferFull()) this.notifySendBufferDrain(); } private assertSendOrdering(encodedMsg: EncodedTransportMessage) { diff --git a/transport/sessionStateMachine/common.ts b/transport/sessionStateMachine/common.ts index 3da66aa4..927950c8 100644 --- a/transport/sessionStateMachine/common.ts +++ b/transport/sessionStateMachine/common.ts @@ -146,15 +146,9 @@ export interface SessionOptions { * session close. */ sendBufferHighWaterMark: number; - /** - * Optional per-stream advisory watermark. When set, stream writers use - * this limit instead of the session watermark. This is not a memory cap. - */ + /** Opt-in per-stream advisory watermark instead of the session watermark. */ streamSendBufferHighWaterMark?: number; - /** - * Opt-in ACK batching for inbound non-ACK messages. Outbound messages - * already carry an ACK and cancel any pending batch. - */ + /** Opt-in ACK batching for inbound non-ACK messages. */ receiverAck?: { maxMessages: number; maxDelayMs: number }; /** * The codec to use for encoding/decoding messages over the wire @@ -213,6 +207,11 @@ export interface IdentifiedSessionListeners { ) => void; } +interface StreamSendBuffer { + size: number; + waiter?: PromiseWithResolvers; +} + // all sessions where we know the other side's client id export interface IdentifiedSessionProps extends CommonSessionProps { id: SessionId; @@ -222,7 +221,7 @@ export interface IdentifiedSessionProps extends CommonSessionProps { seqSent: number; sendBuffer: Array; sendBufferDrainWaiter: PromiseWithResolvers | undefined; - streamSendBufferDrainWaiters?: Map>; + streamSendBuffers?: Map; telemetry: TelemetryInfo; protocolVersion: ProtocolVersion; listeners: IdentifiedSessionListeners; @@ -257,12 +256,7 @@ export abstract class IdentifiedSession extends CommonSession { * Carried across session state transitions alongside {@link sendBuffer}. */ sendBufferDrainWaiter: PromiseWithResolvers | undefined; - readonly streamSendBufferDrainWaiters: Map< - string, - PromiseWithResolvers - >; - - private readonly streamSendBufferSizes = new Map(); + readonly streamSendBuffers: Map; constructor(props: IdentifiedSessionProps) { const { @@ -285,11 +279,17 @@ export abstract class IdentifiedSession extends CommonSession { this.ack = ack; this.sendBuffer = sendBuffer; this.sendBufferDrainWaiter = sendBufferDrainWaiter; - this.streamSendBufferDrainWaiters = - props.streamSendBufferDrainWaiters ?? - new Map>(); - for (const msg of sendBuffer) { - this.addStreamBufferEntry(msg); + this.streamSendBuffers = + props.streamSendBuffers ?? new Map(); + if ( + !props.streamSendBuffers && + this.options.streamSendBufferHighWaterMark + ) { + for (const { msg } of sendBuffer) { + const stream = this.streamSendBuffers.get(msg.streamId) ?? { size: 0 }; + stream.size++; + this.streamSendBuffers.set(msg.streamId, stream); + } } this.telemetry = telemetry; this.log = log; @@ -361,7 +361,7 @@ export abstract class IdentifiedSession extends CommonSession { this.options.streamSendBufferHighWaterMark !== undefined ) { return ( - (this.streamSendBufferSizes.get(streamId) ?? 0) >= + (this.streamSendBuffers.get(streamId)?.size ?? 0) >= this.options.streamSendBufferHighWaterMark ); } @@ -380,17 +380,12 @@ export abstract class IdentifiedSession extends CommonSession { return Promise.resolve(); } - if ( - streamId !== undefined && - this.options.streamSendBufferHighWaterMark !== undefined - ) { - let waiter = this.streamSendBufferDrainWaiters.get(streamId); - if (!waiter) { - waiter = createPromiseWithResolvers(); - this.streamSendBufferDrainWaiters.set(streamId, waiter); - } + const stream = + streamId === undefined ? undefined : this.streamSendBuffers.get(streamId); + if (stream) { + stream.waiter ??= createPromiseWithResolvers(); - return waiter.promise; + return stream.waiter.promise; } this.sendBufferDrainWaiter ??= createPromiseWithResolvers(); @@ -399,53 +394,19 @@ export abstract class IdentifiedSession extends CommonSession { } protected notifySendBufferDrain(): void { - if (!this.isSendBufferFull()) { - this.sendBufferDrainWaiter?.resolve(); - this.sendBufferDrainWaiter = undefined; - } - for (const [streamId, waiter] of this.streamSendBufferDrainWaiters) { - if (!this.isSendBufferFull(streamId)) { - waiter.resolve(); - this.streamSendBufferDrainWaiters.delete(streamId); - } - } - } - - private addStreamBufferEntry(msg: EncodedTransportMessage): void { - if (this.options.streamSendBufferHighWaterMark === undefined) { - return; - } - const streamId = msg.msg.streamId; - this.streamSendBufferSizes.set( - streamId, - (this.streamSendBufferSizes.get(streamId) ?? 0) + 1, - ); + this.sendBufferDrainWaiter?.resolve(); + this.sendBufferDrainWaiter = undefined; } protected bufferMessage(msg: EncodedTransportMessage): void { this.sendBuffer.push(msg); - this.addStreamBufferEntry(msg); - } - - protected acknowledgeSendBuffer(ack: number): void { - this.sendBuffer = this.sendBuffer.filter((msg) => { - if (msg.seq >= ack) { - return true; - } - if (this.options.streamSendBufferHighWaterMark === undefined) { - return false; - } - const streamId = msg.msg.streamId; - const remaining = (this.streamSendBufferSizes.get(streamId) ?? 1) - 1; - if (remaining === 0) { - this.streamSendBufferSizes.delete(streamId); - } else { - this.streamSendBufferSizes.set(streamId, remaining); - } - - return false; - }); - this.notifySendBufferDrain(); + if (this.options.streamSendBufferHighWaterMark !== undefined) { + const stream = this.streamSendBuffers.get(msg.msg.streamId) ?? { + size: 0, + }; + stream.size++; + this.streamSendBuffers.set(msg.msg.streamId, stream); + } } send(msg: PartialTransportMessage): SendResult { @@ -469,7 +430,9 @@ export abstract class IdentifiedSession extends CommonSession { _handleClose(): void { // zero out the buffer this.sendBuffer.length = 0; - this.streamSendBufferSizes.clear(); + for (const stream of this.streamSendBuffers.values()) + stream.waiter?.resolve(); + this.streamSendBuffers.clear(); // wake any producers waiting for drain so they don't hang forever, // they should check isWritable/isSendBufferFull before writing again this.notifySendBufferDrain(); diff --git a/transport/sessionStateMachine/flowControl.test.ts b/transport/sessionStateMachine/flowControl.test.ts index d0fccf89..bbd1d94b 100644 --- a/transport/sessionStateMachine/flowControl.test.ts +++ b/transport/sessionStateMachine/flowControl.test.ts @@ -5,10 +5,9 @@ import { Connection } from '../connection'; import { ControlFlags, currentProtocolVersion } from '../message'; import { defaultTransportOptions, ProvidedTransportOptions } from '../options'; import { SessionConnected } from './SessionConnected'; -import { SessionStateGraph } from './transitions'; class RecordingConnection extends Connection { - send = vi.fn(() => true); + send = vi.fn((_data: Uint8Array) => true); close = vi.fn(); } @@ -45,7 +44,7 @@ function connected(options: ProvidedTransportOptions = {}) { onMessageSendFailure: vi.fn(), }, }); - const receive = (controlFlags = 0, ack = 0) => { + const receive = (controlFlags = 0) => { const data = session.codec.toBuffer({ id: `inbound-${session.ack}`, from: 'sender', @@ -54,7 +53,7 @@ function connected(options: ProvidedTransportOptions = {}) { controlFlags, payload: controlFlags ? { type: 'ACK' } : {}, seq: session.ack, - ack, + ack: 0, }); if (!data.ok) throw new Error(data.reason); session.onMessageData(data.value); @@ -75,14 +74,13 @@ describe('receiver-driven ACKs', () => { expect(conn.send).not.toHaveBeenCalled(); receive(); expect(conn.send).toHaveBeenCalledTimes(1); - expect(session.sendBuffer[0].msg).toMatchObject({ - streamId: 'heartbeat', - controlFlags: ControlFlags.AckBit, - payload: { type: 'ACK' }, - }); - expect(session.codec.fromBuffer(session.sendBuffer[0].data)).toMatchObject({ + expect(session.codec.fromBuffer(conn.send.mock.calls[0][0])).toMatchObject({ ok: true, - value: { ack: 3 }, + value: { + ack: 3, + controlFlags: ControlFlags.AckBit, + payload: { type: 'ACK' }, + }, }); vi.advanceTimersByTime(50); expect(conn.send).toHaveBeenCalledTimes(1); @@ -146,39 +144,10 @@ describe('receiver-driven ACKs', () => { session.send({ streamId: 'reply', controlFlags: 0, payload: {} }); vi.advanceTimersByTime(50); expect(conn.send).toHaveBeenCalledTimes(2); - expect(session.codec.fromBuffer(session.sendBuffer[1].data)).toMatchObject({ + expect(session.codec.fromBuffer(conn.send.mock.calls[1][0])).toMatchObject({ ok: true, value: { ack: 2 }, }); session._handleClose(); }); }); - -test('stream waiters drain independently, survive state transitions, and settle on close', async () => { - const { session, receive } = connected({ - sendBufferHighWaterMark: 2, - streamSendBufferHighWaterMark: 2, - }); - for (const streamId of ['light', 'light', 'heavy', 'heavy', 'heavy']) { - session.send({ streamId, controlFlags: 0, payload: {} }); - } - const lightReady = session.waitForSendBufferDrain('light'); - const heavyReady = session.waitForSendBufferDrain('heavy'); - let heavyDrained = false; - void heavyReady.then(() => { - heavyDrained = true; - }); - receive(ControlFlags.AckBit, 2); - await lightReady; - expect(heavyDrained).toBe(false); - expect(session.isSendBufferFull('light')).toBe(false); - expect(session.isSendBufferFull('heavy')).toBe(true); - const disconnected = SessionStateGraph.transition.ConnectedToNoConnection( - session, - { onSessionGracePeriodElapsed: vi.fn(), onMessageSendFailure: vi.fn() }, - ); - expect(disconnected.isSendBufferFull('heavy')).toBe(true); - expect(disconnected.waitForSendBufferDrain('heavy')).toBe(heavyReady); - disconnected._handleClose(); - await heavyReady; -}); diff --git a/transport/sessionStateMachine/transitions.ts b/transport/sessionStateMachine/transitions.ts index fb23c44f..c390ef8a 100644 --- a/transport/sessionStateMachine/transitions.ts +++ b/transport/sessionStateMachine/transitions.ts @@ -60,7 +60,7 @@ function inheritSharedSession( seqSent: session.seqSent, sendBuffer: session.sendBuffer, sendBufferDrainWaiter: session.sendBufferDrainWaiter, - streamSendBufferDrainWaiters: session.streamSendBufferDrainWaiters, + streamSendBuffers: session.streamSendBuffers, telemetry: session.telemetry, options: session.options, log: session.log,