diff --git a/__tests__/drain.test.ts b/__tests__/drain.test.ts index c04ca6be..31075049 100644 --- a/__tests__/drain.test.ts +++ b/__tests__/drain.test.ts @@ -5,9 +5,11 @@ import { createClient, createServer, createServiceSchema, + Err, Middleware, Ok, Procedure, + UNCAUGHT_ERROR_CODE, } from '../router'; import { createMockTransportNetwork, readNextResult } from '../testUtil'; import { waitFor } from '../testUtil/fixtures/cleanup'; @@ -62,6 +64,19 @@ function setup( resWritable.write(Ok(null)); }, }), + claim: Procedure.rpc({ + requestInit: init, + responseData, + responseError: Type.Object({ + code: Type.Literal('SERVER_DRAINING'), + message: Type.String(), + }), + async handler() { + effects.push('claim'); + + return Err({ code: 'SERVER_DRAINING', message: 'upstream' }); + }, + }), }), }; function startServer() { @@ -279,3 +294,18 @@ test('a refused upload is returned without replay', async () => { expect(await f.client.test.upload.upload({}).finalize()).toEqual(refusal); expect(f.effects).toEqual(['hold']); }); + +test('a handler that returns the reserved drain code is not resent', async () => { + // Risk: a handler that already ran is replayed on the next session. + const f = setup(1); + const result = f.client.test.claim.rpc({}); + await waitFor(() => expect(f.effects).toEqual(['claim'])); + await vi.advanceTimersByTimeAsync(10); + f.reconnect(); + await vi.advanceTimersByTimeAsync(20); + expect(f.effects).toEqual(['claim']); + expect(await result).toMatchObject({ + ok: false, + payload: { code: UNCAUGHT_ERROR_CODE, message: 'upstream' }, + }); +}); diff --git a/router/server.ts b/router/server.ts index 827b5690..2ef33b96 100644 --- a/router/server.ts +++ b/router/server.ts @@ -604,6 +604,17 @@ class RiverServer< Result, Static> >({ writeCb: (response) => { + if (!response.ok && response.payload.code === SERVER_DRAINING_CODE) { + this.log?.warn( + `${serviceName}.${procedureName} handler returned the reserved ${SERVER_DRAINING_CODE} code`, + loggingMetadata, + ); + response = Err({ + code: UNCAUGHT_ERROR_CODE, + message: response.payload.message, + }); + } + if (!response.ok) { recordRiverError(span, response.payload); }