[Router] Count open HTTP exchanges until their response body finishes (#39014)

Co-authored-by: Kangyan Zhou <kangyan.zhou@radixark.ai>
Co-authored-by: Claude Opus 5 (1M context) <noreply@anthropic.com>

experimental/sgl-router/Cargo.lock

@@ -3294,6 +3294,7 @@ dependencies = [
  "dynamo-tokenizers",
  "futures",
  "hf-hub",
+ "http-body",
  "http-body-util",
  "hyper",
  "hyper-util",

experimental/sgl-router/Cargo.toml

@@ -50,6 +50,9 @@ tokio = { version = "1.42", features = ["full"] }
 # so inbound h2c no longer rides on a gRPC dependency it has nothing to do with,
 # and it enables HTTP/2 extended CONNECT for h2 websockets.
 axum = { version = "0.8", features = ["macros", "tracing", "http2"] }
+# Named directly (not just via axum) to implement a response-body wrapper:
+# `Frame`/`SizeHint` are not re-exported through `axum::body`.
+http-body = "1"
 tower = { version = "0.5", features = ["full"] }
 tower-http = { version = "0.6", features = ["trace", "compression-gzip", "cors", "timeout", "request-id"] }
 reqwest = { version = "0.12", features = ["stream", "json", "rustls-tls", "http2"], default-features = false }

experimental/sgl-router/src/server/app.rs

@@ -16,6 +16,11 @@ use std::sync::Arc;
 /// status_code}` on exit (incl. early-exit 400/413/503). Their difference =
 /// received-but-not-answered, invisible to post-dispatch `worker_requests_total`.
 /// `route` is the matched template (not raw URI) to bound label cardinality.
+///
+/// Also the only place that sees every HTTP exchange on every route, so it is
+/// where `inflight_http` is taken — the count the termination drain reports on.
+/// The guard rides the response body rather than being dropped here: a
+/// streaming completion has barely started when this function returns.
 async fn count_requests(State(ctx): State<Arc<AppContext>>, req: Request, next: Next) -> Response {
     let method = req.method().as_str().to_owned();
     let route = req
@@ -24,10 +29,11 @@ async fn count_requests(State(ctx): State<Arc<AppContext>>, req: Request, next:
         .map(|m| m.as_str().to_owned())
         .unwrap_or_else(|| "unmatched".to_owned());
     ctx.metrics.record_ingress(&route, &method);
+    let inflight = ctx.inflight_http.enter();
     let resp = next.run(req).await;
     ctx.metrics
         .record_response(&route, &method, resp.status().as_u16());
-    resp
+    resp.map(|body| crate::server::inflight::track_body(body, inflight))
 }
 
 /// Middleware: log 413 PAYLOAD_TOO_LARGE responses with the request method

experimental/sgl-router/src/server/app_context.rs

@@ -10,6 +10,7 @@ use crate::policies::kv_events::{BlockSizeOracle, KvIndexMetrics};
 use crate::policies::prefix_provider::RadixTreePrefixProvider;
 use crate::policies::PolicyRegistry;
 use crate::proxy::Proxy;
+use crate::server::inflight::InflightHttp;
 use crate::server::metrics::MetricsRegistry;
 use crate::tokenizer::TokenizerRegistry;
 use crate::workers::WorkerRegistry;
@@ -41,6 +42,9 @@ pub struct AppContext {
     /// Indexer), where those series would all be a structural zero — see
     /// [`crate::policies::kv_events::KvEventIndex::metrics_source`].
     pub kv_metrics: Option<KvIndexMetrics>,
+    /// Open HTTP exchanges, on every route. What axum's graceful shutdown
+    /// waits on — `active_load` sees only the proxied subset.
+    pub inflight_http: Arc<InflightHttp>,
     ready: AtomicBool,
 }
 
@@ -98,6 +102,7 @@ impl AppContext {
             block_size_oracle: BlockSizeOracle::new(),
             kv_metrics: None,
             engine_load: EngineLoadTable::new(),
+            inflight_http: InflightHttp::new(),
             ready: AtomicBool::new(false),
         }
     }
@@ -155,6 +160,7 @@ impl AppContext {
             block_size_oracle: BlockSizeOracle::new(),
             kv_metrics: None,
             engine_load: EngineLoadTable::new(),
+            inflight_http: InflightHttp::new(),
             ready: AtomicBool::new(false),
         }
     }

experimental/sgl-router/src/server/inflight.rs

@@ -0,0 +1,178 @@
+// SPDX-FileCopyrightText: Copyright (c) 2026 The SGLang Authors
+// SPDX-License-Identifier: Apache-2.0
+
+//! In-flight HTTP accounting, for the termination drain to report on.
+//!
+//! [`ActiveLoadRegistry`](crate::policies::active_load::ActiveLoadRegistry)
+//! counts *proxied* requests — what the workers are busy with. Axum's graceful
+//! shutdown waits on something different and larger: every HTTP exchange still
+//! open on an accepted connection, on any route, until its response body has
+//! finished streaming. The two diverge exactly where the drain gets stuck — a
+//! stalled SSE consumer, a held `/metrics` scrape, a request on a non-proxied
+//! route — so a heartbeat reporting only the first says "0 in flight" about a
+//! pod that is minutes from being SIGKILLed with work outstanding.
+//!
+//! The count is taken at the edge middleware and released when the response
+//! **body** completes, not when the handler returns: for a streaming
+//! completion the handler returns as soon as the headers are ready, which is
+//! the beginning of the wait, not the end of it.
+
+use axum::body::{Body, Bytes, HttpBody};
+use http_body::{Frame, SizeHint};
+use std::pin::Pin;
+use std::sync::atomic::{AtomicUsize, Ordering};
+use std::sync::Arc;
+use std::task::{Context, Poll};
+
+/// Open HTTP exchanges, as axum's graceful shutdown counts them.
+#[derive(Debug, Default)]
+pub struct InflightHttp {
+    open: AtomicUsize,
+}
+
+impl InflightHttp {
+    pub fn new() -> Arc<Self> {
+        Arc::new(Self::default())
+    }
+
+    /// Exchanges whose response body has not finished. Relaxed throughout:
+    /// this is a diagnostic gauge, and a reader racing an increment sees the
+    /// count one tick later, never a torn value.
+    pub fn count(&self) -> usize {
+        self.open.load(Ordering::Relaxed)
+    }
+
+    /// Count one exchange until the returned guard drops.
+    pub fn enter(self: &Arc<Self>) -> InflightGuard {
+        self.open.fetch_add(1, Ordering::Relaxed);
+        InflightGuard(Arc::clone(self))
+    }
+}
+
+/// Releases its exchange on drop — including when the request future is
+/// cancelled (client gone before the response was built), which is why this is
+/// a guard and not a pair of explicit increment/decrement calls.
+pub struct InflightGuard(Arc<InflightHttp>);
+
+impl Drop for InflightGuard {
+    fn drop(&mut self) {
+        self.0.open.fetch_sub(1, Ordering::Relaxed);
+    }
+}
+
+/// A response body that holds `guard` until the last frame is yielded or the
+/// body is dropped.
+struct TrackedBody {
+    inner: Body,
+    _guard: InflightGuard,
+}
+
+/// Wrap `body` so the exchange stays counted for as long as the client is
+/// still being written to.
+pub fn track_body(body: Body, guard: InflightGuard) -> Body {
+    Body::new(TrackedBody {
+        inner: body,
+        _guard: guard,
+    })
+}
+
+impl HttpBody for TrackedBody {
+    type Data = Bytes;
+    type Error = axum::Error;
+
+    fn poll_frame(
+        self: Pin<&mut Self>,
+        cx: &mut Context<'_>,
+    ) -> Poll<Option<Result<Frame<Self::Data>, Self::Error>>> {
+        // `Body` and `InflightGuard` are both `Unpin`, so the projection is a
+        // plain field borrow rather than anything `pin-project` is needed for.
+        Pin::new(&mut self.get_mut().inner).poll_frame(cx)
+    }
+
+    fn is_end_stream(&self) -> bool {
+        self.inner.is_end_stream()
+    }
+
+    fn size_hint(&self) -> SizeHint {
+        self.inner.size_hint()
+    }
+}
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+    use http_body_util::BodyExt;
+
+    #[test]
+    fn the_guard_releases_on_drop() {
+        let counter = InflightHttp::new();
+        assert_eq!(counter.count(), 0);
+        let guard = counter.enter();
+        assert_eq!(counter.count(), 1);
+        let second = counter.enter();
+        assert_eq!(counter.count(), 2);
+        drop(guard);
+        assert_eq!(counter.count(), 1);
+        drop(second);
+        assert_eq!(counter.count(), 0);
+    }
+
+    /// The property the whole module exists for: a response whose headers are
+    /// ready is NOT finished. A counter released when the handler returns
+    /// reports 0 while a streaming completion is still being written — which is
+    /// precisely the drain the heartbeat is supposed to explain.
+    #[tokio::test]
+    async fn a_streaming_body_stays_counted_until_its_last_frame() {
+        let counter = InflightHttp::new();
+        let (mut tx, rx) = futures::channel::mpsc::channel::<Result<Bytes, axum::Error>>(1);
+        let body = track_body(Body::from_stream(rx), counter.enter());
+        assert_eq!(counter.count(), 1, "entering must count the exchange");
+
+        let mut stream = body.into_data_stream();
+        futures::SinkExt::send(&mut tx, Ok(Bytes::from_static(b"chunk")))
+            .await
+            .expect("the stream must accept a chunk");
+        let frame = futures::StreamExt::next(&mut stream).await;
+        assert!(frame.is_some(), "the wrapper must pass frames through");
+        assert_eq!(
+            counter.count(),
+            1,
+            "a body with frames still to come must stay counted",
+        );
+
+        drop(tx);
+        assert!(
+            futures::StreamExt::next(&mut stream).await.is_none(),
+            "the wrapper must end when the inner body does",
+        );
+        drop(stream);
+        assert_eq!(counter.count(), 0, "a finished body must release the count");
+    }
+
+    /// A client that disconnects mid-stream drops the body rather than
+    /// draining it. That must release the count too, or the gauge ratchets up
+    /// over the life of the process and the drain heartbeat reads permanently
+    /// busy.
+    #[tokio::test]
+    async fn an_abandoned_body_releases_the_count() {
+        let counter = InflightHttp::new();
+        let (_tx, rx) = futures::channel::mpsc::channel::<Result<Bytes, axum::Error>>(1);
+        let body = track_body(Body::from_stream(rx), counter.enter());
+        assert_eq!(counter.count(), 1);
+        drop(body);
+        assert_eq!(counter.count(), 0);
+    }
+
+    #[tokio::test]
+    async fn a_unit_body_passes_through_unchanged() {
+        let counter = InflightHttp::new();
+        let body = track_body(Body::from("payload"), counter.enter());
+        let bytes = body
+            .collect()
+            .await
+            .expect("collecting a tracked body must not error")
+            .to_bytes();
+        assert_eq!(&bytes[..], b"payload");
+        assert_eq!(counter.count(), 0);
+    }
+}

