diff --git a/tonic-xds/src/client/circuit_breaking.rs b/tonic-xds/src/client/circuit_breaking.rs new file mode 100644 index 000000000..616aa6e83 --- /dev/null +++ b/tonic-xds/src/client/circuit_breaking.rs @@ -0,0 +1,1101 @@ +#![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 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 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, +} + +impl Drop for ClusterCircuitBreakersInner { + fn drop(&mut self) { + for state in self.configs.iter() { + if let Some(previous) = state.config.swap(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, + }), + } + } + + #[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 = 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, + 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 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, + config: CircuitBreakerRuntimeConfig, + ) { + let previous = state.current_config(); + if previous.as_deref() == 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.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.swap(None) { + self.deactivate_config(previous); + } + } + + fn deactivate_config(&self, config: Arc) { + let counter_key = config.counter_key.clone(); + drop(config); + self.inner.counters.deactivate(&counter_key); + } + + fn acquire(&self, cluster: &str) -> Result { + self.cluster_breaker(cluster).acquire() + } + + #[cfg(test)] + fn in_flight(&self, cluster: &str) -> u32 { + 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 breaker = self.cluster_breaker(cluster); + self.inner + .counters + .dropped_requests(&breaker.current_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: CounterKey, + 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 {} + +#[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 { + 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 { + Self { + config: ArcSwapOption::empty(), + } + } + + fn current_config(&self) -> Option> { + self.config.load_full() + } +} + +#[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, +} + +#[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: &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, + 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: &CounterKey) -> Arc { + self.inner + .counters + .entry(counter_key.clone()) + .or_insert_with(|| Arc::new(InFlightCounter::default())) + .clone() + } + + 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 + }); + 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)] + fn in_flight(&self, counter_key: &CounterKey) -> u32 { + self.inner + .counters + .get(counter_key) + .map(|counter| counter.in_flight()) + .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() + } +} + +#[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: CounterKey, + 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, + breaker_cache: Arc>>, +} + +impl CircuitBreakingLayer { + pub(crate) fn new(circuit_breakers: ClusterCircuitBreakers) -> Self { + Self { + circuit_breakers, + breaker_cache: Arc::new(DashMap::new()), + } + } +} + +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() + } +} + +impl Layer for CircuitBreakingLayer { + type Service = CircuitBreakingService; + + fn layer(&self, service: S) -> Self::Service { + CircuitBreakingService { + inner: service, + circuit_breakers: self.circuit_breakers.clone(), + breaker_cache: self.breaker_cache.clone(), + } + } +} + +#[derive(Clone)] +pub(crate) struct CircuitBreakingService { + inner: S, + circuit_breakers: ClusterCircuitBreakers, + breaker_cache: Arc>>, +} + +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) + .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> + + 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(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(), + ))) + }); + }; + + 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 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)))) + }) + } +} + +#[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 = CounterKey::new(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); + } + + #[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); + 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, }