diff --git a/config.example.yaml b/config.example.yaml index da904b92..976e8c40 100644 --- a/config.example.yaml +++ b/config.example.yaml @@ -360,6 +360,38 @@ downstream: # connection. 0 disables the heartbeat. # sse_keepalive_interval_secs: 15 +# --------------------------------------------------------------------------- +# Shutdown +# --------------------------------------------------------------------------- +# What happens between SIGINT/SIGTERM and the process exiting. +# +# The signal makes /readyz answer 503 straight away, but a load balancer +# only learns that on its next health check and keeps routing new +# connections until then. So the gateway keeps accepting after the signal +# and only stops once the balancer has had time to withdraw it AND nothing +# is left in flight. In-flight requests then drain with no deadline — an +# inference call or an SSE stream may run for minutes — so the platform +# (Kubernetes terminationGracePeriodSeconds, systemd TimeoutStopSec) is +# the only hard bound. Make sure that bound is generous enough for your +# longest request. +shutdown: + # Minimum seconds to keep accepting new connections after the signal, + # while /readyz already reports 503. + # + # Size it ABOVE the detection latency of whatever load-balances this + # instance: a Kubernetes readiness probe needs periodSeconds x + # failureThreshold, an external balancer its own check interval times + # its retry count. Too low and the listener closes while traffic is + # still being routed here; too high only delays the exit. + # + # It is a minimum, not a deadline — after it elapses the gateway still + # waits for the in-flight count to reach zero, so a balancer slower than + # configured cannot make it close under live traffic. + # + # 0 drops the window entirely and is only correct when nothing routes + # here by health check. + min_drain_secs: 30 + # Models, API keys, provider keys, guardrails, cache policies, and # observability exporters are NOT defined in this file. They come from # the configured resource source: the `resources_file` above, direct diff --git a/config.managed.yaml b/config.managed.yaml index bb1bf985..4bc47738 100644 --- a/config.managed.yaml +++ b/config.managed.yaml @@ -165,3 +165,12 @@ upstream: downstream: idle_timeout_secs: 0 # sse_keepalive_interval_secs: 15 + +# Between SIGTERM and exit the gateway keeps accepting for at least +# `min_drain_secs` while /readyz already reports 503, so the load balancer +# in front has time to withdraw it, then waits for in-flight requests to +# finish before closing the listener. Raise it above the detection latency +# of your balancer's health check. In-flight drain itself is unbounded, so +# terminationGracePeriodSeconds must cover your longest request. +shutdown: + min_drain_secs: 30 diff --git a/crates/aisix-core/src/config.rs b/crates/aisix-core/src/config.rs index 7174f886..458ead8a 100644 --- a/crates/aisix-core/src/config.rs +++ b/crates/aisix-core/src/config.rs @@ -72,6 +72,10 @@ pub struct Config { /// emits a heartbeat — see [`DownstreamConfig`]. #[serde(default)] pub downstream: DownstreamConfig, + /// How the process behaves between the shutdown signal and exit — + /// see [`ShutdownConfig`]. + #[serde(default)] + pub shutdown: ShutdownConfig, /// Optional managed-mode configuration. When `managed.enabled = true` /// the admin API and Playground endpoints are **not** bound — the DP /// is a pure etcd reader driven by the aisix.cloud control plane. @@ -1356,6 +1360,50 @@ impl Default for DownstreamConfig { } } +/// What the gateway does between receiving SIGINT/SIGTERM and exiting. +/// +/// The shutdown signal makes `/readyz` answer 503 immediately, but a load +/// balancer only learns that on its next health check — and keeps sending +/// new connections until then. Closing the listener at signal time would +/// therefore refuse every connection routed inside that blind window. So +/// the gateway keeps serving after the signal, and only stops accepting +/// once the balancer has had time to withdraw it AND nothing is left in +/// flight. +/// +/// Once it does stop accepting, in-flight requests drain without a +/// deadline — an inference call or an SSE stream may run for minutes. +/// The platform caps the whole sequence (Kubernetes +/// `terminationGracePeriodSeconds`, systemd `TimeoutStopSec`), which is +/// the only hard bound. +#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)] +#[serde(deny_unknown_fields, default)] +pub struct ShutdownConfig { + /// Minimum seconds to keep accepting new connections after the + /// shutdown signal, while `/readyz` already answers 503. + /// + /// Size it above the detection latency of whatever load-balances this + /// instance — a Kubernetes readiness probe needs `periodSeconds x + /// failureThreshold`, an external balancer its own check interval + /// times its retry count. Too low and the listener closes while the + /// balancer is still routing to it; too high only delays the exit. + /// + /// The window is a *minimum*, not a deadline: after it elapses the + /// gateway still waits for the in-flight count to reach zero before it + /// stops accepting, so a balancer that is slower than configured + /// cannot make it close under live traffic. + /// + /// `0` drops the window entirely: the gateway stops accepting as soon + /// as nothing is in flight. Only correct when nothing routes to this + /// instance by health check. + pub min_drain_secs: u64, +} + +impl Default for ShutdownConfig { + fn default() -> Self { + Self { min_drain_secs: 30 } + } +} + impl Config { /// Load + merge + validate. /// diff --git a/crates/aisix-proxy/src/health.rs b/crates/aisix-proxy/src/health.rs index 5975bbfe..857e9c66 100644 --- a/crates/aisix-proxy/src/health.rs +++ b/crates/aisix-proxy/src/health.rs @@ -33,6 +33,12 @@ static TEXT_PLAIN_UTF8: HeaderValue = HeaderValue::from_static("text/plain; char #[derive(Debug, Default)] pub struct LivezState { shutting_down: AtomicBool, + /// Requests currently being served on the proxy listener, raised and + /// lowered by the telemetry middleware's RAII guard. Read by the + /// shutdown coordinator to decide when closing the listener can no + /// longer interrupt anything; a streaming response keeps its slot for + /// as long as bytes may still flow. + in_flight: AtomicUsize, } impl LivezState { @@ -44,6 +50,28 @@ impl LivezState { self.shutting_down.store(true, Ordering::Relaxed); } + /// Raise the in-flight count. Pair with [`Self::leave`]; the proxy + /// does that through a `Drop` guard so a cancelled request still + /// lowers the count it raised. + pub fn enter(&self) { + self.in_flight.fetch_add(1, Ordering::Relaxed); + } + + /// Lower the in-flight count, saturating at zero so an unpaired + /// decrement cannot wrap the counter and wedge the drain. + pub fn leave(&self) { + let _ = self + .in_flight + .fetch_update(Ordering::Relaxed, Ordering::Relaxed, |v| { + Some(v.saturating_sub(1)) + }); + } + + /// Requests in flight right now. + pub fn in_flight(&self) -> usize { + self.in_flight.load(Ordering::Relaxed) + } + fn shutdown_check(&self) -> Result<(), &'static str> { if self.shutting_down.load(Ordering::Relaxed) { Err("process is shutting down") diff --git a/crates/aisix-proxy/src/lib.rs b/crates/aisix-proxy/src/lib.rs index 4e581d3d..6b4bdd9a 100644 --- a/crates/aisix-proxy/src/lib.rs +++ b/crates/aisix-proxy/src/lib.rs @@ -416,6 +416,33 @@ impl InFlightGuard { } } +/// Holds the process-wide drain count up for one request. +/// +/// Deliberately separate from [`InFlightGuard`]: that one feeds +/// `aisix_proxy_in_flight_requests` and keeps its established semantics of +/// ending when the handler returns, so the published gauge does not shift +/// meaning. The drain gate has to span the **response body** instead — a +/// streaming body is polled after the middleware returns, so a guard +/// released there would read zero while SSE bytes are still flowing and +/// let the shutdown coordinator close the listener under exactly the +/// traffic the drain window exists to protect. +struct DrainGuard { + livez: std::sync::Arc, +} + +impl DrainGuard { + fn new(livez: std::sync::Arc) -> Self { + livez.enter(); + Self { livez } + } +} + +impl Drop for DrainGuard { + fn drop(&mut self) { + self.livez.leave(); + } +} + impl Drop for InFlightGuard { fn drop(&mut self) { self.metrics @@ -478,6 +505,7 @@ async fn record_request_telemetry( // suffix (or any 404 path) would otherwise let an unauthenticated // caller mint unbounded Prometheus time series (#451). let endpoint = normalize_endpoint_label(request.uri().path()); + let version = request.version(); // The cell the handler fills in as it resolves the model and picks a // target, so this layer can attribute a cancelled request to them // (AISIX-Cloud#1317). Installed here because a cancelled handler @@ -502,9 +530,51 @@ async fn record_request_telemetry( endpoint, inbound_protocol_for_endpoint(endpoint), ); - let response = attribution::scope(attribution, next.run(request)).await; + let drain = DrainGuard::new(state.livez.clone()); + let mut response = attribution::scope(attribution, next.run(request)).await; guard.armed = false; + // Sampled AFTER the handler, not before: a request that arrived just + // ahead of the signal and finished inside the window is riding one of + // the pooled connections that most needs retiring. + if state.livez.is_shutting_down() { + retire_connection(&version, &mut response); + } + hold_until_body_done(response, drain) +} + +/// Move `drain` into the response body so the count stays raised until the +/// body is fully written — or dropped, when the client hangs up mid-stream. +fn hold_until_body_done(response: Response, drain: DrainGuard) -> Response { + use http_body_util::BodyExt; + let (parts, body) = response.into_parts(); + let body = axum::body::Body::new(body.map_frame(move |frame| { + // The closure owns the guard; the mapped body owns the closure. + let _hold = &drain; + frame + })); + Response::from_parts(parts, body) +} + +/// Ask an HTTP/1.1 client to retire this connection once the response is +/// read, by answering `Connection: close`. +/// +/// The gateway keeps accepting through the drain window, so a client that +/// pools connections would otherwise hold idle ones open right up to the +/// moment the listener closes — and a request dispatched onto one of those +/// in that instant dies with no response, which is how a graceful shutdown +/// still surfaces as a 502/503 upstream-reset at the caller. Retiring them +/// as they are used means there is nothing idle left to lose. +/// +/// HTTP/2 has no such header (it is a connection-specific field, forbidden +/// by RFC 9113 §8.2.2); its drain signal is the GOAWAY that hyper emits +/// when the listener does shut down. +fn retire_connection(version: &axum::http::Version, response: &mut Response) { + if *version == axum::http::Version::HTTP_2 || *version == axum::http::Version::HTTP_3 { + return; + } response + .headers_mut() + .insert(header::CONNECTION, HeaderValue::from_static("close")); } struct ClientCancelGuard { diff --git a/crates/aisix-server/src/main.rs b/crates/aisix-server/src/main.rs index 5dc08c30..de311986 100644 --- a/crates/aisix-server/src/main.rs +++ b/crates/aisix-server/src/main.rs @@ -1140,7 +1140,11 @@ async fn run(mut cfg: Config) -> anyhow::Result<()> { // Step 10: shutdown coordinator. Whichever of (signal, proxy, admin) // completes first triggers the rest. - let signal_task = tokio::spawn(wait_for_signal(cancel_tx.clone(), livez_state)); + let signal_task = tokio::spawn(wait_for_signal( + cancel_tx.clone(), + livez_state, + Duration::from_secs(cfg.shutdown.min_drain_secs), + )); proxy_serve .await @@ -1889,9 +1893,19 @@ async fn shutdown_signal(mut cancel: watch::Receiver, label: &'static str) } } +/// How often the drain loop re-reads the in-flight count once the +/// minimum window has elapsed. +const DRAIN_POLL_INTERVAL: Duration = Duration::from_millis(100); + +/// How often the drain loop reports that it is still waiting, so a drain +/// that never finishes is visible in the logs rather than looking like a +/// hung process. +const DRAIN_LOG_INTERVAL: Duration = Duration::from_secs(5); + async fn wait_for_signal( cancel_tx: watch::Sender, livez_state: std::sync::Arc, + min_drain: Duration, ) { let ctrl_c = async { let _ = tokio::signal::ctrl_c().await; @@ -1915,8 +1929,44 @@ async fn wait_for_signal( _ = term => tracing::info!("received SIGTERM"), } + // `/readyz` answers 503 from here on. Everything below decides when + // it is safe to stop accepting, which is deliberately NOT the same + // moment: a balancer only learns about the 503 on its next health + // check, and closing the listener before then refuses every + // connection it routes in between. livez_state.mark_shutting_down(); - tokio::time::sleep(std::time::Duration::from_secs(1)).await; + tracing::info!( + min_drain_secs = min_drain.as_secs(), + in_flight = livez_state.in_flight(), + "draining — /readyz now reports 503, still accepting new connections" + ); + + if !min_drain.is_zero() { + tokio::time::sleep(min_drain).await; + } + + // The window is a minimum, not a deadline. A balancer slower than + // configured is still routing traffic here, and that traffic is + // exactly what the in-flight count shows — so keep serving until it + // reaches zero, at which point closing the listener cannot interrupt + // anything. Unbounded on purpose: an inference call or an SSE stream + // may run for minutes, and the platform (Kubernetes + // `terminationGracePeriodSeconds`, systemd `TimeoutStopSec`) is the + // one hard bound. + let mut last_log = std::time::Instant::now(); + loop { + let in_flight = livez_state.in_flight(); + if in_flight == 0 { + break; + } + if last_log.elapsed() >= DRAIN_LOG_INTERVAL { + tracing::info!(in_flight, "still draining in-flight requests"); + last_log = std::time::Instant::now(); + } + tokio::time::sleep(DRAIN_POLL_INTERVAL).await; + } + + tracing::info!("drain complete — closing listeners"); let _ = cancel_tx.send(true); } diff --git a/tests/e2e/src/cases/graceful-drain-e2e.test.ts b/tests/e2e/src/cases/graceful-drain-e2e.test.ts new file mode 100644 index 00000000..8d26ff28 --- /dev/null +++ b/tests/e2e/src/cases/graceful-drain-e2e.test.ts @@ -0,0 +1,207 @@ +import { createHash } from "node:crypto"; +import { afterAll, beforeAll, describe, expect, test } from "vitest"; +import { + EtcdClient, + SeedClient, + spawnApp, + startOpenAiUpstream, + waitConfigPropagation, + type OpenAiUpstream, + type SpawnedApp, +} from "../harness/index.js"; + +// E2E: what happens between SIGTERM and the process exiting. +// +// A load balancer learns that a replica is going away by polling its +// health check, so it keeps routing new connections for one check +// interval after `/readyz` starts answering 503. Closing the listener at +// signal time refuses every connection routed inside that window — which +// is how a rolling update surfaces at the caller as 502/503 even though +// nothing is actually broken. +// +// So the gateway keeps serving after the signal: `/readyz` flips +// immediately, new connections are still accepted for at least +// `shutdown.min_drain_secs`, and the listener only closes once nothing is +// left in flight. Responses carry `Connection: close` throughout, so a +// pooling client retires its connections as it uses them rather than +// holding idle ones open until the listener disappears. + +const CALLER_PLAINTEXT = "sk-graceful-drain-caller"; +const CALLER_KEY_HASH = createHash("sha256").update(CALLER_PLAINTEXT).digest("hex"); + +const DRAIN_WINDOW_SECS = 5; +/** Longer than the drain window, so it is still running when it elapses. */ +const SLOW_UPSTREAM_MS = 9_000; + +function chatBody(): string { + return JSON.stringify({ + model: "graceful-drain", + messages: [{ role: "user", content: "hi" }], + }); +} + +async function chat(proxyUrl: string): Promise { + return fetch(`${proxyUrl}/v1/chat/completions`, { + method: "POST", + headers: { + authorization: `Bearer ${CALLER_PLAINTEXT}`, + "content-type": "application/json", + }, + body: chatBody(), + }); +} + +async function readyzStatus(proxyUrl: string): Promise { + try { + const res = await fetch(`${proxyUrl}/readyz`); + await res.text(); + return res.status; + } catch { + return "refused"; + } +} + +/** Poll until `check` holds, or fail with `what`. */ +async function waitUntil( + check: () => boolean, + what: string, + timeoutMs = 10_000, +): Promise { + const deadline = Date.now() + timeoutMs; + while (Date.now() < deadline) { + if (check()) return; + await new Promise((r) => setTimeout(r, 25)); + } + throw new Error(`timed out waiting for ${what}`); +} + +/** Poll until `/readyz` reports the expected status, or time out. */ +async function waitReadyz( + proxyUrl: string, + want: number | "refused", + timeoutMs = 5_000, +): Promise { + const deadline = Date.now() + timeoutMs; + let last: number | "refused" = "refused"; + while (Date.now() < deadline) { + last = await readyzStatus(proxyUrl); + if (last === want) return; + await new Promise((r) => setTimeout(r, 50)); + } + throw new Error(`/readyz never reported ${want} (last: ${last})`); +} + +describe("graceful drain e2e: SIGTERM stops readiness, not service", () => { + let app: SpawnedApp | undefined; + let upstream: OpenAiUpstream | undefined; + let etcdReachable = false; + + beforeAll(async () => { + const etcd = new EtcdClient(); + etcdReachable = await etcd.ping(); + if (!etcdReachable) return; + + // Request 1 outlives the drain window; request 2 answers at once. + // `/v1/models` (the propagation gate) never reaches the upstream, so + // the two chat calls below are requests 1 and 2 in arrival order. + upstream = await startOpenAiUpstream({ + scriptedResponses: [{ responseDelayMs: SLOW_UPSTREAM_MS }, {}], + }); + app = await spawnApp({ + // The drain phases are reported at INFO; the suite default is WARN. + logLevel: "info", + extra: { shutdown: { min_drain_secs: DRAIN_WINDOW_SECS } }, + }); + const seed = new SeedClient(etcd, app.etcdPrefix); + + const pk = await seed.createProviderKey({ + display_name: "graceful-drain-pk", + api_key: "sk-mock", + api_base: `${upstream.baseUrl}/v1`, + }); + await seed.createModel({ + display_name: "graceful-drain", + provider: "openai", + model_name: "gpt-4o-mini", + provider_key_id: pk.id, + }); + await seed.createApiKey({ + key_hash: CALLER_KEY_HASH, + allowed_models: ["graceful-drain"], + }); + + // Gate on the caller key authenticating — seeded last, so it implies + // the whole seed set is in the snapshot. + await waitConfigPropagation(async () => { + const res = await fetch(`${app!.proxyUrl}/v1/models`, { + headers: { authorization: `Bearer ${CALLER_PLAINTEXT}` }, + }); + await res.text(); + return res.status === 200; + }); + }, 60_000); + + afterAll(async () => { + await app?.exit(); + await upstream?.close(); + }); + + test( + "keeps serving through the drain window, then exits once nothing is in flight", + async (ctx) => { + if (!etcdReachable || !app) { + ctx.skip(); + return; + } + const proxyUrl = app.proxyUrl; + + expect(await readyzStatus(proxyUrl)).toBe(200); + + // A request that outlives the drain window, started before the + // signal. It must complete rather than be cut off, and it must + // hold the listener open past the window. + const inFlight = chat(proxyUrl); + // Gate on the upstream having actually received it, so the signal + // below lands with the request genuinely in flight rather than after + // a sleep that only usually wins the race. + await waitUntil( + () => upstream!.receivedRequests.length > 0, + "the slow request never reached the upstream", + ); + + const signalledAt = Date.now(); + app.signal("SIGTERM"); + + // Readiness withdraws immediately — this is what the balancer polls. + await waitReadyz(proxyUrl, 503, 3_000); + + // …but the listener is still accepting. Probe well past the point + // where it used to close — about a second after the signal — and + // still comfortably inside the window, which is the interval a + // balancer that has not yet re-checked would route into. Serve a + // fast response so this assertion is about acceptance, not latency. + await new Promise((r) => setTimeout(r, 2_500)); + expect(Date.now() - signalledAt).toBeLessThan(DRAIN_WINDOW_SECS * 1000); + const duringWindow = await chat(proxyUrl); + expect(duringWindow.status).toBe(200); + await duringWindow.text(); + // A pooling client must retire the connection instead of holding it + // idle until the listener goes away. + expect(duringWindow.headers.get("connection")).toBe("close"); + + // The window is a minimum, not a deadline: the slow request is + // still running when it elapses, so the listener stays open for it. + const slow = await inFlight; + expect(slow.status).toBe(200); + await slow.text(); + expect(Date.now() - signalledAt).toBeGreaterThan(DRAIN_WINDOW_SECS * 1000); + + // Nothing in flight now — the process closes the listener and exits + // on its own. No SIGKILL, so the clean-shutdown path runs. + await app.waitForExit(15_000); + expect(await readyzStatus(proxyUrl)).toBe("refused"); + expect(app.output()).toContain("drain complete"); + }, + 60_000, + ); +}); diff --git a/tests/e2e/src/cases/graceful-drain-sse-e2e.test.ts b/tests/e2e/src/cases/graceful-drain-sse-e2e.test.ts new file mode 100644 index 00000000..1f714bc8 --- /dev/null +++ b/tests/e2e/src/cases/graceful-drain-sse-e2e.test.ts @@ -0,0 +1,151 @@ +import { createHash } from "node:crypto"; +import { afterAll, beforeAll, describe, expect, test } from "vitest"; +import { + EtcdClient, + SeedClient, + spawnApp, + startOpenAiUpstream, + waitConfigPropagation, + type OpenAiUpstream, + type SpawnedApp, +} from "../harness/index.js"; + +// E2E: a streaming response counts as in-flight for the whole time bytes +// may still flow, not just until the handler hands the response back. +// +// Axum polls a response body AFTER the middleware that produced it has +// returned, so a drain gate released at that point reads zero while an SSE +// stream is still running. The listener would then close as soon as the +// minimum window elapsed — under exactly the traffic shape an AI gateway +// carries most of, and for exactly the load balancer the window exists to +// outlast: one slower to withdraw the instance than the window is long. + +const CALLER_PLAINTEXT = "sk-graceful-drain-sse-caller"; +const CALLER_KEY_HASH = createHash("sha256").update(CALLER_PLAINTEXT).digest("hex"); + +const DRAIN_WINDOW_SECS = 3; +/** 10 events, one per second: the stream outlives the window several times over. */ +const SSE_EVENT_COUNT = 10; +const SSE_EVENT_DELAY_MS = 1_000; + +async function tcpAccepts(proxyUrl: string): Promise { + try { + const res = await fetch(`${proxyUrl}/readyz`); + await res.text(); + return true; + } catch { + return false; + } +} + +describe("graceful drain e2e: a live SSE stream holds the listener open", () => { + let app: SpawnedApp | undefined; + let upstream: OpenAiUpstream | undefined; + let etcdReachable = false; + + beforeAll(async () => { + const etcd = new EtcdClient(); + etcdReachable = await etcd.ping(); + if (!etcdReachable) return; + + upstream = await startOpenAiUpstream({ + streamEvents: [ + ...Array.from({ length: SSE_EVENT_COUNT }, (_, i) => + JSON.stringify({ + id: `chunk-${i}`, + object: "chat.completion.chunk", + created: 1, + model: "gpt-4o-mini", + choices: [{ index: 0, delta: { content: `t${i}` }, finish_reason: null }], + }), + ), + "[DONE]", + ], + eventDelayMs: SSE_EVENT_DELAY_MS, + }); + app = await spawnApp({ + logLevel: "info", + extra: { shutdown: { min_drain_secs: DRAIN_WINDOW_SECS } }, + }); + const seed = new SeedClient(etcd, app.etcdPrefix); + + const pk = await seed.createProviderKey({ + display_name: "graceful-drain-sse-pk", + api_key: "sk-mock", + api_base: `${upstream.baseUrl}/v1`, + }); + await seed.createModel({ + display_name: "graceful-drain-sse", + provider: "openai", + model_name: "gpt-4o-mini", + provider_key_id: pk.id, + }); + await seed.createApiKey({ + key_hash: CALLER_KEY_HASH, + allowed_models: ["graceful-drain-sse"], + }); + + await waitConfigPropagation(async () => { + const res = await fetch(`${app!.proxyUrl}/v1/models`, { + headers: { authorization: `Bearer ${CALLER_PLAINTEXT}` }, + }); + await res.text(); + return res.status === 200; + }); + }, 60_000); + + afterAll(async () => { + await app?.exit(); + await upstream?.close(); + }); + + test( + "keeps accepting past the window while a stream is still producing", + async (ctx) => { + if (!etcdReachable || !app) { + ctx.skip(); + return; + } + const proxyUrl = app.proxyUrl; + + const res = await fetch(`${proxyUrl}/v1/chat/completions`, { + method: "POST", + headers: { + authorization: `Bearer ${CALLER_PLAINTEXT}`, + "content-type": "application/json", + }, + body: JSON.stringify({ + model: "graceful-drain-sse", + messages: [{ role: "user", content: "hi" }], + stream: true, + }), + }); + expect(res.status).toBe(200); + + // Read the first chunk so the response head is committed and the + // handler has returned — the exact point a body-blind drain gate + // would drop back to zero. + const reader = res.body!.getReader(); + const first = await reader.read(); + expect(first.done).toBe(false); + + const signalledAt = Date.now(); + app.signal("SIGTERM"); + + // Well past the window, with the stream still producing. + await new Promise((r) => setTimeout(r, (DRAIN_WINDOW_SECS + 2) * 1000)); + expect(Date.now() - signalledAt).toBeGreaterThan(DRAIN_WINDOW_SECS * 1000); + expect(await tcpAccepts(proxyUrl)).toBe(true); + + // Drain the rest, then the process may finish and exit on its own. + // eslint-disable-next-line no-constant-condition + while (true) { + const { done } = await reader.read(); + if (done) break; + } + await app.waitForExit(20_000); + expect(app.output()).toContain("drain complete"); + }, + 90_000, + ); +}); diff --git a/tests/e2e/src/harness/app.ts b/tests/e2e/src/harness/app.ts index 3c5343bc..0e2c6fc4 100644 --- a/tests/e2e/src/harness/app.ts +++ b/tests/e2e/src/harness/app.ts @@ -305,6 +305,14 @@ async function spawnAppOnce(overrides: AppOverrides = {}): Promise { tracing: { otlp: { enabled: false, endpoint: "http://127.0.0.1:4317", sample_ratio: 1 } }, }, cache: { backend: "memory" }, + // The gateway ships a 30s drain window so a load balancer can + // withdraw a terminating replica before its listener closes. No + // balancer fronts a spawned test binary, and paying that window on + // every teardown would add 30s per app — the harness would SIGKILL + // at SHUTDOWN_GRACE_MS instead, losing the clean-exit path these + // specs rely on. Drain immediately here; the drain spec sets its own + // window through `extra`. + shutdown: { min_drain_secs: 0 }, ...(overrides.snapshotCachePath ? { managed: { snapshot_cache_path: overrides.snapshotCachePath } } : {}),