From d634320e483640b391d0e693110b65662b6ec064 Mon Sep 17 00:00:00 2001 From: Kangyan-Zhou Date: Wed, 16 Sep 2026 00:04:38 -0700 Subject: [PATCH] [Router] Count open HTTP exchanges until their response body finishes (#39014) Co-authored-by: Kangyan Zhou Co-authored-by: Claude Opus 5 (1M context) --- experimental/sgl-router/Cargo.lock | 1 + experimental/sgl-router/Cargo.toml | 3 + experimental/sgl-router/src/server/app.rs | 8 +- .../sgl-router/src/server/app_context.rs | 6 + .../sgl-router/src/server/inflight.rs | 178 ++++++++++++++++++ experimental/sgl-router/src/server/mod.rs | 1 + .../tests/proxy/graceful_shutdown.rs | 138 ++++++++++++++ .../tests/proxy/pd_protocol_binding.rs | 1 + 8 files changed, 335 insertions(+), 1 deletion(-) create mode 100644 experimental/sgl-router/src/server/inflight.rs diff --git a/experimental/sgl-router/Cargo.lock b/experimental/sgl-router/Cargo.lock index 81b322705..592987e60 100644 --- a/experimental/sgl-router/Cargo.lock +++ b/experimental/sgl-router/Cargo.lock @@ -3294,6 +3294,7 @@ dependencies = [ "dynamo-tokenizers", "futures", "hf-hub", + "http-body", "http-body-util", "hyper", "hyper-util", diff --git a/experimental/sgl-router/Cargo.toml b/experimental/sgl-router/Cargo.toml index b24e9d24b..3f27aec9a 100644 --- a/experimental/sgl-router/Cargo.toml +++ b/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 } diff --git a/experimental/sgl-router/src/server/app.rs b/experimental/sgl-router/src/server/app.rs index a769bc150..a35041f4b 100644 --- a/experimental/sgl-router/src/server/app.rs +++ b/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>, 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>, 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 diff --git a/experimental/sgl-router/src/server/app_context.rs b/experimental/sgl-router/src/server/app_context.rs index 2f8284ac3..cdfe184b6 100644 --- a/experimental/sgl-router/src/server/app_context.rs +++ b/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, + /// Open HTTP exchanges, on every route. What axum's graceful shutdown + /// waits on — `active_load` sees only the proxied subset. + pub inflight_http: Arc, 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), } } diff --git a/experimental/sgl-router/src/server/inflight.rs b/experimental/sgl-router/src/server/inflight.rs new file mode 100644 index 000000000..e8a58a709 --- /dev/null +++ b/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 { + 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) -> 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); + +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, 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::>(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::>(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); + } +} diff --git a/experimental/sgl-router/src/server/mod.rs b/experimental/sgl-router/src/server/mod.rs index 800500718..bbfc67398 100644 --- a/experimental/sgl-router/src/server/mod.rs +++ b/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; diff --git a/experimental/sgl-router/tests/proxy/graceful_shutdown.rs b/experimental/sgl-router/tests/proxy/graceful_shutdown.rs index 843551bbc..650f23302 100644 --- a/experimental/sgl-router/tests/proxy/graceful_shutdown.rs +++ b/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, 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"); +} diff --git a/experimental/sgl-router/tests/proxy/pd_protocol_binding.rs b/experimental/sgl-router/tests/proxy/pd_protocol_binding.rs index d7065e448..6b1bbb2d1 100644 --- a/experimental/sgl-router/tests/proxy/pd_protocol_binding.rs +++ b/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()],