Skip to content

Add drain and refused-call resend to the protobuf router - #433

Closed
Armster15 wants to merge 3 commits into
mainfrom
protobuf-drain-resend
Closed

Armster15 wants to merge 3 commits into
mainfrom
protobuf-drain-resend

Conversation

@Armster15

@Armster15 Armster15 commented Oct 21, 2026 •

Copy link
Copy Markdown
Contributor

Why

The protobuf router lacks the drain and resend added in #426. A server must refuse new work before middleware or handlers run so the client can safely resend it.

What changed

  • Add server.drain({ budgetMs }) to the protobuf router. Refuse new calls with SERVER_DRAINING, wait for admitted unary and client-streaming calls, then close the transport when they finish or the budget expires.
  • Add opt-in resendRefusedCalls: { maxResends, newSessionWaitMs } for protobuf unary calls. Only a drain refusal can be resent, and each resend needs a connected session with a new ID.
  • Share the drain state machine, resend option type and validation, and session waiter with the TypeBox router. Keep its behavior and tests unchanged.
  • Register protobuf calls before middleware runs, and document support in both routers.

Versioning

  • Breaking protocol change
  • Breaking ts/js API change

No wire envelope change or version bump. ClientErrorCode gains SERVER_DRAINING, so exhaustive client error switches need a new case. Handler results now use HandlerError, which excludes SERVER_DRAINING. A handler that returns a value typed as the wider ClientError, for example by forwarding a client call's error, must map it first. The protobuf router is experimental; maintainers can decide whether that type change needs the breaking-ts box.

Tests

  • Full suite: 824 passed, 1 skipped, across 38 files. npm run test:single exited 0.
  • npm run check exited 0: TypeScript, Prettier, and ESLint passed.
  • Ten protobuf tests (nine behavior tests and one compile-time check run by npm run check) use public APIs, fake timers, the existing service descriptor, and the mock transport with ProtoCodec. The existing ten TypeBox drain tests are unchanged and pass.
  • Red proofs, each run separately and restored:
    • Disable admission: the refusal test receives a successful response instead.
    • Count server streams as pending: the drain-completion test times out.
    • Stop counting unary calls: the middleware-drain test finds the transport already closed.
    • Disable the drain deadline callback: the transport is still open at the deadline.
    • Remove the new-session-ID check: the new-session test sees three old-session requests instead of one.
    • Remove the resend limit: the limit test sees four requests instead of three.
    • Disable the session-wait deadline callback: the wait-deadline test times out.
    • Route client-streaming finalization through resend while preserving its public API: the no-resend test times out.
    • Move registration after middleware: the admitted-call test finds the transport already closed.
    • Remove the waiter's abort listener: the waiting-abort test times out.
    • Widen HandlerError to every client code: npm run check fails with an unused @ts-expect-error on the handler that returns SERVER_DRAINING.
Appendix

Decision Log

  • Reuse SERVER_DRAINING directly. The protobuf router already has a separate ProtocolErrorCode union and runtime set for River's built-in errors alongside the gRPC-style codes (protobuf/errors.ts:37-60). No new metadata key, error-detail message, or generic UNAVAILABLE retry rule is needed.
  • Share lifecycle policy in router/drain.ts: two states, the existing promise-resolver helper, bounded timers, explicit waiter cleanup, and one resend option contract. Each router supplies its own pending-call predicate and transport close operation.
  • Share one resend loop, withRefusedCallResends, in router/drain.ts. Each router passes a closure that dispatches one call, so the loop runs only when resend is on and the session waiter stays private.
  • Match the TypeBox rules: validate finite time bounds through MAX_TIMER_DELAY_MS, require a safe-integer resend count, and let non-positive counts disable resend.
  • Abort while waiting returns the refusal. Once dispatched, a resend uses the normal call cancellation path. No extra scheduling step or post-wait abort check.
  • Only the runtime admission refusal may carry SERVER_DRAINING. This is enforced in types, not at runtime: HandlerError excludes it from every handler result, raw handlers included. ProtocolErrorCode does not include it, so request-stream errors cannot carry it, and a client cancel with that code is rejected as on main. Other protocol codes stay allowed because handlers forward request-stream errors such as UNEXPECTED_DISCONNECT and return ctx.cancel(...); those codes never trigger resend.
  • Use one router per draining transport. Drain closes the transport; call server.close() separately to dispose service resources.

Facts verified in code

  • protobuf/server.ts:272-289,458-462,796-811: admission refusal, stream cleanup, and registration before middleware.
  • protobuf/server.ts:185-193,346-348: unary/client-streaming pending predicate and shared drain entry point.
  • protobuf/client.ts:154-178 and router/client.ts:305-320: only unary/RPC dispatch enters the shared loop, and only when resend is configured.
  • protobuf/errors.ts:41-62,93-99,169-180 and protobuf/types.ts:18-21,128-144: client-visible codes, the handler error type, and its use in typed and raw handlers.
  • protobuf/errors.ts:41-60 and protobuf/index.ts: the existing protocol error boundary and the refusal-code export.
  • router/drain.ts:18-30,32-63,65-101,103-150: option validation, the shared resend loop, fresh-session waiting with listener/timer cleanup, and the two-state drain controller.
  • Precedents: router/server.ts:1078-1086 retains the RPC/upload-only predicate; the shared loop keeps the dispatch-then-read-session ordering from the TypeBox resend loop in Add graceful server drain and opt-in resend for refused RPCs #426. transport/promises.ts supplies the existing promise-resolver helper.

