From 36b03763c88d1a1db10abd7c029c2306966a511d Mon Sep 17 00:00:00 2001 From: Michael Ingley Date: Tue, 30 Jun 2026 11:55:31 -0500 Subject: [PATCH 1/6] Add A32 circuit breaking limiter --- tonic-xds/src/client/circuit_breaking.rs | 1000 +++++++++++++++++ tonic-xds/src/client/mod.rs | 1 + .../src/xds/resource/circuit_breaking.rs | 11 +- 3 files changed, 1006 insertions(+), 6 deletions(-) create mode 100644 tonic-xds/src/client/circuit_breaking.rs diff --git a/tonic-xds/src/client/circuit_breaking.rs b/tonic-xds/src/client/circuit_breaking.rs new file mode 100644 index 000000000..4e71b742e --- /dev/null +++ b/tonic-xds/src/client/circuit_breaking.rs @@ -0,0 +1,1000 @@ +#![cfg_attr(not(test), allow(dead_code))] + +use std::fmt; +use std::sync::{ + Arc, OnceLock, + atomic::{AtomicU32, AtomicU64, AtomicUsize, Ordering}, +}; +use std::task::{Context, Poll}; + +use bytes::Bytes; +use dashmap::DashMap; +use http::{Request, Response}; +use http_body::{Body, Frame}; +use pin_project_lite::pin_project; +use tokio::sync::watch; +use tonic::body::Body as TonicBody; +use tower::{BoxError, Layer, Service}; + +use crate::client::route::RouteDecision; +use crate::common::async_util::BoxFuture; +use crate::xds::resource::circuit_breaking::{CircuitBreakingConfig, DEFAULT_MAX_REQUESTS}; + +static GLOBAL_COUNTERS: OnceLock> = OnceLock::new(); + +/// Shared circuit-breaking state for xDS clusters. +#[derive(Clone, Debug)] +pub(crate) struct ClusterCircuitBreakers { + inner: Arc, +} + +#[derive(Debug)] +struct ClusterCircuitBreakersInner { + configs: DashMap>, + counters: ClusterRequestCounters, + default_max_requests: u32, +} + +impl Drop for ClusterCircuitBreakersInner { + fn drop(&mut self) { + for state in self.configs.iter() { + if let Some(previous) = state.config_tx.send_replace(None) { + let counter_key = previous.counter_key.clone(); + drop(previous); + self.counters.deactivate(&counter_key); + } + } + } +} + +impl ClusterCircuitBreakers { + pub(crate) fn new() -> Self { + Self::with_counters(ClusterRequestCounters::global()) + } + + fn with_counters(counters: ClusterRequestCounters) -> Self { + Self { + inner: Arc::new(ClusterCircuitBreakersInner { + configs: DashMap::new(), + counters, + default_max_requests: DEFAULT_MAX_REQUESTS, + }), + } + } + + #[allow(dead_code)] + pub(crate) fn set_config(&self, cluster: impl Into, config: CircuitBreakingConfig) { + let cluster = cluster.into(); + self.set_cluster_config(cluster.clone(), cluster, config); + } + + pub(crate) fn set_cluster_config( + &self, + cluster: impl Into, + eds_service_name: impl Into, + config: CircuitBreakingConfig, + ) { + let cluster = cluster.into(); + let eds_service_name = eds_service_name.into(); + let counter_key = counter_key(&cluster, &eds_service_name); + let counter = self.inner.counters.counter(&counter_key); + let state = self.ensure_state(&cluster); + self.update_state_config( + &state, + CircuitBreakerRuntimeConfig { + max_requests: config.max_requests, + counter_key: Arc::from(counter_key), + counter, + }, + ); + } + + fn ensure_state(&self, cluster: &str) -> Arc { + if let Some(state) = self.inner.configs.get(cluster) { + return state.clone(); + } + + self.inner + .configs + .entry(cluster.to_string()) + .or_insert_with(|| Arc::new(ClusterCircuitBreakerState::new())) + .clone() + } + + fn update_state_config( + &self, + state: &ClusterCircuitBreakerState, + config: CircuitBreakerRuntimeConfig, + ) { + let previous = state.current_config(); + if previous.as_ref() == Some(&config) { + return; + } + + let counter_key_changed = previous + .as_ref() + .is_none_or(|previous| previous.counter_key != config.counter_key); + if counter_key_changed { + self.inner.counters.activate(&config.counter_key); + } + + let previous = state.config_tx.send_replace(Some(config)); + if counter_key_changed && let Some(previous) = previous { + self.deactivate_config(previous); + } + } + + fn clear_state(&self, state: &ClusterCircuitBreakerState) { + if let Some(previous) = state.config_tx.send_replace(None) { + self.deactivate_config(previous); + } + } + + fn deactivate_config(&self, config: CircuitBreakerRuntimeConfig) { + let counter_key = config.counter_key.clone(); + drop(config); + self.inner.counters.deactivate(&counter_key); + } + + fn acquire(&self, cluster: &str) -> Result { + self.acquire_with_config(self.runtime_config_or_default(cluster)) + } + + fn acquire_with_config( + &self, + runtime_config: CircuitBreakerRuntimeConfig, + ) -> Result { + let limit = CircuitBreakerLimit { + max_requests: runtime_config.max_requests, + }; + self.inner + .counters + .acquire( + runtime_config.counter_key, + runtime_config.counter, + limit.max_requests, + ) + .ok_or(limit) + } + + fn runtime_config_or_default(&self, cluster: &str) -> CircuitBreakerRuntimeConfig { + self.inner + .configs + .get(cluster) + .and_then(|state| state.current_config()) + .unwrap_or_else(|| { + let counter_key = counter_key(cluster, cluster); + let counter = self.inner.counters.counter(&counter_key); + CircuitBreakerRuntimeConfig { + max_requests: self.inner.default_max_requests, + counter_key: Arc::from(counter_key), + counter, + } + }) + } + + #[cfg(test)] + fn in_flight(&self, cluster: &str) -> u32 { + let runtime_config = self.runtime_config_or_default(cluster); + self.inner.counters.in_flight(&runtime_config.counter_key) + } + + #[cfg(test)] + fn dropped_requests(&self, cluster: &str) -> u64 { + let runtime_config = self.runtime_config_or_default(cluster); + self.inner + .counters + .dropped_requests(&runtime_config.counter_key) + } + + #[cfg(test)] + fn counter_count(&self) -> usize { + self.inner.counters.counter_count() + } + + #[cfg(test)] + fn clear_cluster_config(&self, cluster: &str) { + if let Some((_, state)) = self.inner.configs.remove(cluster) { + self.clear_state(&state); + } + } + + #[cfg(test)] + pub(crate) fn new_for_test() -> Self { + Self::with_counters(ClusterRequestCounters::isolated()) + } +} + +impl Default for ClusterCircuitBreakers { + fn default() -> Self { + Self::new() + } +} + +#[derive(Clone, Copy, Debug)] +struct CircuitBreakerLimit { + max_requests: u32, +} + +#[derive(Clone, Debug)] +struct CircuitBreakerRuntimeConfig { + max_requests: u32, + counter_key: Arc, + counter: Arc, +} + +impl PartialEq for CircuitBreakerRuntimeConfig { + fn eq(&self, other: &Self) -> bool { + self.max_requests == other.max_requests && self.counter_key == other.counter_key + } +} + +impl Eq for CircuitBreakerRuntimeConfig {} + +fn counter_key(cluster: &str, eds_service_name: &str) -> String { + format!("{cluster}\0{eds_service_name}") +} + +#[derive(Debug)] +struct ClusterCircuitBreakerState { + config_tx: watch::Sender>, +} + +impl ClusterCircuitBreakerState { + fn new() -> Self { + let (config_tx, _) = watch::channel(None); + Self { config_tx } + } + + fn current_config(&self) -> Option { + self.config_tx.borrow().clone() + } +} + +#[derive(Clone, Debug)] +struct ClusterRequestCounters { + inner: Arc, +} + +#[derive(Debug, Default)] +struct ClusterRequestCounterState { + counters: DashMap>, + active_refs: DashMap>, +} + +impl ClusterRequestCounters { + fn global() -> Self { + Self { + inner: GLOBAL_COUNTERS + .get_or_init(|| Arc::new(ClusterRequestCounterState::default())) + .clone(), + } + } + + #[cfg(test)] + fn isolated() -> Self { + Self { + inner: Arc::new(ClusterRequestCounterState::default()), + } + } + + fn activate(&self, counter_key: &str) { + self.inner + .active_refs + .entry(counter_key.to_string()) + .or_insert_with(|| Arc::new(AtomicUsize::new(0))) + .fetch_add(1, Ordering::AcqRel); + } + + fn deactivate(&self, counter_key: &str) { + let should_cleanup = self + .inner + .active_refs + .get(counter_key) + .and_then(|active_refs| { + active_refs + .fetch_update(Ordering::AcqRel, Ordering::Acquire, |count| { + count.checked_sub(1) + }) + .ok() + }) + .is_some_and(|previous| previous <= 1); + + if should_cleanup { + self.cleanup_if_unused(counter_key); + } + } + + fn acquire( + &self, + counter_key: Arc, + counter: Arc, + limit: u32, + ) -> Option { + if counter.try_acquire(limit) { + Some(CircuitBreakerPermit { + counter: Some(counter), + counter_key, + counters: self.clone(), + }) + } else { + counter.record_drop(); + self.cleanup_if_unused(&counter_key); + None + } + } + + fn counter(&self, counter_key: &str) -> Arc { + self.inner + .counters + .entry(counter_key.to_string()) + .or_insert_with(|| Arc::new(InFlightCounter::default())) + .clone() + } + + fn cleanup_if_unused(&self, counter_key: &str) { + let active_refs = self.active_refs(counter_key); + if active_refs != 0 { + return; + } + + self.inner.counters.remove_if(counter_key, |_, counter| { + counter.in_flight() == 0 && Arc::strong_count(counter) == 1 + }); + self.inner + .active_refs + .remove_if(counter_key, |_, refs| refs.load(Ordering::Acquire) == 0); + } + + fn active_refs(&self, counter_key: &str) -> usize { + self.inner + .active_refs + .get(counter_key) + .map(|refs| refs.load(Ordering::Acquire)) + .unwrap_or(0) + } + + #[cfg(test)] + fn in_flight(&self, counter_key: &str) -> u32 { + self.inner + .counters + .get(counter_key) + .map(|counter| counter.in_flight()) + .unwrap_or(0) + } + + #[cfg(test)] + fn dropped_requests(&self, counter_key: &str) -> u64 { + self.inner + .counters + .get(counter_key) + .map(|counter| counter.dropped_requests()) + .unwrap_or(0) + } + + #[cfg(test)] + fn counter_count(&self) -> usize { + self.inner.counters.len() + } +} + +#[derive(Debug, Default)] +struct InFlightCounter { + in_flight: AtomicU32, + /// Local A32 drop accounting, kept with the global counter so future LRS + /// support can export `total_dropped_requests` without changing enforcement. + dropped_requests: AtomicU64, +} + +impl InFlightCounter { + fn try_acquire(&self, limit: u32) -> bool { + loop { + let current = self.in_flight.load(Ordering::Acquire); + if current >= limit { + return false; + } + + if self + .in_flight + .compare_exchange_weak(current, current + 1, Ordering::AcqRel, Ordering::Acquire) + .is_ok() + { + return true; + } + } + } + + fn in_flight(&self) -> u32 { + self.in_flight.load(Ordering::Acquire) + } + + fn record_drop(&self) { + self.dropped_requests.fetch_add(1, Ordering::AcqRel); + } + + #[cfg(test)] + fn dropped_requests(&self) -> u64 { + self.dropped_requests.load(Ordering::Acquire) + } +} + +#[derive(Debug)] +struct CircuitBreakerPermit { + counter: Option>, + counter_key: Arc, + counters: ClusterRequestCounters, +} + +impl Drop for CircuitBreakerPermit { + fn drop(&mut self) { + if let Some(counter) = self.counter.take() { + counter.in_flight.fetch_sub(1, Ordering::AcqRel); + } + self.counters.cleanup_if_unused(&self.counter_key); + } +} + +/// Tower layer that enforces A32 max in-flight requests per xDS cluster. +#[derive(Clone)] +pub(crate) struct CircuitBreakingLayer { + circuit_breakers: ClusterCircuitBreakers, +} + +impl CircuitBreakingLayer { + pub(crate) fn new(circuit_breakers: ClusterCircuitBreakers) -> Self { + Self { circuit_breakers } + } +} + +impl fmt::Debug for CircuitBreakingLayer { + fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { + f.debug_struct("CircuitBreakingLayer") + .field("circuit_breakers", &self.circuit_breakers) + .finish() + } +} + +impl Layer for CircuitBreakingLayer { + type Service = CircuitBreakingService; + + fn layer(&self, service: S) -> Self::Service { + CircuitBreakingService { + inner: service, + circuit_breakers: self.circuit_breakers.clone(), + } + } +} + +#[derive(Clone)] +pub(crate) struct CircuitBreakingService { + inner: S, + circuit_breakers: ClusterCircuitBreakers, +} + +impl fmt::Debug for CircuitBreakingService { + fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { + f.debug_struct("CircuitBreakingService") + .field("inner", &self.inner) + .field("circuit_breakers", &self.circuit_breakers) + .finish() + } +} + +impl Service> for CircuitBreakingService +where + S: Service, Response = Response, Error: Into> + + Clone + + Send + + 'static, + S::Future: Send + 'static, + B: Send + 'static, +{ + type Response = Response; + type Error = BoxError; + type Future = BoxFuture>; + + fn poll_ready(&mut self, _cx: &mut Context<'_>) -> Poll> { + Poll::Ready(Ok(())) + } + + fn call(&mut self, request: Request) -> Self::Future { + let Some(route_decision) = request.extensions().get::().cloned() else { + return Box::pin(async { + Ok(status_response(tonic::Status::internal( + CircuitBreakingError::NoRoutingDecision.to_string(), + ))) + }); + }; + + let cluster = route_decision.cluster; + let circuit_breakers = self.circuit_breakers.clone(); + let clone = self.inner.clone(); + let mut inner = std::mem::replace(&mut self.inner, clone); + Box::pin(async move { + let permit = match circuit_breakers.acquire(&cluster) { + Ok(permit) => permit, + Err(limit) => return Ok(limit_exceeded_response(&cluster, limit)), + }; + + std::future::poll_fn(|cx| inner.poll_ready(cx)) + .await + .map_err(Into::into)?; + let response = inner.call(request).await.map_err(Into::into)?; + Ok(response.map(|body| TonicBody::new(PermitBody::new(body, permit)))) + }) + } +} + +#[derive(Debug, Clone, thiserror::Error)] +enum CircuitBreakingError { + #[error("No routing decision extension from the routing layer available")] + NoRoutingDecision, +} + +fn limit_exceeded_response(cluster: &str, limit: CircuitBreakerLimit) -> Response { + status_response(tonic::Status::unavailable(format!( + "circuit breaker open for cluster '{cluster}': max_requests limit {} reached", + limit.max_requests, + ))) +} + +fn status_response(status: tonic::Status) -> Response { + status.into_http::() +} + +pin_project! { + #[derive(Debug)] + struct PermitBody { + #[pin] + inner: B, + permit: Option, + } +} + +impl PermitBody { + fn new(inner: B, permit: CircuitBreakerPermit) -> Self { + Self { + inner, + permit: Some(permit), + } + } +} + +impl Body for PermitBody +where + B: Body, +{ + type Data = Bytes; + type Error = B::Error; + + fn poll_frame( + self: std::pin::Pin<&mut Self>, + cx: &mut Context<'_>, + ) -> Poll, Self::Error>>> { + let mut this = self.project(); + match this.inner.as_mut().poll_frame(cx) { + Poll::Ready(None) => { + this.permit.take(); + Poll::Ready(None) + } + Poll::Ready(Some(Ok(frame))) => { + if frame.is_trailers() { + this.permit.take(); + } + Poll::Ready(Some(Ok(frame))) + } + Poll::Ready(Some(Err(err))) => { + this.permit.take(); + Poll::Ready(Some(Err(err))) + } + Poll::Pending => Poll::Pending, + } + } + + fn is_end_stream(&self) -> bool { + self.inner.is_end_stream() + } + + fn size_hint(&self) -> http_body::SizeHint { + self.inner.size_hint() + } +} + +#[cfg(test)] +mod tests { + use std::convert::Infallible; + use std::future::Future; + use std::pin::Pin; + use std::sync::Arc; + use std::sync::atomic::{AtomicU32, Ordering}; + use std::task::{Context, Poll}; + + use bytes::Bytes; + use http::{HeaderMap, Request, Response}; + use http_body::{Body, Frame}; + use tonic::Code; + use tower::Layer; + use tower::ServiceExt; + use tower::retry::Policy; + use tower::service_fn; + + use crate::client::retry::RetryLayer; + + use super::*; + + const CLUSTER: &str = "cluster-a"; + + fn request() -> Request { + let mut request = Request::new(TonicBody::empty()); + request.extensions_mut().insert(RouteDecision { + cluster: CLUSTER.to_string(), + request_hash: None, + }); + request + } + + fn configured_breakers(max_requests: u32) -> ClusterCircuitBreakers { + let breakers = ClusterCircuitBreakers::new_for_test(); + breakers.set_config(CLUSTER, CircuitBreakingConfig { max_requests }); + breakers + } + + #[tokio::test] + async fn rejects_requests_when_cluster_limit_is_reached() { + let breakers = configured_breakers(1); + let calls = Arc::new(AtomicU32::new(0)); + let call_counter = calls.clone(); + let service = service_fn(move |_request: Request| { + call_counter.fetch_add(1, Ordering::SeqCst); + async { Ok::<_, BoxError>(Response::new(TonicBody::new(PendingBody))) } + }); + let mut service = CircuitBreakingLayer::new(breakers.clone()).layer(service); + + let first = service + .ready() + .await + .unwrap() + .call(request()) + .await + .unwrap(); + assert_eq!(breakers.in_flight(CLUSTER), 1); + + let second = service + .ready() + .await + .unwrap() + .call(request()) + .await + .unwrap(); + let status = tonic::Status::from_header_map(second.headers()).unwrap(); + assert_eq!(status.code(), Code::Unavailable); + assert!(status.message().contains("max_requests limit 1")); + assert_eq!(calls.load(Ordering::SeqCst), 1); + assert_eq!(breakers.dropped_requests(CLUSTER), 1); + + drop(first); + assert_eq!(breakers.in_flight(CLUSTER), 0); + + let _third = service + .ready() + .await + .unwrap() + .call(request()) + .await + .unwrap(); + assert_eq!(calls.load(Ordering::SeqCst), 2); + } + + #[tokio::test] + async fn releases_permit_when_response_body_reaches_trailers() { + let breakers = configured_breakers(1); + let service = service_fn(|_request: Request| async { + Ok::<_, BoxError>(Response::new(TonicBody::new(DataThenTrailersBody { + state: BodyState::Data, + }))) + }); + let mut service = CircuitBreakingLayer::new(breakers.clone()).layer(service); + + let response = service + .ready() + .await + .unwrap() + .call(request()) + .await + .unwrap(); + assert_eq!(breakers.in_flight(CLUSTER), 1); + + let mut body = response.into_body(); + let data_frame = std::future::poll_fn(|cx| Pin::new(&mut body).poll_frame(cx)).await; + assert!(data_frame.unwrap().unwrap().is_data()); + assert_eq!(breakers.in_flight(CLUSTER), 1); + + let trailers_frame = std::future::poll_fn(|cx| Pin::new(&mut body).poll_frame(cx)).await; + assert!(trailers_frame.unwrap().unwrap().is_trailers()); + assert_eq!(breakers.in_flight(CLUSTER), 0); + } + + #[tokio::test] + async fn releases_permit_when_response_future_is_dropped() { + let breakers = configured_breakers(1); + let service = service_fn(|_request: Request| async { + std::future::pending::, BoxError>>().await + }); + let mut service = CircuitBreakingLayer::new(breakers.clone()).layer(service); + + let mut future = service.ready().await.unwrap().call(request()); + std::future::poll_fn(|cx| match Future::poll(Pin::new(&mut future), cx) { + Poll::Pending => Poll::Ready(()), + Poll::Ready(_) => panic!("inner service should remain pending"), + }) + .await; + assert_eq!(breakers.in_flight(CLUSTER), 1); + + drop(future); + assert_eq!(breakers.in_flight(CLUSTER), 0); + } + + #[tokio::test] + async fn reports_missing_route_decision_as_grpc_status() { + let service = service_fn(|_request: Request| async { + Ok::<_, BoxError>(Response::new(TonicBody::empty())) + }); + let mut service = + CircuitBreakingLayer::new(ClusterCircuitBreakers::new_for_test()).layer(service); + + let response = service + .ready() + .await + .unwrap() + .call(Request::new(TonicBody::empty())) + .await + .unwrap(); + let status = tonic::Status::from_header_map(response.headers()).unwrap(); + assert_eq!(status.code(), Code::Internal); + assert!(status.message().contains("No routing decision")); + } + + #[tokio::test] + async fn limit_responses_do_not_enter_retry_policy() { + let breakers = configured_breakers(1); + let retry_observations = Arc::new(AtomicU32::new(0)); + let policy = CountingUnavailablePolicy { + retry_observations: retry_observations.clone(), + }; + + let service = service_fn( + |_request: Request>| async { + Ok::<_, BoxError>(Response::new(TonicBody::new(PendingBody))) + }, + ); + let mut service = tower::ServiceBuilder::new() + .layer(CircuitBreakingLayer::new(breakers)) + .layer(RetryLayer::new(policy)) + .service(service); + + let _first = service + .ready() + .await + .unwrap() + .call(request()) + .await + .unwrap(); + let second = service + .ready() + .await + .unwrap() + .call(request()) + .await + .unwrap(); + let status = tonic::Status::from_header_map(second.headers()).unwrap(); + + assert_eq!(status.code(), Code::Unavailable); + assert_eq!(retry_observations.load(Ordering::SeqCst), 0); + } + + #[test] + fn eds_service_name_change_uses_independent_counter() { + let breakers = ClusterCircuitBreakers::new_for_test(); + breakers.set_cluster_config(CLUSTER, "eds-a", CircuitBreakingConfig { max_requests: 1 }); + let first = breakers.acquire(CLUSTER).unwrap(); + assert_eq!(breakers.in_flight(CLUSTER), 1); + + breakers.set_cluster_config(CLUSTER, "eds-b", CircuitBreakingConfig { max_requests: 1 }); + assert_eq!(breakers.in_flight(CLUSTER), 0); + let second = breakers.acquire(CLUSTER).unwrap(); + assert_eq!(breakers.in_flight(CLUSTER), 1); + assert!(breakers.acquire(CLUSTER).is_err()); + + drop(second); + assert_eq!(breakers.in_flight(CLUSTER), 0); + drop(first); + assert_eq!(breakers.counter_count(), 1); + } + + #[test] + fn cluster_removal_cleans_up_counter_after_in_flight_requests_finish() { + let breakers = configured_breakers(1); + let permit = breakers.acquire(CLUSTER).unwrap(); + assert_eq!(breakers.counter_count(), 1); + + breakers.clear_cluster_config(CLUSTER); + assert_eq!(breakers.counter_count(), 1); + + drop(permit); + assert_eq!(breakers.counter_count(), 0); + } + + #[test] + fn dropping_breakers_releases_config_counter_ref() { + let counters = ClusterRequestCounters::isolated(); + let breakers = ClusterCircuitBreakers::with_counters(counters.clone()); + breakers.set_config(CLUSTER, CircuitBreakingConfig { max_requests: 1 }); + let permit = breakers.acquire(CLUSTER).unwrap(); + drop(permit); + assert_eq!(counters.counter_count(), 1); + + drop(breakers); + + assert_eq!(counters.counter_count(), 0); + } + + #[test] + fn cleanup_keeps_counter_with_outstanding_clone() { + let counters = ClusterRequestCounters::isolated(); + let counter_key = counter_key(CLUSTER, CLUSTER); + let counter = counters.counter(&counter_key); + + counters.cleanup_if_unused(&counter_key); + assert_eq!(counters.counter_count(), 1); + + drop(counter); + counters.cleanup_if_unused(&counter_key); + assert_eq!(counters.counter_count(), 0); + } + + #[tokio::test] + async fn rejects_over_limit_without_waiting_for_inner_ready() { + let breakers = configured_breakers(1); + let calls = Arc::new(AtomicU32::new(0)); + let service = BackpressuredService { + ready_budget: Arc::new(AtomicU32::new(1)), + calls: calls.clone(), + }; + let mut service = CircuitBreakingLayer::new(breakers.clone()).layer(service); + + let first = service + .ready() + .await + .unwrap() + .call(request()) + .await + .unwrap(); + assert_eq!(breakers.in_flight(CLUSTER), 1); + assert_eq!(calls.load(Ordering::SeqCst), 1); + + let second = tokio::time::timeout( + tokio::time::Duration::from_millis(50), + service.ready().await.unwrap().call(request()), + ) + .await + .expect("over-limit request should not wait for inner readiness") + .unwrap(); + let status = tonic::Status::from_header_map(second.headers()).unwrap(); + assert_eq!(status.code(), Code::Unavailable); + assert_eq!(calls.load(Ordering::SeqCst), 1); + + drop(first); + } + + #[derive(Clone, Debug)] + struct CountingUnavailablePolicy { + retry_observations: Arc, + } + + impl Policy>, Response, BoxError> + for CountingUnavailablePolicy + { + type Future = std::future::Ready<()>; + + fn retry( + &mut self, + _req: &mut Request>, + result: &mut Result, BoxError>, + ) -> Option { + if let Ok(response) = result + && tonic::Status::from_header_map(response.headers()) + .is_some_and(|status| status.code() == Code::Unavailable) + { + self.retry_observations.fetch_add(1, Ordering::SeqCst); + } + None + } + + fn clone_request( + &mut self, + req: &Request>, + ) -> Option>> { + Some(req.clone()) + } + } + + #[derive(Clone, Debug)] + struct BackpressuredService { + ready_budget: Arc, + calls: Arc, + } + + impl Service> for BackpressuredService { + type Response = Response; + type Error = BoxError; + type Future = std::future::Ready>; + + fn poll_ready(&mut self, _cx: &mut Context<'_>) -> Poll> { + if self + .ready_budget + .fetch_update(Ordering::AcqRel, Ordering::Acquire, |remaining| { + remaining.checked_sub(1) + }) + .is_ok() + { + Poll::Ready(Ok(())) + } else { + Poll::Pending + } + } + + fn call(&mut self, _request: Request) -> Self::Future { + self.calls.fetch_add(1, Ordering::SeqCst); + std::future::ready(Ok(Response::new(TonicBody::new(PendingBody)))) + } + } + + #[derive(Debug)] + struct PendingBody; + + impl Body for PendingBody { + type Data = Bytes; + type Error = Infallible; + + fn poll_frame( + self: Pin<&mut Self>, + _cx: &mut Context<'_>, + ) -> Poll, Self::Error>>> { + Poll::Pending + } + } + + #[derive(Debug)] + enum BodyState { + Data, + Trailers, + Done, + } + + #[derive(Debug)] + struct DataThenTrailersBody { + state: BodyState, + } + + impl Body for DataThenTrailersBody { + type Data = Bytes; + type Error = Infallible; + + fn poll_frame( + mut self: Pin<&mut Self>, + _cx: &mut Context<'_>, + ) -> Poll, Self::Error>>> { + match self.state { + BodyState::Data => { + self.state = BodyState::Trailers; + Poll::Ready(Some(Ok(Frame::data(Bytes::from_static(b"hello"))))) + } + BodyState::Trailers => { + self.state = BodyState::Done; + Poll::Ready(Some(Ok(Frame::trailers(HeaderMap::new())))) + } + BodyState::Done => Poll::Ready(None), + } + } + } +} diff --git a/tonic-xds/src/client/mod.rs b/tonic-xds/src/client/mod.rs index 3e02c9b29..d5a7e5a01 100644 --- a/tonic-xds/src/client/mod.rs +++ b/tonic-xds/src/client/mod.rs @@ -1,4 +1,5 @@ pub(crate) mod channel; +pub(crate) mod circuit_breaking; pub(crate) mod cluster; pub(crate) mod endpoint; pub(crate) mod lb; diff --git a/tonic-xds/src/xds/resource/circuit_breaking.rs b/tonic-xds/src/xds/resource/circuit_breaking.rs index 28c7200bc..6cf9934b2 100644 --- a/tonic-xds/src/xds/resource/circuit_breaking.rs +++ b/tonic-xds/src/xds/resource/circuit_breaking.rs @@ -5,9 +5,9 @@ //! because they are connection-pool or retry specific and do not apply to gRPC's //! A32 request limiter. //! -//! This parser intentionally stays detached from `ClusterResource` until -//! enforcement lands; otherwise cluster validation would advertise support before -//! requests are actually limited. +//! Client-side limiter primitives can consume this config; production CDS wiring +//! is added separately so validation only advertises support once enforcement is +//! in the request path. //! //! [gRFC A32]: https://github.com/grpc/proposal/blob/master/A32-xds-circuit-breaking.md @@ -22,9 +22,8 @@ pub(crate) const DEFAULT_MAX_REQUESTS: u32 = 1024; pub(crate) struct CircuitBreakingConfig { /// Maximum number of in-flight requests allowed for the upstream cluster. /// - /// This scaffolds the parsed CDS value only; enforcement is wired in a - /// follow-up change so request-lifetime accounting can be handled correctly - /// for streaming RPCs. + /// The client-side limiter holds requests against this threshold for the + /// full response-body lifetime. pub(crate) max_requests: u32, } From 01c865f8e5c798d9572f62ca8bbbff2b6d5fd5cf Mon Sep 17 00:00:00 2001 From: Michael Ingley Date: Tue, 30 Jun 2026 11:56:31 -0500 Subject: [PATCH 2/6] Wire A32 circuit breaking into xDS channels --- tonic-xds/src/client/channel.rs | 185 ++++++++- tonic-xds/src/client/circuit_breaking.rs | 360 +++++++++++++++++- .../src/client/loadbalance/loadbalancer.rs | 44 ++- tonic-xds/src/client/retry.rs | 6 + tonic-xds/src/xds/cache.rs | 3 + tonic-xds/src/xds/cluster_discovery.rs | 2 + .../src/xds/resource/circuit_breaking.rs | 4 +- tonic-xds/src/xds/resource/cluster.rs | 34 ++ 8 files changed, 609 insertions(+), 29 deletions(-) diff --git a/tonic-xds/src/client/channel.rs b/tonic-xds/src/client/channel.rs index 5132d128a..f9e9d6121 100644 --- a/tonic-xds/src/client/channel.rs +++ b/tonic-xds/src/client/channel.rs @@ -1,4 +1,5 @@ use crate::XdsUri; +use crate::client::circuit_breaking::{CircuitBreakingClusterDiscovery, ClusterCircuitBreakers}; use crate::client::cluster::ClusterClientRegistryGrpc; use crate::client::endpoint::{EndpointAddress, EndpointChannel}; use crate::client::lb::{ClusterDiscovery, XdsLbService}; @@ -227,14 +228,18 @@ impl XdsChannelBuilder { resource_manager: XdsResourceManager, ) -> XdsChannelGrpc { let router: Arc = Arc::new(XdsRouter::new(&cache)); + let circuit_breakers = ClusterCircuitBreakers::default(); #[cfg(feature = "_tls-any")] let discovery: Arc< dyn ClusterDiscovery>, - > = Arc::new(XdsClusterDiscovery::new(cache, cert_provider_registry)); + > = Arc::new(XdsClusterDiscovery::new( + cache.clone(), + cert_provider_registry, + )); #[cfg(not(feature = "_tls-any"))] let discovery: Arc< dyn ClusterDiscovery>, - > = Arc::new(XdsClusterDiscovery::new(cache)); + > = Arc::new(XdsClusterDiscovery::new(cache.clone())); let retry_policy = GrpcRetryPolicy::new(GrpcRetryPolicyConfig::default()); let resources = Arc::new(XdsChannelResources { @@ -243,6 +248,10 @@ impl XdsChannelBuilder { }); let routing_layer = XdsRoutingLayer::new(router, self.authority()); + let discovery = Arc::new( + CircuitBreakingClusterDiscovery::new(discovery, circuit_breakers) + .with_cluster_cache(cache), + ); let retry_layer = RetryLayer::new(retry_policy); let cluster_registry = Arc::new(ClusterClientRegistryGrpc::new()); let lb_service = XdsLbService::new(cluster_registry, discovery); @@ -275,8 +284,28 @@ impl XdsChannelBuilder { router: Arc, discovery: Arc>>, retry_policy: GrpcRetryPolicy, + ) -> XdsChannelGrpc { + self.build_grpc_channel_from_parts_with_circuit_breakers( + router, + discovery, + retry_policy, + ClusterCircuitBreakers::new_for_test(), + ) + } + + #[cfg(test)] + pub(crate) fn build_grpc_channel_from_parts_with_circuit_breakers( + &self, + router: Arc, + discovery: Arc>>, + retry_policy: GrpcRetryPolicy, + circuit_breakers: ClusterCircuitBreakers, ) -> XdsChannelGrpc { let routing_layer = XdsRoutingLayer::new(router, self.authority()); + let discovery = Arc::new(CircuitBreakingClusterDiscovery::new( + discovery, + circuit_breakers, + )); let retry_layer = RetryLayer::new(retry_policy); let cluster_registry = Arc::new(ClusterClientRegistryGrpc::new()); let lb_service = XdsLbService::new(cluster_registry, discovery); @@ -306,6 +335,7 @@ mod tests { use super::{XdsChannelBuilder, XdsChannelConfig}; use crate::XdsUri; use crate::client::channel::XdsChannelGrpc; + use crate::client::circuit_breaking::ClusterCircuitBreakers; use crate::client::endpoint::EndpointAddress; use crate::client::endpoint::EndpointChannel; @@ -323,6 +353,7 @@ mod tests { use crate::testutil::grpc::TestServer; use crate::xds::cache::XdsCache; use crate::xds::resource::EndpointsResource; + use crate::xds::resource::circuit_breaking::CircuitBreakingConfig; use crate::xds::resource::route_config::RouteConfigResource; use std::sync::Arc; use tokio::sync::mpsc; @@ -554,16 +585,140 @@ mod tests { assert_eq!(response.into_inner().message, "retry-server: retry-test"); } + #[tokio::test] + async fn test_xds_channel_enforces_injected_circuit_breaking_limit() { + use crate::client::retry::{GrpcRetryPolicy, GrpcRetryPolicyConfig}; + + let (_, servers) = setup_grpc_servers(1).await; + let xds_manager = Arc::new(MockXdsManager::from_test_servers(&servers)); + let circuit_breakers = ClusterCircuitBreakers::new_for_test(); + circuit_breakers.set_config("test-cluster", CircuitBreakingConfig { max_requests: 0 }); + + let retry_policy = GrpcRetryPolicy::new( + GrpcRetryPolicyConfig::new() + .retry_on(vec![tonic::Code::Unavailable]) + .num_retries(1), + ); + let xds_channel = XdsChannelBuilder::new(test_config()) + .build_grpc_channel_from_parts_with_circuit_breakers( + xds_manager.clone(), + xds_manager.clone(), + retry_policy, + circuit_breakers, + ); + let mut client = GreeterClient::new(xds_channel); + + let error = client + .say_hello(HelloRequest { + name: "limited".to_string(), + }) + .await + .unwrap_err(); + + assert_eq!(error.code(), tonic::Code::Unavailable); + assert!(error.message().contains("max_requests limit 0")); + + for server in servers { + let _ = server.shutdown.send(()); + let _ = server.handle.await; + } + } + + #[tokio::test] + async fn test_xds_channel_uses_cds_circuit_breaking_config() { + let cluster_name = "test-cluster"; + let (_, servers) = setup_grpc_servers(1).await; + + let cache = Arc::new(XdsCache::new()); + cache.update_route_config(make_test_route_config(cluster_name)); + cache.update_cluster( + cluster_name, + make_test_cluster_with_circuit_breaking( + cluster_name, + CircuitBreakingConfig { max_requests: 0 }, + ), + ); + cache.update_endpoints(cluster_name, make_test_endpoints(cluster_name, &servers)); + + let channel = build_xds_channel_from_cache(cache).await; + let mut client = GreeterClient::new(channel); + let error = client + .say_hello(HelloRequest { + name: "cds-limited".to_string(), + }) + .await + .unwrap_err(); + + assert_eq!(error.code(), tonic::Code::Unavailable); + assert!(error.message().contains("max_requests limit 0")); + + for server in servers { + let _ = server.shutdown.send(()); + let _ = server.handle.await; + } + } + + #[tokio::test] + async fn test_xds_channel_waits_for_cds_before_circuit_breaking() { + let cluster_name = "test-cluster"; + let (_, servers) = setup_grpc_servers(1).await; + + let cache = Arc::new(XdsCache::new()); + cache.update_route_config(make_test_route_config(cluster_name)); + cache.update_endpoints(cluster_name, make_test_endpoints(cluster_name, &servers)); + + let channel = build_xds_channel_from_cache(cache.clone()).await; + let mut client = GreeterClient::new(channel); + let mut request = Box::pin(client.say_hello(HelloRequest { + name: "wait-for-cds".to_string(), + })); + + let early = + tokio::time::timeout(tokio::time::Duration::from_millis(20), request.as_mut()).await; + assert!( + early.is_err(), + "request should wait for CDS before acquiring a circuit-breaking permit", + ); + + cache.update_cluster( + cluster_name, + make_test_cluster_with_circuit_breaking( + cluster_name, + CircuitBreakingConfig { max_requests: 0 }, + ), + ); + + let result = tokio::time::timeout(tokio::time::Duration::from_secs(2), request) + .await + .expect("request should complete after CDS update"); + let error = result.unwrap_err(); + assert_eq!(error.code(), tonic::Code::Unavailable); + assert!(error.message().contains("max_requests limit 0")); + + for server in servers { + let _ = server.shutdown.send(()); + let _ = server.handle.await; + } + } + /// Helper: creates a minimal plaintext `ClusterResource` for tests that /// drive `XdsClusterDiscovery`. The cluster watch in `discover_cluster` /// blocks until a cluster is in the cache. fn make_test_cluster(cluster_name: &str) -> Arc { + make_test_cluster_with_circuit_breaking(cluster_name, CircuitBreakingConfig::default()) + } + + fn make_test_cluster_with_circuit_breaking( + cluster_name: &str, + circuit_breaking: CircuitBreakingConfig, + ) -> Arc { use crate::xds::resource::cluster::{ClusterResource, LbPolicy}; Arc::new(ClusterResource { name: cluster_name.to_string(), eds_service_name: None, lb_policy: LbPolicy::RoundRobin, security: None, + circuit_breaking, }) } @@ -614,27 +769,23 @@ mod tests { /// Builds an XdsChannelGrpc using real XdsRouter and XdsClusterDiscovery /// backed by the given cache. async fn build_xds_channel_from_cache(cache: Arc) -> XdsChannelGrpc { - use crate::xds::cluster_discovery::XdsClusterDiscovery; - use crate::xds::routing::XdsRouter; - - let router: Arc = Arc::new(XdsRouter::new(&cache)); + use crate::xds::resource_manager::XdsResourceManager; + let xds_client = xds_client::XdsClient::disconnected(); + let resource_manager = + XdsResourceManager::new(xds_client.clone(), cache.clone(), "test-listener".into()); + let builder = XdsChannelBuilder::new(test_config()); #[cfg(feature = "_tls-any")] - let discovery: Arc< - dyn ClusterDiscovery>, - > = { + { use crate::xds::cert_provider::CertProviderRegistry; let registry = Arc::new(CertProviderRegistry::from_bootstrap(&Default::default()).unwrap()); - Arc::new(XdsClusterDiscovery::new(cache, registry)) - }; + builder.build_from_cache(cache, registry, xds_client, resource_manager) + } #[cfg(not(feature = "_tls-any"))] - let discovery: Arc< - dyn ClusterDiscovery>, - > = Arc::new(XdsClusterDiscovery::new(cache)); - - let builder = XdsChannelBuilder::new(test_config()); - builder.build_grpc_channel_from_parts(router, discovery, GrpcRetryPolicy::default()) + { + builder.build_from_cache(cache, xds_client, resource_manager) + } } /// Tests the full xDS stack (XdsRouter + XdsClusterDiscovery) with a diff --git a/tonic-xds/src/client/circuit_breaking.rs b/tonic-xds/src/client/circuit_breaking.rs index 4e71b742e..68dafc52d 100644 --- a/tonic-xds/src/client/circuit_breaking.rs +++ b/tonic-xds/src/client/circuit_breaking.rs @@ -1,27 +1,43 @@ -#![cfg_attr(not(test), allow(dead_code))] - use std::fmt; use std::sync::{ Arc, OnceLock, - atomic::{AtomicU32, AtomicU64, AtomicUsize, Ordering}, + atomic::{AtomicBool, AtomicU32, AtomicU64, AtomicUsize, Ordering}, }; use std::task::{Context, Poll}; use bytes::Bytes; use dashmap::DashMap; +use futures_util::StreamExt as _; use http::{Request, Response}; use http_body::{Body, Frame}; use pin_project_lite::pin_project; use tokio::sync::watch; use tonic::body::Body as TonicBody; -use tower::{BoxError, Layer, Service}; +#[cfg(test)] +use tower::Layer; +use tower::discover::Change; +use tower::load::Load; +use tower::{BoxError, Service}; +use crate::client::lb::{BoxDiscover, ClusterDiscovery}; +#[cfg(test)] use crate::client::route::RouteDecision; use crate::common::async_util::BoxFuture; +use crate::xds::cache::XdsCache; +use crate::xds::resource::ClusterResource; use crate::xds::resource::circuit_breaking::{CircuitBreakingConfig, DEFAULT_MAX_REQUESTS}; static GLOBAL_COUNTERS: OnceLock> = OnceLock::new(); +/// Response marker for A32 drops. Retry policy checks this marker so +/// circuit-broken `UNAVAILABLE` responses are never retried. +#[derive(Debug, Clone, Copy)] +pub(crate) struct CircuitBreakerDrop; + +pub(crate) fn is_circuit_breaker_drop(response: &Response) -> bool { + response.extensions().get::().is_some() +} + /// Shared circuit-breaking state for xDS clusters. #[derive(Clone, Debug)] pub(crate) struct ClusterCircuitBreakers { @@ -89,6 +105,39 @@ impl ClusterCircuitBreakers { ); } + fn ensure_cluster_watch(&self, cache: Arc, cluster: &str) { + let state = self.ensure_state(cluster); + if state.watch_started.swap(true, Ordering::AcqRel) { + return; + } + + let cluster = cluster.to_string(); + let circuit_breakers = self.clone(); + tokio::spawn(async move { + let mut cluster_watch = cache.watch_cluster(&cluster); + while let Some(cluster_resource) = cluster_watch.next().await { + let config = CircuitBreakerRuntimeConfig::from_cluster( + &cluster_resource, + &circuit_breakers.inner.counters, + ); + circuit_breakers.update_state_config(&state, config); + } + circuit_breakers.remove_cluster_state(&cluster, &state); + }); + } + + pub(crate) fn cluster_breaker( + &self, + cluster: impl Into, + cluster_cache: Option>, + ) -> ClusterCircuitBreaker { + ClusterCircuitBreaker { + cluster: Arc::from(cluster.into()), + circuit_breakers: self.clone(), + cluster_cache, + } + } + fn ensure_state(&self, cluster: &str) -> Arc { if let Some(state) = self.inner.configs.get(cluster) { return state.clone(); @@ -124,6 +173,16 @@ impl ClusterCircuitBreakers { } } + fn remove_cluster_state(&self, cluster: &str, state: &Arc) { + if let Some((_, state)) = self + .inner + .configs + .remove_if(cluster, |_, current| Arc::ptr_eq(current, state)) + { + self.clear_state(&state); + } + } + fn clear_state(&self, state: &ClusterCircuitBreakerState) { if let Some(previous) = state.config_tx.send_replace(None) { self.deactivate_config(previous); @@ -140,6 +199,14 @@ impl ClusterCircuitBreakers { self.acquire_with_config(self.runtime_config_or_default(cluster)) } + async fn acquire_when_ready( + &self, + cluster: &str, + ) -> Option> { + let runtime_config = self.wait_for_config(cluster).await?; + Some(self.acquire_with_config(runtime_config)) + } + fn acquire_with_config( &self, runtime_config: CircuitBreakerRuntimeConfig, @@ -157,6 +224,24 @@ impl ClusterCircuitBreakers { .ok_or(limit) } + async fn wait_for_config(&self, cluster: &str) -> Option { + let state = self.ensure_state(cluster); + if let Some(config) = state.current_config() { + return Some(config); + } + + let mut rx = state.config_tx.subscribe(); + loop { + if rx.changed().await.is_err() { + return None; + } + + if let Some(config) = rx.borrow_and_update().clone() { + return Some(config); + } + } + } + fn runtime_config_or_default(&self, cluster: &str) -> CircuitBreakerRuntimeConfig { self.inner .configs @@ -231,6 +316,18 @@ impl PartialEq for CircuitBreakerRuntimeConfig { impl Eq for CircuitBreakerRuntimeConfig {} +impl CircuitBreakerRuntimeConfig { + fn from_cluster(cluster: &ClusterResource, counters: &ClusterRequestCounters) -> Self { + let counter_key = counter_key(&cluster.name, cluster.eds_service_name()); + let counter = counters.counter(&counter_key); + Self { + max_requests: cluster.circuit_breaking.max_requests, + counter_key: Arc::from(counter_key), + counter, + } + } +} + fn counter_key(cluster: &str, eds_service_name: &str) -> String { format!("{cluster}\0{eds_service_name}") } @@ -238,12 +335,16 @@ fn counter_key(cluster: &str, eds_service_name: &str) -> String { #[derive(Debug)] struct ClusterCircuitBreakerState { config_tx: watch::Sender>, + watch_started: AtomicBool, } impl ClusterCircuitBreakerState { fn new() -> Self { let (config_tx, _) = watch::channel(None); - Self { config_tx } + Self { + config_tx, + watch_started: AtomicBool::new(false), + } } fn current_config(&self) -> Option { @@ -419,7 +520,7 @@ impl InFlightCounter { } #[derive(Debug)] -struct CircuitBreakerPermit { +pub(crate) struct CircuitBreakerPermit { counter: Option>, counter_key: Arc, counters: ClusterRequestCounters, @@ -434,18 +535,192 @@ impl Drop for CircuitBreakerPermit { } } +#[derive(Clone)] +pub(crate) struct CircuitBreakingClusterDiscovery { + inner: Arc>, + circuit_breakers: ClusterCircuitBreakers, + cluster_cache: Option>, +} + +impl CircuitBreakingClusterDiscovery { + pub(crate) fn new( + inner: Arc>, + circuit_breakers: ClusterCircuitBreakers, + ) -> Self { + Self { + inner, + circuit_breakers, + cluster_cache: None, + } + } + + pub(crate) fn with_cluster_cache(mut self, cache: Arc) -> Self { + self.cluster_cache = Some(cache); + self + } +} + +impl fmt::Debug for CircuitBreakingClusterDiscovery { + fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { + f.debug_struct("CircuitBreakingClusterDiscovery") + .field("circuit_breakers", &self.circuit_breakers) + .field("watching_cluster_cache", &self.cluster_cache.is_some()) + .finish() + } +} + +impl ClusterDiscovery> + for CircuitBreakingClusterDiscovery +where + Endpoint: Send + 'static, + S: Send + 'static, +{ + fn discover_cluster( + &self, + cluster_name: &str, + ) -> BoxDiscover> { + let breaker = self + .circuit_breakers + .cluster_breaker(cluster_name.to_string(), self.cluster_cache.clone()); + Box::pin( + self.inner + .discover_cluster(cluster_name) + .map(move |change| { + change.map(|change| match change { + Change::Insert(endpoint, service) => Change::Insert( + endpoint, + CircuitBreakingEndpointService::new(service, breaker.clone()), + ), + Change::Remove(endpoint) => Change::Remove(endpoint), + }) + }), + ) + } +} + +#[derive(Clone)] +pub(crate) struct ClusterCircuitBreaker { + cluster: Arc, + circuit_breakers: ClusterCircuitBreakers, + cluster_cache: Option>, +} + +impl fmt::Debug for ClusterCircuitBreaker { + fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { + f.debug_struct("ClusterCircuitBreaker") + .field("cluster", &self.cluster) + .field("circuit_breakers", &self.circuit_breakers) + .field("watching_cluster_cache", &self.cluster_cache.is_some()) + .finish() + } +} + +impl ClusterCircuitBreaker { + async fn acquire(&self) -> Result> { + if let Some(cache) = self.cluster_cache.clone() { + self.circuit_breakers + .ensure_cluster_watch(cache, &self.cluster); + match self + .circuit_breakers + .acquire_when_ready(&self.cluster) + .await + { + Some(Ok(permit)) => Ok(permit), + Some(Err(limit)) => Err(limit_exceeded_response(&self.cluster, limit)), + None => Err(cluster_unavailable_response(&self.cluster)), + } + } else { + self.circuit_breakers + .acquire(&self.cluster) + .map_err(|limit| limit_exceeded_response(&self.cluster, limit)) + } + } +} + +#[derive(Clone)] +pub(crate) struct CircuitBreakingEndpointService { + inner: S, + breaker: ClusterCircuitBreaker, +} + +impl CircuitBreakingEndpointService { + fn new(inner: S, breaker: ClusterCircuitBreaker) -> Self { + Self { inner, breaker } + } +} + +impl fmt::Debug for CircuitBreakingEndpointService { + fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { + f.debug_struct("CircuitBreakingEndpointService") + .field("inner", &self.inner) + .field("breaker", &self.breaker) + .finish() + } +} + +impl Service> for CircuitBreakingEndpointService +where + S: Service, Response = Response, Error: Into> + + Clone + + Send + + 'static, + S::Future: Send + 'static, + B: Send + 'static, +{ + type Response = Response; + type Error = BoxError; + type Future = BoxFuture>; + + fn poll_ready(&mut self, cx: &mut Context<'_>) -> Poll> { + self.inner.poll_ready(cx).map_err(Into::into) + } + + fn call(&mut self, request: Request) -> Self::Future { + let breaker = self.breaker.clone(); + let clone = self.inner.clone(); + let mut inner = std::mem::replace(&mut self.inner, clone); + Box::pin(async move { + let permit = match breaker.acquire().await { + Ok(permit) => permit, + Err(response) => return Ok(response), + }; + + let response = inner.call(request).await.map_err(Into::into)?; + Ok(hold_permit(response, permit)) + }) + } +} + +impl Load for CircuitBreakingEndpointService { + type Metric = S::Metric; + + fn load(&self) -> Self::Metric { + self.inner.load() + } +} + +pub(crate) fn hold_permit( + response: Response, + permit: CircuitBreakerPermit, +) -> Response { + response.map(|body| TonicBody::new(PermitBody::new(body, permit))) +} + /// Tower layer that enforces A32 max in-flight requests per xDS cluster. +#[cfg(test)] #[derive(Clone)] pub(crate) struct CircuitBreakingLayer { circuit_breakers: ClusterCircuitBreakers, } +#[cfg(test)] impl CircuitBreakingLayer { pub(crate) fn new(circuit_breakers: ClusterCircuitBreakers) -> Self { Self { circuit_breakers } } } +#[cfg(test)] impl fmt::Debug for CircuitBreakingLayer { fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { f.debug_struct("CircuitBreakingLayer") @@ -454,6 +729,7 @@ impl fmt::Debug for CircuitBreakingLayer { } } +#[cfg(test)] impl Layer for CircuitBreakingLayer { type Service = CircuitBreakingService; @@ -466,11 +742,13 @@ impl Layer for CircuitBreakingLayer { } #[derive(Clone)] +#[cfg(test)] pub(crate) struct CircuitBreakingService { inner: S, circuit_breakers: ClusterCircuitBreakers, } +#[cfg(test)] impl fmt::Debug for CircuitBreakingService { fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { f.debug_struct("CircuitBreakingService") @@ -480,6 +758,7 @@ impl fmt::Debug for CircuitBreakingService { } } +#[cfg(test)] impl Service> for CircuitBreakingService where S: Service, Response = Response, Error: Into> @@ -526,15 +805,24 @@ where } #[derive(Debug, Clone, thiserror::Error)] +#[cfg(test)] enum CircuitBreakingError { #[error("No routing decision extension from the routing layer available")] NoRoutingDecision, } fn limit_exceeded_response(cluster: &str, limit: CircuitBreakerLimit) -> Response { - status_response(tonic::Status::unavailable(format!( + let mut response = status_response(tonic::Status::unavailable(format!( "circuit breaker open for cluster '{cluster}': max_requests limit {} reached", limit.max_requests, + ))); + response.extensions_mut().insert(CircuitBreakerDrop); + response +} + +fn cluster_unavailable_response(cluster: &str) -> Response { + status_response(tonic::Status::unavailable(format!( + "cluster '{cluster}' is no longer available", ))) } @@ -792,6 +1080,64 @@ mod tests { assert_eq!(retry_observations.load(Ordering::SeqCst), 0); } + #[tokio::test] + async fn endpoint_waiting_for_ready_does_not_hold_permit() { + let breakers = configured_breakers(1); + let service = BackpressuredService { + ready_budget: Arc::new(AtomicU32::new(0)), + calls: Arc::new(AtomicU32::new(0)), + }; + let mut service = + CircuitBreakingEndpointService::new(service, breakers.cluster_breaker(CLUSTER, None)); + + let early = + tokio::time::timeout(tokio::time::Duration::from_millis(20), service.ready()).await; + + assert!( + early.is_err(), + "endpoint wrapper should wait for inner endpoint readiness", + ); + assert_eq!(breakers.in_flight(CLUSTER), 0); + } + + #[tokio::test] + async fn endpoint_limit_responses_are_not_retried() { + use crate::client::retry::{GrpcRetryPolicy, GrpcRetryPolicyConfig}; + + let breakers = configured_breakers(0); + let calls = Arc::new(AtomicU32::new(0)); + let call_counter = calls.clone(); + let service = service_fn( + move |_request: Request>| { + call_counter.fetch_add(1, Ordering::SeqCst); + async { Ok::<_, BoxError>(Response::new(TonicBody::new(PendingBody))) } + }, + ); + let service = + CircuitBreakingEndpointService::new(service, breakers.cluster_breaker(CLUSTER, None)); + let retry_policy = GrpcRetryPolicy::new( + GrpcRetryPolicyConfig::new() + .retry_on(vec![tonic::Code::Unavailable]) + .num_retries(1), + ); + let mut service = tower::ServiceBuilder::new() + .layer(RetryLayer::new(retry_policy)) + .service(service); + + let response = service + .ready() + .await + .unwrap() + .call(Request::new(TonicBody::empty())) + .await + .unwrap(); + let status = tonic::Status::from_header_map(response.headers()).unwrap(); + + assert_eq!(status.code(), Code::Unavailable); + assert_eq!(calls.load(Ordering::SeqCst), 0); + assert_eq!(breakers.dropped_requests(CLUSTER), 1); + } + #[test] fn eds_service_name_change_uses_independent_counter() { let breakers = ClusterCircuitBreakers::new_for_test(); diff --git a/tonic-xds/src/client/loadbalance/loadbalancer.rs b/tonic-xds/src/client/loadbalance/loadbalancer.rs index 556407619..13eee61d0 100644 --- a/tonic-xds/src/client/loadbalance/loadbalancer.rs +++ b/tonic-xds/src/client/loadbalance/loadbalancer.rs @@ -38,6 +38,7 @@ use tower::discover::{Change, Discover}; use arc_swap::ArcSwap; +use crate::client::circuit_breaking::is_circuit_breaker_drop; use crate::client::endpoint::{Connector, EndpointAddress}; use crate::client::loadbalance::channel_state::{ EjectionConfig, IdleChannel, ReadyChannel, UnejectedChannel, @@ -48,6 +49,26 @@ use crate::client::loadbalance::outlier_detection::{OutlierDetector, OutlierStat use crate::client::loadbalance::pickers::ChannelPicker; use crate::xds::resource::outlier_detection::OutlierDetectionConfig; +trait LbResponseOutcome { + fn is_circuit_breaker_drop(&self) -> bool { + false + } +} + +impl LbResponseOutcome for http::Response { + fn is_circuit_breaker_drop(&self) -> bool { + is_circuit_breaker_drop(self) + } +} + +fn outlier_outcome(result: &Result) -> Option { + match result { + Ok(response) if response.is_circuit_breaker_drop() => None, + Ok(_) => Some(true), + Err(_) => Some(false), + } +} + /// Future returned by [`LoadBalancer::call`]. Either resolves /// immediately with an [`LbError`] or drives the selected channel. pub(crate) enum LbFuture { @@ -316,7 +337,7 @@ where D::Error: Into, C: Connector + Send + Sync + 'static, C::Service: Service + Clone + Send + 'static, - >::Response: Send + 'static, + >::Response: LbResponseOutcome + Send + 'static, >::Error: Into, >::Future: Send + 'static, Req: Send + 'static, @@ -369,7 +390,9 @@ where .await .map_err(|e| LbError::LbChannelPollReadyError(e.into()))?; let result = svc.call(req).await; - svc.record_outcome(result.is_ok()); + if let Some(success) = outlier_outcome(&result) { + svc.record_outcome(success); + } result.map_err(|e| LbError::LbChannelCallError(e.into())) })) } @@ -378,6 +401,7 @@ where #[cfg(test)] mod tests { use super::*; + use crate::client::circuit_breaking::CircuitBreakerDrop; use crate::client::endpoint::Connector; use crate::client::loadbalance::pickers::p2c::P2cPicker; use crate::common::async_util::BoxFuture; @@ -443,6 +467,8 @@ mod tests { } } + impl LbResponseOutcome for &'static str {} + // -- Mock connector -- /// A connector that returns a pending future until signaled via oneshot. @@ -587,6 +613,20 @@ mod tests { Ok(Change::Remove(addr(port))) } + #[test] + fn circuit_breaker_drops_are_not_endpoint_outlier_outcomes() { + let normal: Result, tower::BoxError> = Ok(http::Response::new(())); + assert_eq!(outlier_outcome(&normal), Some(true)); + + let mut dropped_response = http::Response::new(()); + dropped_response.extensions_mut().insert(CircuitBreakerDrop); + let dropped: Result, tower::BoxError> = Ok(dropped_response); + assert_eq!(outlier_outcome(&dropped), None); + + let error: Result, tower::BoxError> = Err("endpoint error".into()); + assert_eq!(outlier_outcome(&error), Some(false)); + } + /// A burst of inserts drained in one poll rebuilds the ring exactly once, /// with the full member set (gRFC A42 ring-hash relies on this). A later /// drain (the remove) is a distinct, single rebuild. diff --git a/tonic-xds/src/client/retry.rs b/tonic-xds/src/client/retry.rs index c8f7c8c6d..194fd5cc9 100644 --- a/tonic-xds/src/client/retry.rs +++ b/tonic-xds/src/client/retry.rs @@ -16,6 +16,8 @@ use tower::retry::Policy; use tower::retry::Retry; use tower::{Layer, Service}; +use crate::client::circuit_breaking::is_circuit_breaker_drop; + /// Check if an error's source chain contains a retryable connection-level error. /// /// These are errors where the request was definitely **not** sent, making it safe to retry. @@ -57,6 +59,10 @@ pub(crate) fn is_retryable( match result { Err(err) => is_retryable_connection_error(err.as_ref()), Ok(response) => { + if is_circuit_breaker_drop(response) { + return false; + } + let status = tonic::Status::from_header_map(response.headers()); match status { Some(status) => is_retryable_grpc_status_code(status.code(), &policy.retry_on), diff --git a/tonic-xds/src/xds/cache.rs b/tonic-xds/src/xds/cache.rs index 14773324f..5846440a6 100644 --- a/tonic-xds/src/xds/cache.rs +++ b/tonic-xds/src/xds/cache.rs @@ -190,11 +190,14 @@ mod tests { } fn make_cluster(name: &str, lb: LbPolicy) -> Arc { + use crate::xds::resource::circuit_breaking::CircuitBreakingConfig; + Arc::new(ClusterResource { name: name.to_string(), eds_service_name: None, lb_policy: lb, security: None, + circuit_breaking: CircuitBreakingConfig::default(), }) } diff --git a/tonic-xds/src/xds/cluster_discovery.rs b/tonic-xds/src/xds/cluster_discovery.rs index 4c9e232a0..b59ef5aae 100644 --- a/tonic-xds/src/xds/cluster_discovery.rs +++ b/tonic-xds/src/xds/cluster_discovery.rs @@ -311,6 +311,7 @@ impl Connector for TlsConnector { #[cfg(test)] mod tests { use super::*; + use crate::xds::resource::circuit_breaking::CircuitBreakingConfig; use crate::xds::resource::cluster::{ClusterResource, LbPolicy}; fn plaintext_cluster() -> ClusterResource { @@ -319,6 +320,7 @@ mod tests { eds_service_name: None, lb_policy: LbPolicy::RoundRobin, security: None, + circuit_breaking: CircuitBreakingConfig::default(), } } diff --git a/tonic-xds/src/xds/resource/circuit_breaking.rs b/tonic-xds/src/xds/resource/circuit_breaking.rs index 6cf9934b2..db9bf8c23 100644 --- a/tonic-xds/src/xds/resource/circuit_breaking.rs +++ b/tonic-xds/src/xds/resource/circuit_breaking.rs @@ -5,9 +5,7 @@ //! because they are connection-pool or retry specific and do not apply to gRPC's //! A32 request limiter. //! -//! Client-side limiter primitives can consume this config; production CDS wiring -//! is added separately so validation only advertises support once enforcement is -//! in the request path. +//! `ClusterResource` carries this config to the client-side limiter. //! //! [gRFC A32]: https://github.com/grpc/proposal/blob/master/A32-xds-circuit-breaking.md diff --git a/tonic-xds/src/xds/resource/cluster.rs b/tonic-xds/src/xds/resource/cluster.rs index a21619e32..7d672a6d1 100644 --- a/tonic-xds/src/xds/resource/cluster.rs +++ b/tonic-xds/src/xds/resource/cluster.rs @@ -6,6 +6,7 @@ use prost::Message; use xds_client::resource::TypeUrl; use xds_client::{Error, Resource}; +use super::circuit_breaking::CircuitBreakingConfig; use super::security::{ClusterSecurityConfig, parse_transport_socket}; /// Validated Cluster resource. @@ -20,6 +21,8 @@ pub(crate) struct ClusterResource { /// TLS security config parsed from `transport_socket`. `None` means the /// cluster uses plaintext connections. pub security: Option, + /// Circuit-breaking config for the cluster. + pub circuit_breaking: CircuitBreakingConfig, } /// Load balancing policies. @@ -67,12 +70,14 @@ impl Resource for ClusterResource { }; let security = parse_transport_socket(message.transport_socket)?; + let circuit_breaking = CircuitBreakingConfig::from_proto(message.circuit_breakers.as_ref()); Ok(ClusterResource { name, eds_service_name, lb_policy, security, + circuit_breaking, }) } } @@ -105,6 +110,7 @@ mod tests { assert_eq!(validated.name, "my-cluster"); assert_eq!(validated.lb_policy, LbPolicy::RoundRobin); assert!(validated.eds_service_name.is_none()); + assert_eq!(validated.circuit_breaking, CircuitBreakingConfig::default()); } #[test] @@ -141,6 +147,34 @@ mod tests { assert_eq!(validated.lb_policy, LbPolicy::LeastRequest); } + #[test] + fn test_circuit_breaking_config() { + use envoy_types::pb::envoy::config::cluster::v3::CircuitBreakers; + use envoy_types::pb::envoy::config::cluster::v3::circuit_breakers::Thresholds; + use envoy_types::pb::envoy::config::core::v3::RoutingPriority; + use envoy_types::pb::google::protobuf::UInt32Value; + + let cluster = Cluster { + name: "cb-cluster".to_string(), + lb_policy: cluster::LbPolicy::RoundRobin as i32, + circuit_breakers: Some(CircuitBreakers { + thresholds: vec![Thresholds { + priority: RoutingPriority::Default as i32, + max_requests: Some(UInt32Value { value: 7 }), + ..Default::default() + }], + ..Default::default() + }), + ..Default::default() + }; + + let validated = ClusterResource::validate(cluster).unwrap(); + assert_eq!( + validated.circuit_breaking, + CircuitBreakingConfig { max_requests: 7 }, + ); + } + #[test] fn test_unsupported_lb_policy_is_rejected() { let cluster = Cluster { From c4559d3e654de9a7028f678ea395e3323449d143 Mon Sep 17 00:00:00 2001 From: Michael Ingley Date: Tue, 30 Jun 2026 14:14:03 -0500 Subject: [PATCH 3/6] Use ArcSwap for A32 limiter config --- tonic-xds/src/client/circuit_breaking.rs | 44 ++++++++++++++---------- 1 file changed, 26 insertions(+), 18 deletions(-) diff --git a/tonic-xds/src/client/circuit_breaking.rs b/tonic-xds/src/client/circuit_breaking.rs index 4e71b742e..258914ec4 100644 --- a/tonic-xds/src/client/circuit_breaking.rs +++ b/tonic-xds/src/client/circuit_breaking.rs @@ -7,12 +7,12 @@ use std::sync::{ }; use std::task::{Context, Poll}; +use arc_swap::ArcSwapOption; use bytes::Bytes; use dashmap::DashMap; use http::{Request, Response}; use http_body::{Body, Frame}; use pin_project_lite::pin_project; -use tokio::sync::watch; use tonic::body::Body as TonicBody; use tower::{BoxError, Layer, Service}; @@ -38,7 +38,7 @@ struct ClusterCircuitBreakersInner { impl Drop for ClusterCircuitBreakersInner { fn drop(&mut self) { for state in self.configs.iter() { - if let Some(previous) = state.config_tx.send_replace(None) { + if let Some(previous) = state.config.swap(None) { let counter_key = previous.counter_key.clone(); drop(previous); self.counters.deactivate(&counter_key); @@ -107,7 +107,7 @@ impl ClusterCircuitBreakers { config: CircuitBreakerRuntimeConfig, ) { let previous = state.current_config(); - if previous.as_ref() == Some(&config) { + if previous.as_deref() == Some(&config) { return; } @@ -118,19 +118,19 @@ impl ClusterCircuitBreakers { self.inner.counters.activate(&config.counter_key); } - let previous = state.config_tx.send_replace(Some(config)); + let previous = state.config.swap(Some(Arc::new(config))); if counter_key_changed && let Some(previous) = previous { self.deactivate_config(previous); } } fn clear_state(&self, state: &ClusterCircuitBreakerState) { - if let Some(previous) = state.config_tx.send_replace(None) { + if let Some(previous) = state.config.swap(None) { self.deactivate_config(previous); } } - fn deactivate_config(&self, config: CircuitBreakerRuntimeConfig) { + fn deactivate_config(&self, config: Arc) { let counter_key = config.counter_key.clone(); drop(config); self.inner.counters.deactivate(&counter_key); @@ -142,7 +142,7 @@ impl ClusterCircuitBreakers { fn acquire_with_config( &self, - runtime_config: CircuitBreakerRuntimeConfig, + runtime_config: Arc, ) -> Result { let limit = CircuitBreakerLimit { max_requests: runtime_config.max_requests, @@ -150,14 +150,14 @@ impl ClusterCircuitBreakers { self.inner .counters .acquire( - runtime_config.counter_key, - runtime_config.counter, + runtime_config.counter_key.clone(), + runtime_config.counter.clone(), limit.max_requests, ) .ok_or(limit) } - fn runtime_config_or_default(&self, cluster: &str) -> CircuitBreakerRuntimeConfig { + fn runtime_config_or_default(&self, cluster: &str) -> Arc { self.inner .configs .get(cluster) @@ -165,11 +165,11 @@ impl ClusterCircuitBreakers { .unwrap_or_else(|| { let counter_key = counter_key(cluster, cluster); let counter = self.inner.counters.counter(&counter_key); - CircuitBreakerRuntimeConfig { + Arc::new(CircuitBreakerRuntimeConfig { max_requests: self.inner.default_max_requests, counter_key: Arc::from(counter_key), counter, - } + }) }) } @@ -235,19 +235,27 @@ fn counter_key(cluster: &str, eds_service_name: &str) -> String { format!("{cluster}\0{eds_service_name}") } -#[derive(Debug)] struct ClusterCircuitBreakerState { - config_tx: watch::Sender>, + config: ArcSwapOption, +} + +impl fmt::Debug for ClusterCircuitBreakerState { + fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { + f.debug_struct("ClusterCircuitBreakerState") + .field("current_config", &self.current_config()) + .finish() + } } impl ClusterCircuitBreakerState { fn new() -> Self { - let (config_tx, _) = watch::channel(None); - Self { config_tx } + Self { + config: ArcSwapOption::empty(), + } } - fn current_config(&self) -> Option { - self.config_tx.borrow().clone() + fn current_config(&self) -> Option> { + self.config.load_full() } } From 4a20f3920419154ea95fcaee7cd1b70bf9723ceb Mon Sep 17 00:00:00 2001 From: Michael Ingley Date: Tue, 30 Jun 2026 16:24:24 -0500 Subject: [PATCH 4/6] Refine A32 circuit breaker state --- tonic-xds/src/client/circuit_breaking.rs | 223 ++++++++++++++++------- 1 file changed, 158 insertions(+), 65 deletions(-) diff --git a/tonic-xds/src/client/circuit_breaking.rs b/tonic-xds/src/client/circuit_breaking.rs index 258914ec4..616aa6e83 100644 --- a/tonic-xds/src/client/circuit_breaking.rs +++ b/tonic-xds/src/client/circuit_breaking.rs @@ -32,7 +32,6 @@ pub(crate) struct ClusterCircuitBreakers { struct ClusterCircuitBreakersInner { configs: DashMap>, counters: ClusterRequestCounters, - default_max_requests: u32, } impl Drop for ClusterCircuitBreakersInner { @@ -57,7 +56,6 @@ impl ClusterCircuitBreakers { inner: Arc::new(ClusterCircuitBreakersInner { configs: DashMap::new(), counters, - default_max_requests: DEFAULT_MAX_REQUESTS, }), } } @@ -76,14 +74,14 @@ impl ClusterCircuitBreakers { ) { let cluster = cluster.into(); let eds_service_name = eds_service_name.into(); - let counter_key = counter_key(&cluster, &eds_service_name); + let counter_key = CounterKey::new(cluster.as_str(), eds_service_name.as_str()); let counter = self.inner.counters.counter(&counter_key); let state = self.ensure_state(&cluster); self.update_state_config( &state, CircuitBreakerRuntimeConfig { max_requests: config.max_requests, - counter_key: Arc::from(counter_key), + counter_key, counter, }, ); @@ -101,6 +99,18 @@ impl ClusterCircuitBreakers { .clone() } + fn cluster_breaker(&self, cluster: &str) -> Arc { + let state = self.ensure_state(cluster); + let cluster: Arc = Arc::from(cluster); + let default_counter_key = CounterKey::same_cluster(cluster.clone()); + Arc::new(ClusterCircuitBreaker { + cluster, + state, + counters: self.inner.counters.clone(), + default_counter_key, + }) + } + fn update_state_config( &self, state: &ClusterCircuitBreakerState, @@ -137,54 +147,23 @@ impl ClusterCircuitBreakers { } fn acquire(&self, cluster: &str) -> Result { - self.acquire_with_config(self.runtime_config_or_default(cluster)) - } - - fn acquire_with_config( - &self, - runtime_config: Arc, - ) -> Result { - let limit = CircuitBreakerLimit { - max_requests: runtime_config.max_requests, - }; - self.inner - .counters - .acquire( - runtime_config.counter_key.clone(), - runtime_config.counter.clone(), - limit.max_requests, - ) - .ok_or(limit) - } - - fn runtime_config_or_default(&self, cluster: &str) -> Arc { - self.inner - .configs - .get(cluster) - .and_then(|state| state.current_config()) - .unwrap_or_else(|| { - let counter_key = counter_key(cluster, cluster); - let counter = self.inner.counters.counter(&counter_key); - Arc::new(CircuitBreakerRuntimeConfig { - max_requests: self.inner.default_max_requests, - counter_key: Arc::from(counter_key), - counter, - }) - }) + self.cluster_breaker(cluster).acquire() } #[cfg(test)] fn in_flight(&self, cluster: &str) -> u32 { - let runtime_config = self.runtime_config_or_default(cluster); - self.inner.counters.in_flight(&runtime_config.counter_key) + let breaker = self.cluster_breaker(cluster); + self.inner + .counters + .in_flight(&breaker.current_counter_key()) } #[cfg(test)] fn dropped_requests(&self, cluster: &str) -> u64 { - let runtime_config = self.runtime_config_or_default(cluster); + let breaker = self.cluster_breaker(cluster); self.inner .counters - .dropped_requests(&runtime_config.counter_key) + .dropped_requests(&breaker.current_counter_key()) } #[cfg(test)] @@ -219,7 +198,7 @@ struct CircuitBreakerLimit { #[derive(Clone, Debug)] struct CircuitBreakerRuntimeConfig { max_requests: u32, - counter_key: Arc, + counter_key: CounterKey, counter: Arc, } @@ -231,8 +210,26 @@ impl PartialEq for CircuitBreakerRuntimeConfig { impl Eq for CircuitBreakerRuntimeConfig {} -fn counter_key(cluster: &str, eds_service_name: &str) -> String { - format!("{cluster}\0{eds_service_name}") +#[derive(Clone, Debug, Eq, Hash, PartialEq)] +struct CounterKey { + cluster: Arc, + eds_service_name: Arc, +} + +impl CounterKey { + fn new(cluster: impl Into>, eds_service_name: impl Into>) -> Self { + Self { + cluster: cluster.into(), + eds_service_name: eds_service_name.into(), + } + } + + fn same_cluster(cluster: Arc) -> Self { + Self { + cluster: cluster.clone(), + eds_service_name: cluster, + } + } } struct ClusterCircuitBreakerState { @@ -259,6 +256,52 @@ impl ClusterCircuitBreakerState { } } +#[derive(Debug)] +struct ClusterCircuitBreaker { + cluster: Arc, + state: Arc, + counters: ClusterRequestCounters, + default_counter_key: CounterKey, +} + +impl ClusterCircuitBreaker { + fn acquire(&self) -> Result { + if let Some(config) = self.state.current_config() { + return self.acquire_with_config( + config.counter_key.clone(), + config.counter.clone(), + config.max_requests, + ); + } + + let counter = self.counters.counter(&self.default_counter_key); + self.acquire_with_config( + self.default_counter_key.clone(), + counter, + DEFAULT_MAX_REQUESTS, + ) + } + + fn acquire_with_config( + &self, + counter_key: CounterKey, + counter: Arc, + max_requests: u32, + ) -> Result { + let limit = CircuitBreakerLimit { max_requests }; + self.counters + .acquire(counter_key, counter, max_requests) + .ok_or(limit) + } + + fn current_counter_key(&self) -> CounterKey { + self.state + .current_config() + .map(|config| config.counter_key.clone()) + .unwrap_or_else(|| self.default_counter_key.clone()) + } +} + #[derive(Clone, Debug)] struct ClusterRequestCounters { inner: Arc, @@ -266,8 +309,8 @@ struct ClusterRequestCounters { #[derive(Debug, Default)] struct ClusterRequestCounterState { - counters: DashMap>, - active_refs: DashMap>, + counters: DashMap>, + active_refs: DashMap>, } impl ClusterRequestCounters { @@ -286,15 +329,15 @@ impl ClusterRequestCounters { } } - fn activate(&self, counter_key: &str) { + fn activate(&self, counter_key: &CounterKey) { self.inner .active_refs - .entry(counter_key.to_string()) + .entry(counter_key.clone()) .or_insert_with(|| Arc::new(AtomicUsize::new(0))) .fetch_add(1, Ordering::AcqRel); } - fn deactivate(&self, counter_key: &str) { + fn deactivate(&self, counter_key: &CounterKey) { let should_cleanup = self .inner .active_refs @@ -315,7 +358,7 @@ impl ClusterRequestCounters { fn acquire( &self, - counter_key: Arc, + counter_key: CounterKey, counter: Arc, limit: u32, ) -> Option { @@ -332,15 +375,15 @@ impl ClusterRequestCounters { } } - fn counter(&self, counter_key: &str) -> Arc { + fn counter(&self, counter_key: &CounterKey) -> Arc { self.inner .counters - .entry(counter_key.to_string()) + .entry(counter_key.clone()) .or_insert_with(|| Arc::new(InFlightCounter::default())) .clone() } - fn cleanup_if_unused(&self, counter_key: &str) { + fn cleanup_if_unused(&self, counter_key: &CounterKey) { let active_refs = self.active_refs(counter_key); if active_refs != 0 { return; @@ -354,7 +397,7 @@ impl ClusterRequestCounters { .remove_if(counter_key, |_, refs| refs.load(Ordering::Acquire) == 0); } - fn active_refs(&self, counter_key: &str) -> usize { + fn active_refs(&self, counter_key: &CounterKey) -> usize { self.inner .active_refs .get(counter_key) @@ -363,7 +406,7 @@ impl ClusterRequestCounters { } #[cfg(test)] - fn in_flight(&self, counter_key: &str) -> u32 { + fn in_flight(&self, counter_key: &CounterKey) -> u32 { self.inner .counters .get(counter_key) @@ -372,7 +415,7 @@ impl ClusterRequestCounters { } #[cfg(test)] - fn dropped_requests(&self, counter_key: &str) -> u64 { + fn dropped_requests(&self, counter_key: &CounterKey) -> u64 { self.inner .counters .get(counter_key) @@ -429,7 +472,7 @@ impl InFlightCounter { #[derive(Debug)] struct CircuitBreakerPermit { counter: Option>, - counter_key: Arc, + counter_key: CounterKey, counters: ClusterRequestCounters, } @@ -446,11 +489,15 @@ impl Drop for CircuitBreakerPermit { #[derive(Clone)] pub(crate) struct CircuitBreakingLayer { circuit_breakers: ClusterCircuitBreakers, + breaker_cache: Arc>>, } impl CircuitBreakingLayer { pub(crate) fn new(circuit_breakers: ClusterCircuitBreakers) -> Self { - Self { circuit_breakers } + Self { + circuit_breakers, + breaker_cache: Arc::new(DashMap::new()), + } } } @@ -458,6 +505,7 @@ impl fmt::Debug for CircuitBreakingLayer { fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { f.debug_struct("CircuitBreakingLayer") .field("circuit_breakers", &self.circuit_breakers) + .field("cached_clusters", &self.breaker_cache.len()) .finish() } } @@ -469,6 +517,7 @@ impl Layer for CircuitBreakingLayer { CircuitBreakingService { inner: service, circuit_breakers: self.circuit_breakers.clone(), + breaker_cache: self.breaker_cache.clone(), } } } @@ -477,6 +526,7 @@ impl Layer for CircuitBreakingLayer { pub(crate) struct CircuitBreakingService { inner: S, circuit_breakers: ClusterCircuitBreakers, + breaker_cache: Arc>>, } impl fmt::Debug for CircuitBreakingService { @@ -484,10 +534,24 @@ impl fmt::Debug for CircuitBreakingService { f.debug_struct("CircuitBreakingService") .field("inner", &self.inner) .field("circuit_breakers", &self.circuit_breakers) + .field("cached_clusters", &self.breaker_cache.len()) .finish() } } +impl CircuitBreakingService { + fn breaker_for_cluster(&self, cluster: &str) -> Arc { + if let Some(breaker) = self.breaker_cache.get(cluster) { + return breaker.clone(); + } + + self.breaker_cache + .entry(cluster.to_string()) + .or_insert_with(|| self.circuit_breakers.cluster_breaker(cluster)) + .clone() + } +} + impl Service> for CircuitBreakingService where S: Service, Response = Response, Error: Into> @@ -506,7 +570,11 @@ where } fn call(&mut self, request: Request) -> Self::Future { - let Some(route_decision) = request.extensions().get::().cloned() else { + let Some(cluster) = request + .extensions() + .get::() + .map(|route_decision| route_decision.cluster.as_str()) + else { return Box::pin(async { Ok(status_response(tonic::Status::internal( CircuitBreakingError::NoRoutingDecision.to_string(), @@ -514,14 +582,13 @@ where }); }; - let cluster = route_decision.cluster; - let circuit_breakers = self.circuit_breakers.clone(); + let breaker = self.breaker_for_cluster(cluster); let clone = self.inner.clone(); let mut inner = std::mem::replace(&mut self.inner, clone); Box::pin(async move { - let permit = match circuit_breakers.acquire(&cluster) { + let permit = match breaker.acquire() { Ok(permit) => permit, - Err(limit) => return Ok(limit_exceeded_response(&cluster, limit)), + Err(limit) => return Ok(limit_exceeded_response(&breaker.cluster, limit)), }; std::future::poll_fn(|cx| inner.poll_ready(cx)) @@ -849,7 +916,7 @@ mod tests { #[test] fn cleanup_keeps_counter_with_outstanding_clone() { let counters = ClusterRequestCounters::isolated(); - let counter_key = counter_key(CLUSTER, CLUSTER); + let counter_key = CounterKey::new(CLUSTER, CLUSTER); let counter = counters.counter(&counter_key); counters.cleanup_if_unused(&counter_key); @@ -860,6 +927,32 @@ mod tests { assert_eq!(counters.counter_count(), 0); } + #[test] + fn structured_counter_keys_do_not_collide_on_embedded_delimiters() { + let counters = ClusterRequestCounters::isolated(); + let first_key = CounterKey::new("cluster\0eds", "service"); + let second_key = CounterKey::new("cluster", "eds\0service"); + + let first = counters.counter(&first_key); + let second = counters.counter(&second_key); + + assert!(!Arc::ptr_eq(&first, &second)); + assert_eq!(counters.counter_count(), 2); + } + + #[test] + fn cached_cluster_breaker_does_not_pin_default_counter() { + let counters = ClusterRequestCounters::isolated(); + let breakers = ClusterCircuitBreakers::with_counters(counters.clone()); + let breaker = breakers.cluster_breaker(CLUSTER); + let permit = breaker.acquire().unwrap(); + assert_eq!(counters.counter_count(), 1); + + drop(permit); + assert_eq!(counters.counter_count(), 0); + assert_eq!(breaker.cluster.as_ref(), CLUSTER); + } + #[tokio::test] async fn rejects_over_limit_without_waiting_for_inner_ready() { let breakers = configured_breakers(1); From c6b89daa833f14454c55ea633398d7fc0c8343f4 Mon Sep 17 00:00:00 2001 From: Michael Ingley Date: Mon, 13 Jul 2026 18:50:05 -0500 Subject: [PATCH 5/6] tonic-xds: correct A32 limiter semantics --- tonic-xds/src/client/circuit_breaking.rs | 525 ++++++++++++++++------- tonic-xds/src/client/retry.rs | 3 + 2 files changed, 363 insertions(+), 165 deletions(-) diff --git a/tonic-xds/src/client/circuit_breaking.rs b/tonic-xds/src/client/circuit_breaking.rs index 616aa6e83..96e9e7ed9 100644 --- a/tonic-xds/src/client/circuit_breaking.rs +++ b/tonic-xds/src/client/circuit_breaking.rs @@ -2,7 +2,7 @@ use std::fmt; use std::sync::{ - Arc, OnceLock, + Arc, Mutex, OnceLock, atomic::{AtomicU32, AtomicU64, AtomicUsize, Ordering}, }; use std::task::{Context, Poll}; @@ -39,8 +39,9 @@ impl Drop for ClusterCircuitBreakersInner { for state in self.configs.iter() { if let Some(previous) = state.config.swap(None) { let counter_key = previous.counter_key.clone(); + previous.counter.deactivate(); drop(previous); - self.counters.deactivate(&counter_key); + self.counters.cleanup_if_unused(&counter_key); } } } @@ -116,6 +117,10 @@ impl ClusterCircuitBreakers { state: &ClusterCircuitBreakerState, config: CircuitBreakerRuntimeConfig, ) { + let _update_guard = state + .update_lock + .lock() + .unwrap_or_else(|poisoned| poisoned.into_inner()); let previous = state.current_config(); if previous.as_deref() == Some(&config) { return; @@ -125,9 +130,10 @@ impl ClusterCircuitBreakers { .as_ref() .is_none_or(|previous| previous.counter_key != config.counter_key); if counter_key_changed { - self.inner.counters.activate(&config.counter_key); + config.counter.activate(); } + drop(previous); let previous = state.config.swap(Some(Arc::new(config))); if counter_key_changed && let Some(previous) = previous { self.deactivate_config(previous); @@ -135,6 +141,10 @@ impl ClusterCircuitBreakers { } fn clear_state(&self, state: &ClusterCircuitBreakerState) { + let _update_guard = state + .update_lock + .lock() + .unwrap_or_else(|poisoned| poisoned.into_inner()); if let Some(previous) = state.config.swap(None) { self.deactivate_config(previous); } @@ -142,8 +152,9 @@ impl ClusterCircuitBreakers { fn deactivate_config(&self, config: Arc) { let counter_key = config.counter_key.clone(); + config.counter.deactivate(); drop(config); - self.inner.counters.deactivate(&counter_key); + self.inner.counters.cleanup_if_unused(&counter_key); } fn acquire(&self, cluster: &str) -> Result { @@ -160,10 +171,7 @@ impl ClusterCircuitBreakers { #[cfg(test)] fn dropped_requests(&self, cluster: &str) -> u64 { - let breaker = self.cluster_breaker(cluster); - self.inner - .counters - .dropped_requests(&breaker.current_counter_key()) + self.ensure_state(cluster).dropped_requests() } #[cfg(test)] @@ -173,7 +181,8 @@ impl ClusterCircuitBreakers { #[cfg(test)] fn clear_cluster_config(&self, cluster: &str) { - if let Some((_, state)) = self.inner.configs.remove(cluster) { + let state = self.inner.configs.get(cluster).map(|state| state.clone()); + if let Some(state) = state { self.clear_state(&state); } } @@ -234,12 +243,18 @@ impl CounterKey { struct ClusterCircuitBreakerState { config: ArcSwapOption, + update_lock: Mutex<()>, + dropped_requests: AtomicU64, } impl fmt::Debug for ClusterCircuitBreakerState { fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { f.debug_struct("ClusterCircuitBreakerState") .field("current_config", &self.current_config()) + .field( + "dropped_requests", + &self.dropped_requests.load(Ordering::Acquire), + ) .finish() } } @@ -248,12 +263,23 @@ impl ClusterCircuitBreakerState { fn new() -> Self { Self { config: ArcSwapOption::empty(), + update_lock: Mutex::new(()), + dropped_requests: AtomicU64::new(0), } } fn current_config(&self) -> Option> { self.config.load_full() } + + fn record_drop(&self) { + self.dropped_requests.fetch_add(1, Ordering::AcqRel); + } + + #[cfg(test)] + fn dropped_requests(&self) -> u64 { + self.dropped_requests.load(Ordering::Acquire) + } } #[derive(Debug)] @@ -289,9 +315,13 @@ impl ClusterCircuitBreaker { max_requests: u32, ) -> Result { let limit = CircuitBreakerLimit { max_requests }; - self.counters - .acquire(counter_key, counter, max_requests) - .ok_or(limit) + match self.counters.acquire(counter_key, counter, max_requests) { + Some(permit) => Ok(permit), + None => { + self.state.record_drop(); + Err(limit) + } + } } fn current_counter_key(&self) -> CounterKey { @@ -310,7 +340,6 @@ struct ClusterRequestCounters { #[derive(Debug, Default)] struct ClusterRequestCounterState { counters: DashMap>, - active_refs: DashMap>, } impl ClusterRequestCounters { @@ -329,33 +358,6 @@ impl ClusterRequestCounters { } } - fn activate(&self, counter_key: &CounterKey) { - self.inner - .active_refs - .entry(counter_key.clone()) - .or_insert_with(|| Arc::new(AtomicUsize::new(0))) - .fetch_add(1, Ordering::AcqRel); - } - - fn deactivate(&self, counter_key: &CounterKey) { - let should_cleanup = self - .inner - .active_refs - .get(counter_key) - .and_then(|active_refs| { - active_refs - .fetch_update(Ordering::AcqRel, Ordering::Acquire, |count| { - count.checked_sub(1) - }) - .ok() - }) - .is_some_and(|previous| previous <= 1); - - if should_cleanup { - self.cleanup_if_unused(counter_key); - } - } - fn acquire( &self, counter_key: CounterKey, @@ -369,8 +371,11 @@ impl ClusterRequestCounters { counters: self.clone(), }) } else { - counter.record_drop(); - self.cleanup_if_unused(&counter_key); + let should_cleanup = counter.is_unused(); + drop(counter); + if should_cleanup { + self.cleanup_if_unused(&counter_key); + } None } } @@ -384,25 +389,9 @@ impl ClusterRequestCounters { } fn cleanup_if_unused(&self, counter_key: &CounterKey) { - let active_refs = self.active_refs(counter_key); - if active_refs != 0 { - return; - } - self.inner.counters.remove_if(counter_key, |_, counter| { - counter.in_flight() == 0 && Arc::strong_count(counter) == 1 + counter.is_unused() && Arc::strong_count(counter) == 1 }); - self.inner - .active_refs - .remove_if(counter_key, |_, refs| refs.load(Ordering::Acquire) == 0); - } - - fn active_refs(&self, counter_key: &CounterKey) -> usize { - self.inner - .active_refs - .get(counter_key) - .map(|refs| refs.load(Ordering::Acquire)) - .unwrap_or(0) } #[cfg(test)] @@ -414,15 +403,6 @@ impl ClusterRequestCounters { .unwrap_or(0) } - #[cfg(test)] - fn dropped_requests(&self, counter_key: &CounterKey) -> u64 { - self.inner - .counters - .get(counter_key) - .map(|counter| counter.dropped_requests()) - .unwrap_or(0) - } - #[cfg(test)] fn counter_count(&self) -> usize { self.inner.counters.len() @@ -432,12 +412,26 @@ impl ClusterRequestCounters { #[derive(Debug, Default)] struct InFlightCounter { in_flight: AtomicU32, - /// Local A32 drop accounting, kept with the global counter so future LRS - /// support can export `total_dropped_requests` without changing enforcement. - dropped_requests: AtomicU64, + active_refs: AtomicUsize, } impl InFlightCounter { + fn activate(&self) { + self.active_refs.fetch_add(1, Ordering::AcqRel); + } + + fn deactivate(&self) { + let result = self + .active_refs + .fetch_update(Ordering::AcqRel, Ordering::Acquire, |count| { + count.checked_sub(1) + }); + assert!( + result.is_ok(), + "attempted to deactivate an inactive circuit breaker counter" + ); + } + fn try_acquire(&self, limit: u32) -> bool { loop { let current = self.in_flight.load(Ordering::Acquire); @@ -459,13 +453,16 @@ impl InFlightCounter { self.in_flight.load(Ordering::Acquire) } - fn record_drop(&self) { - self.dropped_requests.fetch_add(1, Ordering::AcqRel); + fn release(&self) { + let previous = self.in_flight.fetch_sub(1, Ordering::AcqRel); + assert!( + previous > 0, + "attempted to release an inactive circuit breaker permit" + ); } - #[cfg(test)] - fn dropped_requests(&self) -> u64 { - self.dropped_requests.load(Ordering::Acquire) + fn is_unused(&self) -> bool { + self.in_flight() == 0 && self.active_refs.load(Ordering::Acquire) == 0 } } @@ -479,13 +476,20 @@ struct CircuitBreakerPermit { impl Drop for CircuitBreakerPermit { fn drop(&mut self) { if let Some(counter) = self.counter.take() { - counter.in_flight.fetch_sub(1, Ordering::AcqRel); + counter.release(); + let should_cleanup = counter.is_unused(); + drop(counter); + if should_cleanup { + self.counters.cleanup_if_unused(&self.counter_key); + } } - self.counters.cleanup_if_unused(&self.counter_key); } } /// Tower layer that enforces A32 max in-flight requests per xDS cluster. +/// +/// This layer must wrap the ready per-cluster dispatch service inside retries so +/// each admitted call represents one upstream attempt rather than queued work. #[derive(Clone)] pub(crate) struct CircuitBreakingLayer { circuit_breakers: ClusterCircuitBreakers, @@ -565,8 +569,8 @@ where type Error = BoxError; type Future = BoxFuture>; - fn poll_ready(&mut self, _cx: &mut Context<'_>) -> Poll> { - Poll::Ready(Ok(())) + fn poll_ready(&mut self, cx: &mut Context<'_>) -> Poll> { + self.inner.poll_ready(cx).map_err(Into::into) } fn call(&mut self, request: Request) -> Self::Future { @@ -583,17 +587,18 @@ where }; let breaker = self.breaker_for_cluster(cluster); + let permit = match breaker.acquire() { + Ok(permit) => permit, + Err(limit) => { + return Box::pin(std::future::ready(Ok(limit_exceeded_response( + &breaker.cluster, + limit, + )))); + } + }; let clone = self.inner.clone(); let mut inner = std::mem::replace(&mut self.inner, clone); Box::pin(async move { - let permit = match breaker.acquire() { - Ok(permit) => permit, - Err(limit) => return Ok(limit_exceeded_response(&breaker.cluster, limit)), - }; - - std::future::poll_fn(|cx| inner.poll_ready(cx)) - .await - .map_err(Into::into)?; let response = inner.call(request).await.map_err(Into::into)?; Ok(response.map(|body| TonicBody::new(PermitBody::new(body, permit)))) }) @@ -606,11 +611,23 @@ enum CircuitBreakingError { NoRoutingDecision, } +#[derive(Clone, Copy, Debug)] +struct LocalCircuitBreakerDrop; + +pub(crate) fn is_local_circuit_breaker_drop(response: &Response) -> bool { + response + .extensions() + .get::() + .is_some() +} + fn limit_exceeded_response(cluster: &str, limit: CircuitBreakerLimit) -> Response { - status_response(tonic::Status::unavailable(format!( + let mut response = status_response(tonic::Status::unavailable(format!( "circuit breaker open for cluster '{cluster}': max_requests limit {} reached", limit.max_requests, - ))) + ))); + response.extensions_mut().insert(LocalCircuitBreakerDrop); + response } fn status_response(status: tonic::Status) -> Response { @@ -683,6 +700,7 @@ mod tests { use std::sync::Arc; use std::sync::atomic::{AtomicU32, Ordering}; use std::task::{Context, Poll}; + use std::time::Duration; use bytes::Bytes; use http::{HeaderMap, Request, Response}; @@ -690,10 +708,11 @@ mod tests { use tonic::Code; use tower::Layer; use tower::ServiceExt; - use tower::retry::Policy; use tower::service_fn; - use crate::client::retry::RetryLayer; + use crate::client::retry::{ + GrpcRetryBackoffConfig, GrpcRetryPolicy, GrpcRetryPolicyConfig, RetryLayer, + }; use super::*; @@ -760,6 +779,78 @@ mod tests { assert_eq!(calls.load(Ordering::SeqCst), 2); } + #[tokio::test] + async fn cached_service_applies_live_limit_updates_without_resetting_in_flight() { + let breakers = configured_breakers(2); + let calls = Arc::new(AtomicU32::new(0)); + let call_counter = calls.clone(); + let service = service_fn(move |_request: Request| { + call_counter.fetch_add(1, Ordering::SeqCst); + async { Ok::<_, BoxError>(Response::new(TonicBody::new(PendingBody))) } + }); + let mut service = CircuitBreakingLayer::new(breakers.clone()).layer(service); + + let first = service + .ready() + .await + .unwrap() + .call(request()) + .await + .unwrap(); + let second = service + .ready() + .await + .unwrap() + .call(request()) + .await + .unwrap(); + assert_eq!(breakers.in_flight(CLUSTER), 2); + + breakers.set_config(CLUSTER, CircuitBreakingConfig { max_requests: 1 }); + + let third = service + .ready() + .await + .unwrap() + .call(request()) + .await + .unwrap(); + assert_eq!( + tonic::Status::from_header_map(third.headers()) + .unwrap() + .code(), + Code::Unavailable + ); + + drop(first); + assert_eq!(breakers.in_flight(CLUSTER), 1); + let fourth = service + .ready() + .await + .unwrap() + .call(request()) + .await + .unwrap(); + assert_eq!( + tonic::Status::from_header_map(fourth.headers()) + .unwrap() + .code(), + Code::Unavailable + ); + + drop(second); + assert_eq!(breakers.in_flight(CLUSTER), 0); + let _fifth = service + .ready() + .await + .unwrap() + .call(request()) + .await + .unwrap(); + assert_eq!(calls.load(Ordering::SeqCst), 3); + assert_eq!(breakers.dropped_requests(CLUSTER), 2); + } + #[tokio::test] async fn releases_permit_when_response_body_reaches_trailers() { let breakers = configured_breakers(1); @@ -789,6 +880,29 @@ mod tests { assert_eq!(breakers.in_flight(CLUSTER), 0); } + #[tokio::test] + async fn releases_permit_when_response_body_returns_error() { + let breakers = configured_breakers(1); + let service = service_fn(|_request: Request| async { + Ok::<_, BoxError>(Response::new(TonicBody::new(ErrorBody { emitted: false }))) + }); + let mut service = CircuitBreakingLayer::new(breakers.clone()).layer(service); + + let response = service + .ready() + .await + .unwrap() + .call(request()) + .await + .unwrap(); + assert_eq!(breakers.in_flight(CLUSTER), 1); + + let mut body = response.into_body(); + let frame = std::future::poll_fn(|cx| Pin::new(&mut body).poll_frame(cx)).await; + assert!(frame.unwrap().is_err()); + assert_eq!(breakers.in_flight(CLUSTER), 0); + } + #[tokio::test] async fn releases_permit_when_response_future_is_dropped() { let breakers = configured_breakers(1); @@ -829,42 +943,97 @@ mod tests { assert!(status.message().contains("No routing decision")); } + #[tokio::test] + async fn oneshot_honors_config_after_consuming_service() { + let calls = Arc::new(AtomicU32::new(0)); + let call_counter = calls.clone(); + let service = service_fn(move |_request: Request| { + call_counter.fetch_add(1, Ordering::SeqCst); + async { Ok::<_, BoxError>(Response::new(TonicBody::empty())) } + }); + let service = CircuitBreakingLayer::new(configured_breakers(0)).layer(service); + + let response = service.oneshot(request()).await.unwrap(); + let status = tonic::Status::from_header_map(response.headers()).unwrap(); + + assert_eq!(status.code(), Code::Unavailable); + assert_eq!(calls.load(Ordering::SeqCst), 0); + } + #[tokio::test] async fn limit_responses_do_not_enter_retry_policy() { - let breakers = configured_breakers(1); - let retry_observations = Arc::new(AtomicU32::new(0)); - let policy = CountingUnavailablePolicy { - retry_observations: retry_observations.clone(), - }; + let breakers = configured_breakers(0); + let policy = GrpcRetryPolicy::new( + GrpcRetryPolicyConfig::new() + .retry_on(vec![Code::Unavailable]) + .num_retries(4) + .retry_backoff( + GrpcRetryBackoffConfig::new(Duration::from_millis(1)) + .max_interval(Duration::from_millis(1)), + ), + ); + let calls = Arc::new(AtomicU32::new(0)); + let call_counter = calls.clone(); let service = service_fn( - |_request: Request>| async { - Ok::<_, BoxError>(Response::new(TonicBody::new(PendingBody))) + move |_request: Request>| { + call_counter.fetch_add(1, Ordering::SeqCst); + async { Ok::<_, BoxError>(Response::new(TonicBody::empty())) } }, ); let mut service = tower::ServiceBuilder::new() - .layer(CircuitBreakingLayer::new(breakers)) .layer(RetryLayer::new(policy)) + .layer(CircuitBreakingLayer::new(breakers.clone())) .service(service); - let _first = service - .ready() - .await - .unwrap() - .call(request()) - .await - .unwrap(); - let second = service + let response = service .ready() .await .unwrap() .call(request()) .await .unwrap(); - let status = tonic::Status::from_header_map(second.headers()).unwrap(); + let status = tonic::Status::from_header_map(response.headers()).unwrap(); assert_eq!(status.code(), Code::Unavailable); - assert_eq!(retry_observations.load(Ordering::SeqCst), 0); + assert_eq!(breakers.dropped_requests(CLUSTER), 1); + assert_eq!(calls.load(Ordering::SeqCst), 0); + } + + #[test] + fn shared_counters_enforce_process_limit_with_per_channel_drop_counts() { + let counters = ClusterRequestCounters::isolated(); + let first_breakers = ClusterCircuitBreakers::with_counters(counters.clone()); + let second_breakers = ClusterCircuitBreakers::with_counters(counters.clone()); + first_breakers.set_config(CLUSTER, CircuitBreakingConfig { max_requests: 1 }); + second_breakers.set_config(CLUSTER, CircuitBreakingConfig { max_requests: 1 }); + + let first = first_breakers.acquire(CLUSTER).unwrap(); + assert!(second_breakers.acquire(CLUSTER).is_err()); + assert_eq!(first_breakers.dropped_requests(CLUSTER), 0); + assert_eq!(second_breakers.dropped_requests(CLUSTER), 1); + + drop(first); + let second = second_breakers.acquire(CLUSTER).unwrap(); + drop(second); + drop(first_breakers); + drop(second_breakers); + assert_eq!(counters.counter_count(), 0); + } + + #[test] + fn default_limit_rejects_the_1025th_request() { + let breakers = ClusterCircuitBreakers::new_for_test(); + let breaker = breakers.cluster_breaker(CLUSTER); + let permits: Vec<_> = (0..DEFAULT_MAX_REQUESTS) + .map(|_| breaker.acquire().unwrap()) + .collect(); + + assert!(breaker.acquire().is_err()); + assert_eq!(breakers.dropped_requests(CLUSTER), 1); + + drop(permits); + assert_eq!(breakers.counter_count(), 0); } #[test] @@ -886,6 +1055,20 @@ mod tests { assert_eq!(breakers.counter_count(), 1); } + #[test] + fn idle_eds_service_name_change_cleans_up_previous_counter() { + let counters = ClusterRequestCounters::isolated(); + let breakers = ClusterCircuitBreakers::with_counters(counters.clone()); + breakers.set_cluster_config(CLUSTER, "eds-a", CircuitBreakingConfig { max_requests: 1 }); + assert_eq!(counters.counter_count(), 1); + + breakers.set_cluster_config(CLUSTER, "eds-b", CircuitBreakingConfig { max_requests: 1 }); + assert_eq!(counters.counter_count(), 1); + + drop(breakers); + assert_eq!(counters.counter_count(), 0); + } + #[test] fn cluster_removal_cleans_up_counter_after_in_flight_requests_finish() { let breakers = configured_breakers(1); @@ -899,6 +1082,41 @@ mod tests { assert_eq!(breakers.counter_count(), 0); } + #[tokio::test] + async fn cached_breaker_observes_cluster_removal_and_recreation() { + let breakers = configured_breakers(1); + let calls = Arc::new(AtomicU32::new(0)); + let call_counter = calls.clone(); + let service = service_fn(move |_request: Request| { + call_counter.fetch_add(1, Ordering::SeqCst); + async { Ok::<_, BoxError>(Response::new(TonicBody::empty())) } + }); + let mut service = CircuitBreakingLayer::new(breakers.clone()).layer(service); + + let first = service + .ready() + .await + .unwrap() + .call(request()) + .await + .unwrap(); + drop(first); + + breakers.clear_cluster_config(CLUSTER); + breakers.set_config(CLUSTER, CircuitBreakingConfig { max_requests: 0 }); + + let second = service + .ready() + .await + .unwrap() + .call(request()) + .await + .unwrap(); + let status = tonic::Status::from_header_map(second.headers()).unwrap(); + assert_eq!(status.code(), Code::Unavailable); + assert_eq!(calls.load(Ordering::SeqCst), 1); + } + #[test] fn dropping_breakers_releases_config_counter_ref() { let counters = ClusterRequestCounters::isolated(); @@ -954,69 +1172,24 @@ mod tests { } #[tokio::test] - async fn rejects_over_limit_without_waiting_for_inner_ready() { + async fn waiting_for_inner_readiness_does_not_acquire_permit() { let breakers = configured_breakers(1); let calls = Arc::new(AtomicU32::new(0)); let service = BackpressuredService { - ready_budget: Arc::new(AtomicU32::new(1)), + ready_budget: Arc::new(AtomicU32::new(0)), calls: calls.clone(), }; let mut service = CircuitBreakingLayer::new(breakers.clone()).layer(service); - let first = service - .ready() - .await - .unwrap() - .call(request()) - .await - .unwrap(); - assert_eq!(breakers.in_flight(CLUSTER), 1); - assert_eq!(calls.load(Ordering::SeqCst), 1); - - let second = tokio::time::timeout( - tokio::time::Duration::from_millis(50), - service.ready().await.unwrap().call(request()), - ) - .await - .expect("over-limit request should not wait for inner readiness") - .unwrap(); - let status = tonic::Status::from_header_map(second.headers()).unwrap(); - assert_eq!(status.code(), Code::Unavailable); - assert_eq!(calls.load(Ordering::SeqCst), 1); - - drop(first); - } - - #[derive(Clone, Debug)] - struct CountingUnavailablePolicy { - retry_observations: Arc, - } - - impl Policy>, Response, BoxError> - for CountingUnavailablePolicy - { - type Future = std::future::Ready<()>; - - fn retry( - &mut self, - _req: &mut Request>, - result: &mut Result, BoxError>, - ) -> Option { - if let Ok(response) = result - && tonic::Status::from_header_map(response.headers()) - .is_some_and(|status| status.code() == Code::Unavailable) - { - self.retry_observations.fetch_add(1, Ordering::SeqCst); - } - None - } + let mut ready = Box::pin(service.ready()); + std::future::poll_fn(|cx| match ready.as_mut().poll(cx) { + Poll::Pending => Poll::Ready(()), + Poll::Ready(_) => panic!("inner service should remain backpressured"), + }) + .await; - fn clone_request( - &mut self, - req: &Request>, - ) -> Option>> { - Some(req.clone()) - } + assert_eq!(breakers.in_flight(CLUSTER), 0); + assert_eq!(calls.load(Ordering::SeqCst), 0); } #[derive(Clone, Debug)] @@ -1065,6 +1238,28 @@ mod tests { } } + #[derive(Debug)] + struct ErrorBody { + emitted: bool, + } + + impl Body for ErrorBody { + type Data = Bytes; + type Error = tonic::Status; + + fn poll_frame( + mut self: Pin<&mut Self>, + _cx: &mut Context<'_>, + ) -> Poll, Self::Error>>> { + if self.emitted { + Poll::Ready(None) + } else { + self.emitted = true; + Poll::Ready(Some(Err(tonic::Status::internal("body failed")))) + } + } + } + #[derive(Debug)] enum BodyState { Data, diff --git a/tonic-xds/src/client/retry.rs b/tonic-xds/src/client/retry.rs index c8f7c8c6d..0cb2ead25 100644 --- a/tonic-xds/src/client/retry.rs +++ b/tonic-xds/src/client/retry.rs @@ -16,6 +16,8 @@ use tower::retry::Policy; use tower::retry::Retry; use tower::{Layer, Service}; +use crate::client::circuit_breaking::is_local_circuit_breaker_drop; + /// Check if an error's source chain contains a retryable connection-level error. /// /// These are errors where the request was definitely **not** sent, making it safe to retry. @@ -56,6 +58,7 @@ pub(crate) fn is_retryable( ) -> bool { match result { Err(err) => is_retryable_connection_error(err.as_ref()), + Ok(response) if is_local_circuit_breaker_drop(response) => false, Ok(response) => { let status = tonic::Status::from_header_map(response.headers()); match status { From e90cebce9f0c82e794d58e8282fb76212a417f76 Mon Sep 17 00:00:00 2001 From: Michael Ingley Date: Thu, 23 Jul 2026 16:11:14 -0500 Subject: [PATCH 6/6] tonic-xds: address circuit breaker review feedback --- tonic-xds/src/client/circuit_breaking.rs | 44 +++++++++++++----------- 1 file changed, 24 insertions(+), 20 deletions(-) diff --git a/tonic-xds/src/client/circuit_breaking.rs b/tonic-xds/src/client/circuit_breaking.rs index 96e9e7ed9..699090eb8 100644 --- a/tonic-xds/src/client/circuit_breaking.rs +++ b/tonic-xds/src/client/circuit_breaking.rs @@ -24,17 +24,17 @@ static GLOBAL_COUNTERS: OnceLock> = OnceLock::ne /// Shared circuit-breaking state for xDS clusters. #[derive(Clone, Debug)] -pub(crate) struct ClusterCircuitBreakers { - inner: Arc, +pub(crate) struct ClusterCircuitBreakerRegistry { + inner: Arc, } #[derive(Debug)] -struct ClusterCircuitBreakersInner { +struct ClusterCircuitBreakerRegistryInner { configs: DashMap>, counters: ClusterRequestCounters, } -impl Drop for ClusterCircuitBreakersInner { +impl Drop for ClusterCircuitBreakerRegistryInner { fn drop(&mut self) { for state in self.configs.iter() { if let Some(previous) = state.config.swap(None) { @@ -47,14 +47,14 @@ impl Drop for ClusterCircuitBreakersInner { } } -impl ClusterCircuitBreakers { +impl ClusterCircuitBreakerRegistry { pub(crate) fn new() -> Self { Self::with_counters(ClusterRequestCounters::global()) } fn with_counters(counters: ClusterRequestCounters) -> Self { Self { - inner: Arc::new(ClusterCircuitBreakersInner { + inner: Arc::new(ClusterCircuitBreakerRegistryInner { configs: DashMap::new(), counters, }), @@ -193,7 +193,7 @@ impl ClusterCircuitBreakers { } } -impl Default for ClusterCircuitBreakers { +impl Default for ClusterCircuitBreakerRegistry { fn default() -> Self { Self::new() } @@ -492,12 +492,12 @@ impl Drop for CircuitBreakerPermit { /// each admitted call represents one upstream attempt rather than queued work. #[derive(Clone)] pub(crate) struct CircuitBreakingLayer { - circuit_breakers: ClusterCircuitBreakers, + circuit_breakers: ClusterCircuitBreakerRegistry, breaker_cache: Arc>>, } impl CircuitBreakingLayer { - pub(crate) fn new(circuit_breakers: ClusterCircuitBreakers) -> Self { + pub(crate) fn new(circuit_breakers: ClusterCircuitBreakerRegistry) -> Self { Self { circuit_breakers, breaker_cache: Arc::new(DashMap::new()), @@ -529,7 +529,7 @@ impl Layer for CircuitBreakingLayer { #[derive(Clone)] pub(crate) struct CircuitBreakingService { inner: S, - circuit_breakers: ClusterCircuitBreakers, + circuit_breakers: ClusterCircuitBreakerRegistry, breaker_cache: Arc>>, } @@ -611,6 +611,10 @@ enum CircuitBreakingError { NoRoutingDecision, } +/// Marks responses rejected by the local circuit breaker before reaching an endpoint. +/// +/// The retry layer uses this extension to distinguish local `UNAVAILABLE` drops +/// from retryable responses returned by an upstream service. #[derive(Clone, Copy, Debug)] struct LocalCircuitBreakerDrop; @@ -727,8 +731,8 @@ mod tests { request } - fn configured_breakers(max_requests: u32) -> ClusterCircuitBreakers { - let breakers = ClusterCircuitBreakers::new_for_test(); + fn configured_breakers(max_requests: u32) -> ClusterCircuitBreakerRegistry { + let breakers = ClusterCircuitBreakerRegistry::new_for_test(); breakers.set_config(CLUSTER, CircuitBreakingConfig { max_requests }); breakers } @@ -929,7 +933,7 @@ mod tests { Ok::<_, BoxError>(Response::new(TonicBody::empty())) }); let mut service = - CircuitBreakingLayer::new(ClusterCircuitBreakers::new_for_test()).layer(service); + CircuitBreakingLayer::new(ClusterCircuitBreakerRegistry::new_for_test()).layer(service); let response = service .ready() @@ -1003,8 +1007,8 @@ mod tests { #[test] fn shared_counters_enforce_process_limit_with_per_channel_drop_counts() { let counters = ClusterRequestCounters::isolated(); - let first_breakers = ClusterCircuitBreakers::with_counters(counters.clone()); - let second_breakers = ClusterCircuitBreakers::with_counters(counters.clone()); + let first_breakers = ClusterCircuitBreakerRegistry::with_counters(counters.clone()); + let second_breakers = ClusterCircuitBreakerRegistry::with_counters(counters.clone()); first_breakers.set_config(CLUSTER, CircuitBreakingConfig { max_requests: 1 }); second_breakers.set_config(CLUSTER, CircuitBreakingConfig { max_requests: 1 }); @@ -1023,7 +1027,7 @@ mod tests { #[test] fn default_limit_rejects_the_1025th_request() { - let breakers = ClusterCircuitBreakers::new_for_test(); + let breakers = ClusterCircuitBreakerRegistry::new_for_test(); let breaker = breakers.cluster_breaker(CLUSTER); let permits: Vec<_> = (0..DEFAULT_MAX_REQUESTS) .map(|_| breaker.acquire().unwrap()) @@ -1038,7 +1042,7 @@ mod tests { #[test] fn eds_service_name_change_uses_independent_counter() { - let breakers = ClusterCircuitBreakers::new_for_test(); + let breakers = ClusterCircuitBreakerRegistry::new_for_test(); breakers.set_cluster_config(CLUSTER, "eds-a", CircuitBreakingConfig { max_requests: 1 }); let first = breakers.acquire(CLUSTER).unwrap(); assert_eq!(breakers.in_flight(CLUSTER), 1); @@ -1058,7 +1062,7 @@ mod tests { #[test] fn idle_eds_service_name_change_cleans_up_previous_counter() { let counters = ClusterRequestCounters::isolated(); - let breakers = ClusterCircuitBreakers::with_counters(counters.clone()); + let breakers = ClusterCircuitBreakerRegistry::with_counters(counters.clone()); breakers.set_cluster_config(CLUSTER, "eds-a", CircuitBreakingConfig { max_requests: 1 }); assert_eq!(counters.counter_count(), 1); @@ -1120,7 +1124,7 @@ mod tests { #[test] fn dropping_breakers_releases_config_counter_ref() { let counters = ClusterRequestCounters::isolated(); - let breakers = ClusterCircuitBreakers::with_counters(counters.clone()); + let breakers = ClusterCircuitBreakerRegistry::with_counters(counters.clone()); breakers.set_config(CLUSTER, CircuitBreakingConfig { max_requests: 1 }); let permit = breakers.acquire(CLUSTER).unwrap(); drop(permit); @@ -1161,7 +1165,7 @@ mod tests { #[test] fn cached_cluster_breaker_does_not_pin_default_counter() { let counters = ClusterRequestCounters::isolated(); - let breakers = ClusterCircuitBreakers::with_counters(counters.clone()); + let breakers = ClusterCircuitBreakerRegistry::with_counters(counters.clone()); let breaker = breakers.cluster_breaker(CLUSTER); let permit = breaker.acquire().unwrap(); assert_eq!(counters.counter_count(), 1);