Skip to content

Commit 082197a

Browse files
igalshilmanclaude
andauthored
[restate-sdk-tunnel] Client-initiated graceful shutdown (#749)
* [restate-sdk-tunnel] Client-initiated graceful shutdown Adds connection.shutdown() and opt-in signal handling so an in-process tunnel deployment can drain on SIGTERM without dropping invocations during a rolling restart. On shutdown the client advertises `supports-client-drain` at handshake and, while draining, refuses each NEW forwarded invocation with `503` + `x-restate-tunnel-draining: true` WITHOUT running the handler — the server then deselects this connection and retries the invocation on a healthy replica (no double-execution, since the handler never ran). In-flight invocations keep running and are awaited (bounded by `drainGraceMs`, default 120s) before teardown; an idle connection just closes. API: - TunnelConnection.shutdown({ graceMs? }) — drain, then stop; wire it to SIGTERM. Degrades to an abrupt close() if supportsClientDrain was not advertised. - gracefulShutdown option (off by default) — opt in to auto-installing one-shot signal handlers that call shutdown() then process.exit(0). - supportsClientDrain option (default true) — advertise the capability. Requires the matching tunnel-server support (restate-cloud). Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com> * [restate-sdk-tunnel] Model connection lifecycle as an explicit state machine Replace the settled/serving/openedAt/handshakePromise/detachForDrain flag set on ConnectionAttempt with one AttemptState discriminated union (connecting -> handshaking -> serving -> draining{trigger} -> closed), each variant carrying only its phase's data, so illegal combinations (e.g. serving without a session) cannot be represented. Server- and client-initiated drains become the same explicit `draining` state distinguished by `trigger`, making the essential difference visible rather than implicit: a server drain detaches the session to the DrainingRegistry and redials (the zero-drop property), while a client drain refuses new invocations and finishes in place with no redial. run()-resolution is decoupled from teardown (resolve-once) to support the detach handover — previously expressed implicitly by leaving `serving` true after settle. The engine gains an activeConnections registry so shutdown() drives each live connection's beginClientDrain(). No behavior change; tsc, eslint, and 92/92 vitest cases pass. Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com> * [restate-sdk-tunnel] Decompose ConnectionAttempt into focused units Pull the cross-cutting concerns out of the ConnectionAttempt class so it reads as an explicit stage pipeline holding almost no shared mutable state: - dial() — connect + TLS + ALPN + connect-timeout + abort, owning its own timer/listeners and the in-flight socket - Completion — resolve run() exactly once (was resolved + resolveRun) - Watchdog — liveness ping, owning its interval + miss counter - classifyRequest — pure request -> discriminated TunnelRequest, so routing is a switch rather than a path-sniffing if-else chain ConnectionAttempt drops from ~10 mutable fields to 4 and orchestrates dial -> establish, with handshake/serve event-driven from the router. No behavior change; tsc, eslint, and 92/92 vitest pass. Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com> * [restate-sdk-tunnel] Extract the slot supervisor from connectTunnel connectTunnel() had ~12 mutable bindings shared across eight nested closures. Pull the cohesive concerns into their own units so the engine is a thin assembler holding only lifecycle flags + public-handle state: - Supervisor (supervisor.ts) — the slot map, per-server reconnect loops, and the resolve/reconcile loop; exposes done / fatalError / stopResolving() / abortAll(). Replaces slots + runSlot + startSlot + stopAllSlots + the supervisor IIFE + the two ad-hoc AbortControllers. A fatal outcome is reported via a single onFatal hook instead of reaching into engine state. - InflightTracker — in-flight count + whenDrained(graceMs); replaces globalInflight + inflightDrained + waitForInflightDrain. - Deferred<T> (util.ts) — resolve/reject-once for readiness; replaces the readyResolve/readyReject + catch dance. connectTunnel now builds the deps, creates the Supervisor, and owns only close()/shutdown()/signal wiring and the handle. No behavior change; tsc, eslint, and 92/92 vitest pass. Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com> * [restate-sdk-tunnel] Replace engine lifecycle booleans with one disjoint phase The engine tracked its lifecycle in three booleans (stopped / toreDown / shuttingDown) whose combinations could express nonsense. Collapse them into a single disjoint state — running | draining | closed: - isShuttingDown() -> phase === "draining" - teardown() guard -> phase === "closed" (then sets it) - shutdown() re-entry guard -> phase !== "running" close() is running|draining -> closed; shutdown() is running -> draining -> closed. One source of truth, illegal combinations unrepresentable. No behavior change; tsc, eslint, and 92/92 vitest pass. Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com> * [restate-sdk-tunnel] Carry per-phase data on the engine lifecycle state Make the engine state a data-carrying disjoint union instead of a bare marker: running — no phase-local data (the registries, supervisor, and the handle's learned info are lifetime-scoped — used across every phase and exposed by the handle even after close — so they stay in stable fields, not the state). draining — holds `completed`, the in-progress graceful-shutdown promise, so a re-entrant shutdown()/close() coalesces onto the same drain instead of starting another. The promise is recorded on the state before the first await, so isShuttingDown() holds for the whole drain. closed — terminal. The drain work moves into drainGracefully(); shutdown() is now just the state transition. No behavior change; tsc, eslint, and 92/92 vitest pass. Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com> * [restate-sdk-tunnel] Hold live machinery in the engine state, not closures Apply the data-in-state principle to the engine: phase-owned resources live on the state, and each transition narrows then destructures what it needs, instead of reaching for shared closure vars. - running/draining carry `active` (the supervisor + the event-loop anchor); draining additionally carries its in-flight `completed` drain promise. - closed carries nothing — the live machinery is provably gone, so teardown (which destructures `{ supervisor, keepAlive } = state.active`) and drainGracefully can only touch it in a live phase. Two kinds of state stay out of the union by design, with comments saying why: the registries (activeSockets/activeConnections/draining/inflight) are injected into every connection via ConnectionDeps, so they are shared infrastructure; and connectionCount/lastInfo/fatalError are the handle's observable output, readable in every phase including after close, so they live in one Output record. No behavior change; tsc, eslint, and 92/92 vitest pass. Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com> * [restate-sdk-tunnel] Move phase-owned resources onto the connection state Finish applying the data-in-state principle to ConnectionAttempt: the handshake timer and the liveness watchdog were nullable instance fields reachable from any phase; move them onto the phases that own them so methods narrow then destructure instead of touching fields that may be undefined. - handshaking now carries `firstRequestTimer` (the bound on waiting for the cloud to open /_/start-tunnel). - serving and draining carry the running `watchdog`. - teardown reads them via `stopPhaseResources()` off the narrowed state; a server drain destructures `{ session, openedAt, watchdog } = this.state`. `socket` and `completion` stay instance fields — genuinely lifetime-scoped (socket is destroyed in every teardown path and handed to the registry on a server drain; completion resolves run() once). No behavior change; tsc, eslint, and 92/92 vitest pass. Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com> * [restate-sdk-tunnel] Address PR review on graceful shutdown - shutdown() honors supportsClientDrain: false. The docs promise it degrades to an abrupt close() when the capability was not advertised, but it still entered draining and emitted the sentinel — which the server ignores for non-capable clients, so refused requests just keep getting routed back (a refusal loop, not a drain). shutdown() now returns close() in that case. Added a regression test: with supportsClientDrain: false, shutdown() tears down immediately despite an in-flight invocation instead of waiting the grace. - gracefulShutdown signal handlers are removed on teardown. We installed process.once(signal, ...) but never unregistered the not-yet-fired handler, so a closed connection could later intercept a signal and call process.exit(0) on the host process. Track the handlers and removeListener them in teardown(). No behavior change to the happy path; tsc, eslint, and 93/93 vitest pass. Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com> * [restate-sdk-tunnel] Enable graceful shutdown on SIGTERM by default gracefulShutdown was opt-in, so out of the box a SIGTERM killed the process abruptly (in-flight invocations cut) even though the drain capability was advertised. Flip it on by default: an unset (or `true`) gracefulShutdown now installs a SIGTERM handler that drains then exits, so an operator-managed in-process deployment gets zero-dropped-invocation rollouts with no wiring. Only an explicit `gracefulShutdown: false` opts out (manage signals/exit yourself and call shutdown() by hand). Handlers are still removed on close, so this doesn't leak signal listeners. Updated the option doc and the default test. tsc, eslint, and 95/95 vitest pass. Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com> * [restate-sdk-tunnel] Don't leak a signal handler on the aborted-signal path connectTunnel installed the gracefulShutdown handlers AFTER the already-aborted-signal check, whose synchronous close() ran teardown() while signalHandlers was still empty — so an already-aborted signal (now reachable by default, since gracefulShutdown is on) returned a closed connection that still held a live SIGTERM handler able to later call process.exit(0). Register the handlers before the aborted-signal handling so that path's close()/teardown() removes them like every other teardown. Adds a regression test asserting process.listenerCount("SIGTERM") is unchanged after constructing with an already-aborted signal + gracefulShutdown. (Codex finding on #749.) tsc, eslint, and 96/96 vitest pass. Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com> --------- Co-authored-by: Claude Opus 4.8 <noreply@anthropic.com>
1 parent e765029 commit 082197a

10 files changed

Lines changed: 1375 additions & 469 deletions

File tree

packages/libs/restate-sdk-tunnel/src/connect.ts

Lines changed: 229 additions & 202 deletions
Large diffs are not rendered by default.

packages/libs/restate-sdk-tunnel/src/connection.ts

Lines changed: 518 additions & 264 deletions
Large diffs are not rendered by default.

packages/libs/restate-sdk-tunnel/src/handshake.ts

Lines changed: 8 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -63,6 +63,13 @@ export interface HandshakeCredentials {
6363
* obliges us to open a replacement connection on drain.
6464
*/
6565
supportsDrain: boolean;
66+
/**
67+
* Advertise `supports-client-drain: true`. Tells the server that on
68+
* shutdown we refuse new invocations with the `x-restate-tunnel-draining`
69+
* sentinel (rather than dropping them); only then does the server trust
70+
* that sentinel to deselect this connection.
71+
*/
72+
supportsClientDrain: boolean;
6673
}
6774

6875
export const START_TUNNEL_PATH = "/_/start-tunnel";
@@ -166,6 +173,7 @@ export function performHandshake(
166173
"environment-id": creds.environmentId,
167174
"tunnel-name": creds.tunnelName,
168175
...(creds.supportsDrain && { "supports-drain": "true" }),
176+
...(creds.supportsClientDrain && { "supports-client-drain": "true" }),
169177
});
170178
res.end();
171179
});

