From 696c8f4fb149550688bb2f5b462f4aac7043cbbc Mon Sep 17 00:00:00 2001 From: Armaan Aggarwal <49625265+Armster15@users.noreply.github.com> Date: Wed, 21 Oct 2026 22:00:52 +0000 Subject: [PATCH 1/5] Add graceful server drain and opt-in refused RPC resend Assisted-by: Replit --- README.md | 39 ++ .../__snapshots__/serialize.test.ts.snap | 432 ++++++++++++++++++ __tests__/drain.test.ts | 256 +++++++++++ router/client.ts | 84 +++- router/errors.ts | 5 + router/index.ts | 1 + router/server.ts | 59 ++- 7 files changed, 863 insertions(+), 13 deletions(-) create mode 100644 __tests__/drain.test.ts diff --git a/README.md b/README.md index eb674811..f8bfe0c9 100644 --- a/README.md +++ b/README.md @@ -170,6 +170,45 @@ return Ok({ result: 42 }); return Err({ code: 'INVALID_INPUT', message: 'Value must be positive' }); ``` +#### Graceful server drain + +The TypeBox router server exposes `server.drain({ budgetMs: number }): Promise`. +The budget must be finite and non-negative. Drain immediately stops accepting new +procedure invocations. RPCs, uploads, subscriptions, and streams receive an error +result with the reserved code `SERVER_DRAINING` (exported as +`SERVER_DRAINING_CODE`). This code guarantees that the handler and middleware did +not run. Application handlers must not emit this reserved code. + +Existing calls may finish until the budget expires. Drain then closes the entire +server transport and deletes its sessions; these sessions cannot be resumed. +The promise resolves when the transport has closed. Repeated calls share the +first drain's budget and promise. An idle server closes immediately. +Closing the transport cancels unfinished calls but cannot force JavaScript +handlers to stop; handlers should observe their cancellation signal. +Call `server.close()` separately to dispose service and extended-context resources. +Do not share a draining transport with another router. + +Clients can opt in to resending refused unary RPCs: + +```ts +const client = createClient(transport, 'SERVER', { + resendRefusedCalls: { maxResends: 3, newSessionWaitMs: 30_000 }, +}); +``` + +This option is off by default. Only `SERVER_DRAINING` results are retried. +Each retry waits for a connected session with a different session ID, not just +a new connection on the old session. The transport's existing reconnection policy +still applies; route reconnects to a healthy server. Each wait has its own timeout. +The original refusal is returned if the retry budget or wait expires, the +transport closes, or the caller aborts while waiting. Normal call cancellation +applies during an attempt. Limits must be finite and non-negative, and +`maxResends` must be a safe integer. + +Uploads, subscriptions, and streams are never replayed. There is no unbounded +backoff after exhaustion. These APIs do not currently apply to the experimental +protobuf router. + #### Custom Error Types You can define custom error schemas for your procedures: diff --git a/__tests__/__snapshots__/serialize.test.ts.snap b/__tests__/__snapshots__/serialize.test.ts.snap index 72114ad8..a24ec354 100644 --- a/__tests__/__snapshots__/serialize.test.ts.snap +++ b/__tests__/__snapshots__/serialize.test.ts.snap @@ -19,6 +19,22 @@ exports[`serialize server to jsonschema > serialize entire service schema 1`] = "add": { "errors": { "anyOf": [ + { + "properties": { + "code": { + "const": "SERVER_DRAINING", + "type": "string", + }, + "message": { + "type": "string", + }, + }, + "required": [ + "code", + "message", + ], + "type": "object", + }, { "properties": { "code": { @@ -142,6 +158,22 @@ exports[`serialize server to jsonschema > serialize entire service schema 1`] = "array": { "errors": { "anyOf": [ + { + "properties": { + "code": { + "const": "SERVER_DRAINING", + "type": "string", + }, + "message": { + "type": "string", + }, + }, + "required": [ + "code", + "message", + ], + "type": "object", + }, { "properties": { "code": { @@ -260,6 +292,22 @@ exports[`serialize server to jsonschema > serialize entire service schema 1`] = "arrayStream": { "errors": { "anyOf": [ + { + "properties": { + "code": { + "const": "SERVER_DRAINING", + "type": "string", + }, + "message": { + "type": "string", + }, + }, + "required": [ + "code", + "message", + ], + "type": "object", + }, { "properties": { "code": { @@ -382,6 +430,22 @@ exports[`serialize server to jsonschema > serialize entire service schema 1`] = "echo": { "errors": { "anyOf": [ + { + "properties": { + "code": { + "const": "SERVER_DRAINING", + "type": "string", + }, + "message": { + "type": "string", + }, + }, + "required": [ + "code", + "message", + ], + "type": "object", + }, { "properties": { "code": { @@ -514,6 +578,22 @@ exports[`serialize server to jsonschema > serialize entire service schema 1`] = "description": "Echos back whatever we sent", "errors": { "anyOf": [ + { + "properties": { + "code": { + "const": "SERVER_DRAINING", + "type": "string", + }, + "message": { + "type": "string", + }, + }, + "required": [ + "code", + "message", + ], + "type": "object", + }, { "properties": { "code": { @@ -675,6 +755,22 @@ exports[`serialize server to jsonschema > serialize entire service schema 1`] = "echoWithPrefix": { "errors": { "anyOf": [ + { + "properties": { + "code": { + "const": "SERVER_DRAINING", + "type": "string", + }, + "message": { + "type": "string", + }, + }, + "required": [ + "code", + "message", + ], + "type": "object", + }, { "properties": { "code": { @@ -813,6 +909,22 @@ exports[`serialize server to jsonschema > serialize entire service schema 1`] = "unimplementedSubscription": { "errors": { "anyOf": [ + { + "properties": { + "code": { + "const": "SERVER_DRAINING", + "type": "string", + }, + "message": { + "type": "string", + }, + }, + "required": [ + "code", + "message", + ], + "type": "object", + }, { "properties": { "code": { @@ -922,6 +1034,22 @@ exports[`serialize server to jsonschema > serialize entire service schema 1`] = "unimplementedUpload": { "errors": { "anyOf": [ + { + "properties": { + "code": { + "const": "SERVER_DRAINING", + "type": "string", + }, + "message": { + "type": "string", + }, + }, + "required": [ + "code", + "message", + ], + "type": "object", + }, { "properties": { "code": { @@ -1044,6 +1172,22 @@ exports[`serialize service to jsonschema > serialize backwards compatible with v "add": { "errors": { "anyOf": [ + { + "properties": { + "code": { + "const": "SERVER_DRAINING", + "type": "string", + }, + "message": { + "type": "string", + }, + }, + "required": [ + "code", + "message", + ], + "type": "object", + }, { "properties": { "code": { @@ -1167,6 +1311,22 @@ exports[`serialize service to jsonschema > serialize backwards compatible with v "array": { "errors": { "anyOf": [ + { + "properties": { + "code": { + "const": "SERVER_DRAINING", + "type": "string", + }, + "message": { + "type": "string", + }, + }, + "required": [ + "code", + "message", + ], + "type": "object", + }, { "properties": { "code": { @@ -1285,6 +1445,22 @@ exports[`serialize service to jsonschema > serialize backwards compatible with v "arrayStream": { "errors": { "anyOf": [ + { + "properties": { + "code": { + "const": "SERVER_DRAINING", + "type": "string", + }, + "message": { + "type": "string", + }, + }, + "required": [ + "code", + "message", + ], + "type": "object", + }, { "properties": { "code": { @@ -1407,6 +1583,22 @@ exports[`serialize service to jsonschema > serialize backwards compatible with v "echo": { "errors": { "anyOf": [ + { + "properties": { + "code": { + "const": "SERVER_DRAINING", + "type": "string", + }, + "message": { + "type": "string", + }, + }, + "required": [ + "code", + "message", + ], + "type": "object", + }, { "properties": { "code": { @@ -1539,6 +1731,22 @@ exports[`serialize service to jsonschema > serialize backwards compatible with v "description": "Echos back whatever we sent", "errors": { "anyOf": [ + { + "properties": { + "code": { + "const": "SERVER_DRAINING", + "type": "string", + }, + "message": { + "type": "string", + }, + }, + "required": [ + "code", + "message", + ], + "type": "object", + }, { "properties": { "code": { @@ -1700,6 +1908,22 @@ exports[`serialize service to jsonschema > serialize backwards compatible with v "echoWithPrefix": { "errors": { "anyOf": [ + { + "properties": { + "code": { + "const": "SERVER_DRAINING", + "type": "string", + }, + "message": { + "type": "string", + }, + }, + "required": [ + "code", + "message", + ], + "type": "object", + }, { "properties": { "code": { @@ -1838,6 +2062,22 @@ exports[`serialize service to jsonschema > serialize backwards compatible with v "unimplementedSubscription": { "errors": { "anyOf": [ + { + "properties": { + "code": { + "const": "SERVER_DRAINING", + "type": "string", + }, + "message": { + "type": "string", + }, + }, + "required": [ + "code", + "message", + ], + "type": "object", + }, { "properties": { "code": { @@ -1947,6 +2187,22 @@ exports[`serialize service to jsonschema > serialize backwards compatible with v "unimplementedUpload": { "errors": { "anyOf": [ + { + "properties": { + "code": { + "const": "SERVER_DRAINING", + "type": "string", + }, + "message": { + "type": "string", + }, + }, + "required": [ + "code", + "message", + ], + "type": "object", + }, { "properties": { "code": { @@ -2067,6 +2323,22 @@ exports[`serialize service to jsonschema > serialize basic service 1`] = ` "add": { "errors": { "anyOf": [ + { + "properties": { + "code": { + "const": "SERVER_DRAINING", + "type": "string", + }, + "message": { + "type": "string", + }, + }, + "required": [ + "code", + "message", + ], + "type": "object", + }, { "properties": { "code": { @@ -2190,6 +2462,22 @@ exports[`serialize service to jsonschema > serialize basic service 1`] = ` "array": { "errors": { "anyOf": [ + { + "properties": { + "code": { + "const": "SERVER_DRAINING", + "type": "string", + }, + "message": { + "type": "string", + }, + }, + "required": [ + "code", + "message", + ], + "type": "object", + }, { "properties": { "code": { @@ -2308,6 +2596,22 @@ exports[`serialize service to jsonschema > serialize basic service 1`] = ` "arrayStream": { "errors": { "anyOf": [ + { + "properties": { + "code": { + "const": "SERVER_DRAINING", + "type": "string", + }, + "message": { + "type": "string", + }, + }, + "required": [ + "code", + "message", + ], + "type": "object", + }, { "properties": { "code": { @@ -2430,6 +2734,22 @@ exports[`serialize service to jsonschema > serialize basic service 1`] = ` "echo": { "errors": { "anyOf": [ + { + "properties": { + "code": { + "const": "SERVER_DRAINING", + "type": "string", + }, + "message": { + "type": "string", + }, + }, + "required": [ + "code", + "message", + ], + "type": "object", + }, { "properties": { "code": { @@ -2562,6 +2882,22 @@ exports[`serialize service to jsonschema > serialize basic service 1`] = ` "description": "Echos back whatever we sent", "errors": { "anyOf": [ + { + "properties": { + "code": { + "const": "SERVER_DRAINING", + "type": "string", + }, + "message": { + "type": "string", + }, + }, + "required": [ + "code", + "message", + ], + "type": "object", + }, { "properties": { "code": { @@ -2723,6 +3059,22 @@ exports[`serialize service to jsonschema > serialize basic service 1`] = ` "echoWithPrefix": { "errors": { "anyOf": [ + { + "properties": { + "code": { + "const": "SERVER_DRAINING", + "type": "string", + }, + "message": { + "type": "string", + }, + }, + "required": [ + "code", + "message", + ], + "type": "object", + }, { "properties": { "code": { @@ -2861,6 +3213,22 @@ exports[`serialize service to jsonschema > serialize basic service 1`] = ` "unimplementedSubscription": { "errors": { "anyOf": [ + { + "properties": { + "code": { + "const": "SERVER_DRAINING", + "type": "string", + }, + "message": { + "type": "string", + }, + }, + "required": [ + "code", + "message", + ], + "type": "object", + }, { "properties": { "code": { @@ -2970,6 +3338,22 @@ exports[`serialize service to jsonschema > serialize basic service 1`] = ` "unimplementedUpload": { "errors": { "anyOf": [ + { + "properties": { + "code": { + "const": "SERVER_DRAINING", + "type": "string", + }, + "message": { + "type": "string", + }, + }, + "required": [ + "code", + "message", + ], + "type": "object", + }, { "properties": { "code": { @@ -3090,6 +3474,22 @@ exports[`serialize service to jsonschema > serialize service with binary 1`] = ` "getFile": { "errors": { "anyOf": [ + { + "properties": { + "code": { + "const": "SERVER_DRAINING", + "type": "string", + }, + "message": { + "type": "string", + }, + }, + "required": [ + "code", + "message", + ], + "type": "object", + }, { "properties": { "code": { @@ -3264,6 +3664,22 @@ exports[`serialize service to jsonschema > serialize service with errors 1`] = ` ], "type": "object", }, + { + "properties": { + "code": { + "const": "SERVER_DRAINING", + "type": "string", + }, + "message": { + "type": "string", + }, + }, + "required": [ + "code", + "message", + ], + "type": "object", + }, { "properties": { "code": { @@ -3407,6 +3823,22 @@ exports[`serialize service to jsonschema > serialize service with errors 1`] = ` ], "type": "object", }, + { + "properties": { + "code": { + "const": "SERVER_DRAINING", + "type": "string", + }, + "message": { + "type": "string", + }, + }, + "required": [ + "code", + "message", + ], + "type": "object", + }, { "properties": { "code": { diff --git a/__tests__/drain.test.ts b/__tests__/drain.test.ts new file mode 100644 index 00000000..ebe7fa27 --- /dev/null +++ b/__tests__/drain.test.ts @@ -0,0 +1,256 @@ +import { afterEach, expect, test, vi } from 'vitest'; +import { Type } from 'typebox'; +import { + ClientOptions, + createClient, + createServer, + createServiceSchema, + Ok, + Procedure, + SERVER_DRAINING_CODE, +} from '../router'; +import { createMockTransportNetwork, readNextResult } from '../testUtil'; +import { waitFor } from '../testUtil/fixtures/cleanup'; + +const cleanups: Array<() => void | Promise> = []; +afterEach(async () => { + for (const cleanup of cleanups.splice(0).reverse()) await cleanup(); +}); + +async function setup(options: Partial = {}) { + const network = createMockTransportNetwork(); + cleanups.push(() => network.cleanup()); + const clientTransport = network.getClientTransport('client'); + cleanups.push(() => clientTransport.close()); + const ran = vi.fn(); + let release!: () => void; + const gate = new Promise((resolve) => { + release = resolve; + }); + const services = { + test: createServiceSchema().define({ + hold: Procedure.rpc({ + requestInit: Type.Object({}), + responseData: Type.Null(), + async handler() { + await gate; + + return Ok(null); + }, + }), + rpc: Procedure.rpc({ + requestInit: Type.Object({}), + responseData: Type.Null(), + async handler() { + ran(); + + return Ok(null); + }, + }), + upload: Procedure.upload({ + requestInit: Type.Object({}), + requestData: Type.Null(), + responseData: Type.Null(), + async handler() { + ran(); + + return Ok(null); + }, + }), + subscription: Procedure.subscription({ + requestInit: Type.Object({}), + responseData: Type.Null(), + async handler({ resWritable }) { + ran(); + resWritable.close(); + }, + }), + stream: Procedure.stream({ + requestInit: Type.Object({}), + requestData: Type.Null(), + responseData: Type.Null(), + async handler({ resWritable }) { + ran(); + resWritable.close(); + }, + }), + }), + }; + function startServer() { + const transport = network.getServerTransport(); + const server = createServer(transport, services); + cleanups.push(async () => { + transport.close(); + await server.close(); + }); + + return { server, transport }; + } + const { server, transport } = startServer(); + const client = createClient( + clientTransport, + 'SERVER', + options, + ); + const held = client.test.hold.rpc({}); + await waitFor(() => expect(server.streams.size).toBe(1)); + cleanups.push(release); + + return { + client, + clientTransport, + server, + transport, + release, + held, + ran, + startServer, + }; +} + +const refusal = { + ok: false, + payload: { + code: SERVER_DRAINING_CODE, + message: 'Server is draining; the procedure was not run', + }, +}; + +test.each(['rpc', 'upload', 'subscription', 'stream'] as const)( + 'drain refuses new %s calls without running the handler', + async (kind) => { + const f = await setup(); + const draining = f.server.drain({ budgetMs: 10_000 }); + switch (kind) { + case 'rpc': + expect(await f.client.test.rpc.rpc({})).toEqual(refusal); + break; + case 'upload': + expect(await f.client.test.upload.upload({}).finalize()).toEqual( + refusal, + ); + break; + case 'subscription': + expect( + await readNextResult( + f.client.test.subscription.subscribe({}).resReadable, + ), + ).toEqual(refusal); + break; + case 'stream': + expect( + await readNextResult(f.client.test.stream.stream({}).resReadable), + ).toEqual(refusal); + break; + } + expect(f.ran).not.toHaveBeenCalled(); + f.release(); + await draining; + }, +); + +test('drain lets an existing RPC finish and removes its session', async () => { + const f = await setup(); + const done = vi.fn(); + const draining = f.server.drain({ budgetMs: 10_000 }).then(done); + await vi.advanceTimersByTimeAsync(50); + expect(done).not.toHaveBeenCalled(); + f.release(); + expect(await f.held).toEqual(Ok(null)); + await draining; + expect(f.transport.sessions.size).toBe(0); + expect(f.transport.getStatus()).toBe('closed'); +}); + +test('drain closes at the budget even when a handler is still running', async () => { + const f = await setup(); + const draining = f.server.drain({ budgetMs: 1_000 }); + await vi.advanceTimersByTimeAsync(1_000); + await draining; + expect(f.transport.getStatus()).toBe('closed'); + expect(f.transport.sessions.size).toBe(0); +}); + +test('repeated drains share the same completion', async () => { + const f = await setup(); + const first = f.server.drain({ budgetMs: 1_000 }); + expect(f.server.drain({ budgetMs: 0 })).toBe(first); + f.release(); + await first; + expect(f.server.drain({ budgetMs: 0 })).toBe(first); +}); + +test('opt-in RPC resend waits for a different session', async () => { + const f = await setup({ + resendRefusedCalls: { maxResends: 1, newSessionWaitMs: 30_000 }, + }); + const previous = f.clientTransport.sessions.get('SERVER')?.id; + const draining = f.server.drain({ budgetMs: 10_000 }); + const done = vi.fn(); + const result = f.client.test.rpc.rpc({}).then((value) => { + done(); + + return value; + }); + await vi.advanceTimersByTimeAsync(100); + expect(done).not.toHaveBeenCalled(); + f.release(); + await draining; + f.startServer(); + expect(await result).toEqual(Ok(null)); + expect(f.clientTransport.sessions.get('SERVER')?.id).not.toBe(previous); + expect(f.ran).toHaveBeenCalledTimes(1); +}); + +test('resend wait timeout surfaces the original refusal', async () => { + const f = await setup({ + resendRefusedCalls: { maxResends: 1, newSessionWaitMs: 100 }, + }); + void f.server.drain({ budgetMs: 10_000 }); + const done = vi.fn(); + const result = f.client.test.rpc.rpc({}).then((value) => { + done(); + + return value; + }); + await vi.advanceTimersByTimeAsync(30); + expect(done).not.toHaveBeenCalled(); + expect(await result).toEqual(refusal); + expect(f.ran).not.toHaveBeenCalled(); +}); + +test('zero resend budget surfaces the refusal without waiting', async () => { + const f = await setup({ + resendRefusedCalls: { maxResends: 0, newSessionWaitMs: 30_000 }, + }); + void f.server.drain({ budgetMs: 10_000 }); + expect(await f.client.test.rpc.rpc({})).toEqual(refusal); +}); + +test('abort stops waiting for a new session and preserves the refusal', async () => { + const f = await setup({ + resendRefusedCalls: { maxResends: 1, newSessionWaitMs: 30_000 }, + }); + void f.server.drain({ budgetMs: 10_000 }); + const controller = new AbortController(); + const done = vi.fn(); + const result = f.client.test.rpc + .rpc({}, { signal: controller.signal }) + .then((value) => { + done(); + + return value; + }); + await vi.advanceTimersByTimeAsync(100); + expect(done).not.toHaveBeenCalled(); + controller.abort(); + expect(await result).toEqual(refusal); +}); + +test('opt-in does not replay an upload refusal', async () => { + const f = await setup({ + resendRefusedCalls: { maxResends: 1, newSessionWaitMs: 30_000 }, + }); + void f.server.drain({ budgetMs: 10_000 }); + expect(await f.client.test.upload.upload({}).finalize()).toEqual(refusal); +}); diff --git a/router/client.ts b/router/client.ts index e64966cc..d58ae98f 100644 --- a/router/client.ts +++ b/router/client.ts @@ -34,10 +34,12 @@ import { PayloadType, ValidProcType } from './procedures'; import { BaseErrorSchemaType, CANCEL_CODE, + SERVER_DRAINING_CODE, ReaderErrorResultSchema, UNEXPECTED_DISCONNECT_CODE, validationErrorToRiverErrors, } from './errors'; +import { SessionState } from '../transport/sessionStateMachine/common'; export interface CallOptions { signal?: AbortSignal; @@ -205,6 +207,7 @@ function _createRecursiveProxy( } export interface ClientOptions { + resendRefusedCalls?: { maxResends: number; newSessionWaitMs: number }; connectOnInvoke: boolean; eagerlyConnect: boolean; /** @@ -260,6 +263,16 @@ export function createClient< } const clientOptions = { ...defaultClientOptions, ...providedClientOptions }; + const resend = clientOptions.resendRefusedCalls; + if ( + resend && + (!Number.isSafeInteger(resend.maxResends) || + resend.maxResends < 0 || + !Number.isFinite(resend.newSessionWaitMs) || + resend.newSessionWaitMs < 0) + ) { + throw new RangeError('Invalid resendRefusedCalls limits'); + } if (clientOptions.eagerlyConnect) { transport.connect(serverId); } @@ -294,15 +307,68 @@ export function createClient< callOptions as CallOptions | undefined, ); - return handleProc( - procMethod === 'subscribe' ? 'subscription' : procMethod, - transport, - serverId, - init, - serviceName, - procName, - merged.signal, - ); + const invoke = () => + handleProc( + procMethod === 'subscribe' ? 'subscription' : procMethod, + transport, + serverId, + init, + serviceName, + procName, + merged.signal, + ); + if (procMethod !== 'rpc' || !resend) return invoke(); + + return (async () => { + for (let attempt = 0; ; attempt++) { + const response = invoke() as ReturnType>; + const sessionId = transport.sessions.get(serverId)?.id; + const result = await response; + const waitForNewSession = (): Promise => { + return new Promise((resolve) => { + const finish = (ready: boolean) => { + clearTimeout(timer); + transport.removeEventListener('sessionTransition', check); + transport.removeEventListener('transportStatus', check); + merged.signal?.removeEventListener('abort', aborted); + resolve(ready); + }; + const aborted = () => finish(false); + const check = () => { + const session = transport.sessions.get(serverId); + if ( + merged.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), + resend.newSessionWaitMs, + ); + transport.addEventListener('sessionTransition', check); + transport.addEventListener('transportStatus', check); + merged.signal?.addEventListener('abort', aborted, { once: true }); + check(); + }); + }; + if ( + result.ok || + result.payload.code !== SERVER_DRAINING_CODE || + attempt >= resend.maxResends || + !(await waitForNewSession()) + ) { + return result; + } + } + })(); }, []) as Client; } diff --git a/router/errors.ts b/router/errors.ts index 376b0299..1b76bd7e 100644 --- a/router/errors.ts +++ b/router/errors.ts @@ -35,6 +35,7 @@ export const INVALID_REQUEST_CODE = 'INVALID_REQUEST'; * {@link CANCEL_CODE} is the code used when either server or client cancels the stream. */ export const CANCEL_CODE = 'CANCEL'; +export const SERVER_DRAINING_CODE = 'SERVER_DRAINING'; type TLiteralString = TLiteral; @@ -122,6 +123,10 @@ export const CancelResultSchema = ErrResultSchema(CancelErrorSchema); * on the client). */ export const ReaderErrorSchema = Type.Union([ + Type.Object({ + code: Type.Literal(SERVER_DRAINING_CODE), + message: Type.String(), + }), Type.Object({ code: Type.Literal(UNCAUGHT_ERROR_CODE), message: Type.String(), diff --git a/router/index.ts b/router/index.ts index f748101b..c816d912 100644 --- a/router/index.ts +++ b/router/index.ts @@ -39,6 +39,7 @@ export { UNEXPECTED_DISCONNECT_CODE, INVALID_REQUEST_CODE, CANCEL_CODE, + SERVER_DRAINING_CODE, ReaderErrorSchema, BaseErrorSchemaType, } from './errors'; diff --git a/router/server.ts b/router/server.ts index f29f9552..3f55be1f 100644 --- a/router/server.ts +++ b/router/server.ts @@ -6,6 +6,7 @@ import { UNEXPECTED_DISCONNECT_CODE, CANCEL_CODE, INVALID_REQUEST_CODE, + SERVER_DRAINING_CODE, BaseErrorSchemaType, ValidationErrors, castTypeboxValueErrors, @@ -72,6 +73,7 @@ export interface Server< streams: Map>; close: () => Promise; + drain: (options: { budgetMs: number }) => Promise; } interface StreamInitProps { @@ -146,6 +148,10 @@ class RiverServer< >; private unregisterTransportListeners: () => void; + private drainState: + | { status: 'accepting' } + | { status: 'draining'; promise: Promise; finish: () => void } + | { status: 'drained'; promise: Promise } = { status: 'accepting' }; constructor( transport: ServerTransport< @@ -231,6 +237,21 @@ class RiverServer< return; } + if (this.drainState.status !== 'accepting') { + const session = newStreamProps.initialSession; + this.cancelStream( + message.from, + transport.getSessionBoundSendFn(message.from, session.id), + streamId, + Err({ + code: SERVER_DRAINING_CODE, + message: 'Server is draining; the procedure was not run', + }), + ); + + return; + } + // if its not a cancelled stream, validate and create a new stream createHandlerSpan( transport.tracer, @@ -543,6 +564,9 @@ class RiverServer< const cleanup = () => { finishedController.abort(); this.streams.delete(streamId); + if (this.streams.size === 0 && this.drainState.status === 'draining') { + this.drainState.finish(); + } void runDeferredCleanups(); }; @@ -776,6 +800,10 @@ class RiverServer< } }; + if (!finishedController.signal.aborted) { + this.streams.set(streamId, procStream); + } + // Start the middleware chain, which will eventually call runProcedureHandler this.middlewares.reduceRight( (next: () => void, middleware: Middleware) => { @@ -791,10 +819,6 @@ class RiverServer< void runProcedureHandler(); }, )(); - - if (!finishedController.signal.aborted) { - this.streams.set(streamId, procStream); - } } private getContext(service: AnyService, serviceName: string) { @@ -1052,6 +1076,33 @@ class RiverServer< sessionScopedSend(msg); } + drain({ budgetMs }: { budgetMs: number }): Promise { + if (!Number.isFinite(budgetMs) || budgetMs < 0) { + throw new RangeError('budgetMs must be a finite non-negative number'); + } + switch (this.drainState.status) { + case 'draining': + case 'drained': + return this.drainState.promise; + case 'accepting': { + let finish!: () => void; + const waiting = new Promise((resolve) => { + finish = resolve; + }); + const timer = setTimeout(finish, budgetMs); + const promise = waiting.then(() => { + clearTimeout(timer); + this.drainState = { status: 'drained', promise }; + this.transport.close(); + }); + this.drainState = { status: 'draining', promise, finish }; + if (this.streams.size === 0) finish(); + + return promise; + } + } + } + async close() { this.unregisterTransportListeners(); From 75b77b0b84cc64c223621c549a40e0b2d86de216 Mon Sep 17 00:00:00 2001 From: Armaan Aggarwal <49625265+Armster15@users.noreply.github.com> Date: Wed, 21 Oct 2026 22:29:09 +0000 Subject: [PATCH 2/5] server: wait only for single responses during drain Long-lived procedures must not delay transport shutdown. Track admitted calls before middleware and finish when RPCs and uploads end, with a bounded timer as the deadline. Assisted-by: Replit --- __tests__/server-drain.test.ts | 134 +++++++++++++++++++++++++++++++++ router/server.ts | 33 +++++--- 2 files changed, 155 insertions(+), 12 deletions(-) create mode 100644 __tests__/server-drain.test.ts diff --git a/__tests__/server-drain.test.ts b/__tests__/server-drain.test.ts new file mode 100644 index 00000000..14434377 --- /dev/null +++ b/__tests__/server-drain.test.ts @@ -0,0 +1,134 @@ +import { afterEach, expect, test, vi } from 'vitest'; +import { Type } from 'typebox'; +import { + createClient, + createServer, + createServiceSchema, + Middleware, + Ok, + Procedure, +} from '../router'; +import { createMockTransportNetwork, readNextResult } from '../testUtil'; +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(middlewares: Array = []) { + const network = createMockTransportNetwork(); + const gate = createPromiseWithResolvers(); + const effects: Array = []; + const services = { + test: createServiceSchema().define({ + hold: Procedure.rpc({ + requestInit: Type.Object({}), + responseData: Type.Null(), + async handler() { + effects.push('hold'); + await gate.promise; + + return Ok(null); + }, + }), + rpc: Procedure.rpc({ + requestInit: Type.Object({}), + responseData: Type.Null(), + async handler() { + effects.push('rpc'); + + return Ok(null); + }, + }), + subscription: Procedure.subscription({ + requestInit: Type.Object({}), + responseData: Type.Null(), + async handler({ resWritable }) { + resWritable.write(Ok(null)); + }, + }), + }), + }; + const transport = network.getServerTransport(); + const clientTransport = network.getClientTransport('client'); + const server = createServer(transport, services, { middlewares }); + const client = createClient(clientTransport, 'SERVER'); + cleanups.push(async () => { + gate.resolve(); + clientTransport.close(); + transport.close(); + await server.close(); + await network.cleanup(); + }); + + return { server, transport, clientTransport, client, gate, effects }; +} + +test('drain refuses calls before middleware or handler runs', async () => { + // Risk: a refused call that ran would be unsafe to resend. + const effects: Array = []; + const f = setup([ + ({ next }) => { + effects.push('middleware'); + next(); + }, + ]); + void f.client.test.hold.rpc({}); + await waitFor(() => expect(f.effects).toContain('hold')); + effects.length = 0; + void f.server.drain({ budgetMs: 1000 }); + expect(await f.client.test.rpc.rpc({})).toMatchObject({ + ok: false, + payload: { code: 'SERVER_DRAINING' }, + }); + expect(effects).toEqual([]); + expect(f.effects).toEqual(['hold']); +}); + +test('drain inside middleware waits for its admitted rpc result', async () => { + // Risk: an accepted call is closed as if the server were idle. + const f = setup([ + ({ next }) => { + void f.server.drain({ budgetMs: 1000 }); + next(); + }, + ]); + const result = f.client.test.hold.rpc({}); + await waitFor(() => expect(f.effects).toContain('hold')); + expect(f.transport.getStatus()).not.toBe('closed'); + f.gate.resolve(); + expect(await result).toEqual(Ok(null)); + await f.server.drain({ budgetMs: 1000 }); + expect(f.transport.getStatus()).toBe('closed'); +}); + +test('an open subscription does not hold drain or preserve the old session', async () => { + // Risk: long streams keep clients on the old server for the whole budget. + const f = setup(); + const subscription = f.client.test.subscription.subscribe({}); + expect(await readNextResult(subscription.resReadable)).toEqual(Ok(null)); + const result = f.client.test.hold.rpc({}); + await waitFor(() => expect(f.effects).toContain('hold')); + const draining = f.server.drain({ budgetMs: 1000 }); + expect(f.transport.getStatus()).not.toBe('closed'); + f.gate.resolve(); + expect(await result).toEqual(Ok(null)); + expect(f.transport.getStatus()).toBe('closed'); + await draining; + expect(f.transport.getStatus()).toBe('closed'); + expect(f.transport.sessions.size).toBe(0); + expect(f.server.streams.size).toBe(0); +}); + +test('the drain budget closes a server whose rpc never returns', async () => { + // Risk: shutdown hangs. + const f = setup(); + void f.client.test.hold.rpc({}); + await waitFor(() => expect(f.effects).toContain('hold')); + const draining = f.server.drain({ budgetMs: 100 }); + await vi.advanceTimersByTimeAsync(100); + expect(f.transport.getStatus()).toBe('closed'); + await draining; +}); diff --git a/router/server.ts b/router/server.ts index 3f55be1f..dbab2f28 100644 --- a/router/server.ts +++ b/router/server.ts @@ -51,6 +51,7 @@ import { ServerTransport } from '../transport/server'; import { ReadableImpl, WritableImpl } from './streams'; import { IdentifiedSession } from '../transport/sessionStateMachine/common'; import { SessionBoundSendFn } from '../transport/transport'; +import { createPromiseWithResolvers } from '../transport/promises'; type StreamId = string; @@ -150,8 +151,9 @@ class RiverServer< private unregisterTransportListeners: () => void; private drainState: | { status: 'accepting' } - | { status: 'draining'; promise: Promise; finish: () => void } - | { status: 'drained'; promise: Promise } = { status: 'accepting' }; + | { status: 'draining'; promise: Promise; finish: () => void } = { + status: 'accepting', + }; constructor( transport: ServerTransport< @@ -564,7 +566,7 @@ class RiverServer< const cleanup = () => { finishedController.abort(); this.streams.delete(streamId); - if (this.streams.size === 0 && this.drainState.status === 'draining') { + if (!this.hasPendingResponse() && this.drainState.status === 'draining') { this.drainState.finish(); } void runDeferredCleanups(); @@ -1076,27 +1078,34 @@ class RiverServer< sessionScopedSend(msg); } + private hasPendingResponse(): boolean { + return [...this.streams.values()].some( + ({ procedure }) => + procedure.type === 'rpc' || procedure.type === 'upload', + ); + } + drain({ budgetMs }: { budgetMs: number }): Promise { - if (!Number.isFinite(budgetMs) || budgetMs < 0) { - throw new RangeError('budgetMs must be a finite non-negative number'); + if ( + !Number.isFinite(budgetMs) || + budgetMs < 0 || + budgetMs > 2_147_483_647 + ) { + throw new RangeError('budgetMs must be between 0 and 2147483647'); } switch (this.drainState.status) { case 'draining': - case 'drained': return this.drainState.promise; case 'accepting': { - let finish!: () => void; - const waiting = new Promise((resolve) => { - finish = resolve; - }); + const { promise: waiting, resolve: finish } = + createPromiseWithResolvers(); const timer = setTimeout(finish, budgetMs); const promise = waiting.then(() => { clearTimeout(timer); - this.drainState = { status: 'drained', promise }; this.transport.close(); }); this.drainState = { status: 'draining', promise, finish }; - if (this.streams.size === 0) finish(); + if (!this.hasPendingResponse()) finish(); return promise; } From 8d4f5a5d84a44a3842840a46a67b9dada7a96708 Mon Sep 17 00:00:00 2001 From: Armaan Aggarwal <49625265+Armster15@users.noreply.github.com> Date: Wed, 21 Oct 2026 22:37:16 +0000 Subject: [PATCH 3/5] client: isolate refused RPC resends and bound session waits Only refused RPCs can be sent again on a fresh connected session. Keep waiting and cleanup in one named operation, and check cancellation before dispatching each resend. Assisted-by: Replit --- __tests__/client-drain.test.ts | 176 +++++++++++++++++++++++++++++++++ router/client.ts | 165 +++++++++++++++++++------------ router/server.ts | 2 +- 3 files changed, 279 insertions(+), 64 deletions(-) create mode 100644 __tests__/client-drain.test.ts diff --git a/__tests__/client-drain.test.ts b/__tests__/client-drain.test.ts new file mode 100644 index 00000000..88a0650c --- /dev/null +++ b/__tests__/client-drain.test.ts @@ -0,0 +1,176 @@ +import { afterEach, expect, test, vi } from 'vitest'; +import { Type } from 'typebox'; +import { + createClient, + createServer, + createServiceSchema, + Ok, + Procedure, +} from '../router'; +import { createMockTransportNetwork } from '../testUtil'; +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(); +}); + +async function setup(maxResends = 2, newSessionWaitMs = 1000) { + const network = createMockTransportNetwork(); + const gate = createPromiseWithResolvers(); + const effects: Array = []; + const attempts: Array = []; + const services = { + test: createServiceSchema().define({ + hold: Procedure.rpc({ + requestInit: Type.Object({}), + responseData: Type.Null(), + async handler() { + effects.push('hold'); + await gate.promise; + + return Ok(null); + }, + }), + rpc: Procedure.rpc({ + requestInit: Type.Object({}), + responseData: Type.Null(), + async handler() { + effects.push('rpc'); + + return Ok(null); + }, + }), + upload: Procedure.upload({ + requestInit: Type.Object({}), + requestData: Type.Null(), + responseData: Type.Null(), + async handler() { + effects.push('upload'); + + return Ok(null); + }, + }), + }), + }; + function startServer() { + const transport = network.getServerTransport(); + const server = createServer(transport, services); + transport.addEventListener('message', (message) => { + if (message.procedureName === 'rpc') + 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(); + blocker.close(); + clientTransport.close(); + }); + const { server, transport } = startServer(); + const client = createClient(clientTransport, 'SERVER', { + resendRefusedCalls: { maxResends, newSessionWaitMs }, + }); + void createClient(blocker, 'SERVER').test.hold.rpc({}); + await waitFor(() => expect(effects).toEqual(['hold'])); + const draining = server.drain({ budgetMs: 10000 }); + + return { + client, + clientTransport, + gate, + effects, + attempts, + draining, + transport, + startServer, + }; +} + +const refusal = { + ok: false, + payload: { + code: 'SERVER_DRAINING', + message: 'Server is draining; the procedure was not run', + }, +}; + +test('a refused rpc waits for a connected session with a new id', async () => { + // Risk: replaying work on the draining server. + const f = await setup(); + const result = f.client.test.rpc.rpc({}); + await waitFor(() => expect(f.attempts).toHaveLength(1)); + await vi.advanceTimersByTimeAsync(20); + expect(f.attempts).toHaveLength(1); + f.gate.resolve(); + await f.draining; + f.startServer(); + f.clientTransport.hardDisconnect(); + f.clientTransport.connect('SERVER'); + expect(await result).toEqual(Ok(null)); + expect(f.attempts).toHaveLength(2); + expect(f.attempts[1]).not.toBe(f.attempts[0]); + expect(f.effects).toEqual(['hold', 'rpc']); +}); + +test('successive refusals stop after a positive resend budget', async () => { + // Risk: unbounded replay. + const f = await setup(2); + const result = f.client.test.rpc.rpc({}); + for (let attempt = 1; attempt <= 3; attempt++) { + await waitFor(() => expect(f.attempts).toHaveLength(attempt)); + await vi.advanceTimersByTimeAsync(10); + f.clientTransport.hardDisconnect(); + f.clientTransport.connect('SERVER'); + } + expect(await result).toEqual(refusal); + await vi.advanceTimersByTimeAsync(20); + expect(f.attempts).toHaveLength(3); + expect(f.effects).toEqual(['hold']); +}); + +test('the session wait deadline returns the original refusal', async () => { + // Risk: the rpc stays pending forever. + const f = await setup(2, 100); + const result = f.client.test.rpc.rpc({}); + await waitFor(() => expect(f.attempts).toHaveLength(1)); + await vi.advanceTimersByTimeAsync(100); + expect(await result).toEqual(refusal); + expect(f.effects).toEqual(['hold']); +}); + +test('abort while waiting preserves the refusal and removes the waiter', async () => { + // Risk: cancelled work still runs. + const f = await setup(2, 30_000); + const waiting = new AbortController(); + const first = f.client.test.rpc.rpc({}, { signal: waiting.signal }); + await waitFor(() => expect(f.attempts).toHaveLength(1)); + await vi.advanceTimersByTimeAsync(10); + waiting.abort(); + expect(await first).toEqual(refusal); + f.gate.resolve(); + await f.draining; + f.startServer(); + f.clientTransport.hardDisconnect(); + f.clientTransport.connect('SERVER'); + await vi.advanceTimersByTimeAsync(500); + expect(f.effects).toEqual(['hold']); + expect(f.attempts).toHaveLength(1); +}); + +test('a refused upload is returned without replay', async () => { + // Risk: replaying a consumed upload. + const f = await setup(); + expect(await f.client.test.upload.upload({}).finalize()).toEqual(refusal); + expect(f.effects).toEqual(['hold']); +}); diff --git a/router/client.ts b/router/client.ts index d58ae98f..5c46bc5e 100644 --- a/router/client.ts +++ b/router/client.ts @@ -269,7 +269,8 @@ export function createClient< (!Number.isSafeInteger(resend.maxResends) || resend.maxResends < 0 || !Number.isFinite(resend.newSessionWaitMs) || - resend.newSessionWaitMs < 0) + resend.newSessionWaitMs < 0 || + resend.newSessionWaitMs > 2_147_483_647) ) { throw new RangeError('Invalid resendRefusedCalls limits'); } @@ -277,6 +278,75 @@ export function createClient< 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, + 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( + sessionId, + limits.newSessionWaitMs, + signal, + ); + if (!ready || signal?.aborted) return result; + } + } + return _createRecursiveProxy((opts) => { const [serviceName, procName, procMethod] = [...opts.path]; if (!(serviceName && procName && procMethod)) { @@ -307,68 +377,19 @@ export function createClient< callOptions as CallOptions | undefined, ); - const invoke = () => - handleProc( - procMethod === 'subscribe' ? 'subscription' : procMethod, - transport, - serverId, - init, - serviceName, - procName, - merged.signal, - ); - if (procMethod !== 'rpc' || !resend) return invoke(); - - return (async () => { - for (let attempt = 0; ; attempt++) { - const response = invoke() as ReturnType>; - const sessionId = transport.sessions.get(serverId)?.id; - const result = await response; - const waitForNewSession = (): Promise => { - return new Promise((resolve) => { - const finish = (ready: boolean) => { - clearTimeout(timer); - transport.removeEventListener('sessionTransition', check); - transport.removeEventListener('transportStatus', check); - merged.signal?.removeEventListener('abort', aborted); - resolve(ready); - }; - const aborted = () => finish(false); - const check = () => { - const session = transport.sessions.get(serverId); - if ( - merged.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), - resend.newSessionWaitMs, - ); - transport.addEventListener('sessionTransition', check); - transport.addEventListener('transportStatus', check); - merged.signal?.addEventListener('abort', aborted, { once: true }); - check(); - }); - }; - if ( - result.ok || - result.payload.code !== SERVER_DRAINING_CODE || - attempt >= resend.maxResends || - !(await waitForNewSession()) - ) { - return result; - } - } - })(); + if (procMethod === 'rpc' && resend) { + return rpcWithResends(init, serviceName, procName, resend, merged.signal); + } + + return handleProc( + procMethod === 'subscribe' ? 'subscription' : procMethod, + transport, + serverId, + init, + serviceName, + procName, + merged.signal, + ); }, []) as Client; } @@ -388,6 +409,24 @@ type AnyProcReturn = | ReturnType> | ReturnType>; +function handleProc( + procType: 'rpc', + transport: ClientTransport, + serverId: TransportClientId, + init: Static, + serviceName: string, + procedureName: string, + abortSignal?: AbortSignal, +): ReturnType>; +function handleProc( + procType: ValidProcType, + transport: ClientTransport, + serverId: TransportClientId, + init: Static, + serviceName: string, + procedureName: string, + abortSignal?: AbortSignal, +): AnyProcReturn; function handleProc( procType: ValidProcType, transport: ClientTransport, diff --git a/router/server.ts b/router/server.ts index dbab2f28..3b18252a 100644 --- a/router/server.ts +++ b/router/server.ts @@ -566,7 +566,7 @@ class RiverServer< const cleanup = () => { finishedController.abort(); this.streams.delete(streamId); - if (!this.hasPendingResponse() && this.drainState.status === 'draining') { + if (this.drainState.status === 'draining' && !this.hasPendingResponse()) { this.drainState.finish(); } void runDeferredCleanups(); From 6079d4a28cdcf499f5308c00f13c766e6c78014a Mon Sep 17 00:00:00 2001 From: Armaan Aggarwal <49625265+Armster15@users.noreply.github.com> Date: Wed, 21 Oct 2026 22:43:09 +0000 Subject: [PATCH 4/5] tests: consolidate drain coverage and document cancellation Keep one behavior test for each drain and resend risk. Share the transport fixture and state when cancellation preserves a refusal instead of cancelling an in-flight call. Assisted-by: Replit --- README.md | 43 ++--- __tests__/client-drain.test.ts | 176 ----------------- __tests__/drain.test.ts | 336 +++++++++++++++++---------------- __tests__/server-drain.test.ts | 134 ------------- 4 files changed, 190 insertions(+), 499 deletions(-) delete mode 100644 __tests__/client-drain.test.ts delete mode 100644 __tests__/server-drain.test.ts diff --git a/README.md b/README.md index f8bfe0c9..ac1340fe 100644 --- a/README.md +++ b/README.md @@ -172,23 +172,12 @@ return Err({ code: 'INVALID_INPUT', message: 'Value must be positive' }); #### Graceful server drain -The TypeBox router server exposes `server.drain({ budgetMs: number }): Promise`. -The budget must be finite and non-negative. Drain immediately stops accepting new -procedure invocations. RPCs, uploads, subscriptions, and streams receive an error -result with the reserved code `SERVER_DRAINING` (exported as -`SERVER_DRAINING_CODE`). This code guarantees that the handler and middleware did -not run. Application handlers must not emit this reserved code. - -Existing calls may finish until the budget expires. Drain then closes the entire -server transport and deletes its sessions; these sessions cannot be resumed. -The promise resolves when the transport has closed. Repeated calls share the -first drain's budget and promise. An idle server closes immediately. -Closing the transport cancels unfinished calls but cannot force JavaScript -handlers to stop; handlers should observe their cancellation signal. -Call `server.close()` separately to dispose service and extended-context resources. -Do not share a draining transport with another router. - -Clients can opt in to resending refused unary RPCs: +`server.drain({ budgetMs })` refuses all new procedures with `SERVER_DRAINING` +(`SERVER_DRAINING_CODE`) before middleware or handlers run. Handlers must not emit it. +Drain waits only for admitted RPCs and uploads, until they finish or the budget expires. +It then closes the whole transport, ends streams, and deletes sessions so they cannot resume. +The promise resolves after close. Use one router per draining transport. +Call `server.close()` separately to dispose services. ```ts const client = createClient(transport, 'SERVER', { @@ -196,18 +185,14 @@ const client = createClient(transport, 'SERVER', { }); ``` -This option is off by default. Only `SERVER_DRAINING` results are retried. -Each retry waits for a connected session with a different session ID, not just -a new connection on the old session. The transport's existing reconnection policy -still applies; route reconnects to a healthy server. Each wait has its own timeout. -The original refusal is returned if the retry budget or wait expires, the -transport closes, or the caller aborts while waiting. Normal call cancellation -applies during an attempt. Limits must be finite and non-negative, and -`maxResends` must be a safe integer. - -Uploads, subscriptions, and streams are never replayed. There is no unbounded -backoff after exhaustion. These APIs do not currently apply to the experimental -protobuf router. +Resend is off by default and applies only to RPCs refused with `SERVER_DRAINING`. +Each resend waits for a connected session with a new ID. Other errors and disconnects are not retried. +Resends stop at `maxResends`. Each wait stops at `newSessionWaitMs`. +An abort before resend dispatch returns the original refusal without resending, while an abort after dispatch follows normal call cancellation. +Wait expiry or transport close also returns the refusal. +`maxResends` must be a non-negative safe integer. Both time limits must be finite, non-negative, and at most 2147483647 ms. +The protobuf router is excluded. +`ReaderErrorSchema` gains a code. Exhaustive TypeScript error switches need a new case. #### Custom Error Types diff --git a/__tests__/client-drain.test.ts b/__tests__/client-drain.test.ts deleted file mode 100644 index 88a0650c..00000000 --- a/__tests__/client-drain.test.ts +++ /dev/null @@ -1,176 +0,0 @@ -import { afterEach, expect, test, vi } from 'vitest'; -import { Type } from 'typebox'; -import { - createClient, - createServer, - createServiceSchema, - Ok, - Procedure, -} from '../router'; -import { createMockTransportNetwork } from '../testUtil'; -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(); -}); - -async function setup(maxResends = 2, newSessionWaitMs = 1000) { - const network = createMockTransportNetwork(); - const gate = createPromiseWithResolvers(); - const effects: Array = []; - const attempts: Array = []; - const services = { - test: createServiceSchema().define({ - hold: Procedure.rpc({ - requestInit: Type.Object({}), - responseData: Type.Null(), - async handler() { - effects.push('hold'); - await gate.promise; - - return Ok(null); - }, - }), - rpc: Procedure.rpc({ - requestInit: Type.Object({}), - responseData: Type.Null(), - async handler() { - effects.push('rpc'); - - return Ok(null); - }, - }), - upload: Procedure.upload({ - requestInit: Type.Object({}), - requestData: Type.Null(), - responseData: Type.Null(), - async handler() { - effects.push('upload'); - - return Ok(null); - }, - }), - }), - }; - function startServer() { - const transport = network.getServerTransport(); - const server = createServer(transport, services); - transport.addEventListener('message', (message) => { - if (message.procedureName === 'rpc') - 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(); - blocker.close(); - clientTransport.close(); - }); - const { server, transport } = startServer(); - const client = createClient(clientTransport, 'SERVER', { - resendRefusedCalls: { maxResends, newSessionWaitMs }, - }); - void createClient(blocker, 'SERVER').test.hold.rpc({}); - await waitFor(() => expect(effects).toEqual(['hold'])); - const draining = server.drain({ budgetMs: 10000 }); - - return { - client, - clientTransport, - gate, - effects, - attempts, - draining, - transport, - startServer, - }; -} - -const refusal = { - ok: false, - payload: { - code: 'SERVER_DRAINING', - message: 'Server is draining; the procedure was not run', - }, -}; - -test('a refused rpc waits for a connected session with a new id', async () => { - // Risk: replaying work on the draining server. - const f = await setup(); - const result = f.client.test.rpc.rpc({}); - await waitFor(() => expect(f.attempts).toHaveLength(1)); - await vi.advanceTimersByTimeAsync(20); - expect(f.attempts).toHaveLength(1); - f.gate.resolve(); - await f.draining; - f.startServer(); - f.clientTransport.hardDisconnect(); - f.clientTransport.connect('SERVER'); - expect(await result).toEqual(Ok(null)); - expect(f.attempts).toHaveLength(2); - expect(f.attempts[1]).not.toBe(f.attempts[0]); - expect(f.effects).toEqual(['hold', 'rpc']); -}); - -test('successive refusals stop after a positive resend budget', async () => { - // Risk: unbounded replay. - const f = await setup(2); - const result = f.client.test.rpc.rpc({}); - for (let attempt = 1; attempt <= 3; attempt++) { - await waitFor(() => expect(f.attempts).toHaveLength(attempt)); - await vi.advanceTimersByTimeAsync(10); - f.clientTransport.hardDisconnect(); - f.clientTransport.connect('SERVER'); - } - expect(await result).toEqual(refusal); - await vi.advanceTimersByTimeAsync(20); - expect(f.attempts).toHaveLength(3); - expect(f.effects).toEqual(['hold']); -}); - -test('the session wait deadline returns the original refusal', async () => { - // Risk: the rpc stays pending forever. - const f = await setup(2, 100); - const result = f.client.test.rpc.rpc({}); - await waitFor(() => expect(f.attempts).toHaveLength(1)); - await vi.advanceTimersByTimeAsync(100); - expect(await result).toEqual(refusal); - expect(f.effects).toEqual(['hold']); -}); - -test('abort while waiting preserves the refusal and removes the waiter', async () => { - // Risk: cancelled work still runs. - const f = await setup(2, 30_000); - const waiting = new AbortController(); - const first = f.client.test.rpc.rpc({}, { signal: waiting.signal }); - await waitFor(() => expect(f.attempts).toHaveLength(1)); - await vi.advanceTimersByTimeAsync(10); - waiting.abort(); - expect(await first).toEqual(refusal); - f.gate.resolve(); - await f.draining; - f.startServer(); - f.clientTransport.hardDisconnect(); - f.clientTransport.connect('SERVER'); - await vi.advanceTimersByTimeAsync(500); - expect(f.effects).toEqual(['hold']); - expect(f.attempts).toHaveLength(1); -}); - -test('a refused upload is returned without replay', async () => { - // Risk: replaying a consumed upload. - const f = await setup(); - expect(await f.client.test.upload.upload({}).finalize()).toEqual(refusal); - expect(f.effects).toEqual(['hold']); -}); diff --git a/__tests__/drain.test.ts b/__tests__/drain.test.ts index ebe7fa27..7a8acfd0 100644 --- a/__tests__/drain.test.ts +++ b/__tests__/drain.test.ts @@ -5,80 +5,72 @@ import { createClient, createServer, createServiceSchema, + Middleware, Ok, Procedure, - SERVER_DRAINING_CODE, } from '../router'; import { createMockTransportNetwork, readNextResult } from '../testUtil'; 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(); }); -async function setup(options: Partial = {}) { - const network = createMockTransportNetwork(); - cleanups.push(() => network.cleanup()); - const clientTransport = network.getClientTransport('client'); - cleanups.push(() => clientTransport.close()); - const ran = vi.fn(); - let release!: () => void; - const gate = new Promise((resolve) => { - release = resolve; +function setup( + maxResends?: number, + newSessionWaitMs = 1000, + middlewares: Array = [], +) { + const network = createMockTransportNetwork({ + client: { baseIntervalMs: 1, maxJitterMs: 0 }, }); + const gate = createPromiseWithResolvers(); + const uploadGate = createPromiseWithResolvers(); + const effects: Array = []; + const attempts: Array = []; + const init = Type.Object({}); + const responseData = Type.Null(); const services = { test: createServiceSchema().define({ - hold: Procedure.rpc({ - requestInit: Type.Object({}), - responseData: Type.Null(), - async handler() { - await gate; - - return Ok(null); - }, - }), rpc: Procedure.rpc({ - requestInit: Type.Object({}), - responseData: Type.Null(), - async handler() { - ran(); + requestInit: Type.Boolean(), + responseData, + async handler({ reqInit }) { + effects.push(reqInit ? 'hold' : 'rpc'); + if (reqInit) await gate.promise; return Ok(null); }, }), upload: Procedure.upload({ - requestInit: Type.Object({}), + requestInit: init, requestData: Type.Null(), - responseData: Type.Null(), + responseData, async handler() { - ran(); + effects.push('upload'); + await uploadGate.promise; return Ok(null); }, }), subscription: Procedure.subscription({ - requestInit: Type.Object({}), - responseData: Type.Null(), - async handler({ resWritable }) { - ran(); - resWritable.close(); - }, - }), - stream: Procedure.stream({ - requestInit: Type.Object({}), - requestData: Type.Null(), - responseData: Type.Null(), + requestInit: init, + responseData, async handler({ resWritable }) { - ran(); - resWritable.close(); + resWritable.write(Ok(null)); }, }), }), }; function startServer() { const transport = network.getServerTransport(); - const server = createServer(transport, services); + const server = createServer(transport, services, { middlewares }); + transport.addEventListener('message', (message) => { + if (message.procedureName === 'rpc' && message.payload === false) + attempts.push(transport.sessions.get(message.from)?.id); + }); cleanups.push(async () => { transport.close(); await server.close(); @@ -86,171 +78,195 @@ async function setup(options: Partial = {}) { 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 options: Partial = {}; + if (maxResends !== undefined) + options.resendRefusedCalls = { maxResends, newSessionWaitMs }; const client = createClient( clientTransport, 'SERVER', options, ); - const held = client.test.hold.rpc({}); - await waitFor(() => expect(server.streams.size).toBe(1)); - cleanups.push(release); + async function block() { + void createClient(blocker, 'SERVER').test.rpc.rpc(true); + await waitFor(() => expect(effects).toContain('hold')); + } + function reconnect() { + clientTransport.hardDisconnect(); + clientTransport.connect('SERVER'); + } return { - client, - clientTransport, server, transport, - release, - held, - ran, + client, + clientTransport, + gate, + uploadGate, + effects, + attempts, startServer, + block, + reconnect, }; } const refusal = { ok: false, payload: { - code: SERVER_DRAINING_CODE, + code: 'SERVER_DRAINING', message: 'Server is draining; the procedure was not run', }, }; -test.each(['rpc', 'upload', 'subscription', 'stream'] as const)( - 'drain refuses new %s calls without running the handler', - async (kind) => { - const f = await setup(); - const draining = f.server.drain({ budgetMs: 10_000 }); - switch (kind) { - case 'rpc': - expect(await f.client.test.rpc.rpc({})).toEqual(refusal); - break; - case 'upload': - expect(await f.client.test.upload.upload({}).finalize()).toEqual( - refusal, - ); - break; - case 'subscription': - expect( - await readNextResult( - f.client.test.subscription.subscribe({}).resReadable, - ), - ).toEqual(refusal); - break; - case 'stream': - expect( - await readNextResult(f.client.test.stream.stream({}).resReadable), - ).toEqual(refusal); - break; - } - expect(f.ran).not.toHaveBeenCalled(); - f.release(); - await draining; - }, -); +test('drain refuses calls before middleware or handler runs', async () => { + // Risk: a refused call that ran would be unsafe to resend. + const middlewareEffects: Array = []; + const f = setup(undefined, 1000, [ + ({ next }) => { + middlewareEffects.push('middleware'); + next(); + }, + ]); + await f.block(); + middlewareEffects.length = 0; + void f.server.drain({ budgetMs: 10000 }); + expect(await f.client.test.rpc.rpc(false)).toEqual(refusal); + expect(middlewareEffects).toEqual([]); + expect(f.effects).toEqual(['hold']); +}); -test('drain lets an existing RPC finish and removes its session', async () => { - const f = await setup(); - const done = vi.fn(); - const draining = f.server.drain({ budgetMs: 10_000 }).then(done); - await vi.advanceTimersByTimeAsync(50); - expect(done).not.toHaveBeenCalled(); - f.release(); - expect(await f.held).toEqual(Ok(null)); - await draining; - expect(f.transport.sessions.size).toBe(0); +test('drain inside middleware waits for its admitted rpc result', async () => { + // Risk: an accepted call is closed as if the server were idle. + const f = setup(undefined, 1000, [ + ({ next }) => { + void f.server.drain({ budgetMs: 10000 }); + next(); + }, + ]); + const result = f.client.test.rpc.rpc(true); + await waitFor(() => expect(f.effects).toContain('hold')); + expect(f.transport.getStatus()).not.toBe('closed'); + f.gate.resolve(); + expect(await result).toEqual(Ok(null)); + await f.server.drain({ budgetMs: 10000 }); expect(f.transport.getStatus()).toBe('closed'); }); -test('drain closes at the budget even when a handler is still running', async () => { - const f = await setup(); - const draining = f.server.drain({ budgetMs: 1_000 }); - await vi.advanceTimersByTimeAsync(1_000); - await draining; +test('a subscription does not hold drain after rpc and upload finish', async () => { + // Risk: long streams keep clients on the old server for the whole budget. + const f = setup(); + expect( + await readNextResult(f.client.test.subscription.subscribe({}).resReadable), + ).toEqual(Ok(null)); + const rpc = f.client.test.rpc.rpc(true); + const upload = f.client.test.upload.upload({}).finalize(); + await waitFor(() => expect(f.effects).toEqual(['hold', 'upload'])); + const draining = f.server.drain({ budgetMs: 10000 }); + f.gate.resolve(); + expect(await rpc).toEqual(Ok(null)); + expect(f.transport.getStatus()).not.toBe('closed'); + f.uploadGate.resolve(); + expect(await upload).toEqual(Ok(null)); expect(f.transport.getStatus()).toBe('closed'); + await draining; expect(f.transport.sessions.size).toBe(0); + expect(f.server.streams.size).toBe(0); }); -test('repeated drains share the same completion', async () => { - const f = await setup(); - const first = f.server.drain({ budgetMs: 1_000 }); - expect(f.server.drain({ budgetMs: 0 })).toBe(first); - f.release(); - await first; - expect(f.server.drain({ budgetMs: 0 })).toBe(first); +test('the drain budget closes a server whose rpc 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('opt-in RPC resend waits for a different session', async () => { - const f = await setup({ - resendRefusedCalls: { maxResends: 1, newSessionWaitMs: 30_000 }, - }); - const previous = f.clientTransport.sessions.get('SERVER')?.id; - const draining = f.server.drain({ budgetMs: 10_000 }); - const done = vi.fn(); - const result = f.client.test.rpc.rpc({}).then((value) => { - done(); - - return value; - }); - await vi.advanceTimersByTimeAsync(100); - expect(done).not.toHaveBeenCalled(); - f.release(); +test('a refused rpc waits for a connected session with a new id', async () => { + // Risk: replaying work on the draining server. + const f = setup(2); + await f.block(); + const draining = f.server.drain({ budgetMs: 10000 }); + const result = f.client.test.rpc.rpc(false); + 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).toEqual(Ok(null)); - expect(f.clientTransport.sessions.get('SERVER')?.id).not.toBe(previous); - expect(f.ran).toHaveBeenCalledTimes(1); + expect(f.attempts).toHaveLength(2); + expect(f.attempts[1]).not.toBe(f.attempts[0]); + expect(f.effects).toEqual(['hold', 'rpc']); }); -test('resend wait timeout surfaces the original refusal', async () => { - const f = await setup({ - resendRefusedCalls: { maxResends: 1, newSessionWaitMs: 100 }, - }); - void f.server.drain({ budgetMs: 10_000 }); - const done = vi.fn(); - const result = f.client.test.rpc.rpc({}).then((value) => { - done(); - - return value; - }); - await vi.advanceTimersByTimeAsync(30); - expect(done).not.toHaveBeenCalled(); +test('successive refusals stop after a positive resend budget', async () => { + // Risk: unbounded replay. + const f = setup(2); + await f.block(); + void f.server.drain({ budgetMs: 10000 }); + const result = f.client.test.rpc.rpc(false); + 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).toEqual(refusal); - expect(f.ran).not.toHaveBeenCalled(); + expect(f.effects).toEqual(['hold']); }); -test('zero resend budget surfaces the refusal without waiting', async () => { - const f = await setup({ - resendRefusedCalls: { maxResends: 0, newSessionWaitMs: 30_000 }, - }); - void f.server.drain({ budgetMs: 10_000 }); - expect(await f.client.test.rpc.rpc({})).toEqual(refusal); +test('the session wait deadline returns the original refusal', async () => { + // Risk: the rpc stays pending forever. + const f = setup(2, 100); + await f.block(); + void f.server.drain({ budgetMs: 10000 }); + const result = f.client.test.rpc.rpc(false); + await waitFor(() => expect(f.attempts).toHaveLength(1)); + await vi.advanceTimersByTimeAsync(100); + expect(await result).toEqual(refusal); + expect(f.effects).toEqual(['hold']); }); -test('abort stops waiting for a new session and preserves the refusal', async () => { - const f = await setup({ - resendRefusedCalls: { maxResends: 1, newSessionWaitMs: 30_000 }, - }); - void f.server.drain({ budgetMs: 10_000 }); +test('abort while waiting preserves the refusal and removes the waiter', async () => { + // Risk: cancelled work still runs. + const f = setup(2, 30000); + await f.block(); + const draining = f.server.drain({ budgetMs: 10000 }); const controller = new AbortController(); - const done = vi.fn(); - const result = f.client.test.rpc - .rpc({}, { signal: controller.signal }) - .then((value) => { - done(); - - return value; - }); - await vi.advanceTimersByTimeAsync(100); - expect(done).not.toHaveBeenCalled(); + const result = f.client.test.rpc.rpc(false, { signal: controller.signal }); + await waitFor(() => expect(f.attempts).toHaveLength(1)); + await vi.advanceTimersByTimeAsync(10); controller.abort(); expect(await result).toEqual(refusal); + f.gate.resolve(); + await draining; + f.startServer(); + f.reconnect(); + await vi.advanceTimersByTimeAsync(50); + expect(f.effects).toEqual(['hold']); + expect(f.attempts).toHaveLength(1); }); -test('opt-in does not replay an upload refusal', async () => { - const f = await setup({ - resendRefusedCalls: { maxResends: 1, newSessionWaitMs: 30_000 }, - }); - void f.server.drain({ budgetMs: 10_000 }); +test('a refused upload is returned without replay', async () => { + // Risk: replaying a consumed upload. + const f = setup(2); + await f.block(); + void f.server.drain({ budgetMs: 10000 }); expect(await f.client.test.upload.upload({}).finalize()).toEqual(refusal); + expect(f.effects).toEqual(['hold']); }); diff --git a/__tests__/server-drain.test.ts b/__tests__/server-drain.test.ts deleted file mode 100644 index 14434377..00000000 --- a/__tests__/server-drain.test.ts +++ /dev/null @@ -1,134 +0,0 @@ -import { afterEach, expect, test, vi } from 'vitest'; -import { Type } from 'typebox'; -import { - createClient, - createServer, - createServiceSchema, - Middleware, - Ok, - Procedure, -} from '../router'; -import { createMockTransportNetwork, readNextResult } from '../testUtil'; -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(middlewares: Array = []) { - const network = createMockTransportNetwork(); - const gate = createPromiseWithResolvers(); - const effects: Array = []; - const services = { - test: createServiceSchema().define({ - hold: Procedure.rpc({ - requestInit: Type.Object({}), - responseData: Type.Null(), - async handler() { - effects.push('hold'); - await gate.promise; - - return Ok(null); - }, - }), - rpc: Procedure.rpc({ - requestInit: Type.Object({}), - responseData: Type.Null(), - async handler() { - effects.push('rpc'); - - return Ok(null); - }, - }), - subscription: Procedure.subscription({ - requestInit: Type.Object({}), - responseData: Type.Null(), - async handler({ resWritable }) { - resWritable.write(Ok(null)); - }, - }), - }), - }; - const transport = network.getServerTransport(); - const clientTransport = network.getClientTransport('client'); - const server = createServer(transport, services, { middlewares }); - const client = createClient(clientTransport, 'SERVER'); - cleanups.push(async () => { - gate.resolve(); - clientTransport.close(); - transport.close(); - await server.close(); - await network.cleanup(); - }); - - return { server, transport, clientTransport, client, gate, effects }; -} - -test('drain refuses calls before middleware or handler runs', async () => { - // Risk: a refused call that ran would be unsafe to resend. - const effects: Array = []; - const f = setup([ - ({ next }) => { - effects.push('middleware'); - next(); - }, - ]); - void f.client.test.hold.rpc({}); - await waitFor(() => expect(f.effects).toContain('hold')); - effects.length = 0; - void f.server.drain({ budgetMs: 1000 }); - expect(await f.client.test.rpc.rpc({})).toMatchObject({ - ok: false, - payload: { code: 'SERVER_DRAINING' }, - }); - expect(effects).toEqual([]); - expect(f.effects).toEqual(['hold']); -}); - -test('drain inside middleware waits for its admitted rpc result', async () => { - // Risk: an accepted call is closed as if the server were idle. - const f = setup([ - ({ next }) => { - void f.server.drain({ budgetMs: 1000 }); - next(); - }, - ]); - const result = f.client.test.hold.rpc({}); - await waitFor(() => expect(f.effects).toContain('hold')); - expect(f.transport.getStatus()).not.toBe('closed'); - f.gate.resolve(); - expect(await result).toEqual(Ok(null)); - await f.server.drain({ budgetMs: 1000 }); - expect(f.transport.getStatus()).toBe('closed'); -}); - -test('an open subscription does not hold drain or preserve the old session', async () => { - // Risk: long streams keep clients on the old server for the whole budget. - const f = setup(); - const subscription = f.client.test.subscription.subscribe({}); - expect(await readNextResult(subscription.resReadable)).toEqual(Ok(null)); - const result = f.client.test.hold.rpc({}); - await waitFor(() => expect(f.effects).toContain('hold')); - const draining = f.server.drain({ budgetMs: 1000 }); - expect(f.transport.getStatus()).not.toBe('closed'); - f.gate.resolve(); - expect(await result).toEqual(Ok(null)); - expect(f.transport.getStatus()).toBe('closed'); - await draining; - expect(f.transport.getStatus()).toBe('closed'); - expect(f.transport.sessions.size).toBe(0); - expect(f.server.streams.size).toBe(0); -}); - -test('the drain budget closes a server whose rpc never returns', async () => { - // Risk: shutdown hangs. - const f = setup(); - void f.client.test.hold.rpc({}); - await waitFor(() => expect(f.effects).toContain('hold')); - const draining = f.server.drain({ budgetMs: 100 }); - await vi.advanceTimersByTimeAsync(100); - expect(f.transport.getStatus()).toBe('closed'); - await draining; -}); From 97c3eeec08720494c0910e1d57b97ffc41d4d5af Mon Sep 17 00:00:00 2001 From: Armaan Aggarwal <49625265+Armster15@users.noreply.github.com> Date: Wed, 21 Oct 2026 22:58:54 +0000 Subject: [PATCH 5/5] router: simplify drain limits and resend cancellation Share the timer limit already used by the transport and let negative resend counts disable resend. Keep abort handling in the session waiter, remove unused fixture output, and shorten the README to the public contract. Assisted-by: Replit --- README.md | 16 ++-------------- __tests__/drain.test.ts | 11 ++++++++++- router/client.ts | 6 +++--- router/server.ts | 4 ++-- transport/transport.ts | 4 +++- 5 files changed, 20 insertions(+), 21 deletions(-) diff --git a/README.md b/README.md index ac1340fe..60bbc794 100644 --- a/README.md +++ b/README.md @@ -172,12 +172,7 @@ return Err({ code: 'INVALID_INPUT', message: 'Value must be positive' }); #### Graceful server drain -`server.drain({ budgetMs })` refuses all new procedures with `SERVER_DRAINING` -(`SERVER_DRAINING_CODE`) before middleware or handlers run. Handlers must not emit it. -Drain waits only for admitted RPCs and uploads, until they finish or the budget expires. -It then closes the whole transport, ends streams, and deletes sessions so they cannot resume. -The promise resolves after close. Use one router per draining transport. -Call `server.close()` separately to dispose services. +`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. ```ts const client = createClient(transport, 'SERVER', { @@ -185,14 +180,7 @@ const client = createClient(transport, 'SERVER', { }); ``` -Resend is off by default and applies only to RPCs refused with `SERVER_DRAINING`. -Each resend waits for a connected session with a new ID. Other errors and disconnects are not retried. -Resends stop at `maxResends`. Each wait stops at `newSessionWaitMs`. -An abort before resend dispatch returns the original refusal without resending, while an abort after dispatch follows normal call cancellation. -Wait expiry or transport close also returns the refusal. -`maxResends` must be a non-negative safe integer. Both time limits must be finite, non-negative, and at most 2147483647 ms. -The protobuf router is excluded. -`ReaderErrorSchema` gains a code. Exhaustive TypeScript error switches need a new case. +Resend is off by default. It only resends RPCs refused with `SERVER_DRAINING` on a new session, bounded by `maxResends` and `newSessionWaitMs`. `ReaderErrorSchema` gains a code. Exhaustive TypeScript error switches need a new case. #### Custom Error Types diff --git a/__tests__/drain.test.ts b/__tests__/drain.test.ts index 7a8acfd0..c04ca6be 100644 --- a/__tests__/drain.test.ts +++ b/__tests__/drain.test.ts @@ -110,7 +110,6 @@ function setup( server, transport, client, - clientTransport, gate, uploadGate, effects, @@ -230,6 +229,16 @@ test('successive refusals stop after a positive resend budget', async () => { expect(f.effects).toEqual(['hold']); }); +test('a negative resend budget returns the first refusal', async () => { + // Risk: a finite resend budget is rejected instead of disabling resend. + const f = setup(-1); + await f.block(); + void f.server.drain({ budgetMs: 10000 }); + expect(await f.client.test.rpc.rpc(false)).toEqual(refusal); + expect(f.attempts).toHaveLength(1); + expect(f.effects).toEqual(['hold']); +}); + test('the session wait deadline returns the original refusal', async () => { // Risk: the rpc stays pending forever. const f = setup(2, 100); diff --git a/router/client.ts b/router/client.ts index 5c46bc5e..2a59f9d6 100644 --- a/router/client.ts +++ b/router/client.ts @@ -40,6 +40,7 @@ import { validationErrorToRiverErrors, } from './errors'; import { SessionState } from '../transport/sessionStateMachine/common'; +import { MAX_TIMER_DELAY_MS } from '../transport/transport'; export interface CallOptions { signal?: AbortSignal; @@ -267,10 +268,9 @@ export function createClient< if ( resend && (!Number.isSafeInteger(resend.maxResends) || - resend.maxResends < 0 || !Number.isFinite(resend.newSessionWaitMs) || resend.newSessionWaitMs < 0 || - resend.newSessionWaitMs > 2_147_483_647) + resend.newSessionWaitMs > MAX_TIMER_DELAY_MS) ) { throw new RangeError('Invalid resendRefusedCalls limits'); } @@ -343,7 +343,7 @@ export function createClient< limits.newSessionWaitMs, signal, ); - if (!ready || signal?.aborted) return result; + if (!ready) return result; } } diff --git a/router/server.ts b/router/server.ts index 3b18252a..827b5690 100644 --- a/router/server.ts +++ b/router/server.ts @@ -50,7 +50,7 @@ import { Connection } from '../transport/connection'; import { ServerTransport } from '../transport/server'; import { ReadableImpl, WritableImpl } from './streams'; import { IdentifiedSession } from '../transport/sessionStateMachine/common'; -import { SessionBoundSendFn } from '../transport/transport'; +import { MAX_TIMER_DELAY_MS, SessionBoundSendFn } from '../transport/transport'; import { createPromiseWithResolvers } from '../transport/promises'; type StreamId = string; @@ -1089,7 +1089,7 @@ class RiverServer< if ( !Number.isFinite(budgetMs) || budgetMs < 0 || - budgetMs > 2_147_483_647 + budgetMs > MAX_TIMER_DELAY_MS ) { throw new RangeError('budgetMs must be between 0 and 2147483647'); } diff --git a/transport/transport.ts b/transport/transport.ts index 122b89cd..2b4bf502 100644 --- a/transport/transport.ts +++ b/transport/transport.ts @@ -37,6 +37,8 @@ import { SessionId } from './sessionStateMachine/common'; import { Tracer } from '@opentelemetry/api'; import { getTracer } from '../tracing'; +export const MAX_TIMER_DELAY_MS = 2_147_483_647; + /** * Represents the possible states of a transport. * @property {'open'} open - The transport is open and operational (note that this doesn't mean it is actively connected) @@ -134,7 +136,7 @@ export abstract class Transport< receiverAck.maxMessages < 1 || !Number.isFinite(receiverAck.maxDelayMs) || receiverAck.maxDelayMs <= 0 || - receiverAck.maxDelayMs > 2_147_483_647) + receiverAck.maxDelayMs > MAX_TIMER_DELAY_MS) ) { throw new Error('receiverAck requires positive message and delay limits'); }