From f9910d59ae6f3a4f904e3eb1b6c0df2204fd6eaf Mon Sep 17 00:00:00 2001 From: Armaan Aggarwal <49625265+Armster15@users.noreply.github.com> Date: Wed, 21 Oct 2026 23:25:40 +0000 Subject: [PATCH 1/3] Add drain and refused-call resend to the protobuf router Share drain lifecycle and session waiting with the TypeBox router. Add protobuf behavior coverage and document both router surfaces. Assisted-by: Replit --- README.md | 2 + __tests__/protobufDrain.test.ts | 245 ++++++++++++++++++++++++++++++++ protobuf/client.ts | 58 +++++++- protobuf/errors.ts | 3 + protobuf/index.ts | 1 + protobuf/server.ts | 40 +++++- router/client.ts | 55 ++----- router/drain.ts | 116 +++++++++++++++ router/server.ts | 44 ++---- 9 files changed, 475 insertions(+), 89 deletions(-) create mode 100644 __tests__/protobufDrain.test.ts create mode 100644 router/drain.ts diff --git a/README.md b/README.md index 60bbc794..d639b218 100644 --- a/README.md +++ b/README.md @@ -174,6 +174,8 @@ return Err({ code: 'INVALID_INPUT', message: 'Value must be positive' }); `server.drain({ budgetMs })` refuses new calls with `SERVER_DRAINING` before handlers run. It waits for in-flight RPCs and uploads up to `budgetMs`, then closes the transport. +Both the TypeBox and protobuf routers support drain and opt-in refused-call resend. For protobuf, drain waits for unary and client-streaming calls; only unary calls can be resent. + ```ts const client = createClient(transport, 'SERVER', { resendRefusedCalls: { maxResends: 3, newSessionWaitMs: 30_000 }, diff --git a/__tests__/protobufDrain.test.ts b/__tests__/protobufDrain.test.ts new file mode 100644 index 00000000..c1efe366 --- /dev/null +++ b/__tests__/protobufDrain.test.ts @@ -0,0 +1,245 @@ +import { afterEach, expect, test, vi } from 'vitest'; +import { + createClient, + createProtoService, + createServer, + Middleware, + Ok, + ProtoCodec, +} from '../protobuf'; +import { createMockTransportNetwork, readNextResult } from '../testUtil'; +import { TestService } from '../testUtil/fixtures/protobuf'; +import { waitFor } from '../testUtil/fixtures/cleanup'; +import { createPromiseWithResolvers } from '../transport/promises'; + +const cleanups: Array<() => void | Promise> = []; +afterEach(async () => { + for (const cleanup of cleanups.splice(0).reverse()) await cleanup(); +}); + +function setup( + maxResends?: number, + newSessionWaitMs = 1000, + middlewares: Array = [], +) { + const network = createMockTransportNetwork({ + client: { codec: ProtoCodec, baseIntervalMs: 1, maxJitterMs: 0 }, + server: { codec: ProtoCodec }, + }); + const gate = createPromiseWithResolvers(); + const uploadGate = createPromiseWithResolvers(); + const effects: Array = []; + const attempts: Array = []; + const service = createProtoService().define(TestService, { + echo: async ({ text }) => { + effects.push(text); + if (text === 'hold') await gate.promise; + + return Ok({ text: 'ok' }); + }, + countUp: ({ resWritable }) => { + resWritable.write(Ok({ value: 1 })); + }, + sum: async () => { + effects.push('sum'); + await uploadGate.promise; + + return Ok({ total: 1 }); + }, + }); + function startServer() { + const transport = network.getServerTransport(); + const server = createServer(transport, [service], { middlewares }); + transport.addEventListener('message', (message) => { + if (message.from === 'client' && message.procedureName === 'Echo') + attempts.push(transport.sessions.get(message.from)?.id); + }); + cleanups.push(async () => { + transport.close(); + await server.close(); + }); + + return { server, transport }; + } + cleanups.push(() => network.cleanup()); + const blocker = network.getClientTransport('blocker'); + const clientTransport = network.getClientTransport('client'); + clientTransport.reconnectOnConnectionDrop = false; + cleanups.push(() => { + gate.resolve(); + uploadGate.resolve(); + blocker.close(); + clientTransport.close(); + }); + const { server, transport } = startServer(); + const client = createClient( + TestService, + clientTransport, + 'SERVER', + maxResends === undefined + ? {} + : { resendRefusedCalls: { maxResends, newSessionWaitMs } }, + ); + async function block() { + void createClient(TestService, blocker, 'SERVER').echo({ text: 'hold' }); + await waitFor(() => expect(effects).toContain('hold')); + } + function reconnect() { + clientTransport.hardDisconnect(); + clientTransport.connect('SERVER'); + } + + return { + server, + transport, + client, + gate, + uploadGate, + effects, + attempts, + startServer, + block, + reconnect, + }; +} + +const refusal = { ok: false, payload: { code: 'SERVER_DRAINING' } }; +const success = { ok: true, payload: { text: 'ok' } }; + +test('drain refuses calls before middleware and handler run', async () => { + // Risk: a refused call that ran would be unsafe to resend. + const middlewareEffects: Array = []; + const f = setup(undefined, 1000, [ + ({ next }) => { + middlewareEffects.push('ran'); + next(); + }, + ]); + await f.block(); + middlewareEffects.length = 0; + void f.server.drain({ budgetMs: 10000 }); + expect(await f.client.echo({ text: 'run' })).toMatchObject(refusal); + expect(middlewareEffects).toEqual([]); + expect(f.effects).toEqual(['hold']); +}); + +test('drain waits for unary and client-streaming calls but not a server stream', async () => { + // Risk: drain drops accepted calls or waits for a long stream. + const f = setup(); + await readNextResult(f.client.countUp({ limit: 1 })); + const unary = f.client.echo({ text: 'hold' }); + const upload = f.client.sum().finalize(); + await waitFor(() => expect(f.effects).toEqual(['hold', 'sum'])); + const draining = f.server.drain({ budgetMs: 10000 }); + expect(f.server.drain({ budgetMs: 1 })).toBe(draining); + f.gate.resolve(); + expect(await unary).toMatchObject(success); + expect(f.transport.getStatus()).not.toBe('closed'); + f.uploadGate.resolve(); + expect(await upload).toMatchObject({ ok: true, payload: { total: 1 } }); + await draining; + expect(f.transport.getStatus()).toBe('closed'); +}); + +test('the drain budget closes a server whose unary call never returns', async () => { + // Risk: shutdown hangs. + const f = setup(); + await f.block(); + const draining = f.server.drain({ budgetMs: 100 }); + await vi.advanceTimersByTimeAsync(100); + expect(f.transport.getStatus()).toBe('closed'); + await draining; +}); + +test('a refused unary call succeeds only on a new session', async () => { + // Risk: replaying work on the draining session. + const f = setup(2); + await f.block(); + const draining = f.server.drain({ budgetMs: 10000 }); + const result = f.client.echo({ text: 'run' }); + await waitFor(() => expect(f.attempts).toHaveLength(1)); + await vi.advanceTimersByTimeAsync(20); + expect(f.attempts).toHaveLength(1); + f.gate.resolve(); + await draining; + f.startServer(); + f.reconnect(); + expect(await result).toMatchObject(success); + expect(f.attempts).toHaveLength(2); + expect(f.attempts[1]).not.toBe(f.attempts[0]); + expect(f.effects).toEqual(['hold', 'run']); +}); + +test('successive refusals stop at maxResends', async () => { + // Risk: unbounded replay. + const f = setup(2); + await f.block(); + void f.server.drain({ budgetMs: 10000 }); + const result = f.client.echo({ text: 'run' }); + for (let attempt = 1; attempt <= 3; attempt++) { + await waitFor(() => expect(f.attempts).toHaveLength(attempt)); + await vi.advanceTimersByTimeAsync(10); + f.reconnect(); + } + await vi.advanceTimersByTimeAsync(20); + expect(f.attempts).toHaveLength(3); + expect(await result).toMatchObject(refusal); + expect(f.effects).toEqual(['hold']); +}); + +test('the session wait deadline returns the refusal', async () => { + // Risk: a refused call stays pending forever. + const f = setup(2, 100); + await f.block(); + void f.server.drain({ budgetMs: 10000 }); + const result = f.client.echo({ text: 'run' }); + await waitFor(() => expect(f.attempts).toHaveLength(1)); + await vi.advanceTimersByTimeAsync(100); + expect(await result).toMatchObject(refusal); + expect(f.effects).toEqual(['hold']); +}); + +test('a refused client-streaming call is never resent', async () => { + // Risk: replaying a consumed request stream. + const f = setup(2); + await f.block(); + void f.server.drain({ budgetMs: 10000 }); + expect(await f.client.sum().finalize()).toMatchObject(refusal); + expect(f.effects).toEqual(['hold']); +}); + +test('drain called from middleware waits for its admitted unary call', async () => { + // Risk: drain closes a call before middleware admits its handler. + const f = setup(undefined, 1000, [ + ({ next }) => { + void f.server.drain({ budgetMs: 10000 }); + next(); + }, + ]); + const result = f.client.echo({ text: 'hold' }); + await waitFor(() => expect(f.effects).toContain('hold')); + expect(f.transport.getStatus()).not.toBe('closed'); + f.gate.resolve(); + expect(await result).toMatchObject(success); + await f.server.drain({ budgetMs: 10000 }); +}); + +test('abort while waiting preserves the refusal without a later resend', async () => { + // Risk: cancelled work still runs after reconnect. + const f = setup(2, 30000); + await f.block(); + const draining = f.server.drain({ budgetMs: 10000 }); + const controller = new AbortController(); + const result = f.client.echo({ text: 'run' }, { signal: controller.signal }); + await waitFor(() => expect(f.attempts).toHaveLength(1)); + await vi.advanceTimersByTimeAsync(10); + controller.abort(); + expect(await result).toMatchObject(refusal); + f.gate.resolve(); + await draining; + f.startServer(); + f.reconnect(); + await vi.advanceTimersByTimeAsync(50); + expect(f.effects).toEqual(['hold']); + expect(f.attempts).toHaveLength(1); +}); diff --git a/protobuf/client.ts b/protobuf/client.ts index 91fcc74c..5f56bce0 100644 --- a/protobuf/client.ts +++ b/protobuf/client.ts @@ -24,6 +24,11 @@ import { } from '../transport/message'; import { generateId } from '../transport/id'; import { ClientHandshakeOptions } from '../router/handshake'; +import { + ResendRefusedCallsOptions, + validateResendLimits, + waitForNewSession, +} from '../router/drain'; import { Err, Ok, type Result } from '../router/result'; import { Readable, @@ -37,6 +42,7 @@ import { type ClientError, isSerializedClientErrorResult } from './errors'; import { CANCEL_CODE, INVALID_REQUEST_CODE, + SERVER_DRAINING_CODE, UNEXPECTED_DISCONNECT_CODE, } from '../router/errors'; import { @@ -59,6 +65,7 @@ import type { export interface ClientOptions { readonly connectOnInvoke: boolean; readonly eagerlyConnect: boolean; + readonly resendRefusedCalls?: ResendRefusedCallsOptions; } const defaultClientOptions: ClientOptions = { @@ -95,6 +102,7 @@ export function createClient< } const clientOptions = { ...defaultClientOptions, ...providedClientOptions }; + validateResendLimits(clientOptions.resendRefusedCalls); if (clientOptions.eagerlyConnect) { transport.connect(serverId); } @@ -142,17 +150,16 @@ function createMethodCaller< } connectOnInvokeIfNeeded(clientOptions, transport, serverId); - const { resReadable } = startMethodCall( + + return unaryWithResends( service, method, transport, serverId, encodeMessageBytes(method.input, request), - true, + clientOptions.resendRefusedCalls, options?.signal, ); - - return getSingleMessage(resReadable, transport.log); }) as unknown as ClientMethod; case 'server_streaming': @@ -243,6 +250,49 @@ function createMethodCaller< } } +async function unaryWithResends< + Method extends DescMethod, + RejectionCodeSchema extends CustomHandshakeErrorCodeSchema, +>( + service: DescService, + method: Method, + transport: ClientTransport, + serverId: TransportClientId, + request: Uint8Array, + limits: ResendRefusedCallsOptions | undefined, + signal?: AbortSignal, +): Promise, ClientError>> { + for (let attempt = 0; ; attempt++) { + const { resReadable } = startMethodCall( + service, + method, + transport, + serverId, + request, + true, + signal, + ); + const sessionId = transport.sessions.get(serverId)?.id; + const result = await getSingleMessage(resReadable, transport.log); + if ( + !limits || + result.ok || + result.payload.code !== SERVER_DRAINING_CODE || + attempt >= limits.maxResends + ) { + return result; + } + const ready = await waitForNewSession( + transport, + serverId, + sessionId, + limits.newSessionWaitMs, + signal, + ); + if (!ready) return result; + } +} + function connectOnInvokeIfNeeded< RejectionCodeSchema extends CustomHandshakeErrorCodeSchema, >( diff --git a/protobuf/errors.ts b/protobuf/errors.ts index 7d45d115..adaf6d9b 100644 --- a/protobuf/errors.ts +++ b/protobuf/errors.ts @@ -1,6 +1,7 @@ import { CANCEL_CODE, INVALID_REQUEST_CODE, + SERVER_DRAINING_CODE, UNCAUGHT_ERROR_CODE, UNEXPECTED_DISCONNECT_CODE, } from '../router/errors'; @@ -40,6 +41,7 @@ export enum RiverErrorCode { export type ProtocolErrorCode = | typeof CANCEL_CODE | typeof INVALID_REQUEST_CODE + | typeof SERVER_DRAINING_CODE | typeof UNCAUGHT_ERROR_CODE | typeof UNEXPECTED_DISCONNECT_CODE; @@ -52,6 +54,7 @@ const riverErrorCodeSet = new Set(Object.values(RiverErrorCode)); const protocolErrorCodeSet = new Set([ CANCEL_CODE, INVALID_REQUEST_CODE, + SERVER_DRAINING_CODE, UNCAUGHT_ERROR_CODE, UNEXPECTED_DISCONNECT_CODE, ]); diff --git a/protobuf/index.ts b/protobuf/index.ts index c31e0eba..d1bc94ed 100644 --- a/protobuf/index.ts +++ b/protobuf/index.ts @@ -59,6 +59,7 @@ export { ReadableBrokenError } from '../router/streams'; export { CANCEL_CODE, INVALID_REQUEST_CODE, + SERVER_DRAINING_CODE, UNCAUGHT_ERROR_CODE, UNEXPECTED_DISCONNECT_CODE, } from '../router/errors'; diff --git a/protobuf/server.ts b/protobuf/server.ts index 898f50c5..9fc7e201 100644 --- a/protobuf/server.ts +++ b/protobuf/server.ts @@ -9,6 +9,7 @@ import { Value } from 'typebox/value'; import { context as otelContext, trace, type Span } from '@opentelemetry/api'; import { Logger } from '../logging'; import { ServerHandshakeOptions } from '../router/handshake'; +import { DrainController } from '../router/drain'; import { Err, Ok, type Result } from '../router/result'; import { ReadableImpl, WritableImpl } from '../router/streams'; import { Connection } from '../transport/connection'; @@ -45,6 +46,7 @@ import { import { CANCEL_CODE, INVALID_REQUEST_CODE, + SERVER_DRAINING_CODE, UNCAUGHT_ERROR_CODE, UNEXPECTED_DISCONNECT_CODE, } from '../router/errors'; @@ -101,6 +103,7 @@ interface ProcStream { */ export interface Server { readonly streams: Map; + drain(options: { budgetMs: number }): Promise; close(): Promise; } @@ -179,6 +182,15 @@ class ProtobufServer< private readonly maxCancelledStreamTombstonesPerSession: number; private unregisterTransportListeners: () => void; + private readonly drainController = new DrainController( + () => + [...this.streams.values()].some( + ({ method }) => + method.methodKind === 'unary' || + method.methodKind === 'client_streaming', + ), + () => this.transport.close(), + ); constructor( transport: ServerTransport< @@ -260,6 +272,23 @@ class ProtobufServer< return; } + if (!this.drainController.accepting) { + this.cancelStream( + message.from, + transport.getSessionBoundSendFn( + message.from, + newStreamProps.initialSession.id, + ), + message.streamId, + { + code: SERVER_DRAINING_CODE, + message: 'Server is draining; the procedure was not run', + }, + ); + + return; + } + createHandlerSpan( transport.tracer, newStreamProps.initialSession, @@ -314,6 +343,10 @@ class ProtobufServer< this.transport.addEventListener('transportStatus', handleTransportStatus); } + drain(options: { budgetMs: number }): Promise { + return this.drainController.drain(options); + } + async close() { this.unregisterTransportListeners(); @@ -425,6 +458,7 @@ class ProtobufServer< const cleanup = () => { finishedController.abort(); this.streams.delete(streamId); + this.drainController.responseFinished(); void runDeferredCleanups(); }; @@ -759,6 +793,8 @@ class ProtobufServer< }; /* eslint-enable @typescript-eslint/no-unsafe-assignment, @typescript-eslint/no-explicit-any */ + this.streams.set(streamId, procStream); + this.middlewares.reduceRight( (next: () => void, middleware: Middleware) => { return () => { @@ -773,10 +809,6 @@ class ProtobufServer< void runProcedureHandler(); }, )(); - - if (!finishedController.signal.aborted) { - this.streams.set(streamId, procStream); - } } private validateNewProcStream( diff --git a/router/client.ts b/router/client.ts index 2a59f9d6..038241a8 100644 --- a/router/client.ts +++ b/router/client.ts @@ -39,8 +39,11 @@ import { UNEXPECTED_DISCONNECT_CODE, validationErrorToRiverErrors, } from './errors'; -import { SessionState } from '../transport/sessionStateMachine/common'; -import { MAX_TIMER_DELAY_MS } from '../transport/transport'; +import { + ResendRefusedCallsOptions, + validateResendLimits, + waitForNewSession, +} from './drain'; export interface CallOptions { signal?: AbortSignal; @@ -208,7 +211,7 @@ function _createRecursiveProxy( } export interface ClientOptions { - resendRefusedCalls?: { maxResends: number; newSessionWaitMs: number }; + resendRefusedCalls?: ResendRefusedCallsOptions; connectOnInvoke: boolean; eagerlyConnect: boolean; /** @@ -265,53 +268,11 @@ export function createClient< const clientOptions = { ...defaultClientOptions, ...providedClientOptions }; const resend = clientOptions.resendRefusedCalls; - if ( - resend && - (!Number.isSafeInteger(resend.maxResends) || - !Number.isFinite(resend.newSessionWaitMs) || - resend.newSessionWaitMs < 0 || - resend.newSessionWaitMs > MAX_TIMER_DELAY_MS) - ) { - throw new RangeError('Invalid resendRefusedCalls limits'); - } + validateResendLimits(resend); if (clientOptions.eagerlyConnect) { transport.connect(serverId); } - function waitForNewSession( - sessionId: string | undefined, - timeoutMs: number, - signal?: AbortSignal, - ): Promise { - return new Promise((resolve) => { - const finish = (ready: boolean) => { - clearTimeout(timer); - transport.removeEventListener('sessionTransition', check); - transport.removeEventListener('transportStatus', check); - signal?.removeEventListener('abort', aborted); - resolve(ready); - }; - const aborted = () => finish(false); - const check = () => { - const session = transport.sessions.get(serverId); - if (signal?.aborted === true || transport.getStatus() === 'closed') { - finish(false); - } else if ( - session && - session.id !== sessionId && - session.state === SessionState.Connected - ) { - finish(true); - } - }; - const timer = setTimeout(() => finish(false), timeoutMs); - transport.addEventListener('sessionTransition', check); - transport.addEventListener('transportStatus', check); - signal?.addEventListener('abort', aborted, { once: true }); - check(); - }); - } - async function rpcWithResends( init: Static, serviceName: string, @@ -339,6 +300,8 @@ export function createClient< return result; } const ready = await waitForNewSession( + transport, + serverId, sessionId, limits.newSessionWaitMs, signal, diff --git a/router/drain.ts b/router/drain.ts new file mode 100644 index 00000000..848ac7e9 --- /dev/null +++ b/router/drain.ts @@ -0,0 +1,116 @@ +import { ClientTransport } from '../transport/client'; +import { Connection } from '../transport/connection'; +import { + CustomHandshakeErrorCodeSchema, + TransportClientId, +} from '../transport/message'; +import { createPromiseWithResolvers } from '../transport/promises'; +import { SessionState } from '../transport/sessionStateMachine/common'; +import { MAX_TIMER_DELAY_MS } from '../transport/transport'; + +export interface ResendRefusedCallsOptions { + maxResends: number; + newSessionWaitMs: number; +} + +export function validateResendLimits( + limits: ResendRefusedCallsOptions | undefined, +): void { + if ( + limits && + (!Number.isSafeInteger(limits.maxResends) || + !Number.isFinite(limits.newSessionWaitMs) || + limits.newSessionWaitMs < 0 || + limits.newSessionWaitMs > MAX_TIMER_DELAY_MS) + ) { + throw new RangeError('Invalid resendRefusedCalls limits'); + } +} + +export function waitForNewSession< + RejectionCodeSchema extends CustomHandshakeErrorCodeSchema, +>( + transport: ClientTransport, + serverId: TransportClientId, + sessionId: string | undefined, + timeoutMs: number, + signal?: AbortSignal, +): Promise { + return new Promise((resolve) => { + const finish = (ready: boolean) => { + clearTimeout(timer); + transport.removeEventListener('sessionTransition', check); + transport.removeEventListener('transportStatus', check); + signal?.removeEventListener('abort', aborted); + resolve(ready); + }; + const aborted = () => finish(false); + const check = () => { + const session = transport.sessions.get(serverId); + if (signal?.aborted === true || transport.getStatus() === 'closed') { + finish(false); + } else if ( + session && + session.id !== sessionId && + session.state === SessionState.Connected + ) { + finish(true); + } + }; + const timer = setTimeout(() => finish(false), timeoutMs); + transport.addEventListener('sessionTransition', check); + transport.addEventListener('transportStatus', check); + signal?.addEventListener('abort', aborted, { once: true }); + check(); + }); +} + +export class DrainController { + private state: + | { status: 'accepting' } + | { status: 'draining'; promise: Promise; finish: () => void } = { + status: 'accepting', + }; + + constructor( + private readonly hasPendingResponse: () => boolean, + private readonly closeTransport: () => void, + ) {} + + get accepting(): boolean { + return this.state.status === 'accepting'; + } + + responseFinished(): void { + if (this.state.status === 'draining' && !this.hasPendingResponse()) { + this.state.finish(); + } + } + + drain({ budgetMs }: { budgetMs: number }): Promise { + if ( + !Number.isFinite(budgetMs) || + budgetMs < 0 || + budgetMs > MAX_TIMER_DELAY_MS + ) { + throw new RangeError('budgetMs must be between 0 and 2147483647'); + } + switch (this.state.status) { + case 'draining': + return this.state.promise; + case 'accepting': { + const { promise: waiting, resolve: finish } = + createPromiseWithResolvers(); + const timer = setTimeout(finish, budgetMs); + const promise = waiting.then(() => { + clearTimeout(timer); + this.closeTransport(); + }); + this.state = { status: 'draining', promise, finish }; + this.responseFinished(); + + return promise; + } + } + } +} diff --git a/router/server.ts b/router/server.ts index 827b5690..db3c46da 100644 --- a/router/server.ts +++ b/router/server.ts @@ -50,8 +50,8 @@ import { Connection } from '../transport/connection'; import { ServerTransport } from '../transport/server'; import { ReadableImpl, WritableImpl } from './streams'; import { IdentifiedSession } from '../transport/sessionStateMachine/common'; -import { MAX_TIMER_DELAY_MS, SessionBoundSendFn } from '../transport/transport'; -import { createPromiseWithResolvers } from '../transport/promises'; +import { SessionBoundSendFn } from '../transport/transport'; +import { DrainController } from './drain'; type StreamId = string; @@ -149,11 +149,10 @@ class RiverServer< >; private unregisterTransportListeners: () => void; - private drainState: - | { status: 'accepting' } - | { status: 'draining'; promise: Promise; finish: () => void } = { - status: 'accepting', - }; + private readonly drainController = new DrainController( + () => this.hasPendingResponse(), + () => this.transport.close(), + ); constructor( transport: ServerTransport< @@ -239,7 +238,7 @@ class RiverServer< return; } - if (this.drainState.status !== 'accepting') { + if (!this.drainController.accepting) { const session = newStreamProps.initialSession; this.cancelStream( message.from, @@ -566,9 +565,7 @@ class RiverServer< const cleanup = () => { finishedController.abort(); this.streams.delete(streamId); - if (this.drainState.status === 'draining' && !this.hasPendingResponse()) { - this.drainState.finish(); - } + this.drainController.responseFinished(); void runDeferredCleanups(); }; @@ -1086,30 +1083,7 @@ class RiverServer< } drain({ budgetMs }: { budgetMs: number }): Promise { - if ( - !Number.isFinite(budgetMs) || - budgetMs < 0 || - budgetMs > MAX_TIMER_DELAY_MS - ) { - throw new RangeError('budgetMs must be between 0 and 2147483647'); - } - switch (this.drainState.status) { - case 'draining': - return this.drainState.promise; - case 'accepting': { - const { promise: waiting, resolve: finish } = - createPromiseWithResolvers(); - const timer = setTimeout(finish, budgetMs); - const promise = waiting.then(() => { - clearTimeout(timer); - this.transport.close(); - }); - this.drainState = { status: 'draining', promise, finish }; - if (!this.hasPendingResponse()) finish(); - - return promise; - } - } + return this.drainController.drain({ budgetMs }); } async close() { From 28f8a6e4a86f2c708eec1567858c3f0c7fa25806 Mon Sep 17 00:00:00 2001 From: Armaan Aggarwal <49625265+Armster15@users.noreply.github.com> Date: Wed, 21 Oct 2026 23:33:15 +0000 Subject: [PATCH 2/3] Keep protobuf handler errors from claiming SERVER_DRAINING Only the runtime admission refusal may carry the reserved code, so a handler that already ran is never resent. Remap a handler-returned SERVER_DRAINING to UNCAUGHT_ERROR and log a warning. Assisted-by: Replit --- __tests__/protobufDrain.test.ts | 20 ++++++++++++++++++++ protobuf/server.ts | 11 +++++++++++ 2 files changed, 31 insertions(+) diff --git a/__tests__/protobufDrain.test.ts b/__tests__/protobufDrain.test.ts index c1efe366..50657750 100644 --- a/__tests__/protobufDrain.test.ts +++ b/__tests__/protobufDrain.test.ts @@ -3,9 +3,12 @@ import { createClient, createProtoService, createServer, + Err, Middleware, Ok, ProtoCodec, + SERVER_DRAINING_CODE, + UNCAUGHT_ERROR_CODE, } from '../protobuf'; import { createMockTransportNetwork, readNextResult } from '../testUtil'; import { TestService } from '../testUtil/fixtures/protobuf'; @@ -34,6 +37,8 @@ function setup( echo: async ({ text }) => { effects.push(text); if (text === 'hold') await gate.promise; + if (text === 'claim') + return Err({ code: SERVER_DRAINING_CODE, message: 'upstream' }); return Ok({ text: 'ok' }); }, @@ -208,6 +213,21 @@ test('a refused client-streaming call is never resent', async () => { expect(f.effects).toEqual(['hold']); }); +test('a handler that returns the reserved drain code is not resent', async () => { + // Risk: a handler that already ran is replayed on the next session. + const f = setup(1); + const result = f.client.echo({ text: 'claim' }); + await waitFor(() => expect(f.attempts).toHaveLength(1)); + await vi.advanceTimersByTimeAsync(10); + f.reconnect(); + await vi.advanceTimersByTimeAsync(20); + expect(f.effects).toEqual(['claim']); + expect(await result).toMatchObject({ + ok: false, + payload: { code: UNCAUGHT_ERROR_CODE, message: 'upstream' }, + }); +}); + test('drain called from middleware waits for its admitted unary call', async () => { // Risk: drain closes a call before middleware admits its handler. const f = setup(undefined, 1000, [ diff --git a/protobuf/server.ts b/protobuf/server.ts index 9fc7e201..8a747e7a 100644 --- a/protobuf/server.ts +++ b/protobuf/server.ts @@ -482,6 +482,17 @@ class ProtobufServer< const resWritable = new WritableImpl({ writeCb: (response) => { + if (!response.ok && response.payload.code === SERVER_DRAINING_CODE) { + this.log?.warn( + `${service.typeName}.${method.name} handler returned the reserved ${SERVER_DRAINING_CODE} code`, + loggingMetadata, + ); + response = Err({ + code: UNCAUGHT_ERROR_CODE, + message: response.payload.message, + }); + } + let payload = response.ok ? response.payload : Err(response.payload); if (response.ok && raw === undefined) { payload = encodeMessageBytes( From d4c4c2cbc3ffb4d9a936f0040031246e8787cc6f Mon Sep 17 00:00:00 2001 From: Armaan Aggarwal <49625265+Armster15@users.noreply.github.com> Date: Thu, 22 Oct 2026 00:39:43 +0000 Subject: [PATCH 3/3] Address review: reserved drain code in types, shared resend loop Replace the runtime remap with a handler error type that excludes SERVER_DRAINING; only the runtime refusal can send it. Share one resend loop between both routers so it exists only when resend is on. Use an options object in the protobuf drain test setup and drop the Risk lines. Assisted-by: Replit --- __tests__/protobuf.test.ts | 7 ++- __tests__/protobufDrain.test.ts | 97 +++++++++++++++------------------ protobuf/client.ts | 73 ++++++++----------------- protobuf/errors.ts | 37 +++++++++++-- protobuf/index.ts | 1 + protobuf/server.ts | 11 ---- protobuf/types.ts | 12 ++-- router/client.ts | 57 ++++++------------- router/drain.ts | 37 ++++++++++++- 9 files changed, 162 insertions(+), 170 deletions(-) diff --git a/__tests__/protobuf.test.ts b/__tests__/protobuf.test.ts index 12e31c38..0919e917 100644 --- a/__tests__/protobuf.test.ts +++ b/__tests__/protobuf.test.ts @@ -4,6 +4,7 @@ import { Type } from 'typebox'; import { BinaryCodec, NaiveJsonCodec } from '../codec'; import { type ClientError, + type ProtocolError, Err, type Middleware, Ok, @@ -39,9 +40,9 @@ const protobufRouterMatrix = transports.flatMap((transport) => ].map((codec) => ({ transport, codec })), ); -function coerceReadableError( - error: ClientError | typeof ReadableBrokenError, -): ClientError { +function coerceReadableError( + error: E | typeof ReadableBrokenError, +): E | ProtocolError { return error.code === ReadableBrokenError.code ? { code: UNEXPECTED_DISCONNECT_CODE, message: error.message } : error; diff --git a/__tests__/protobufDrain.test.ts b/__tests__/protobufDrain.test.ts index 50657750..21046cb2 100644 --- a/__tests__/protobufDrain.test.ts +++ b/__tests__/protobufDrain.test.ts @@ -7,8 +7,8 @@ import { Middleware, Ok, ProtoCodec, + type ProtobufClientOptions, SERVER_DRAINING_CODE, - UNCAUGHT_ERROR_CODE, } from '../protobuf'; import { createMockTransportNetwork, readNextResult } from '../testUtil'; import { TestService } from '../testUtil/fixtures/protobuf'; @@ -20,11 +20,13 @@ afterEach(async () => { for (const cleanup of cleanups.splice(0).reverse()) await cleanup(); }); -function setup( - maxResends?: number, - newSessionWaitMs = 1000, - middlewares: Array = [], -) { +function setup({ + resendRefusedCalls, + middlewares = [], +}: { + resendRefusedCalls?: ResendRefusedCalls; + middlewares?: Array; +} = {}) { const network = createMockTransportNetwork({ client: { codec: ProtoCodec, baseIntervalMs: 1, maxJitterMs: 0 }, server: { codec: ProtoCodec }, @@ -37,8 +39,6 @@ function setup( echo: async ({ text }) => { effects.push(text); if (text === 'hold') await gate.promise; - if (text === 'claim') - return Err({ code: SERVER_DRAINING_CODE, message: 'upstream' }); return Ok({ text: 'ok' }); }, @@ -77,14 +77,9 @@ function setup( clientTransport.close(); }); const { server, transport } = startServer(); - const client = createClient( - TestService, - clientTransport, - 'SERVER', - maxResends === undefined - ? {} - : { resendRefusedCalls: { maxResends, newSessionWaitMs } }, - ); + const client = createClient(TestService, clientTransport, 'SERVER', { + resendRefusedCalls, + }); async function block() { void createClient(TestService, blocker, 'SERVER').echo({ text: 'hold' }); await waitFor(() => expect(effects).toContain('hold')); @@ -108,18 +103,22 @@ function setup( }; } +type ResendRefusedCalls = ProtobufClientOptions['resendRefusedCalls']; +const resendTwice = { maxResends: 2, newSessionWaitMs: 1000 }; + const refusal = { ok: false, payload: { code: 'SERVER_DRAINING' } }; const success = { ok: true, payload: { text: 'ok' } }; test('drain refuses calls before middleware and handler run', async () => { - // Risk: a refused call that ran would be unsafe to resend. const middlewareEffects: Array = []; - const f = setup(undefined, 1000, [ - ({ next }) => { - middlewareEffects.push('ran'); - next(); - }, - ]); + const f = setup({ + middlewares: [ + ({ next }) => { + middlewareEffects.push('ran'); + next(); + }, + ], + }); await f.block(); middlewareEffects.length = 0; void f.server.drain({ budgetMs: 10000 }); @@ -129,7 +128,6 @@ test('drain refuses calls before middleware and handler run', async () => { }); test('drain waits for unary and client-streaming calls but not a server stream', async () => { - // Risk: drain drops accepted calls or waits for a long stream. const f = setup(); await readNextResult(f.client.countUp({ limit: 1 })); const unary = f.client.echo({ text: 'hold' }); @@ -147,7 +145,6 @@ test('drain waits for unary and client-streaming calls but not a server stream', }); test('the drain budget closes a server whose unary call never returns', async () => { - // Risk: shutdown hangs. const f = setup(); await f.block(); const draining = f.server.drain({ budgetMs: 100 }); @@ -157,8 +154,7 @@ test('the drain budget closes a server whose unary call never returns', async () }); test('a refused unary call succeeds only on a new session', async () => { - // Risk: replaying work on the draining session. - const f = setup(2); + const f = setup({ resendRefusedCalls: resendTwice }); await f.block(); const draining = f.server.drain({ budgetMs: 10000 }); const result = f.client.echo({ text: 'run' }); @@ -176,8 +172,7 @@ test('a refused unary call succeeds only on a new session', async () => { }); test('successive refusals stop at maxResends', async () => { - // Risk: unbounded replay. - const f = setup(2); + const f = setup({ resendRefusedCalls: resendTwice }); await f.block(); void f.server.drain({ budgetMs: 10000 }); const result = f.client.echo({ text: 'run' }); @@ -193,8 +188,9 @@ test('successive refusals stop at maxResends', async () => { }); test('the session wait deadline returns the refusal', async () => { - // Risk: a refused call stays pending forever. - const f = setup(2, 100); + const f = setup({ + resendRefusedCalls: { maxResends: 2, newSessionWaitMs: 100 }, + }); await f.block(); void f.server.drain({ budgetMs: 10000 }); const result = f.client.echo({ text: 'run' }); @@ -205,37 +201,29 @@ test('the session wait deadline returns the refusal', async () => { }); test('a refused client-streaming call is never resent', async () => { - // Risk: replaying a consumed request stream. - const f = setup(2); + const f = setup({ resendRefusedCalls: resendTwice }); await f.block(); void f.server.drain({ budgetMs: 10000 }); expect(await f.client.sum().finalize()).toMatchObject(refusal); expect(f.effects).toEqual(['hold']); }); -test('a handler that returns the reserved drain code is not resent', async () => { - // Risk: a handler that already ran is replayed on the next session. - const f = setup(1); - const result = f.client.echo({ text: 'claim' }); - await waitFor(() => expect(f.attempts).toHaveLength(1)); - await vi.advanceTimersByTimeAsync(10); - f.reconnect(); - await vi.advanceTimersByTimeAsync(20); - expect(f.effects).toEqual(['claim']); - expect(await result).toMatchObject({ - ok: false, - payload: { code: UNCAUGHT_ERROR_CODE, message: 'upstream' }, +test('handlers cannot return the reserved drain code', () => { + createProtoService().define(TestService, { + // @ts-expect-error only the runtime may send SERVER_DRAINING + echo: () => Err({ code: SERVER_DRAINING_CODE, message: 'upstream' }), }); }); test('drain called from middleware waits for its admitted unary call', async () => { - // Risk: drain closes a call before middleware admits its handler. - const f = setup(undefined, 1000, [ - ({ next }) => { - void f.server.drain({ budgetMs: 10000 }); - next(); - }, - ]); + const f = setup({ + middlewares: [ + ({ next }) => { + void f.server.drain({ budgetMs: 10000 }); + next(); + }, + ], + }); const result = f.client.echo({ text: 'hold' }); await waitFor(() => expect(f.effects).toContain('hold')); expect(f.transport.getStatus()).not.toBe('closed'); @@ -245,8 +233,9 @@ test('drain called from middleware waits for its admitted unary call', async () }); test('abort while waiting preserves the refusal without a later resend', async () => { - // Risk: cancelled work still runs after reconnect. - const f = setup(2, 30000); + const f = setup({ + resendRefusedCalls: { maxResends: 2, newSessionWaitMs: 30000 }, + }); await f.block(); const draining = f.server.drain({ budgetMs: 10000 }); const controller = new AbortController(); diff --git a/protobuf/client.ts b/protobuf/client.ts index 5f56bce0..22e115d1 100644 --- a/protobuf/client.ts +++ b/protobuf/client.ts @@ -27,7 +27,7 @@ import { ClientHandshakeOptions } from '../router/handshake'; import { ResendRefusedCallsOptions, validateResendLimits, - waitForNewSession, + withRefusedCallResends, } from '../router/drain'; import { Err, Ok, type Result } from '../router/result'; import { @@ -42,7 +42,6 @@ import { type ClientError, isSerializedClientErrorResult } from './errors'; import { CANCEL_CODE, INVALID_REQUEST_CODE, - SERVER_DRAINING_CODE, UNEXPECTED_DISCONNECT_CODE, } from '../router/errors'; import { @@ -151,13 +150,30 @@ function createMethodCaller< connectOnInvokeIfNeeded(clientOptions, transport, serverId); - return unaryWithResends( - service, - method, + const requestBytes = encodeMessageBytes(method.input, request); + const call = () => + getSingleMessage( + startMethodCall( + service, + method, + transport, + serverId, + requestBytes, + true, + options?.signal, + ).resReadable, + transport.log, + ); + const resend = clientOptions.resendRefusedCalls; + if (!resend) { + return call(); + } + + return withRefusedCallResends( + call, transport, serverId, - encodeMessageBytes(method.input, request), - clientOptions.resendRefusedCalls, + resend, options?.signal, ); }) as unknown as ClientMethod; @@ -250,49 +266,6 @@ function createMethodCaller< } } -async function unaryWithResends< - Method extends DescMethod, - RejectionCodeSchema extends CustomHandshakeErrorCodeSchema, ->( - service: DescService, - method: Method, - transport: ClientTransport, - serverId: TransportClientId, - request: Uint8Array, - limits: ResendRefusedCallsOptions | undefined, - signal?: AbortSignal, -): Promise, ClientError>> { - for (let attempt = 0; ; attempt++) { - const { resReadable } = startMethodCall( - service, - method, - transport, - serverId, - request, - true, - signal, - ); - const sessionId = transport.sessions.get(serverId)?.id; - const result = await getSingleMessage(resReadable, transport.log); - if ( - !limits || - result.ok || - result.payload.code !== SERVER_DRAINING_CODE || - attempt >= limits.maxResends - ) { - return result; - } - const ready = await waitForNewSession( - transport, - serverId, - sessionId, - limits.newSessionWaitMs, - signal, - ); - if (!ready) return result; - } -} - function connectOnInvokeIfNeeded< RejectionCodeSchema extends CustomHandshakeErrorCodeSchema, >( diff --git a/protobuf/errors.ts b/protobuf/errors.ts index adaf6d9b..68318b45 100644 --- a/protobuf/errors.ts +++ b/protobuf/errors.ts @@ -41,20 +41,22 @@ export enum RiverErrorCode { export type ProtocolErrorCode = | typeof CANCEL_CODE | typeof INVALID_REQUEST_CODE - | typeof SERVER_DRAINING_CODE | typeof UNCAUGHT_ERROR_CODE | typeof UNEXPECTED_DISCONNECT_CODE; /** - * Error codes visible from the protobuf client surface. + * Error codes visible from the protobuf client surface. `SERVER_DRAINING` is + * sent only by the runtime, before middleware or handlers run. */ -export type ClientErrorCode = RiverErrorCode | ProtocolErrorCode; +export type ClientErrorCode = + | RiverErrorCode + | ProtocolErrorCode + | typeof SERVER_DRAINING_CODE; const riverErrorCodeSet = new Set(Object.values(RiverErrorCode)); const protocolErrorCodeSet = new Set([ CANCEL_CODE, INVALID_REQUEST_CODE, - SERVER_DRAINING_CODE, UNCAUGHT_ERROR_CODE, UNEXPECTED_DISCONNECT_CODE, ]); @@ -88,6 +90,14 @@ export interface ProtocolError extends ClientError { readonly code: ProtocolErrorCode; } +/** + * An error a handler may return. Excludes `SERVER_DRAINING`, which tells + * clients the call never ran and is safe to resend. + */ +export interface HandlerError extends ClientError { + readonly code: RiverErrorCode | ProtocolErrorCode; +} + /** * Returns true when the given value matches the wire shape of a client error * with a canonical RPC error code. @@ -123,7 +133,11 @@ export function isProtocolError(value: unknown): value is ProtocolError { * Returns true when the given value is any client-visible error. */ export function isClientError(value: unknown): value is ClientError { - return isRiverError(value) || isProtocolError(value); + return ( + isRiverError(value) || + isProtocolError(value) || + isServerDrainingError(value) + ); } /** @@ -152,6 +166,19 @@ function isRiverErrorCode(value: unknown): value is RiverErrorCode { return typeof value === 'string' && riverErrorCodeSet.has(value); } +function isServerDrainingError(value: unknown): value is ClientError { + if (!(typeof value === 'object' && value !== null)) { + return false; + } + + const candidate = value as Partial; + + return ( + candidate.code === SERVER_DRAINING_CODE && + typeof candidate.message === 'string' + ); +} + function isProtocolErrorCode(value: unknown): value is ProtocolErrorCode { return typeof value === 'string' && protocolErrorCodeSet.has(value); } diff --git a/protobuf/index.ts b/protobuf/index.ts index d1bc94ed..e988c6b6 100644 --- a/protobuf/index.ts +++ b/protobuf/index.ts @@ -13,6 +13,7 @@ export { export type { ClientError, ClientErrorCode, + HandlerError, ProtocolError, ProtocolErrorCode, RiverErrorDetail, diff --git a/protobuf/server.ts b/protobuf/server.ts index 8a747e7a..9fc7e201 100644 --- a/protobuf/server.ts +++ b/protobuf/server.ts @@ -482,17 +482,6 @@ class ProtobufServer< const resWritable = new WritableImpl({ writeCb: (response) => { - if (!response.ok && response.payload.code === SERVER_DRAINING_CODE) { - this.log?.warn( - `${service.typeName}.${method.name} handler returned the reserved ${SERVER_DRAINING_CODE} code`, - loggingMetadata, - ); - response = Err({ - code: UNCAUGHT_ERROR_CODE, - message: response.payload.message, - }); - } - let payload = response.ok ? response.payload : Err(response.payload); if (response.ok && raw === undefined) { payload = encodeMessageBytes( diff --git a/protobuf/types.ts b/protobuf/types.ts index 71ae7d3e..66d03d4b 100644 --- a/protobuf/types.ts +++ b/protobuf/types.ts @@ -11,13 +11,13 @@ import type { import type { Result } from '../router/result'; import type { Readable, Writable } from '../router/streams'; import type { ProtobufHandlerContext } from './context'; -import type { ClientError, ProtocolError } from './errors'; +import type { ClientError, HandlerError, ProtocolError } from './errors'; type Awaitable = T | PromiseLike; type HandlerResult = Result< MessageInitShape, - ClientError + HandlerError >; /** @@ -125,23 +125,23 @@ type RawMethodImpl< ? ( request: Request, ctx: HandlerContext, - ) => Awaitable> + ) => Awaitable> : Method extends DescMethodServerStreaming ? (param: { readonly request: Request; readonly ctx: HandlerContext; - readonly resWritable: Writable>; + readonly resWritable: Writable>; }) => Awaitable : Method extends DescMethodClientStreaming ? (param: { readonly ctx: HandlerContext; readonly reqReadable: Readable; - }) => Awaitable> + }) => Awaitable> : Method extends DescMethodBiDiStreaming ? (param: { readonly ctx: HandlerContext; readonly reqReadable: Readable; - readonly resWritable: Writable>; + readonly resWritable: Writable>; }) => Awaitable : never; diff --git a/router/client.ts b/router/client.ts index 038241a8..cc949f80 100644 --- a/router/client.ts +++ b/router/client.ts @@ -34,7 +34,6 @@ import { PayloadType, ValidProcType } from './procedures'; import { BaseErrorSchemaType, CANCEL_CODE, - SERVER_DRAINING_CODE, ReaderErrorResultSchema, UNEXPECTED_DISCONNECT_CODE, validationErrorToRiverErrors, @@ -42,7 +41,7 @@ import { import { ResendRefusedCallsOptions, validateResendLimits, - waitForNewSession, + withRefusedCallResends, } from './drain'; export interface CallOptions { @@ -273,43 +272,6 @@ export function createClient< transport.connect(serverId); } - async function rpcWithResends( - init: Static, - serviceName: string, - procedureName: string, - limits: NonNullable, - signal?: AbortSignal, - ) { - for (let attempt = 0; ; attempt++) { - const response = handleProc( - 'rpc', - transport, - serverId, - init, - serviceName, - procedureName, - signal, - ); - const sessionId = transport.sessions.get(serverId)?.id; - const result = await response; - if ( - result.ok || - result.payload.code !== SERVER_DRAINING_CODE || - attempt >= limits.maxResends - ) { - return result; - } - const ready = await waitForNewSession( - transport, - serverId, - sessionId, - limits.newSessionWaitMs, - signal, - ); - if (!ready) return result; - } - } - return _createRecursiveProxy((opts) => { const [serviceName, procName, procMethod] = [...opts.path]; if (!(serviceName && procName && procMethod)) { @@ -341,7 +303,22 @@ export function createClient< ); if (procMethod === 'rpc' && resend) { - return rpcWithResends(init, serviceName, procName, resend, merged.signal); + return withRefusedCallResends( + () => + handleProc( + 'rpc', + transport, + serverId, + init, + serviceName, + procName, + merged.signal, + ), + transport, + serverId, + resend, + merged.signal, + ); } return handleProc( diff --git a/router/drain.ts b/router/drain.ts index 848ac7e9..bc152a1b 100644 --- a/router/drain.ts +++ b/router/drain.ts @@ -7,6 +7,8 @@ import { import { createPromiseWithResolvers } from '../transport/promises'; import { SessionState } from '../transport/sessionStateMachine/common'; import { MAX_TIMER_DELAY_MS } from '../transport/transport'; +import { SERVER_DRAINING_CODE } from './errors'; +import { type ErrorPayload, type Result } from './result'; export interface ResendRefusedCallsOptions { maxResends: number; @@ -27,7 +29,40 @@ export function validateResendLimits( } } -export function waitForNewSession< +export async function withRefusedCallResends< + T, + E extends ErrorPayload, + RejectionCodeSchema extends CustomHandshakeErrorCodeSchema, +>( + call: () => Promise>, + transport: ClientTransport, + serverId: TransportClientId, + limits: ResendRefusedCallsOptions, + signal?: AbortSignal, +): Promise> { + for (let attempt = 0; ; attempt++) { + const response = call(); + const sessionId = transport.sessions.get(serverId)?.id; + const result = await response; + if ( + result.ok || + result.payload.code !== SERVER_DRAINING_CODE || + attempt >= limits.maxResends + ) { + return result; + } + const ready = await waitForNewSession( + transport, + serverId, + sessionId, + limits.newSessionWaitMs, + signal, + ); + if (!ready) return result; + } +} + +function waitForNewSession< RejectionCodeSchema extends CustomHandshakeErrorCodeSchema, >( transport: ClientTransport,