From 4c936140e976fe905ad9d0f7c9d9934efc48c1e6 Mon Sep 17 00:00:00 2001 From: Jarvis Date: Wed, 19 Aug 2026 08:09:42 +0000 Subject: [PATCH 1/2] feat(server): keep serving through a drain window before closing the listener MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit A load balancer learns that an instance is withdrawing on its next health check, not the moment the instance decides to. Between those two points it keeps routing new connections. The gateway closed its listener about a second after SIGTERM, so every connection routed inside that interval was refused and callers saw gateway errors during an ordinary rolling update or scale-down. Separate the two events. On the shutdown signal the gateway now: - answers /readyz and /livez with 503 immediately, as before; - keeps accepting new connections for at least `shutdown.min_drain_secs` (new, defaults to 30s); - adds `Connection: close` to HTTP/1.1 responses, so a pooling client retires its connections as it uses them instead of holding idle ones open until the listener disappears — a request dispatched onto one of those in the closing instant dies with no response, which is how a graceful shutdown still surfaces as an upstream reset at the caller; - stops accepting only once that window has elapsed AND nothing is left in flight, so a balancer slower than configured cannot make it close under live traffic; - drains the remaining in-flight requests without a deadline of its own, as before, leaving `terminationGracePeriodSeconds` / `TimeoutStopSec` as the one hard bound. The in-flight count is a new process-wide counter raised by the telemetry middleware's RAII guard. The `aisix_proxy_in_flight_requests` gauge next to it is sliced by endpoint and protocol and lives behind the metrics registry's lock — the right shape for a dashboard, the wrong one for a drain gate. The e2e harness pins `min_drain_secs: 0`: no balancer fronts a spawned test binary, and paying the window on every teardown would cost 30s per app and lose the clean-exit path the specs rely on. --- config.example.yaml | 32 +++ config.managed.yaml | 9 + crates/aisix-core/src/config.rs | 48 +++++ crates/aisix-proxy/src/health.rs | 28 +++ crates/aisix-proxy/src/lib.rs | 40 +++- crates/aisix-server/src/main.rs | 54 ++++- .../e2e/src/cases/graceful-drain-e2e.test.ts | 185 ++++++++++++++++++ tests/e2e/src/harness/app.ts | 8 + 8 files changed, 401 insertions(+), 3 deletions(-) create mode 100644 tests/e2e/src/cases/graceful-drain-e2e.test.ts 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..742e945f 100644 --- a/crates/aisix-proxy/src/lib.rs +++ b/crates/aisix-proxy/src/lib.rs @@ -393,6 +393,12 @@ fn inbound_protocol_for_endpoint(endpoint: &str) -> &'static str { struct InFlightGuard { metrics: std::sync::Arc, + /// Same edges as the metric, but as one process-wide count the + /// shutdown coordinator can read. The gauge next to it is sliced by + /// endpoint and protocol and lives behind the metrics registry's + /// lock — the right shape for a dashboard, the wrong one for a + /// hot-path drain gate. + livez: std::sync::Arc, /// Bounded route template + protocol family — both `'static` by /// construction (`normalize_endpoint_label` / /// `inbound_protocol_for_endpoint`), so the guard owns no @@ -404,12 +410,15 @@ struct InFlightGuard { impl InFlightGuard { fn new( metrics: std::sync::Arc, + livez: std::sync::Arc, endpoint: &'static str, inbound_protocol: &'static str, ) -> Self { metrics.increment_proxy_in_flight(endpoint, inbound_protocol); + livez.enter(); Self { metrics, + livez, endpoint, inbound_protocol, } @@ -420,6 +429,7 @@ impl Drop for InFlightGuard { fn drop(&mut self) { self.metrics .decrement_proxy_in_flight(self.endpoint, self.inbound_protocol); + self.livez.leave(); } } @@ -478,6 +488,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 @@ -499,12 +510,39 @@ async fn record_request_telemetry( }; let _in_flight = InFlightGuard::new( state.metrics.clone(), + state.livez.clone(), endpoint, inbound_protocol_for_endpoint(endpoint), ); - let response = attribution::scope(attribution, next.run(request)).await; + let draining = state.livez.is_shutting_down(); + let mut response = attribution::scope(attribution, next.run(request)).await; guard.armed = false; + if draining { + retire_connection(&version, &mut response); + } + response +} + +/// 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..0cd3a420 --- /dev/null +++ b/tests/e2e/src/cases/graceful-drain-e2e.test.ts @@ -0,0 +1,185 @@ +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 `/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 () => { + if (!etcdReachable) 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); + // Let it reach the gateway before the signal lands. + await new Promise((r) => setTimeout(r, 300)); + + 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/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 } } : {}), From d85f5867c40cba0157becad0e8f4a9135e24d0aa Mon Sep 17 00:00:00 2001 From: Jarvis Date: Wed, 19 Aug 2026 14:47:46 +0000 Subject: [PATCH 2/2] fix(proxy): hold the drain count until the response body is done MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Axum polls a response body after the middleware that produced it returns, so a drain gate released there reads zero while an SSE stream is still running. The shutdown coordinator would then close the listener as soon as the minimum window elapsed — under the traffic shape an AI gateway carries most of, and for exactly the balancer the window exists to outlast: one slower to withdraw the instance than the window is long. Split the count off from `InFlightGuard` into its own guard moved into the response body, so it falls only when the body is fully written or dropped. `InFlightGuard` keeps feeding `aisix_proxy_in_flight_requests` with its established semantics; folding the body into that published gauge would shift what every existing dashboard reads. Also sample the draining flag AFTER the handler rather than before it: a request that arrived just ahead of the signal and finished inside the window rides one of the pooled connections that most needs retiring, and was missing its `Connection: close`. Tests: a second spec drives a live SSE stream across the signal and pins that the listener still accepts well past the window. The first spec now gates on the upstream having received the slow request instead of sleeping 300ms, and both use the suite's `ctx.skip()` idiom for an unreachable etcd. --- crates/aisix-proxy/src/lib.rs | 60 +++++-- .../e2e/src/cases/graceful-drain-e2e.test.ts | 38 ++++- .../src/cases/graceful-drain-sse-e2e.test.ts | 151 ++++++++++++++++++ 3 files changed, 227 insertions(+), 22 deletions(-) create mode 100644 tests/e2e/src/cases/graceful-drain-sse-e2e.test.ts diff --git a/crates/aisix-proxy/src/lib.rs b/crates/aisix-proxy/src/lib.rs index 742e945f..6b4bdd9a 100644 --- a/crates/aisix-proxy/src/lib.rs +++ b/crates/aisix-proxy/src/lib.rs @@ -393,12 +393,6 @@ fn inbound_protocol_for_endpoint(endpoint: &str) -> &'static str { struct InFlightGuard { metrics: std::sync::Arc, - /// Same edges as the metric, but as one process-wide count the - /// shutdown coordinator can read. The gauge next to it is sliced by - /// endpoint and protocol and lives behind the metrics registry's - /// lock — the right shape for a dashboard, the wrong one for a - /// hot-path drain gate. - livez: std::sync::Arc, /// Bounded route template + protocol family — both `'static` by /// construction (`normalize_endpoint_label` / /// `inbound_protocol_for_endpoint`), so the guard owns no @@ -410,26 +404,49 @@ struct InFlightGuard { impl InFlightGuard { fn new( metrics: std::sync::Arc, - livez: std::sync::Arc, endpoint: &'static str, inbound_protocol: &'static str, ) -> Self { metrics.increment_proxy_in_flight(endpoint, inbound_protocol); - livez.enter(); Self { metrics, - livez, endpoint, inbound_protocol, } } } +/// 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 .decrement_proxy_in_flight(self.endpoint, self.inbound_protocol); - self.livez.leave(); } } @@ -510,17 +527,32 @@ async fn record_request_telemetry( }; let _in_flight = InFlightGuard::new( state.metrics.clone(), - state.livez.clone(), endpoint, inbound_protocol_for_endpoint(endpoint), ); - let draining = state.livez.is_shutting_down(); + let drain = DrainGuard::new(state.livez.clone()); let mut response = attribution::scope(attribution, next.run(request)).await; guard.armed = false; - if draining { + // 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); } - 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 diff --git a/tests/e2e/src/cases/graceful-drain-e2e.test.ts b/tests/e2e/src/cases/graceful-drain-e2e.test.ts index 0cd3a420..8d26ff28 100644 --- a/tests/e2e/src/cases/graceful-drain-e2e.test.ts +++ b/tests/e2e/src/cases/graceful-drain-e2e.test.ts @@ -61,6 +61,20 @@ async function readyzStatus(proxyUrl: string): Promise { } } +/** 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, @@ -134,9 +148,12 @@ describe("graceful drain e2e: SIGTERM stops readiness, not service", () => { test( "keeps serving through the drain window, then exits once nothing is in flight", - async () => { - if (!etcdReachable) return; - const proxyUrl = app!.proxyUrl; + async (ctx) => { + if (!etcdReachable || !app) { + ctx.skip(); + return; + } + const proxyUrl = app.proxyUrl; expect(await readyzStatus(proxyUrl)).toBe(200); @@ -144,11 +161,16 @@ describe("graceful drain e2e: SIGTERM stops readiness, not service", () => { // signal. It must complete rather than be cut off, and it must // hold the listener open past the window. const inFlight = chat(proxyUrl); - // Let it reach the gateway before the signal lands. - await new Promise((r) => setTimeout(r, 300)); + // 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"); + app.signal("SIGTERM"); // Readiness withdraws immediately — this is what the balancer polls. await waitReadyz(proxyUrl, 503, 3_000); @@ -176,9 +198,9 @@ describe("graceful drain e2e: SIGTERM stops readiness, not service", () => { // 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); + await app.waitForExit(15_000); expect(await readyzStatus(proxyUrl)).toBe("refused"); - expect(app!.output()).toContain("drain complete"); + 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, + ); +});