packages/libs/restate-sdk-tunnel/src/options.ts

Lines changed: 30 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -46,6 +46,9 @@ export interface ResolvedOptions {
4646
resolveIntervalMs: number;
4747
supportsDrain: boolean;
4848
drainGraceMs: number;
49+
supportsClientDrain: boolean;
50+
/** Set when auto signal-handling is opted into; undefined leaves signals alone. */
51+
gracefulShutdown?: { signals: NodeJS.Signals[]; graceMs: number };
4952
connectTimeoutMs: number;
5053
handshakeTimeoutMs: number;
5154
reconnectInitialMs: number;
@@ -245,6 +248,7 @@ export function resolveOptions(options: ConnectTunnelOptions): ResolvedOptions {
245248
10_000,
246249
"pingTimeoutMs"
247250
);
251+
const drainGraceMs = positive(options.drainGraceMs, 120_000, "drainGraceMs");
248252

249253
return {
250254
srvName: hasRegion
@@ -264,7 +268,12 @@ export function resolveOptions(options: ConnectTunnelOptions): ResolvedOptions {
264268
"resolveIntervalMs"
265269
),
266270
supportsDrain: options.supportsDrain ?? true,
267-
drainGraceMs: positive(options.drainGraceMs, 120_000, "drainGraceMs"),
271+
drainGraceMs,
272+
supportsClientDrain: options.supportsClientDrain ?? true,
273+
gracefulShutdown: resolveGracefulShutdown(
274+
options.gracefulShutdown,
275+
drainGraceMs
276+
),
268277
connectTimeoutMs: positive(
269278
options.connectTimeoutMs,
270279
5_000,
@@ -305,6 +314,26 @@ export function resolveOptions(options: ConnectTunnelOptions): ResolvedOptions {
305314
};
306315
}
307316

317+
/** Resolve the opt-in auto signal-handling config (undefined = leave signals alone). */
318+
function resolveGracefulShutdown(
319+
option: boolean | { signals?: NodeJS.Signals[]; graceMs?: number } | undefined,
320+
drainGraceMs: number
321+
): { signals: NodeJS.Signals[]; graceMs: number } | undefined {
322+
// On by default: only an explicit `false` opts out.
323+
if (option === false) return undefined;
324+
if (option === undefined || option === true) {
325+
return { signals: ["SIGTERM"], graceMs: drainGraceMs };
326+
}
327+
const signals = option.signals ?? ["SIGTERM"];
328+
if (signals.length === 0) {
329+
throw new Error("tunnel: gracefulShutdown.signals must not be empty");
330+
}
331+
return {
332+
signals,
333+
graceMs: positive(option.graceMs, drainGraceMs, "gracefulShutdown.graceMs"),
334+
};
335+
}
336+
308337
/**
309338
* Build the `tls.connect` options for a tunnel target, or `undefined` for a
310339
* plaintext connection.
Lines changed: 201 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,201 @@
1+
/*
2+
* Copyright (c) 2023-2026 - Restate Software, Inc., Restate GmbH
3+
*
4+
* This file is part of the Restate SDK for Node.js/TypeScript,
5+
* which is released under the MIT license.
6+
*
7+
* You can find a copy of the license in file LICENSE in the root
8+
* directory of this repository or package, or at
9+
* https://github.com/restatedev/sdk-typescript/blob/main/LICENSE
10+
*/
11+
12+
// The slot supervisor.
13+
// =============================================================================
14+
//
15+
// Multi-homing — one tunnel connection per resolved tunnel server (like the
16+
// Rust client; the slot set IS the resolved set, it is not configurable). The
17+
// supervisor resolves the server set, reconciles the slot map against it
18+
// (starting connections to servers that appear, tearing down ones that
19+
// vanish), and re-resolves every `resolveIntervalMs` for SRV discovery. Each
20+
// slot runs its own reconnect loop with fatal-vs-retryable classification.
21+
//
22+
// Invariants:
23+
// E1. A FATAL outcome (unauthorized / bad-tunnel-name / name mismatch) on
24+
// ANY slot stops the WHOLE tunnel — the credentials are shared, so every
25+
// other slot would hit the same wall. It aborts every slot and reports
26+
// via `hooks.onFatal`; the engine surfaces it on `error`/`ready`.
27+
// E2. Backoff resets only after a connection held for
28+
// MIN_UPTIME_FOR_BACKOFF_RESET_MS; a drain only skips the backoff sleep
29+
// under the same guard (drain-spam must compound).
30+
// E3. Teardown is prompt: `abortAll()` aborts in-flight dials via per-slot
31+
// signals and wakes the resolve loop out of any sleep; the (un-abortable)
32+
// DNS work is raced against the wake signal, never awaited.
33+
34+
import type { ResolvedOptions } from "./options.js";
35+
import { resolveTargets, targetKey, type Target } from "./targets.js";
36+
import { runConnection, type ConnectionDeps } from "./connection.js";
37+
import { Backoff, MIN_UPTIME_FOR_BACKOFF_RESET_MS } from "./backoff.js";
38+
import { delay, raceAbortable } from "./util.js";
39+
40+
/** A running per-server connection loop. */
41+
interface Slot {
42+
ctl: AbortController;
43+
done: Promise<void>;
44+
}
45+
46+
export interface SupervisorHooks {
47+
/** A slot hit a non-retryable failure; the whole tunnel must stop (E1). */
48+
onFatal: (err: Error) => void;
49+
}
50+
51+
export class Supervisor {
52+
private readonly slots = new Map<string, Slot>();
53+
/** Aborting this cascades to every slot (each slot chains its ctl to it). */
54+
private readonly stopSignal = new AbortController();
55+
/** Wakes the resolve loop out of a sleep / DNS race (E3). */
56+
private readonly wake = new AbortController();
57+
private stopping = false;
58+
private fatal: Error | undefined;
59+
60+
/** Resolves when the resolve loop has exited AND every slot has settled. */
61+
readonly done: Promise<void>;
62+
63+
constructor(
64+
private readonly opts: ResolvedOptions,
65+
private readonly deps: ConnectionDeps,
66+
private readonly hooks: SupervisorHooks,
67+
private readonly log: (message: string) => void
68+
) {
69+
this.stopSignal.signal.addEventListener("abort", () => this.wake.abort(), {
70+
once: true,
71+
});
72+
this.done = this.supervise();
73+
}
74+
75+
get fatalError(): Error | undefined {
76+
return this.fatal;
77+
}
78+
79+
/**
80+
* Stop starting/resolving new connections; existing slots keep running so
81+
* their connections can finish draining in place (client-initiated drain).
82+
*/
83+
stopResolving(): void {
84+
this.stopping = true;
85+
this.wake.abort();
86+
}
87+
88+
/** Abort every slot and the resolve loop (engine teardown). */
89+
abortAll(): void {
90+
this.stopping = true;
91+
this.stopSignal.abort();
92+
}
93+
94+
private startSlot(key: string, target: Target): void {
95+
const ctl = new AbortController();
96+
// Chain to the global stop so abortAll() cascades; self-detaching.
97+
this.stopSignal.signal.addEventListener("abort", () => ctl.abort(), {
98+
once: true,
99+
signal: ctl.signal,
100+
});
101+
const slot: Slot = { ctl, done: Promise.resolve() };
102+
slot.done = this.runSlot(target, ctl).finally(() => {
103+
// Guarded: this key may have vanished and re-appeared, in which case a
104+
// NEWER slot owns it — don't delete someone else's registration.
105+
if (this.slots.get(key) === slot) this.slots.delete(key);
106+
});
107+
this.slots.set(key, slot);
108+
}
109+
110+
/** The per-server loop: dial → serve → classify outcome → backoff → redial. */
111+
private async runSlot(target: Target, ctl: AbortController): Promise<void> {
112+
const backoff = new Backoff(
113+
this.opts.reconnectInitialMs,
114+
this.opts.reconnectFactor,
115+
this.opts.reconnectMaxMs
116+
);
117+
118+
while (!this.stopping && !ctl.signal.aborted && this.fatal === undefined) {
119+
const outcome = await runConnection(target, ctl.signal, this.deps);
120+
if (this.stopping || ctl.signal.aborted) break;
121+
if (outcome.kind === "fatal") {
122+
// E1: shared credentials — stop everything.
123+
this.fatal = new Error(`tunnel: ${outcome.reason}`);
124+
this.log(`tunnel: FATAL — ${outcome.reason}; stopping all connections`);
125+
this.hooks.onFatal(this.fatal);
126+
this.stopSignal.abort();
127+
break;
128+
}
129+
if (outcome.kind === "served" || outcome.kind === "drained") {
130+
// E2: only a connection that actually held resets the backoff.
131+
const heldLongEnough =
132+
outcome.uptimeMs >= MIN_UPTIME_FOR_BACKOFF_RESET_MS;
133+
if (heldLongEnough) backoff.reset();
134+
if (outcome.kind === "drained" && heldLongEnough) {
135+
// A stable connection was asked to rotate and the server is holding
136+
// the old one open for us — replace it NOW.
137+
this.log("tunnel: draining — reconnecting immediately");
138+
continue;
139+
}
140+
this.log(
141+
outcome.kind === "drained"
142+
? "tunnel: drained shortly after connecting — reconnecting with backoff"
143+
: "tunnel: connection ended — reconnecting"
144+
);
145+
} else {
146+
this.log(`tunnel: ${outcome.reason} — reconnecting`);
147+
}
148+
await delay(backoff.next(), ctl.signal);
149+
}
150+
}
151+
152+
/** Resolve the server set, reconcile slots, repeat. For SRV discovery the
153+
* set is re-resolved every resolveIntervalMs; an explicit set is fixed. */
154+
private async supervise(): Promise<void> {
155+
while (!this.stopping && this.fatal === undefined) {
156+
let targets: Target[];
157+
try {
158+
// E3: race the (un-abortable) DNS work against the wake signal so
159+
// teardown/fatal don't block on a slow resolver — a late result is
160+
// discarded by the stopping/fatal check below.
161+
const resolution = resolveTargets(this.opts);
162+
resolution.catch(() => {}); // a late rejection must not be unhandled
163+
const raced = await raceAbortable(resolution, this.wake.signal);
164+
if (raced === null) break; // woken: stopping or fatal
165+
targets = raced;
166+
} catch (err) {
167+
// Keep whatever slots exist serving; retry the resolution later
168+
// (the Rust client does the same on SRV failures).
169+
this.log(
170+
`tunnel: target resolution failed: ${err instanceof Error ? err.message : String(err)} — retrying`
171+
);
172+
await delay(
173+
Math.min(5_000, this.opts.resolveIntervalMs),
174+
this.wake.signal
175+
);
176+
continue;
177+
}
178+
if (this.stopping || this.fatal !== undefined) break;
179+
180+
const desired = new Map(targets.map((t) => [targetKey(t), t] as const));
181+
for (const [key, target] of desired) {
182+
if (!this.slots.has(key)) {
183+
this.log(`tunnel: starting connection to ${key}`);
184+
this.startSlot(key, target);
185+
}
186+
}
187+
for (const [key, slot] of this.slots) {
188+
if (!desired.has(key)) {
189+
this.log(`tunnel: ${key} no longer resolves — tearing down`);
190+
slot.ctl.abort();
191+
}
192+
}
193+
194+
if (this.opts.srvName === undefined) break; // explicit servers: fixed set
195+
await delay(this.opts.resolveIntervalMs, this.wake.signal);
196+
}
197+
// Slots still in the map are live; evicted ones have already settled. No
198+
// slot can start after the loop exits (stopping/fatal both gate startSlot).
199+
await Promise.all([...this.slots.values()].map((s) => s.done));
200+
}
201+
}

packages/libs/restate-sdk-tunnel/src/types.ts

Lines changed: 34 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -171,6 +171,29 @@ export interface ConnectTunnelOptions {
171171
* standalone tunnel client).
172172
*/
173173
drainGraceMs?: number;
174+
/**
175+
* Advertise client-initiated graceful drain (`supports-client-drain: true`)
176+
* in the handshake. Default `true`. When enabled, {@link
177+
* TunnelConnection.shutdown} (or the opt-in {@link gracefulShutdown} signal
178+
* handler) refuses new invocations with a drain sentinel so Restate Cloud
179+
* stops routing new work to this process while its in-flight invocations
180+
* finish — the basis for zero-dropped-invocation rollouts. Specific to this
181+
* in-process client; the standalone Rust client does not implement it.
182+
*/
183+
supportsClientDrain?: boolean;
184+
/**
185+
* Automatic graceful shutdown on process signals. **On by default**: the
186+
* engine installs a one-shot handler for each signal (default `SIGTERM`)
187+
* that calls {@link TunnelConnection.shutdown} and then `process.exit(0)`
188+
* once draining completes (or the grace elapses) — so an operator-managed
189+
* deployment gets zero-dropped-invocation rollouts with no wiring. The
190+
* handlers are removed when the connection closes.
191+
*
192+
* Pass `false` to opt out entirely — e.g. to manage signals and process
193+
* exit yourself and call {@link TunnelConnection.shutdown} by hand. Pass an
194+
* object to choose the signals and grace, or `true` for the defaults.
195+
*/
196+
gracefulShutdown?: boolean | { signals?: NodeJS.Signals[]; graceMs?: number };
174197
/**
175198
* Reconnect backoff: initial delay in milliseconds. Default 10.
176199
* The delay grows by `reconnectFactor` per failed attempt (with jitter)
@@ -240,6 +263,17 @@ export interface ConnectTunnelOptions {
240263
export interface TunnelConnection {
241264
/** Stop reconnecting, close the current connection, and wait for teardown. */
242265
close(): Promise<void>;
266+
/**
267+
* Gracefully drain, then stop. Stops accepting new invocations (Restate
268+
* Cloud deselects this process via the drain sentinel), lets in-flight
269+
* invocations finish — bounded by `graceMs` (default {@link
270+
* ConnectTunnelOptions.drainGraceMs}) — and then tears down. Resolves once
271+
* drained or the grace elapsed. Wire this to `SIGTERM` for
272+
* zero-dropped-invocation rollouts; {@link close} is the abrupt alternative.
273+
* Requires {@link ConnectTunnelOptions.supportsClientDrain} (the default) to
274+
* have been advertised; otherwise it degrades to an abrupt {@link close}.
275+
*/
276+
shutdown(opts?: { graceMs?: number }): Promise<void>;
243277
/** Number of successful tunnel handshakes since `connectTunnel` was called. */
244278
readonly connectionCount: number;
245279
/**

packages/libs/restate-sdk-tunnel/src/util.ts

Lines changed: 27 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -52,3 +52,30 @@ export async function raceAbortable<T>(
5252
signal.removeEventListener("abort", onAbort);
5353
}
5454
}
55+
56+
/** A promise whose resolve/reject are exposed and fire at most once. */
57+
export class Deferred<T> {
58+
private settled = false;
59+
readonly promise: Promise<T>;
60+
private resolveFn!: (value: T) => void;
61+
private rejectFn!: (err: Error) => void;
62+
63+
constructor() {
64+
this.promise = new Promise<T>((resolve, reject) => {
65+
this.resolveFn = resolve;
66+
this.rejectFn = reject;
67+
});
68+
}
69+
70+
resolve(value: T): void {
71+
if (this.settled) return;
72+
this.settled = true;
73+
this.resolveFn(value);
74+
}
75+
76+
reject(err: Error): void {
77+
if (this.settled) return;
78+
this.settled = true;
79+
this.rejectFn(err);
80+
}
81+
}

packages/libs/restate-sdk-tunnel/test/fake-cloud.ts

Lines changed: 12 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -213,17 +213,27 @@ export function startFakeCloud(options: FakeCloudOptions): Promise<FakeCloud> {
213213
export function roundtrip(
214214
session: http2.ClientHttp2Session,
215215
headers: http2.OutgoingHttpHeaders
216-
): Promise<{ status: number; body: string }> {
216+
): Promise<{
217+
status: number;
218+
headers: http2.IncomingHttpHeaders;
219+
body: string;
220+
}> {
217221
return new Promise((resolve, reject) => {
218222
const req = session.request(headers);
219223
let status = 0;
224+
let responseHeaders: http2.IncomingHttpHeaders = {};
220225
const chunks: Buffer[] = [];
221226
req.on("response", (h) => {
222227
status = Number(h[":status"]);
228+
responseHeaders = h;
223229
});
224230
req.on("data", (c: Buffer) => chunks.push(c));
225231
req.on("end", () =>
226-
resolve({ status, body: Buffer.concat(chunks).toString() })
232+
resolve({
233+
status,
234+
headers: responseHeaders,
235+
body: Buffer.concat(chunks).toString(),
236+
})
227237
);
228238
req.on("error", reject);
229239
req.end();

0 commit comments

Comments
 (0)