Add drain and refused-call resend to the protobuf router

README.md

@@ -174,6 +174,8 @@ return Err({ code: 'INVALID_INPUT', message: 'Value must be positive' });
 
 `server.drain({ budgetMs })` refuses new calls with `SERVER_DRAINING` before handlers run. It waits for in-flight RPCs and uploads up to `budgetMs`, then closes the transport.
 
+Both the TypeBox and protobuf routers support drain and opt-in refused-call resend. For protobuf, drain waits for unary and client-streaming calls; only unary calls can be resent.
+
 ```ts
 const client = createClient<ServiceSurface>(transport, 'SERVER', {
   resendRefusedCalls: { maxResends: 3, newSessionWaitMs: 30_000 },

__tests__/protobuf.test.ts

@@ -4,6 +4,7 @@ import { Type } from 'typebox';
 import { BinaryCodec, NaiveJsonCodec } from '../codec';
 import {
   type ClientError,
+  type ProtocolError,
   Err,
   type Middleware,
   Ok,
@@ -39,9 +40,9 @@ const protobufRouterMatrix = transports.flatMap((transport) =>
   ].map((codec) => ({ transport, codec })),
 );
 
-function coerceReadableError(
-  error: ClientError | typeof ReadableBrokenError,
-): ClientError {
+function coerceReadableError<E extends ClientError>(
+  error: E | typeof ReadableBrokenError,
+): E | ProtocolError {
   return error.code === ReadableBrokenError.code
     ? { code: UNEXPECTED_DISCONNECT_CODE, message: error.message }
     : error;

__tests__/protobufDrain.test.ts

@@ -0,0 +1,254 @@
+import { afterEach, expect, test, vi } from 'vitest';
+import {
+  createClient,
+  createProtoService,
+  createServer,
+  Err,
+  Middleware,
+  Ok,
+  ProtoCodec,
+  type ProtobufClientOptions,
+  SERVER_DRAINING_CODE,
+} from '../protobuf';
+import { createMockTransportNetwork, readNextResult } from '../testUtil';
+import { TestService } from '../testUtil/fixtures/protobuf';
+import { waitFor } from '../testUtil/fixtures/cleanup';
+import { createPromiseWithResolvers } from '../transport/promises';
+
+const cleanups: Array<() => void | Promise<void>> = [];
+afterEach(async () => {
+  for (const cleanup of cleanups.splice(0).reverse()) await cleanup();
+});
+
+function setup({
+  resendRefusedCalls,
+  middlewares = [],
+}: {
+  resendRefusedCalls?: ResendRefusedCalls;
+  middlewares?: Array<Middleware>;
+} = {}) {
+  const network = createMockTransportNetwork({
+    client: { codec: ProtoCodec, baseIntervalMs: 1, maxJitterMs: 0 },
+    server: { codec: ProtoCodec },
+  });
+  const gate = createPromiseWithResolvers<void>();
+  const uploadGate = createPromiseWithResolvers<void>();
+  const effects: Array<string> = [];
+  const attempts: Array<string | undefined> = [];
+  const service = createProtoService().define(TestService, {
+    echo: async ({ text }) => {
+      effects.push(text);
+      if (text === 'hold') await gate.promise;
+
+      return Ok({ text: 'ok' });
+    },
+    countUp: ({ resWritable }) => {
+      resWritable.write(Ok({ value: 1 }));
+    },
+    sum: async () => {
+      effects.push('sum');
+      await uploadGate.promise;
+
+      return Ok({ total: 1 });
+    },
+  });
+  function startServer() {
+    const transport = network.getServerTransport();
+    const server = createServer(transport, [service], { middlewares });
+    transport.addEventListener('message', (message) => {
+      if (message.from === 'client' && message.procedureName === 'Echo')
+        attempts.push(transport.sessions.get(message.from)?.id);
+    });
+    cleanups.push(async () => {
+      transport.close();
+      await server.close();
+    });
+
+    return { server, transport };
+  }
+  cleanups.push(() => network.cleanup());
+  const blocker = network.getClientTransport('blocker');
+  const clientTransport = network.getClientTransport('client');
+  clientTransport.reconnectOnConnectionDrop = false;
+  cleanups.push(() => {
+    gate.resolve();
+    uploadGate.resolve();
+    blocker.close();
+    clientTransport.close();
+  });
+  const { server, transport } = startServer();
+  const client = createClient(TestService, clientTransport, 'SERVER', {
+    resendRefusedCalls,
+  });
+  async function block() {
+    void createClient(TestService, blocker, 'SERVER').echo({ text: 'hold' });
+    await waitFor(() => expect(effects).toContain('hold'));
+  }
+  function reconnect() {
+    clientTransport.hardDisconnect();
+    clientTransport.connect('SERVER');
+  }
+
+  return {
+    server,
+    transport,
+    client,
+    gate,
+    uploadGate,
+    effects,
+    attempts,
+    startServer,
+    block,
+    reconnect,
+  };
+}
+
+type ResendRefusedCalls = ProtobufClientOptions['resendRefusedCalls'];
+const resendTwice = { maxResends: 2, newSessionWaitMs: 1000 };
+
+const refusal = { ok: false, payload: { code: 'SERVER_DRAINING' } };
+const success = { ok: true, payload: { text: 'ok' } };
+
+test('drain refuses calls before middleware and handler run', async () => {
+  const middlewareEffects: Array<string> = [];
+  const f = setup({
+    middlewares: [
+      ({ next }) => {
+        middlewareEffects.push('ran');
+        next();
+      },
+    ],
+  });
+  await f.block();
+  middlewareEffects.length = 0;
+  void f.server.drain({ budgetMs: 10000 });
+  expect(await f.client.echo({ text: 'run' })).toMatchObject(refusal);
+  expect(middlewareEffects).toEqual([]);
+  expect(f.effects).toEqual(['hold']);
+});
+
+test('drain waits for unary and client-streaming calls but not a server stream', async () => {
+  const f = setup();
+  await readNextResult(f.client.countUp({ limit: 1 }));
+  const unary = f.client.echo({ text: 'hold' });
+  const upload = f.client.sum().finalize();
+  await waitFor(() => expect(f.effects).toEqual(['hold', 'sum']));
+  const draining = f.server.drain({ budgetMs: 10000 });
+  expect(f.server.drain({ budgetMs: 1 })).toBe(draining);
+  f.gate.resolve();
+  expect(await unary).toMatchObject(success);
+  expect(f.transport.getStatus()).not.toBe('closed');
+  f.uploadGate.resolve();
+  expect(await upload).toMatchObject({ ok: true, payload: { total: 1 } });
+  await draining;
+  expect(f.transport.getStatus()).toBe('closed');
+});
+
+test('the drain budget closes a server whose unary call never returns', async () => {
+  const f = setup();
+  await f.block();
+  const draining = f.server.drain({ budgetMs: 100 });
+  await vi.advanceTimersByTimeAsync(100);
+  expect(f.transport.getStatus()).toBe('closed');
+  await draining;
+});
+
+test('a refused unary call succeeds only on a new session', async () => {
+  const f = setup({ resendRefusedCalls: resendTwice });
+  await f.block();
+  const draining = f.server.drain({ budgetMs: 10000 });
+  const result = f.client.echo({ text: 'run' });
+  await waitFor(() => expect(f.attempts).toHaveLength(1));
+  await vi.advanceTimersByTimeAsync(20);
+  expect(f.attempts).toHaveLength(1);
+  f.gate.resolve();
+  await draining;
+  f.startServer();
+  f.reconnect();
+  expect(await result).toMatchObject(success);
+  expect(f.attempts).toHaveLength(2);
+  expect(f.attempts[1]).not.toBe(f.attempts[0]);
+  expect(f.effects).toEqual(['hold', 'run']);
+});
+
+test('successive refusals stop at maxResends', async () => {
+  const f = setup({ resendRefusedCalls: resendTwice });
+  await f.block();
+  void f.server.drain({ budgetMs: 10000 });
+  const result = f.client.echo({ text: 'run' });
+  for (let attempt = 1; attempt <= 3; attempt++) {
+    await waitFor(() => expect(f.attempts).toHaveLength(attempt));
+    await vi.advanceTimersByTimeAsync(10);
+    f.reconnect();
+  }
+  await vi.advanceTimersByTimeAsync(20);
+  expect(f.attempts).toHaveLength(3);
+  expect(await result).toMatchObject(refusal);
+  expect(f.effects).toEqual(['hold']);
+});
+
+test('the session wait deadline returns the refusal', async () => {
+  const f = setup({
+    resendRefusedCalls: { maxResends: 2, newSessionWaitMs: 100 },
+  });
+  await f.block();
+  void f.server.drain({ budgetMs: 10000 });
+  const result = f.client.echo({ text: 'run' });
+  await waitFor(() => expect(f.attempts).toHaveLength(1));
+  await vi.advanceTimersByTimeAsync(100);
+  expect(await result).toMatchObject(refusal);
+  expect(f.effects).toEqual(['hold']);
+});
+
+test('a refused client-streaming call is never resent', async () => {
+  const f = setup({ resendRefusedCalls: resendTwice });
+  await f.block();
+  void f.server.drain({ budgetMs: 10000 });
+  expect(await f.client.sum().finalize()).toMatchObject(refusal);
+  expect(f.effects).toEqual(['hold']);
+});
+
+test('handlers cannot return the reserved drain code', () => {
+  createProtoService().define(TestService, {
+    // @ts-expect-error only the runtime may send SERVER_DRAINING
+    echo: () => Err({ code: SERVER_DRAINING_CODE, message: 'upstream' }),
+  });
+});
+
+test('drain called from middleware waits for its admitted unary call', async () => {
+  const f = setup({
+    middlewares: [
+      ({ next }) => {
+        void f.server.drain({ budgetMs: 10000 });
+        next();
+      },
+    ],
+  });
+  const result = f.client.echo({ text: 'hold' });
+  await waitFor(() => expect(f.effects).toContain('hold'));
+  expect(f.transport.getStatus()).not.toBe('closed');
+  f.gate.resolve();
+  expect(await result).toMatchObject(success);
+  await f.server.drain({ budgetMs: 10000 });
+});
+
+test('abort while waiting preserves the refusal without a later resend', async () => {
+  const f = setup({
+    resendRefusedCalls: { maxResends: 2, newSessionWaitMs: 30000 },
+  });
+  await f.block();
+  const draining = f.server.drain({ budgetMs: 10000 });
+  const controller = new AbortController();
+  const result = f.client.echo({ text: 'run' }, { signal: controller.signal });
+  await waitFor(() => expect(f.attempts).toHaveLength(1));
+  await vi.advanceTimersByTimeAsync(10);
+  controller.abort();
+  expect(await result).toMatchObject(refusal);
+  f.gate.resolve();
+  await draining;
+  f.startServer();
+  f.reconnect();
+  await vi.advanceTimersByTimeAsync(50);
+  expect(f.effects).toEqual(['hold']);
+  expect(f.attempts).toHaveLength(1);
+});

protobuf/client.ts

@@ -24,6 +24,11 @@ import {
 } from '../transport/message';
 import { generateId } from '../transport/id';
 import { ClientHandshakeOptions } from '../router/handshake';
+import {
+  ResendRefusedCallsOptions,
+  validateResendLimits,
+  withRefusedCallResends,
+} from '../router/drain';
 import { Err, Ok, type Result } from '../router/result';
 import {
   Readable,
@@ -59,6 +64,7 @@ import type {
 export interface ClientOptions {
   readonly connectOnInvoke: boolean;
   readonly eagerlyConnect: boolean;
+  readonly resendRefusedCalls?: ResendRefusedCallsOptions;
 }
 
 const defaultClientOptions: ClientOptions = {
@@ -95,6 +101,7 @@ export function createClient<
   }
 
   const clientOptions = { ...defaultClientOptions, ...providedClientOptions };
+  validateResendLimits(clientOptions.resendRefusedCalls);
   if (clientOptions.eagerlyConnect) {
     transport.connect(serverId);
   }
@@ -142,17 +149,33 @@ function createMethodCaller<
         }
 
         connectOnInvokeIfNeeded(clientOptions, transport, serverId);
-        const { resReadable } = startMethodCall(
-          service,
-          method,
+
+        const requestBytes = encodeMessageBytes(method.input, request);
+        const call = () =>
+          getSingleMessage(
+            startMethodCall(
+              service,
+              method,
+              transport,
+              serverId,
+              requestBytes,
+              true,
+              options?.signal,
+            ).resReadable,
+            transport.log,
+          );
+        const resend = clientOptions.resendRefusedCalls;
+        if (!resend) {
+          return call();
+        }
+
+        return withRefusedCallResends(
+          call,
           transport,
           serverId,
-          encodeMessageBytes(method.input, request),
-          true,
+          resend,
           options?.signal,
         );
-
-        return getSingleMessage(resReadable, transport.log);
       }) as unknown as ClientMethod<Method>;
 
     case 'server_streaming':