diff --git a/experimental/sgl-router/src/health/circuit_breaker.rs b/experimental/sgl-router/src/health/circuit_breaker.rs index 26d8c73f0..476ee2709 100644 --- a/experimental/sgl-router/src/health/circuit_breaker.rs +++ b/experimental/sgl-router/src/health/circuit_breaker.rs @@ -49,6 +49,7 @@ enum State { struct Inner { state: State, consecutive_failures: u32, + probe_generation: u64, } #[derive(Debug)] @@ -67,6 +68,7 @@ impl CircuitBreaker { inner: Mutex::new(Inner { state: State::Closed, consecutive_failures: 0, + probe_generation: 0, }), config, } @@ -119,33 +121,34 @@ impl CircuitBreaker { CircuitSnapshot { admit, state_code } } - /// True if a request may proceed. Mutates state when transitioning - /// from Open → HalfOpen. + /// Admit a caller that will explicitly record its outcome. pub fn allow(&self) -> bool { + self.acquire().map(|permit| permit.disarm()).is_some() + } + + /// Claim admission, releasing an unfinished recovery probe on cancellation. + pub fn acquire(&self) -> Option> { let mut g = self.inner.lock().unwrap(); - match g.state { - State::Closed => true, - State::Open { opened_at } => { - if opened_at.elapsed() >= self.config.cool_down { - g.state = State::HalfOpen { - probe_in_flight: true, - }; - true - } else { - false - } + let generation = match g.state { + State::Closed => None, + State::Open { opened_at } if opened_at.elapsed() < self.config.cool_down => { + return None; } - State::HalfOpen { probe_in_flight } => { - if probe_in_flight { - false - } else { - g.state = State::HalfOpen { - probe_in_flight: true, - }; - true - } + State::HalfOpen { + probe_in_flight: true, + } => return None, + _ => { + g.probe_generation = g.probe_generation.wrapping_add(1); + g.state = State::HalfOpen { + probe_in_flight: true, + }; + Some(g.probe_generation) } - } + }; + Some(CircuitPermit { + breaker: self, + generation, + }) } pub fn record_success(&self) { @@ -203,6 +206,33 @@ impl CircuitBreaker { } } +/// Releases only the recovery probe claimed by this admission. +pub struct CircuitPermit<'a> { + breaker: &'a CircuitBreaker, + generation: Option, +} + +impl CircuitPermit<'_> { + /// Leave outcome accounting to the caller or streaming completion hook. + pub fn disarm(mut self) { + self.generation = None; + } +} + +impl Drop for CircuitPermit<'_> { + fn drop(&mut self) { + let Some(generation) = self.generation else { + return; + }; + let mut g = self.breaker.inner.lock().unwrap(); + if g.probe_generation == generation { + if let State::HalfOpen { probe_in_flight } = &mut g.state { + *probe_in_flight = false; + } + } + } +} + impl Default for CircuitBreaker { fn default() -> Self { Self::new() @@ -220,6 +250,22 @@ mod tests { }) } + #[test] + fn cancelled_permits_release_only_their_own_probe() { + let b = cb(1, 0); + let closed = b.acquire().unwrap(); + b.record_failure(); + let old_probe = b.acquire().unwrap(); + b.record_success(); + b.record_failure(); + let current_probe = b.acquire().unwrap(); + drop((closed, old_probe)); + assert!(!b.would_allow()); + drop(current_probe); + assert!(b.would_allow()); + assert_eq!(b.snapshot().state_code, 2); + } + #[test] fn state_code_is_closed_by_default() { assert_eq!(CircuitBreaker::new().snapshot().state_code, 0); diff --git a/experimental/sgl-router/src/proxy/mod.rs b/experimental/sgl-router/src/proxy/mod.rs index be687cea0..e23d9063f 100644 --- a/experimental/sgl-router/src/proxy/mod.rs +++ b/experimental/sgl-router/src/proxy/mod.rs @@ -193,7 +193,7 @@ impl Proxy { } } - /// Breaker-gated JSON POST: checks `breaker.allow()` first, classifies the + /// Breaker-gated JSON POST: acquires a cancellation-safe permit first, classifies the /// response status through [`breaker_outcome`] (success / failure / /// backpressure), and returns `ApiError::BreakerOpen` immediately when the /// breaker is Open. @@ -212,11 +212,9 @@ impl Proxy { headers: &HeaderMap, body: Bytes, ) -> Result, ApiError> { - if !breaker.allow() { - return Err(ApiError::BreakerOpen { - worker: worker_url.to_string(), - }); - } + let permit = breaker.acquire().ok_or_else(|| ApiError::BreakerOpen { + worker: worker_url.to_string(), + })?; let worker_url = parse_worker_url(worker_url, breaker)?; let url = worker_url.join(path).map_err(|e| { ApiError::Internal(anyhow::Error::new(e).context(format!("join worker path {path}"))) @@ -268,6 +266,7 @@ impl Proxy { // worker that answers a probe with 503 isn't wedged shut. BreakerOutcome::Neutral => breaker.record_backpressure(), } + permit.disarm(); let mut out = Response::new(Body::from(bytes)); *out.status_mut() = status; out.headers_mut().insert( @@ -277,7 +276,7 @@ impl Proxy { Ok(out) } - /// Breaker-gated streaming POST: checks `breaker.allow()` first, classifies + /// Breaker-gated streaming POST: acquires a cancellation-safe permit first, classifies /// the response status through [`breaker_outcome`], and returns /// `ApiError::BreakerOpen` when Open. /// @@ -307,11 +306,9 @@ impl Proxy { on_stream_end: Option>, expiration: Option, ) -> Result, ApiError> { - if !breaker.allow() { - return Err(ApiError::BreakerOpen { - worker: worker_url.to_string(), - }); - } + let permit = breaker.acquire().ok_or_else(|| ApiError::BreakerOpen { + worker: worker_url.to_string(), + })?; let worker_url = parse_worker_url(worker_url, breaker)?; let url = worker_url.join(path).map_err(|e| { ApiError::Internal(anyhow::Error::new(e).context(format!("join worker path {path}"))) @@ -387,6 +384,7 @@ impl Proxy { } else { None }; + permit.disarm(); let body = sse::bytes_stream_to_body( resp.bytes_stream(), stream_guards, @@ -446,6 +444,49 @@ mod tests { )); } + #[tokio::test] + async fn cancelled_forwards_release_half_open_probes() { + use futures::FutureExt; + + let listener = TcpListener::bind("127.0.0.1:0").await.unwrap(); + let url = format!("http://{}", listener.local_addr().unwrap()); + let proxy = Proxy::new(Duration::from_secs(60)).unwrap(); + let breaker = Arc::new(CircuitBreaker::with_config(CircuitBreakerConfig { + threshold: NonZeroU32::new(1).unwrap(), + cool_down: Duration::ZERO, + })); + breaker.record_failure(); + let headers = HeaderMap::new(); + assert!(proxy + .forward_json_to( + &url, + WireProtocol::Http1, + &breaker, + "/chat", + &headers, + Bytes::new(), + ) + .now_or_never() + .is_none()); + assert!(breaker.would_allow()); + assert!(proxy + .forward_streaming_to( + &url, + WireProtocol::Http1, + &breaker, + "/chat", + &headers, + Bytes::new(), + None, + None, + None, + None, + ) + .now_or_never() + .is_none()); + assert!(breaker.would_allow()); + } + #[test] fn breaker_outcome_treats_backpressure_as_neutral() { // Backpressure: healthy but busy — must not touch the breaker.