experimental/sgl-router/src/server/mod.rs

@@ -5,5 +5,6 @@ pub mod app;
 pub mod app_context;
 pub mod error;
 pub mod header_utils;
+pub mod inflight;
 pub mod metrics;
 pub mod routes;

experimental/sgl-router/tests/proxy/graceful_shutdown.rs

@@ -233,3 +233,141 @@ async fn shutdown_with_no_inflight_returns_promptly() {
         "idle shutdown took too long: {elapsed:?}"
     );
 }
+
+/// Poll until `inflight_http` settles on `want`, so the assertions below do not
+/// race the guard drop that happens on the server task after the client has
+/// already seen the last byte.
+async fn wait_for_inflight_http(ctx: &Arc<AppContext>, want: usize) {
+    let deadline = Instant::now() + Duration::from_secs(5);
+    while Instant::now() < deadline {
+        if ctx.inflight_http.count() == want {
+            return;
+        }
+        tokio::time::sleep(Duration::from_millis(10)).await;
+    }
+    panic!(
+        "inflight_http stayed at {} instead of settling to {want}",
+        ctx.inflight_http.count(),
+    );
+}
+
+/// `inflight_http` is what the drain heartbeat reports, and it is only worth
+/// reporting if it tracks what axum's graceful shutdown actually waits on: the
+/// response BODY finishing, not the handler returning. A streaming completion
+/// hands back its headers immediately, so a count released at handler exit
+/// would read 0 for the entire window the heartbeat exists to explain — the
+/// same blind spot `active_load.inflight_count()` has, reproduced in the
+/// replacement.
+#[tokio::test(flavor = "multi_thread", worker_threads = 4)]
+async fn inflight_http_counts_a_streaming_response_until_its_body_finishes() {
+    let worker = crate::common::mock_worker::MockWorker::start_slow_stream(
+        SLOW_CHUNKS.to_vec(),
+        Duration::from_millis(60),
+    )
+    .await;
+    let ctx = build_ctx_with_worker(&worker.url);
+    let app = build_router(ctx.clone());
+    let listener = TcpListener::bind("127.0.0.1:0").await.unwrap();
+    let addr = listener.local_addr().unwrap();
+    let (stop_tx, stop_rx) = oneshot::channel::<()>();
+    let server = tokio::spawn(async move {
+        axum::serve(listener, app)
+            .with_graceful_shutdown(async move {
+                let _ = stop_rx.await;
+            })
+            .await
+            .unwrap();
+    });
+
+    assert_eq!(ctx.inflight_http.count(), 0, "idle router counts nothing");
+
+    let client = reqwest::Client::builder()
+        .timeout(Duration::from_secs(10))
+        .build()
+        .unwrap();
+    let resp = client
+        .post(format!("http://{addr}/v1/chat/completions"))
+        .json(&serde_json::json!({
+            "model": "tiny",
+            "messages": [{"role": "user", "content": "hi"}],
+            "stream": true,
+        }))
+        .send()
+        .await
+        .unwrap();
+    assert!(
+        resp.status().is_success(),
+        "stream started: {}",
+        resp.status()
+    );
+
+    // Headers are in, ~480 ms of chunks are not. This is precisely the state a
+    // SIGTERM lands in, and the count has to see it.
+    assert_eq!(
+        ctx.inflight_http.count(),
+        1,
+        "a streaming response whose body is still being written must stay counted",
+    );
+
+    let body = resp.bytes().await.unwrap();
+    assert!(
+        String::from_utf8_lossy(&body).contains("data: [DONE]"),
+        "the stream must have run to completion for this to say anything",
+    );
+    wait_for_inflight_http(&ctx, 0).await;
+
+    stop_tx.send(()).unwrap();
+    tokio::time::timeout(Duration::from_secs(5), server)
+        .await
+        .expect("server resolves")
+        .expect("server task joined cleanly");
+}
+
+/// Every route is instrumented, not only the proxied ones. `/metrics`,
+/// `/readyz` and a 404 are exchanges axum's drain waits on too, and they are
+/// exactly the traffic `active_load` cannot see — so a guard that leaked on a
+/// non-proxied route would leave the heartbeat permanently busy and turn the
+/// drain report back into noise.
+#[tokio::test(flavor = "multi_thread", worker_threads = 2)]
+async fn inflight_http_returns_to_zero_after_non_proxied_routes() {
+    let worker = crate::common::mock_worker::MockWorker::start_slow_stream(
+        SLOW_CHUNKS.to_vec(),
+        Duration::from_millis(1),
+    )
+    .await;
+    let ctx = build_ctx_with_worker(&worker.url);
+    let app = build_router(ctx.clone());
+    let listener = TcpListener::bind("127.0.0.1:0").await.unwrap();
+    let addr = listener.local_addr().unwrap();
+    let (stop_tx, stop_rx) = oneshot::channel::<()>();
+    let server = tokio::spawn(async move {
+        axum::serve(listener, app)
+            .with_graceful_shutdown(async move {
+                let _ = stop_rx.await;
+            })
+            .await
+            .unwrap();
+    });
+
+    let client = reqwest::Client::builder()
+        .timeout(Duration::from_secs(10))
+        .build()
+        .unwrap();
+    for path in ["/metrics", "/readyz", "/healthz", "/v1/models", "/nope"] {
+        let resp = client
+            .get(format!("http://{addr}{path}"))
+            .send()
+            .await
+            .unwrap_or_else(|e| panic!("GET {path} failed: {e}"));
+        // Body consumed, not just headers: an unread body is an unfinished
+        // exchange and would make this assert nothing.
+        let _ = resp.bytes().await.unwrap();
+    }
+    wait_for_inflight_http(&ctx, 0).await;
+
+    stop_tx.send(()).unwrap();
+    tokio::time::timeout(Duration::from_secs(5), server)
+        .await
+        .expect("server resolves")
+        .expect("server task joined cleanly");
+}

experimental/sgl-router/tests/proxy/pd_protocol_binding.rs

@@ -141,6 +141,7 @@ fn config() -> Config {
             affinity: None,
             fused: None,
             eligibility: None,
+            sampling_overrides: Default::default(),
         },
         discovery: DiscoveryBackend::StaticUrls(StaticUrlsDiscoveryConfig {
             urls: vec!["http://placeholder:0".into()],