From 36b03763c88d1a1db10abd7c029c2306966a511d Mon Sep 17 00:00:00 2001 From: Michael Ingley Date: Tue, 30 Jun 2026 11:55:31 -0500 Subject: [PATCH 1/5] 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 c4559d3e654de9a7028f678ea395e3323449d143 Mon Sep 17 00:00:00 2001 From: Michael Ingley Date: Tue, 30 Jun 2026 14:14:03 -0500 Subject: [PATCH 2/5] 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 3/5] 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 4/5] 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 5/5] 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);