Review dispositions

  • Pass A — SHIP. No refused call reaches middleware or handlers. Admitted calls are registered before middleware, finite-response calls hold drain, long streams do not, and the deadline closes stuck work. The waiter removes both transport listeners, its abort listener, and its timer on every completion path. Resend remains opt-in and unary-only. No new casts or any were added. Existing TypeBox tests and the full protobuf matrix pass.
  • Pass B — SHIP after one cleanup pass. Share the option type as well as its validation, and fix the required blank lines. The shared state machine removes copied lifecycle code; dispatch stays local. No file crosses 1,000 lines: the protobuf server was already 1,053 lines. No unrelated refactor or test-only parameter was added.
  • Review finding: a handler could return SERVER_DRAINING after running, and the client would resend it. 28f8a6e fixed it with a runtime remap. Per code review, d4c4c2c replaces that remap with HandlerError, so the code is rejected at compile time.
  • Code review, d4c4c2c: share the resend loop and branch at the call site; use an options object in test setup; drop the // Risk: comments. Pass A and Pass B were rerun on this diff: SHIP. The full suite and npm run check pass.
  • Test size: 254 lines for ten tests, including the middleware-registration and abort cases beyond the seven requested core cases. Keep those two tests because they cover integration risks at the new protobuf entry points.
  • Size: production +290/-138 lines; tests +258/-3 lines; README +2 lines.
  • Limits: no real-network rollout or mixed-version compatibility test was run. Listener cleanup is checked in code and by the abort/no-later-resend behavior, not by private listener counts.

~ written by ⠕ Replit

Share drain lifecycle and session waiting with the TypeBox router.
Add protobuf behavior coverage and document both router surfaces.

Assisted-by: Replit
@Armster15
Armster15 marked this pull request as ready for review October 21, 2026 23:31
@Armster15
Armster15 requested a review from a team as a code owner October 21, 2026 23:31
@Armster15
Armster15 requested review from wernst and removed request for a team October 21, 2026 23:31
Only the runtime admission refusal may carry the reserved code, so a
handler that already ran is never resent. Remap a handler-returned
SERVER_DRAINING to UNCAUGHT_ERROR and log a warning.

Assisted-by: Replit
Comment thread protobuf/client.ts Outdated
request: Uint8Array,
limits: ResendRefusedCallsOptions | undefined,
signal?: AbortSignal,
): Promise<Result<MessageShape<Method['output']>, ClientError>> {

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

7 positional params that get forwarded straight into startMethodCall, plus the !limits branch inside. should this branch at the call site like the typebox client (resend ? rpcWithResends : handleProc) so the loop only exists when resend is on?

Copy link
Copy Markdown
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Done in d4c4c2c. The loop is now one shared withRefusedCallResends(call, transport, serverId, limits, signal) in router/drain.ts, used by both routers. The protobuf unary case builds a call closure and only enters the loop when resendRefusedCalls is set, matching the TypeBox call site. unaryWithResends, the !limits branch, and the TypeBox rpcWithResends copy are gone.

Comment thread protobuf/server.ts Outdated

const resWritable = new WritableImpl<HandlerResponse>({
writeCb: (response) => {
if (!response.ok && response.payload.code === SERVER_DRAINING_CODE) {

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

feels like a type thing not a runtime thing. can the handler's error type exclude ProtocolErrorCode instead of remapping this one code? a handler returning CANCEL or UNEXPECTED_DISCONNECT still goes through untouched here

Copy link
Copy Markdown
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Agreed, moved it to the type level in d4c4c2c and dropped the runtime remap. Handler results (typed and raw) now use HandlerError, which excludes SERVER_DRAINING; a test with @ts-expect-error fails npm run check if the type is widened.

I kept the other protocol codes allowed on purpose. Handlers return ctx.cancel(...) (CANCEL) and forward request-stream errors like UNEXPECTED_DISCONNECT; the existing protobuf tests do both. Those codes never trigger resend. SERVER_DRAINING is also out of ProtocolErrorCode now, so request-stream errors cannot carry it and a client cancel with that code is rejected as on main.

Comment thread __tests__/protobufDrain.test.ts Outdated
test('drain refuses calls before middleware and handler run', async () => {
// Risk: a refused call that ran would be unsafe to resend.
const middlewareEffects: Array<string> = [];
const f = setup(undefined, 1000, [

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

setup(undefined, 1000, [...]) 👀 options object?

Copy link
Copy Markdown
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Switched to setup({ resendRefusedCalls, middlewares }) in d4c4c2c.

Comment thread __tests__/protobufDrain.test.ts Outdated
const success = { ok: true, payload: { text: 'ok' } };

test('drain refuses calls before middleware and handler run', async () => {
// Risk: a refused call that ran would be unsafe to resend.

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

nit: these // Risk: lines read like the PR description, drop? (same for the rest of the file)

Copy link
Copy Markdown
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Dropped in d4c4c2c.

Replace the runtime remap with a handler error type that excludes
SERVER_DRAINING; only the runtime refusal can send it. Share one resend
loop between both routers so it exists only when resend is on. Use an
options object in the protobuf drain test setup and drop the Risk lines.

Assisted-by: Replit
@Armster15 Armster15 closed this Oct 22, 2026
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

3 participants