feat(server): keep serving through a drain window before closing the listener

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

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

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.
     ///

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")

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<health::LivezState>,
+}
+
+impl DrainGuard {
+    fn new(livez: std::sync::Arc<health::LivezState>) -> 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 {

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<bool>, 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<bool>,
     livez_state: std::sync::Arc<aisix_proxy::LivezState>,
+    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);
 }