From 008470abd80a3c50c95c16aa41ac97d51586819e Mon Sep 17 00:00:00 2001 From: Kan Wu Date: Mon, 21 Sep 2026 10:07:07 -0700 Subject: [PATCH] [sgl-router] Bound streaming lifetimes and release guards on idle disconnect (#40391) Co-authored-by: Shangming Cai --- experimental/sgl-router/src/config/cli.rs | 9 +- experimental/sgl-router/src/config/mod.rs | 4 + experimental/sgl-router/src/config/types.rs | 3 + experimental/sgl-router/src/main.rs | 3 +- experimental/sgl-router/src/proxy/mod.rs | 182 ++++++++- experimental/sgl-router/src/proxy/sse.rs | 358 +++++++++++------- experimental/sgl-router/src/server/metrics.rs | 54 ++- .../src/server/routes/chat/forward.rs | 14 +- .../sgl-router/tests/proxy/chat_routing.rs | 1 + .../sgl-router/tests/proxy/h2c_forward.rs | 1 + .../tests/proxy/pd_pool_isolation.rs | 106 ++++++ 11 files changed, 571 insertions(+), 164 deletions(-) diff --git a/experimental/sgl-router/src/config/cli.rs b/experimental/sgl-router/src/config/cli.rs index 30d5c34cf..ee1d4e54b 100644 --- a/experimental/sgl-router/src/config/cli.rs +++ b/experimental/sgl-router/src/config/cli.rs @@ -114,9 +114,13 @@ pub struct ServerArgs { /// Per-request upstream timeout in seconds. #[arg(long, default_value_t = default_proxy_request_timeout_secs())] pub request_timeout_secs: u64, + /// Maximum silence between upstream stream chunks, in seconds. + #[arg(long, default_value_t = ProxyConfig::default().stream_idle_timeout_secs)] + pub stream_idle_timeout_secs: u64, - /// Max lifetime of an in-flight request entry before the janitor - /// reaps it (returns 504 `stale_request_expired`). + /// Maximum in-flight request lifetime in seconds, including streaming responses. + /// Expiry returns 504 `stale_request_expired` before response headers are sent; + /// after streaming starts, it aborts the body without changing the HTTP status. #[arg(long, default_value_t = default_stale_request_timeout_secs())] pub stale_request_timeout_secs: u64, @@ -389,6 +393,7 @@ impl Cli { discovery, proxy: ProxyConfig { request_timeout_secs: self.server.request_timeout_secs, + stream_idle_timeout_secs: self.server.stream_idle_timeout_secs, }, router_inflight_load: InflightLoadConfig { stale_request_timeout_secs: self.server.stale_request_timeout_secs, diff --git a/experimental/sgl-router/src/config/mod.rs b/experimental/sgl-router/src/config/mod.rs index f086f4607..67def2ec2 100644 --- a/experimental/sgl-router/src/config/mod.rs +++ b/experimental/sgl-router/src/config/mod.rs @@ -48,6 +48,10 @@ impl Config { validate_bucket_config(bucket_config)?; } self.model.sampling_overrides.validate()?; + ensure!( + self.proxy.stream_idle_timeout_secs > 0, + "stream_idle_timeout_secs must be greater than zero" + ); ensure!( self.server.shutdown_drain_secs <= MAX_SHUTDOWN_DRAIN_SECS, "shutdown_drain_secs must be at most {MAX_SHUTDOWN_DRAIN_SECS} (got {}); \ diff --git a/experimental/sgl-router/src/config/types.rs b/experimental/sgl-router/src/config/types.rs index 58bae48da..7c5feb13c 100644 --- a/experimental/sgl-router/src/config/types.rs +++ b/experimental/sgl-router/src/config/types.rs @@ -19,6 +19,8 @@ pub struct Config { pub struct ProxyConfig { /// Timeout for upstream response headers and body. Counts as a circuit-breaker failure. pub request_timeout_secs: u64, + /// Maximum silence between streamed upstream chunks before the stream fails. + pub stream_idle_timeout_secs: u64, } pub fn default_proxy_request_timeout_secs() -> u64 { @@ -29,6 +31,7 @@ impl Default for ProxyConfig { fn default() -> Self { Self { request_timeout_secs: default_proxy_request_timeout_secs(), + stream_idle_timeout_secs: 180, } } } diff --git a/experimental/sgl-router/src/main.rs b/experimental/sgl-router/src/main.rs index dd2c132e9..b41c0b4c3 100644 --- a/experimental/sgl-router/src/main.rs +++ b/experimental/sgl-router/src/main.rs @@ -276,7 +276,8 @@ fn build_app_context( let block_size_oracle = engine_state.block_size_oracle(); let proxy = Arc::new( Proxy::new(Duration::from_secs(config.proxy.request_timeout_secs)) - .context("build proxy client")?, + .context("build proxy client")? + .with_stream_idle_timeout(Duration::from_secs(config.proxy.stream_idle_timeout_secs)), ); let mut app_context = AppContext::with_router_inflight_load( diff --git a/experimental/sgl-router/src/proxy/mod.rs b/experimental/sgl-router/src/proxy/mod.rs index 227af7948..be687cea0 100644 --- a/experimental/sgl-router/src/proxy/mod.rs +++ b/experimental/sgl-router/src/proxy/mod.rs @@ -16,6 +16,7 @@ use bytes::Bytes; use reqwest::{Client, Url}; use std::sync::Arc; use std::time::Duration; +use tokio_util::sync::CancellationToken; /// Parse a worker URL emitted by discovery. On failure, trip the worker's /// circuit breaker so the malformed worker drops out of subsequent @@ -38,14 +39,14 @@ enum BreakerOutcome { /// The worker was responsive — a 2xx, or a 4xx it answered cleanly (a /// client's bad request says nothing about worker health). The non-streaming /// arm records success immediately; the streaming arm defers to the pump's - /// completion hook, which records success or failure by - /// [`sse::StreamEnd::transport_ok`], since a 2xx head can still be followed + /// completion hook, which classifies [`sse::StreamEnd::reason`], + /// since a 2xx head can still be followed /// by a body that never completes. Success, /// A real fault (5xx other than backpressure) → `record_failure`: count /// toward opening. Failure, - /// Backpressure (the worker is responsive but at capacity) → + /// Backpressure or router-side stream expiry → /// `record_backpressure`: never opens the breaker and, while Closed, leaves /// an in-progress failure streak intact — but still resolves a half-open /// probe so a recovered-but-busy worker isn't wedged shut. @@ -81,6 +82,20 @@ fn breaker_outcome(status: reqwest::StatusCode) -> BreakerOutcome { } } +/// Router-side expiry says nothing about worker health. Preserve the existing +/// treatment of completed streams and client disconnects; upstream faults, +/// idle timeouts, and pump panics remain failures. +fn stream_breaker_outcome(end: sse::StreamEnd) -> BreakerOutcome { + use sse::StreamEndReason; + match end.reason { + StreamEndReason::Expired => BreakerOutcome::Neutral, + StreamEndReason::Completed | StreamEndReason::ClientDisconnect => BreakerOutcome::Success, + StreamEndReason::UpstreamError + | StreamEndReason::IdleTimeout + | StreamEndReason::PumpPanicked => BreakerOutcome::Failure, + } +} + #[derive(Debug)] pub struct Proxy { /// The negotiating client: HTTP/1.1 in cleartext, and ALPN `h2, http/1.1` @@ -94,6 +109,8 @@ pub struct Proxy { /// Wall-clock timeout applied to non-streaming upstream requests. Streaming /// requests deliberately do not use this (long generations are valid). pub request_timeout: Duration, + /// Maximum silence between streamed upstream chunks; `None` waits forever. + pub stream_idle_timeout: Option, } /// Build a forwarding client for `protocol`, sharing pool/connect tuning @@ -127,9 +144,15 @@ impl Proxy { default_client: build_client(WireProtocol::Http1)?, h2c_client: build_client(WireProtocol::H2c)?, request_timeout, + stream_idle_timeout: None, }) } + pub fn with_stream_idle_timeout(mut self, timeout: Duration) -> Self { + self.stream_idle_timeout = Some(timeout); + self + } + /// The forwarding client for `protocol`, taken from the selected worker's /// [`crate::workers::Worker::protocol`]. fn client_for(&self, protocol: WireProtocol) -> &Client { @@ -282,6 +305,7 @@ impl Proxy { stream_guards: Option>, on_first_byte: Option>, on_stream_end: Option>, + expiration: Option, ) -> Result, ApiError> { if !breaker.allow() { return Err(ApiError::BreakerOpen { @@ -345,10 +369,10 @@ impl Proxy { BreakerOutcome::Success => { let breaker_for_hook = Arc::clone(breaker); Some(Box::new(move |end| { - if end.transport_ok { - breaker_for_hook.record_success(); - } else { - breaker_for_hook.record_failure(); + match stream_breaker_outcome(end) { + BreakerOutcome::Success => breaker_for_hook.record_success(), + BreakerOutcome::Failure => breaker_for_hook.record_failure(), + BreakerOutcome::Neutral => breaker_for_hook.record_backpressure(), } if let Some(hook) = caller_end_hook { hook(end); @@ -368,6 +392,10 @@ impl Proxy { stream_guards, on_complete, first_byte_hook, + sse::StreamLimits { + idle_timeout: self.stream_idle_timeout, + expiration, + }, ); let mut out = Response::new(body); *out.status_mut() = status; @@ -467,6 +495,145 @@ mod tests { (format!("http://127.0.0.1:{port}"), tx) } + async fn spawn_pending_stream_worker() -> (String, oneshot::Sender<()>) { + use futures::StreamExt; + + let listener = TcpListener::bind("127.0.0.1:0").await.unwrap(); + let address = listener.local_addr().unwrap(); + let app = Router::new().route( + "/v1/chat/completions", + post(|| async { + Body::from_stream( + futures::stream::iter([Ok::<_, std::io::Error>(Bytes::from_static( + b"data: chunk\n\n", + ))]) + .chain(futures::stream::pending()), + ) + }), + ); + let (tx, rx) = oneshot::channel(); + tokio::spawn(async move { + axum::serve(listener, app) + .with_graceful_shutdown(async move { + let _ = rx.await; + }) + .await + .unwrap(); + }); + (format!("http://{address}"), tx) + } + + async fn pending_stream_body( + proxy: &Proxy, + url: &str, + breaker: &Arc, + expiration: Option, + ) -> Body { + use http_body_util::BodyExt; + + let response = proxy + .forward_streaming_to( + url, + WireProtocol::Http1, + breaker, + "/v1/chat/completions", + &HeaderMap::new(), + Bytes::from_static(b"{}"), + None, + None, + None, + expiration, + ) + .await + .unwrap(); + assert_eq!(response.status(), StatusCode::OK); + let mut body = response.into_body(); + tokio::time::timeout(Duration::from_secs(2), body.frame()) + .await + .unwrap() + .unwrap() + .unwrap(); + body + } + + #[tokio::test] + async fn stream_expiry_preserves_breaker_failure_streak() { + use http_body_util::BodyExt; + + let (url, _shutdown) = spawn_pending_stream_worker().await; + let proxy = Proxy::new(Duration::from_secs(5)).unwrap(); + let breaker = Arc::new(CircuitBreaker::new()); + breaker.record_failure(); + breaker.record_failure(); + for _ in 0..6 { + let expiration = CancellationToken::new(); + let body = pending_stream_body(&proxy, &url, &breaker, Some(expiration.clone())).await; + expiration.cancel(); + let error = tokio::time::timeout(Duration::from_secs(2), body.collect()) + .await + .unwrap() + .unwrap_err(); + assert!(error.to_string().contains("stale_request_timeout")); + assert_eq!( + breaker.snapshot().state_code, + 0, + "expiry must not add a failure" + ); + } + breaker.record_failure(); + assert_eq!( + breaker.snapshot().state_code, + 1, + "expiry must not reset prior failures" + ); + } + + #[tokio::test] + async fn stream_expiry_resolves_half_open_probe() { + use http_body_util::BodyExt; + + let (url, _shutdown) = spawn_pending_stream_worker().await; + let proxy = Proxy::new(Duration::from_secs(5)).unwrap(); + let breaker = Arc::new(CircuitBreaker::with_config(CircuitBreakerConfig { + threshold: NonZeroU32::new(1).unwrap(), + cool_down: Duration::ZERO, + })); + breaker.record_failure(); + let expiration = CancellationToken::new(); + let body = pending_stream_body(&proxy, &url, &breaker, Some(expiration.clone())).await; + assert_eq!(breaker.snapshot().state_code, 2); + assert!(!breaker.would_allow()); + expiration.cancel(); + let error = tokio::time::timeout(Duration::from_secs(2), body.collect()) + .await + .unwrap() + .unwrap_err(); + assert!(error.to_string().contains("stale_request_timeout")); + assert_eq!(breaker.snapshot().state_code, 0); + assert!(breaker.would_allow()); + } + + #[tokio::test] + async fn stream_idle_timeout_still_trips_breaker() { + use http_body_util::BodyExt; + + let (url, _shutdown) = spawn_pending_stream_worker().await; + let mut proxy = Proxy::new(Duration::from_secs(5)).unwrap(); + proxy.stream_idle_timeout = Some(Duration::from_millis(20)); + let breaker = Arc::new(CircuitBreaker::with_config(CircuitBreakerConfig { + threshold: NonZeroU32::new(1).unwrap(), + cool_down: Duration::from_secs(30), + })); + let body = pending_stream_body(&proxy, &url, &breaker, None).await; + let error = tokio::time::timeout(Duration::from_secs(2), body.collect()) + .await + .unwrap() + .unwrap_err(); + assert!(error.to_string().contains("idle timeout")); + assert_eq!(breaker.snapshot().state_code, 1); + assert!(!breaker.would_allow()); + } + /// A saturated engine's own queue-full 503s must not trip the router's /// circuit breaker. Dispatch far past any plausible failure threshold and /// assert the breaker stays Closed and admitting. @@ -610,6 +777,7 @@ mod tests { None, None, None, + None, ) .await .expect("streaming dispatch should reach the worker"); diff --git a/experimental/sgl-router/src/proxy/sse.rs b/experimental/sgl-router/src/proxy/sse.rs index 3ce5e8ea0..695af12af 100644 --- a/experimental/sgl-router/src/proxy/sse.rs +++ b/experimental/sgl-router/src/proxy/sse.rs @@ -4,22 +4,32 @@ //! SSE passthrough — bridges a reqwest `bytes_stream()` into an axum Body. use std::panic::AssertUnwindSafe; -use std::sync::Arc; +use std::time::Duration; use axum::body::Body; use bytes::Bytes; use futures::{FutureExt, StreamExt}; use tokio_stream::wrappers::ReceiverStream; +use tokio_util::sync::CancellationToken; + +/// Why the SSE pump stopped, independently of any SSE error event it observed. +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +pub enum StreamEndReason { + Completed, + UpstreamError, + IdleTimeout, + /// The router's stale-request deadline expired, regardless of worker health. + Expired, + ClientDisconnect, + PumpPanicked, +} /// How the SSE pump ended, reported to the `on_complete` hook. #[derive(Debug, Clone, Copy)] pub struct StreamEnd { - /// No upstream stream error and no pump panic. - pub transport_ok: bool, + pub reason: StreamEndReason, /// An SSE error event (`data: {"error"...}`) rode the stream. pub saw_error_event: bool, - /// The client dropped the response body before upstream finished. - pub client_disconnect: bool, } /// A `data:` line whose payload's first JSON key is `error` — tolerant of @@ -59,140 +69,126 @@ impl ErrorEventScanner { } } +/// Bounds on a streaming response beyond what the upstream stream itself provides. +#[derive(Debug, Clone, Default)] +pub struct StreamLimits { + /// Maximum silence between upstream chunks. `None` waits indefinitely. + pub idle_timeout: Option, + /// Fires when the stale-request janitor expires the request. + pub expiration: Option, +} + /// Bridge a byte stream into an axum Body that streams chunks unchanged. /// -/// Spawns one tokio task per stream so the handler can return immediately. -/// Uses a **bounded** 64-slot channel so `tx.send().await` naturally -/// backpressures the upstream read when the client (axum Body consumer) falls -/// behind — an unbounded channel would buffer hundreds of MB for a slow client -/// receiving a long completion. +/// One tokio task pumps upstream chunks through a bounded 64-slot channel so a +/// slow client backpressures the upstream read. The pump stops as soon as the +/// client disconnects, even while upstream is silent, when `limits.idle_timeout` +/// elapses between chunks, or when `limits.expiration` fires. /// -/// # Backpressure note -/// The channel bound of 64 absorbs short bursts while still limiting -/// worst-case outstanding bytes to 64 × chunk_size (typically a few MB). +/// The terminal result travels on a separate channel and is chained after the +/// data, so a full queue cannot block cleanup or turn a failed stream into a +/// clean EOF. A pump panic is reported the same way. /// -/// # Client disconnect -/// When the axum Body is dropped the receiver is closed; `tx.send()` then -/// returns `Err`, which breaks the loop — no upstream bytes are read after the -/// client disconnects. -/// -/// # Panic safety -/// The pump future is wrapped in `AssertUnwindSafe(..).catch_unwind()`. If the -/// upstream stream panics, we surface a loud `io::Error` to the client; without -/// this, the body would EOF cleanly and clients couldn't distinguish that from -/// success — the worst failure class (truncated output that looks complete). -/// -/// # Stream guards -/// When `stream_guards` is `Some`, the value is **moved into the spawned task** -/// and held for the entire body lifetime. It is dropped only when the SSE -/// pump finishes (stream exhausted, client disconnects, or upstream errors). -/// The opaque `Box` accepts any drop-only payload — most -/// commonly a tuple of [`crate::workers::LoadGuard`] and -/// [`crate::state::load_monitor::router_inflight_load::RouterInflightLoadGuard`]. The proxy does not -/// inspect the value; it relies entirely on `Drop` semantics, so callers can -/// pack arbitrary cleanup state in. Pass `None` for callers that manage the -/// guard externally (e.g. non-streaming paths where the handler itself is the -/// guard scope). -/// -/// # Completion hook -/// When `on_complete` is `Some`, it runs exactly once when the pump task -/// finishes with the transport, SSE error-event, and client-disconnect state. -/// -/// # First-byte hook -/// When `on_first_byte` is `Some`, the closure runs exactly once, the moment -/// the first `Ok` chunk is read from the upstream stream — i.e. time to first -/// token. It does NOT fire if the stream ends or errors before any `Ok` chunk -/// arrives. `forward_streaming_to` passes a closure that records -/// `sgl_router_ttft_seconds` for successful streaming responses. +/// `guards` is held until the pump finishes; `on_first_byte` runs on the first +/// `Ok` chunk; `on_complete` runs exactly once with the final [`StreamEnd`]. pub fn bytes_stream_to_body( - stream: S, - stream_guards: Option>, + mut stream: S, + guards: Option>, on_complete: Option>, - on_first_byte: Option>, + mut on_first_byte: Option>, + limits: StreamLimits, ) -> Body where S: futures::Stream> + Send + Unpin + 'static, E: std::fmt::Display + Send + Sync + 'static, { - let (tx, rx) = tokio::sync::mpsc::channel(64); + let (tx, rx) = tokio::sync::mpsc::channel::(64); + let (terminal_tx, terminal_rx) = tokio::sync::oneshot::channel(); tokio::spawn(async move { - let tx_for_panic = tx.clone(); - let outcome = Arc::new(parking_lot::Mutex::new(StreamEnd { - transport_ok: true, + let mut end = StreamEnd { + reason: StreamEndReason::Completed, saw_error_event: false, - client_disconnect: false, - })); - let outcome_setter = Arc::clone(&outcome); - // `None` once an error event is found — the scan is done for good. - let mut scanner = Some(ErrorEventScanner::default()); - let pump = AssertUnwindSafe(async move { - // Hold the guards for the task's lifetime — dropped when this - // block exits (stream done or client disconnect). Leading - // underscore suppresses the "unused variable" lint while - // keeping intent explicit. - let _hold = stream_guards; - let mut on_first_byte = on_first_byte; - let mut s = stream; - while let Some(chunk) = s.next().await { - let item: Result = chunk.map_err(|e| { - let msg = e.to_string(); - tracing::warn!(error = %msg, "upstream SSE stream errored mid-flight"); - std::io::Error::other(msg) - }); - let is_err_chunk = item.is_err(); - match &item { - Ok(bytes) => { - // TTFT hook: at most once (`take()`); an error-first - // stream never produced a token, so it stays unfired. + }; + let mut scanner = ErrorEventScanner::default(); + let idle = limits.idle_timeout.unwrap_or(Duration::MAX); + let expired = async { + match limits.expiration { + Some(token) => token.cancelled().await, + None => std::future::pending().await, + } + }; + // Disconnect and expiration race the whole forwarding loop, so they + // fire while `send` waits on a full queue as well as while upstream is silent. + let pump = async { + tokio::select! { + biased; + _ = tx.closed() => { + end.reason = StreamEndReason::ClientDisconnect; + Ok(()) + } + _ = expired => { + end.reason = StreamEndReason::Expired; + Err(std::io::Error::other("SSE stream exceeded stale_request_timeout")) + } + result = async { + loop { + let bytes = match tokio::time::timeout(idle, stream.next()).await { + Ok(None) => return (StreamEndReason::Completed, Ok(())), + Ok(Some(Ok(bytes))) => bytes, + Ok(Some(Err(e))) => return ( + StreamEndReason::UpstreamError, + Err(std::io::Error::other(e.to_string())), + ), + Err(_) => return ( + StreamEndReason::IdleTimeout, + Err(std::io::Error::other("SSE upstream idle timeout")), + ), + }; if let Some(hook) = on_first_byte.take() { hook(); } - if scanner.as_mut().is_some_and(|scanner| scanner.feed(bytes)) { - outcome_setter.lock().saw_error_event = true; - scanner = None; + if !end.saw_error_event { + end.saw_error_event = scanner.feed(&bytes); + } + if tx.send(bytes).await.is_err() { + // Receiver gone; the `tx.closed()` arm reports the disconnect. + std::future::pending::<()>().await; } } - Err(_) => outcome_setter.lock().transport_ok = false, - } - if tx.send(item).await.is_err() { - // Receiver dropped. If we were about to ship an upstream - // error there's nothing left to report; otherwise this is - // a clean client-side disconnect — log at debug since it's - // not a router-side fault. - if !is_err_chunk { - tracing::debug!("SSE client disconnected mid-stream"); - outcome_setter.lock().client_disconnect = true; - } - break; - } - if is_err_chunk { - // Surfaced upstream error to client; stop reading. - break; + } => { + end.reason = result.0; + result.1 } } - }); - let pump_result = pump.catch_unwind().await; - let panicked = pump_result.is_err(); - if let Err(panic_payload) = pump_result { - let msg = panic_payload - .downcast_ref::<&'static str>() - .map(|s| (*s).to_string()) - .or_else(|| panic_payload.downcast_ref::().cloned()) - .unwrap_or_else(|| "".to_string()); - tracing::error!(error = %msg, "SSE pump task panicked"); - let _ = tx_for_panic - .send(Err(std::io::Error::other(format!( - "SSE pump panicked: {msg}" - )))) - .await; - } + }; + let result = match AssertUnwindSafe(pump).catch_unwind().await { + Ok(result) => result, + Err(payload) => { + end.reason = StreamEndReason::PumpPanicked; + let message = payload + .downcast_ref::<&str>() + .copied() + .or_else(|| payload.downcast_ref::().map(String::as_str)) + .unwrap_or(""); + Err(std::io::Error::other(format!( + "SSE pump panicked: {message}" + ))) + } + }; if let Some(hook) = on_complete { - let mut end = *outcome.lock(); - end.transport_ok &= !panicked; hook(end); } + drop(guards); + let _ = terminal_tx.send(result); }); - Body::from_stream(ReceiverStream::new(rx)) + let terminal = futures::stream::once(terminal_rx).filter_map(|result| { + futures::future::ready(match result { + Ok(Ok(())) => None, + Ok(Err(error)) => Some(Err(error)), + Err(_) => Some(Err(std::io::Error::other("SSE pump cancelled"))), + }) + }); + Body::from_stream(ReceiverStream::new(rx).map(Ok).chain(terminal)) } #[cfg(test)] @@ -202,6 +198,108 @@ mod tests { use futures::stream; use http_body_util::BodyExt; + fn limited_body( + stream: S, + limits: StreamLimits, + ) -> (Body, tokio::sync::oneshot::Receiver) + where + S: futures::Stream> + Send + Unpin + 'static, + { + let (tx, rx) = tokio::sync::oneshot::channel(); + let body = bytes_stream_to_body( + stream, + None, + Some(Box::new(move |end| { + let _ = tx.send(end); + })), + None, + limits, + ); + (body, rx) + } + + #[tokio::test(start_paused = true)] + async fn idle_disconnect_releases_guards_without_waiting_for_upstream() { + struct Release(Option>); + impl Drop for Release { + fn drop(&mut self) { + let _ = self.0.take().unwrap().send(()); + } + } + let (tx, rx) = tokio::sync::oneshot::channel(); + let body = bytes_stream_to_body( + stream::pending::>(), + Some(Box::new(Release(Some(tx)))), + None, + None, + StreamLimits::default(), + ); + tokio::task::yield_now().await; + drop(body); + tokio::time::timeout(Duration::from_millis(1), rx) + .await + .unwrap() + .unwrap(); + } + + #[tokio::test(start_paused = true)] + async fn idle_timeout_is_a_visible_upstream_failure() { + let (body, end) = limited_body( + stream::pending(), + StreamLimits { + idle_timeout: Some(Duration::from_secs(1)), + expiration: None, + }, + ); + assert!(body + .collect() + .await + .unwrap_err() + .to_string() + .contains("idle timeout")); + let end = end.await.unwrap(); + assert_eq!(end.reason, StreamEndReason::IdleTimeout); + } + + #[tokio::test(start_paused = true)] + async fn expiration_releases_guards_while_queue_is_full() { + struct Release(Option>); + impl Drop for Release { + fn drop(&mut self) { + let _ = self.0.take().unwrap().send(()); + } + } + let token = CancellationToken::new(); + let (release_tx, release_rx) = tokio::sync::oneshot::channel(); + let (end_tx, end_rx) = tokio::sync::oneshot::channel(); + let body = bytes_stream_to_body( + stream::repeat_with(|| Ok::<_, std::io::Error>(Bytes::from_static(b"chunk"))), + Some(Box::new(Release(Some(release_tx)))), + Some(Box::new(move |end| { + let _ = end_tx.send(end); + })), + None, + StreamLimits { + idle_timeout: None, + expiration: Some(token.clone()), + }, + ); + tokio::task::yield_now().await; + token.cancel(); + // Cleanup must finish before the client frees any channel capacity. + tokio::time::timeout(Duration::from_millis(1), release_rx) + .await + .expect("expiration must release guards while the queue remains full") + .unwrap(); + assert_eq!(end_rx.await.unwrap().reason, StreamEndReason::Expired); + // The queue is full, so the failure must ride the terminal channel. + assert!(body + .collect() + .await + .unwrap_err() + .to_string() + .contains("stale_request_timeout")); + } #[tokio::test] async fn passes_through_a_simple_byte_stream() { let chunks = vec![ @@ -209,7 +307,7 @@ mod tests { Ok(Bytes::from_static(b"world")), ]; let s = stream::iter(chunks); - let body = bytes_stream_to_body(s, None, None, None); + let body = bytes_stream_to_body(s, None, None, None, StreamLimits::default()); let bytes = body.collect().await.unwrap().to_bytes(); assert_eq!(&bytes[..], b"hello world"); } @@ -233,6 +331,7 @@ mod tests { Some(Box::new(move || { fired_c.fetch_add(1, Ordering::SeqCst); })), + StreamLimits::default(), ); let _ = body.collect().await.unwrap(); assert_eq!( @@ -260,6 +359,7 @@ mod tests { Some(Box::new(move || { fired_c.fetch_add(1, Ordering::SeqCst); })), + StreamLimits::default(), ); let _ = body.collect().await; assert_eq!( @@ -276,7 +376,7 @@ mod tests { Err(std::io::Error::other("upstream blew up mid-stream")), ]; let s = stream::iter(chunks); - let body = bytes_stream_to_body(s, None, None, None); + let body = bytes_stream_to_body(s, None, None, None, StreamLimits::default()); // Collecting a body that terminates with an error must return Err. let result = body.collect().await; assert!( @@ -338,8 +438,9 @@ mod tests { // that arm, the closure unwrap-or-elses would panic itself or // produce an empty message, which this test catches. let s = PanicAnyOnSecondPoll { polls: 0 }; - let body = bytes_stream_to_body(s, None, None, None); + let (body, end) = limited_body(s, StreamLimits::default()); let result = body.collect().await; + assert_eq!(end.await.unwrap().reason, StreamEndReason::PumpPanicked); assert!( result.is_err(), "expected body collect to surface non-string panic as Err, got Ok" @@ -361,7 +462,7 @@ mod tests { // The pump task panics mid-stream. The client must see a loud Err, // NOT a silently-truncated success. let s = PanicOnSecondPoll { polls: 0 }; - let body = bytes_stream_to_body(s, None, None, None); + let body = bytes_stream_to_body(s, None, None, None, StreamLimits::default()); let result = body.collect().await; assert!( result.is_err(), @@ -420,7 +521,7 @@ mod tests { yielded: 0, max: 1000, // way more than we'll let it consume }; - let body = bytes_stream_to_body(stream, None, None, None); + let body = bytes_stream_to_body(stream, None, None, None, StreamLimits::default()); // Read exactly one frame, then drop the body to simulate client disconnect. let mut data_stream = body.into_data_stream(); @@ -431,7 +532,7 @@ mod tests { // Give the pump generous time to make additional polls if its break is // broken. Healthy code: pump fills the 64-slot channel, then on the // next iteration tx.send().await detects receiver-drop and breaks. - tokio::time::sleep(std::time::Duration::from_millis(200)).await; + tokio::time::sleep(Duration::from_millis(200)).await; let final_polls = polls.load(Ordering::SeqCst); assert!( final_polls <= 70, @@ -503,6 +604,7 @@ mod tests { let _ = tx.send(end); })), None, + StreamLimits::default(), ); (body, rx) } @@ -520,9 +622,8 @@ mod tests { let (body, completion) = body_with_completion(chunks); let _ = body.collect().await.unwrap(); let end = stream_end(completion).await; - assert!(end.transport_ok); + assert_eq!(end.reason, StreamEndReason::Completed); assert!(end.saw_error_event); - assert!(!end.client_disconnect); } #[tokio::test] @@ -536,7 +637,7 @@ mod tests { let (body, completion) = body_with_completion(chunks); let _ = body.collect().await; let end = stream_end(completion).await; - assert!(!end.transport_ok); + assert_eq!(end.reason, StreamEndReason::UpstreamError); assert!(end.saw_error_event); } @@ -546,9 +647,8 @@ mod tests { let (body, completion) = body_with_completion(chunks); let _ = body.collect().await; let end = stream_end(completion).await; - assert!(!end.transport_ok); + assert_eq!(end.reason, StreamEndReason::UpstreamError); assert!(!end.saw_error_event); - assert!(!end.client_disconnect); } #[tokio::test] @@ -559,9 +659,8 @@ mod tests { let (body, completion) = body_with_completion(chunks); let _ = body.collect().await.unwrap(); let end = stream_end(completion).await; - assert!(end.transport_ok); + assert_eq!(end.reason, StreamEndReason::Completed); assert!(!end.saw_error_event); - assert!(!end.client_disconnect); } #[tokio::test] @@ -576,7 +675,6 @@ mod tests { let _ = stream.next().await; drop(stream); let end = stream_end(completion).await; - assert!(end.transport_ok); - assert!(end.client_disconnect); + assert_eq!(end.reason, StreamEndReason::ClientDisconnect); } } diff --git a/experimental/sgl-router/src/server/metrics.rs b/experimental/sgl-router/src/server/metrics.rs index d153c0b61..9a7b00679 100644 --- a/experimental/sgl-router/src/server/metrics.rs +++ b/experimental/sgl-router/src/server/metrics.rs @@ -92,7 +92,7 @@ //! The exposition is text/plain; version=0.0.4 per the Prometheus spec. use crate::config::PolicyKind; -use crate::proxy::sse::StreamEnd; +use crate::proxy::sse::{StreamEnd, StreamEndReason}; use parking_lot::Mutex; use std::collections::HashMap; use std::sync::atomic::{AtomicI64, AtomicU64, Ordering}; @@ -236,14 +236,19 @@ pub enum StreamOutcome { UpstreamError, /// The client disconnected before the stream finished. ClientDisconnect, + /// The router's stale-request deadline aborted the stream. + Expired, } pub(crate) fn classify_stream_end(end: StreamEnd) -> StreamOutcome { - match (end.transport_ok, end.saw_error_event, end.client_disconnect) { - (false, _, _) => StreamOutcome::UpstreamError, - (_, true, _) => StreamOutcome::StreamErrorEvent, - (_, _, true) => StreamOutcome::ClientDisconnect, - _ => StreamOutcome::Ok, + match end.reason { + StreamEndReason::Expired => StreamOutcome::Expired, + StreamEndReason::UpstreamError + | StreamEndReason::IdleTimeout + | StreamEndReason::PumpPanicked => StreamOutcome::UpstreamError, + _ if end.saw_error_event => StreamOutcome::StreamErrorEvent, + StreamEndReason::ClientDisconnect => StreamOutcome::ClientDisconnect, + StreamEndReason::Completed => StreamOutcome::Ok, } } @@ -254,6 +259,7 @@ impl StreamOutcome { Self::StreamErrorEvent => "stream_error_event", Self::UpstreamError => "upstream_error", Self::ClientDisconnect => "client_disconnect", + Self::Expired => "expired", } } } @@ -1500,22 +1506,26 @@ mod tests { fn stream_outcome_precedence() { use StreamOutcome::*; - for (transport_ok, saw_error_event, client_disconnect, expected) in [ - (false, false, false, UpstreamError), - (false, false, true, UpstreamError), - (false, true, false, UpstreamError), - (false, true, true, UpstreamError), - (true, false, false, Ok), - (true, false, true, ClientDisconnect), - (true, true, false, StreamErrorEvent), - (true, true, true, StreamErrorEvent), + for (reason, expected) in [ + (StreamEndReason::Completed, Ok), + (StreamEndReason::ClientDisconnect, ClientDisconnect), + (StreamEndReason::UpstreamError, UpstreamError), + (StreamEndReason::IdleTimeout, UpstreamError), + (StreamEndReason::PumpPanicked, UpstreamError), + (StreamEndReason::Expired, Expired), ] { - let end = StreamEnd { - transport_ok, - saw_error_event, - client_disconnect, - }; - assert_eq!(classify_stream_end(end), expected, "{end:?}"); + for saw_error_event in [false, true] { + let end = StreamEnd { + reason, + saw_error_event, + }; + let expected = if saw_error_event && matches!(expected, Ok | ClientDisconnect) { + StreamErrorEvent + } else { + expected + }; + assert_eq!(classify_stream_end(end), expected, "{end:?}"); + } } } @@ -1527,12 +1537,14 @@ mod tests { reg.record_stream_outcome("http://w:30000", "tiny", StreamOutcome::StreamErrorEvent); reg.record_stream_outcome("http://w:30000", "tiny", StreamOutcome::UpstreamError); reg.record_stream_outcome("http://w:30000", "tiny", StreamOutcome::ClientDisconnect); + reg.record_stream_outcome("http://w:30000", "tiny", StreamOutcome::Expired); let out = reg.render(); for expected in [ r#"sgl_router_stream_outcome_total{worker_url="http://w:30000",model_id="tiny",outcome="ok"} 2"#, r#"sgl_router_stream_outcome_total{worker_url="http://w:30000",model_id="tiny",outcome="stream_error_event"} 1"#, r#"sgl_router_stream_outcome_total{worker_url="http://w:30000",model_id="tiny",outcome="upstream_error"} 1"#, r#"sgl_router_stream_outcome_total{worker_url="http://w:30000",model_id="tiny",outcome="client_disconnect"} 1"#, + r#"sgl_router_stream_outcome_total{worker_url="http://w:30000",model_id="tiny",outcome="expired"} 1"#, ] { assert_metric_line(&out, expected); } diff --git a/experimental/sgl-router/src/server/routes/chat/forward.rs b/experimental/sgl-router/src/server/routes/chat/forward.rs index f07fa64e9..562310ab3 100644 --- a/experimental/sgl-router/src/server/routes/chat/forward.rs +++ b/experimental/sgl-router/src/server/routes/chat/forward.rs @@ -5,7 +5,7 @@ use super::preparation::{generate_room_id, BootstrapFields, PreparedChatRequest}; use crate::discovery::WorkerMode; -use crate::proxy::sse::StreamEnd; +use crate::proxy::sse::{StreamEnd, StreamEndReason}; use crate::server::app_context::AppContext; use crate::server::error::ApiError; use crate::server::metrics::{ @@ -20,6 +20,7 @@ use axum::response::IntoResponse; use bytes::Bytes; use std::sync::Arc; use std::time::Instant; +use tokio_util::sync::CancellationToken; const CHAT_PATH: &str = "/v1/chat/completions"; // Expose the selected decode worker to both PD workers and the client. @@ -64,8 +65,6 @@ pub(super) async fn forward_chat_request( request.input_token_count, 0, ); - // PD requests keep using the prefill expiration token after dispatching decode. - let expiration_token = active_request_guard.cancel_token().clone(); // Attribute the outcome to the worker supplying the client-visible response. let metrics = DispatchMetrics::new( ctx, @@ -105,6 +104,9 @@ pub(super) async fn forward_chat_request( (prefill, prefill_load_guards) }; + // In PD mode, prefill can finish before decode. Watch the registration + // held by the response so expiration remains live for its full lifetime. + let expiration_token = response_load_guards.1.cancel_token().clone(); let response_future = forward_to_response_worker( ctx, &response_worker, @@ -112,6 +114,7 @@ pub(super) async fn forward_chat_request( body, response_load_guards, &metrics, + expiration_token.clone(), ); // A ready response wins if request expiration fires in the same poll. let result = tokio::select! { @@ -192,6 +195,7 @@ async fn forward_to_response_worker( body: Bytes, load_guards: LoadGuards, metrics: &DispatchMetrics, + expiration: CancellationToken, ) -> Result, ApiError> { if metrics.streaming { // Load and duration guards live until the SSE pump ends, not just until headers arrive. @@ -208,6 +212,7 @@ async fn forward_to_response_worker( Some(stream_guards), Some(metrics.first_byte_callback()), Some(metrics.stream_end_callback(worker.url.clone())), + Some(expiration), ) .await } else { @@ -279,6 +284,9 @@ impl DispatchMetrics { let metrics = Arc::clone(&self.registry); let model = self.model.clone(); Box::new(move |end| { + if end.reason == StreamEndReason::Expired { + metrics.record_stale_request(StaleRequestOutcome::Expired); + } metrics.record_stream_outcome(&response_worker_url, &model, classify_stream_end(end)); }) } diff --git a/experimental/sgl-router/tests/proxy/chat_routing.rs b/experimental/sgl-router/tests/proxy/chat_routing.rs index 28bf7c718..d3218c2f8 100644 --- a/experimental/sgl-router/tests/proxy/chat_routing.rs +++ b/experimental/sgl-router/tests/proxy/chat_routing.rs @@ -1129,6 +1129,7 @@ async fn forward_streaming_to_records_failure_on_mid_stream_drop() { None, None, None, + None, ) .await; diff --git a/experimental/sgl-router/tests/proxy/h2c_forward.rs b/experimental/sgl-router/tests/proxy/h2c_forward.rs index 0cbca58ab..983fce2b4 100644 --- a/experimental/sgl-router/tests/proxy/h2c_forward.rs +++ b/experimental/sgl-router/tests/proxy/h2c_forward.rs @@ -168,6 +168,7 @@ async fn h2c_client_streams_sse_from_http2_only_worker() { flag.store(true, std::sync::atomic::Ordering::SeqCst); })), None, + None, ) .await .expect("h2c client must stream from an HTTP/2-only worker"); diff --git a/experimental/sgl-router/tests/proxy/pd_pool_isolation.rs b/experimental/sgl-router/tests/proxy/pd_pool_isolation.rs index 23ed74189..3fbe944b2 100644 --- a/experimental/sgl-router/tests/proxy/pd_pool_isolation.rs +++ b/experimental/sgl-router/tests/proxy/pd_pool_isolation.rs @@ -92,6 +92,112 @@ fn chat_request() -> Request { .unwrap() } +#[tokio::test] +async fn pd_decode_stream_expires_after_prefill_completes() { + use sgl_router::state::load_monitor::router_inflight_load::{ + MockClock, RouterInflightLoadRegistry, + }; + use std::time::Instant; + + let prefill = crate::common::mock_worker::MockWorker::start(vec![]).await; + let decode = crate::common::mock_worker::MockWorker::start_slow_stream( + vec!["data: chunk\n\n"; 1000], + Duration::from_millis(10), + ) + .await; + let cfg = config(); + let tokenizers = Arc::new(TokenizerRegistry::load_from_config(&cfg).unwrap()); + let registry = Arc::new(WorkerRegistry::default()); + registry + .add(WorkerSpec { + id: WorkerId("p1".into()), + url: prefill.url.clone(), + mode: WorkerMode::Prefill, + model_ids: vec![ModelId("tiny".into())], + bootstrap_port: Some(8997), + }) + .unwrap(); + registry + .add(WorkerSpec { + id: WorkerId("d1".into()), + url: decode.url.clone(), + mode: WorkerMode::Decode, + model_ids: vec![ModelId("tiny".into())], + bootstrap_port: None, + }) + .unwrap(); + let prefill_worker = registry.get(&WorkerId("p1".into())).unwrap(); + let decode_worker = registry.get(&WorkerId("d1".into())).unwrap(); + let clock = Arc::new(MockClock::new(Instant::now())); + let inflight = RouterInflightLoadRegistry::new(clock.clone(), Duration::from_secs(10)); + let policies = Arc::new(build_registry_with_defaults(&cfg).unwrap()); + let proxy = Arc::new(Proxy::new(Duration::from_secs(5)).unwrap()); + let ctx = Arc::new(AppContext::with_router_inflight_load( + cfg, + tokenizers, + proxy, + registry, + policies, + inflight.clone(), + )); + let request = Request::builder() + .method("POST") + .uri("/v1/chat/completions") + .header("content-type", "application/json") + .body(Body::from( + serde_json::json!({ + "model": "tiny", + "messages": [{"role": "user", "content": "hi"}], + "stream": true, + }) + .to_string(), + )) + .unwrap(); + // Two prior faults make any accidental expiry failure trip the default breaker. + decode_worker.breaker.record_failure(); + decode_worker.breaker.record_failure(); + let response = build_router(ctx.clone()).oneshot(request).await.unwrap(); + assert_eq!(response.status(), StatusCode::OK); + let mut body = response.into_body(); + tokio::time::timeout(Duration::from_secs(2), body.frame()) + .await + .unwrap() + .unwrap() + .unwrap(); + tokio::time::timeout(Duration::from_secs(2), async { + while inflight.inflight_count() != 1 || prefill_worker.router_inflight_load() != 0 { + tokio::task::yield_now().await; + } + }) + .await + .expect("prefill must complete before expiring decode"); + assert_eq!(decode_worker.router_inflight_load(), 1); + clock.advance(Duration::from_secs(11)); + assert_eq!(inflight.sweep_stale(), 1); + let error = tokio::time::timeout(Duration::from_secs(2), body.collect()) + .await + .expect("decode stream must stop when its registration expires") + .unwrap_err(); + assert!(error.to_string().contains("stale_request_timeout")); + assert_eq!(decode_worker.router_inflight_load(), 0); + assert_eq!(decode_worker.breaker.snapshot().state_code, 0); + let metrics = ctx.metrics.render(); + assert!(metrics + .lines() + .any(|line| line == r#"sgl_router_stale_requests_total{outcome="expired"} 1"#)); + let expected = format!( + r#"sgl_router_stream_outcome_total{{worker_url="{}",model_id="tiny",outcome="expired"}} 1"#, + decode.url, + ); + assert!(metrics.lines().any(|line| line == expected)); + assert!(!metrics + .lines() + .any(|line| line.starts_with("sgl_router_stream_outcome_total{") + && line.contains(r#"outcome="upstream_error""#))); + decode_worker.breaker.record_failure(); + assert_eq!(decode_worker.breaker.snapshot().state_code, 1); +} + /// Gap closer #1: PD mode with only decode workers → 503 with /// `no_prefill_workers_available`. The chat route is a prefill /// dispatch, so a decode-only pool means partial failure.