[sgl-router] Release cancelled circuit-breaker probes (#40603)

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<CircuitPermit<'_>> {
         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<u64>,
+}
+
+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);

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<Response<Body>, 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<Box<dyn FnOnce(sse::StreamEnd) + Send + 'static>>,
         expiration: Option<CancellationToken>,
     ) -> Result<Response<Body>, 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.