diff --git a/README.md b/README.md index eb674811..60bbc794 100644 --- a/README.md +++ b/README.md @@ -170,6 +170,18 @@ return Ok({ result: 42 }); return Err({ code: 'INVALID_INPUT', message: 'Value must be positive' }); ``` +#### Graceful server drain + +`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', { + resendRefusedCalls: { maxResends: 3, newSessionWaitMs: 30_000 }, +}); +``` + +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 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..c04ca6be --- /dev/null +++ b/__tests__/drain.test.ts @@ -0,0 +1,281 @@ +import { afterEach, expect, test, vi } from 'vitest'; +import { Type } from 'typebox'; +import { + ClientOptions, + 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( + 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({ + rpc: Procedure.rpc({ + requestInit: Type.Boolean(), + responseData, + async handler({ reqInit }) { + effects.push(reqInit ? 'hold' : 'rpc'); + if (reqInit) await gate.promise; + + return Ok(null); + }, + }), + upload: Procedure.upload({ + requestInit: init, + requestData: Type.Null(), + responseData, + async handler() { + effects.push('upload'); + await uploadGate.promise; + + return Ok(null); + }, + }), + subscription: Procedure.subscription({ + requestInit: init, + responseData, + async handler({ resWritable }) { + resWritable.write(Ok(null)); + }, + }), + }), + }; + function startServer() { + const transport = network.getServerTransport(); + 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(); + }); + + 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, + ); + 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 { + server, + transport, + client, + gate, + uploadGate, + effects, + attempts, + startServer, + block, + reconnect, + }; +} + +const refusal = { + ok: false, + payload: { + code: 'SERVER_DRAINING', + message: 'Server is draining; the procedure was not run', + }, +}; + +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 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('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('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('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.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 = 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.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); + 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 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 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('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/router/client.ts b/router/client.ts index e64966cc..2a59f9d6 100644 --- a/router/client.ts +++ b/router/client.ts @@ -34,10 +34,13 @@ 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'; +import { MAX_TIMER_DELAY_MS } from '../transport/transport'; export interface CallOptions { signal?: AbortSignal; @@ -205,6 +208,7 @@ function _createRecursiveProxy( } export interface ClientOptions { + resendRefusedCalls?: { maxResends: number; newSessionWaitMs: number }; connectOnInvoke: boolean; eagerlyConnect: boolean; /** @@ -260,10 +264,89 @@ 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'); + } 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, + 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) return result; + } + } + return _createRecursiveProxy((opts) => { const [serviceName, procName, procMethod] = [...opts.path]; if (!(serviceName && procName && procMethod)) { @@ -294,6 +377,10 @@ export function createClient< callOptions as CallOptions | undefined, ); + if (procMethod === 'rpc' && resend) { + return rpcWithResends(init, serviceName, procName, resend, merged.signal); + } + return handleProc( procMethod === 'subscribe' ? 'subscription' : procMethod, transport, @@ -322,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/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..827b5690 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, @@ -49,7 +50,8 @@ 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; @@ -72,6 +74,7 @@ export interface Server< streams: Map>; close: () => Promise; + drain: (options: { budgetMs: number }) => Promise; } interface StreamInitProps { @@ -146,6 +149,11 @@ class RiverServer< >; private unregisterTransportListeners: () => void; + private drainState: + | { status: 'accepting' } + | { status: 'draining'; promise: Promise; finish: () => void } = { + status: 'accepting', + }; constructor( transport: ServerTransport< @@ -231,6 +239,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 +566,9 @@ class RiverServer< const cleanup = () => { finishedController.abort(); this.streams.delete(streamId); + if (this.drainState.status === 'draining' && !this.hasPendingResponse()) { + this.drainState.finish(); + } void runDeferredCleanups(); }; @@ -776,6 +802,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 +821,6 @@ class RiverServer< void runProcedureHandler(); }, )(); - - if (!finishedController.signal.aborted) { - this.streams.set(streamId, procStream); - } } private getContext(service: AnyService, serviceName: string) { @@ -1052,6 +1078,40 @@ 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 || + 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; + } + } + } + async close() { this.unregisterTransportListeners(); 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'); }