From c1d28a038e099e4db4751ce42579266288d5c3ee Mon Sep 17 00:00:00 2001 From: YutaoMa Date: Thu, 6 Aug 2026 14:44:29 -0700 Subject: [PATCH 1/2] fix(tonic-xds): one config load per request, 3 fewer hot-path allocs ## Motivation The routing layer loaded the config twice per request with the header-mutation hook in between, so the hook and the matching that consumes its output could see different versions and the request would silently take the default route. The same split also skipped the hook until the first config arrived. Matching was async only to host that wait, so every request paid a nested boxed future and the owned copies a `'static` future forces; the second load handed the hook an owned metadata clone, and the config `Arc` stayed alive for the whole request rather than just the decision. ## Solution One required trait method returning either an available config or a future for the first one. Matching borrows it and turns synchronous, dropping a nested boxed future, an authority `String`, a `HeaderMap` clone and a metadata deep clone; block-scoped so the `Arc` releases before the inner call. --- tonic-xds/src/client/channel.rs | 16 +- tonic-xds/src/client/route.rs | 219 ++++++++++++++++++--- tonic-xds/src/xds/resource/route_config.rs | 22 ++- tonic-xds/src/xds/routing.rs | 70 +++---- 4 files changed, 263 insertions(+), 64 deletions(-) diff --git a/tonic-xds/src/client/channel.rs b/tonic-xds/src/client/channel.rs index 8700437a4..d0272099a 100644 --- a/tonic-xds/src/client/channel.rs +++ b/tonic-xds/src/client/channel.rs @@ -459,7 +459,6 @@ mod tests { use crate::client::route::RouteDecision; use crate::client::route::RouteInput; use crate::client::route::Router; - use crate::common::async_util::BoxFuture; use crate::testutil::grpc::GreeterClient; use crate::testutil::grpc::HelloRequest; use crate::testutil::grpc::TestServer; @@ -513,15 +512,18 @@ mod tests { } impl Router for MockXdsManager { + fn acquire(&self) -> crate::client::route::AcquiredConfig { + crate::client::route::AcquiredConfig::Ready(Arc::new(Default::default())) + } + fn route( &self, _input: &RouteInput<'_>, - ) -> BoxFuture> { - Box::pin(async move { - Ok(RouteDecision { - cluster: "test-cluster".to_string(), - request_hash: None, - }) + _config: &crate::xds::resource::route_config::RouteConfigResource, + ) -> Result { + Ok(RouteDecision { + cluster: "test-cluster".to_string(), + request_hash: None, }) } } diff --git a/tonic-xds/src/client/route.rs b/tonic-xds/src/client/route.rs index d3f47afe0..0e8ff99e9 100644 --- a/tonic-xds/src/client/route.rs +++ b/tonic-xds/src/client/route.rs @@ -23,7 +23,7 @@ */ use crate::common::async_util::BoxFuture; -use crate::xds::resource::route_config::RouteConfigMetadata; +use crate::xds::resource::route_config::{RouteConfigMetadata, RouteConfigResource}; use crate::xds::routing::RoutingError; use http::Request; use std::sync::Arc; @@ -64,19 +64,44 @@ pub trait PreRouteInterceptor: Send + Sync + 'static { fn on_request(&self, headers: &mut http::HeaderMap, metadata: &RouteConfigMetadata); } +/// A route config obtained in one step, to serve a single request with. +/// +/// Two cases so the common one -- a config is already in effect -- costs +/// neither an allocation nor an await, while a request racing startup can still +/// wait for the first config. +pub(crate) enum AcquiredConfig { + /// A config is already in effect. + Ready(Arc), + /// None has arrived yet; await this for the first one, bounded by an + /// implementation-defined timeout. + Pending(BoxFuture, RoutingError>>), +} + +impl AcquiredConfig { + /// Resolves to the config, awaiting only when one is not already available. + pub(crate) async fn get(self) -> Result, RoutingError> { + match self { + Self::Ready(config) => Ok(config), + Self::Pending(wait) => wait.await, + } + } +} + /// Trait for routing requests to clusters. /// /// Implementations resolve a request's authority and headers into a target /// cluster name. The xDS-backed implementation is /// [`XdsRouter`](crate::xds::routing::XdsRouter). pub(crate) trait Router: Send + Sync + 'static { - fn route(&self, input: &RouteInput<'_>) -> BoxFuture>; + /// Obtains the route config to serve one request with. + fn acquire(&self) -> AcquiredConfig; - /// Current route-config metadata, if available, used to feed a - /// [`PreRouteInterceptor`]. Defaults to `None` (e.g. for mock routers). - fn metadata(&self) -> Option { - None - } + /// Resolves `input` against `config`. + fn route( + &self, + input: &RouteInput<'_>, + config: &RouteConfigResource, + ) -> Result; } /// Tower service for routing requests to the appropriate cluster. @@ -115,17 +140,19 @@ where let authority = self.authority.clone(); let mut inner_service = self.inner.clone(); Box::pin(async move { - if let Some(interceptor) = interceptor.as_ref() - && let Some(metadata) = router.metadata() - { - interceptor.on_request(request.headers_mut(), &metadata); - } - let headers = &request.headers(); - let route_input = RouteInput { - authority: &authority, - headers, + // Scoped so the config `Arc` is dropped before the inner call: + // to avoid holding that config version for the life of a long-running stream. + let route_decision = { + let config = router.acquire().get().await?; + if let Some(interceptor) = interceptor.as_ref() { + interceptor.on_request(request.headers_mut(), &config.metadata); + } + let route_input = RouteInput { + authority: &authority, + headers: request.headers(), + }; + router.route(&route_input, &config)? }; - let route_decision = router.route(&route_input).await?; request.extensions_mut().insert(route_decision); inner_service.call(request).await.map_err(Into::into) }) @@ -187,19 +214,165 @@ mod tests { } impl Router for CaptureAuthorityRouter { - fn route(&self, input: &RouteInput<'_>) -> BoxFuture> { + fn acquire(&self) -> AcquiredConfig { + AcquiredConfig::Ready(Arc::new(RouteConfigResource::default())) + } + + fn route( + &self, + input: &RouteInput<'_>, + _config: &RouteConfigResource, + ) -> Result { *self.captured.lock().unwrap() = Some(input.authority.to_string()); - Box::pin(async move { + Ok(RouteDecision { + cluster: "test-cluster".to_string(), + request_hash: None, + }) + } + } + + #[tokio::test] + async fn interceptor_runs_when_config_arrives_late() { + use crate::xds::cache::XdsCache; + use crate::xds::resource::route_config::{ + PathSpecifierConfig, RouteConfig, RouteConfigAction, RouteConfigMatch, + VirtualHostConfig, + }; + use crate::xds::routing::XdsRouter; + + struct MarkPartition; + impl PreRouteInterceptor for MarkPartition { + fn on_request(&self, headers: &mut http::HeaderMap, metadata: &RouteConfigMetadata) { + let partition = metadata + .filter_metadata("partition") + .expect("interceptor ran without config metadata"); + headers.insert( + "x-partition", + http::HeaderValue::from_bytes(&partition).expect("valid header"), + ); + } + } + + // Start with no route config at all. + let cache = XdsCache::new(); + let xds_router = XdsRouter::new(&cache); + assert!( + xds_router.snapshot().is_none(), + "precondition: no config yet" + ); + let router: Arc = Arc::new(xds_router); + + let seen: Arc>>> = Arc::new(Mutex::new(Vec::new())); + let sink = seen.clone(); + let inner = service_fn(move |req: Request<()>| { + let sink = sink.clone(); + async move { + sink.lock().unwrap().push( + req.headers() + .get("x-partition") + .map(|v| v.to_str().expect("utf8").to_string()), + ); + Ok::<_, BoxError>(http::Response::new(())) + } + }); + let svc = XdsRoutingLayer::new(router, Some(Arc::new(MarkPartition)), Arc::from("svc")) + .layer(inner); + + let call = tokio::spawn(svc.oneshot(Request::builder().uri("/pkg.S/M").body(()).unwrap())); + + tokio::time::sleep(std::time::Duration::from_millis(50)).await; + assert!( + seen.lock().unwrap().is_empty(), + "request was served before any config existed", + ); + + let mut filter_metadata = std::collections::HashMap::new(); + filter_metadata.insert("partition".to_string(), bytes::Bytes::from_static(b"7")); + cache.update_route_config(Arc::new(RouteConfigResource { + name: "rc".into(), + virtual_hosts: vec![VirtualHostConfig { + name: "vh".into(), + domains: vec!["svc".into()], + routes: vec![RouteConfig { + match_criteria: RouteConfigMatch { + path_specifier: PathSpecifierConfig::Prefix("/".into()), + headers: vec![], + case_sensitive: true, + match_fraction: None, + }, + action: RouteConfigAction::Cluster("c".into()), + }], + }], + metadata: RouteConfigMetadata::from_encoded( + filter_metadata, + std::collections::HashMap::new(), + ), + })); + + call.await.expect("task").expect("request"); + let seen = seen.lock().unwrap(); + assert_eq!(seen.len(), 1, "expected exactly one request to be served"); + assert_eq!( + seen[0].as_deref(), + Some("7"), + "interceptor did not run for a request that predated the first config", + ); + } + + #[tokio::test] + async fn releases_the_route_config_before_calling_the_inner_service() { + struct SharedConfigRouter { + config: Arc, + } + + impl Router for SharedConfigRouter { + fn acquire(&self) -> AcquiredConfig { + AcquiredConfig::Ready(self.config.clone()) + } + + fn route( + &self, + _input: &RouteInput<'_>, + _config: &RouteConfigResource, + ) -> Result { Ok(RouteDecision { - cluster: "test-cluster".to_string(), + cluster: "c".to_string(), request_hash: None, }) - }) + } } + + let config = Arc::new(RouteConfigResource::default()); + let router: Arc = Arc::new(SharedConfigRouter { + config: config.clone(), + }); + + let baseline = Arc::strong_count(&config); + + let observed = Arc::new(Mutex::new(None)); + let sink = observed.clone(); + let probe = Arc::downgrade(&config); + let inner = service_fn(move |_req: Request<()>| { + let sink = sink.clone(); + let probe = probe.clone(); + async move { + *sink.lock().unwrap() = Some(probe.strong_count()); + Ok::<_, BoxError>(http::Response::new(())) + } + }); + let svc = XdsRoutingLayer::new(router, None, Arc::from("svc")).layer(inner); + + svc.oneshot(Request::builder().uri("/pkg.S/M").body(()).unwrap()) + .await + .expect("request"); + + assert_eq!( + observed.lock().unwrap().expect("inner service ran"), + baseline, + "the routing layer still held the route config while the request ran", + ); } - /// Verifies the routing layer always sources `authority` from its layer - /// config, not from the request URI. #[tokio::test] async fn uses_layer_authority_regardless_of_request_uri() { let captured = Arc::new(Mutex::new(None)); diff --git a/tonic-xds/src/xds/resource/route_config.rs b/tonic-xds/src/xds/resource/route_config.rs index fd201a87c..68a09b53e 100644 --- a/tonic-xds/src/xds/resource/route_config.rs +++ b/tonic-xds/src/xds/resource/route_config.rs @@ -91,6 +91,26 @@ impl RouteConfigMetadata { self.filter_metadata.is_empty() && self.typed_filter_metadata.is_empty() } + /// Constructs a `RouteConfigMetadata` directly from pre-encoded bytes. + /// + /// Lets downstream crates build realistic metadata to unit-test a + /// [`PreRouteInterceptor`](crate::PreRouteInterceptor) without standing up + /// a control plane. + /// + /// Takes already-encoded `google.protobuf.Struct` bytes so that this does + /// not expose proto binding types to the public API surface. + #[cfg(any(test, feature = "testutil"))] + #[must_use] + pub fn from_encoded( + filter_metadata: HashMap, + typed_filter_metadata: HashMap, + ) -> Self { + Self { + filter_metadata, + typed_filter_metadata, + } + } + /// Builds the view from an xDS `Metadata`, pre-encoding each namespace's /// `Struct`/`Any` to bytes. pub(crate) fn from_proto(metadata: Metadata) -> Self { @@ -120,7 +140,7 @@ impl RouteConfigMetadata { } /// Validated RouteConfiguration. -#[derive(Debug, Clone)] +#[derive(Debug, Clone, Default)] pub(crate) struct RouteConfigResource { pub name: String, pub virtual_hosts: Vec, diff --git a/tonic-xds/src/xds/routing.rs b/tonic-xds/src/xds/routing.rs index d06761b78..eed0adf06 100644 --- a/tonic-xds/src/xds/routing.rs +++ b/tonic-xds/src/xds/routing.rs @@ -43,14 +43,13 @@ use std::time::Duration; use arc_swap::ArcSwapOption; use tokio::sync::watch; -use crate::client::route::{RouteDecision, RouteInput, Router}; -use crate::common::async_util::{AbortOnDrop, BoxFuture}; +use crate::client::route::{AcquiredConfig, RouteDecision, RouteInput, Router}; +use crate::common::async_util::AbortOnDrop; use crate::xds::cache::XdsCache; use crate::xds::resource::hash_policy::HashPolicyConfig; use crate::xds::resource::route_config::{ HeaderMatchSpecifierConfig, HeaderMatcherConfig, PathSpecifierConfig, RouteConfig, - RouteConfigAction, RouteConfigMatch, RouteConfigMetadata, RouteConfigResource, - VirtualHostConfig, WeightedCluster, + RouteConfigAction, RouteConfigMatch, RouteConfigResource, VirtualHostConfig, WeightedCluster, }; /// Default timeout for waiting for the initial route config (matches gRFC A57 @@ -99,35 +98,37 @@ impl XdsRouter { _watch_task: AbortOnDrop(handle), } } + + /// The route config currently in effect, or `None` if none has arrived yet. + pub(crate) fn snapshot(&self) -> Option> { + self.route_config.load_full() + } } impl Router for XdsRouter { - fn route(&self, input: &RouteInput<'_>) -> BoxFuture> { - let authority = input.authority.to_string(); - let headers = input.headers.clone(); - - // Fast path: config already available, no cloning needed. - if let Some(rc) = self.route_config.load_full() { - return Box::pin(async move { resolve_route(&rc, &authority, &headers) }); + fn acquire(&self) -> AcquiredConfig { + if let Some(config) = self.snapshot() { + return AcquiredConfig::Ready(config); } - - // Slow path: wait for the initial route config, matching standard - // gRPC behavior where RPCs block until the resolver provides the - // first update. + // Wait for the initial route config, matching standard gRPC behavior + // where RPCs block until the resolver provides the first update. let route_config_ref = self.route_config.clone(); let mut ready_rx = self.ready_rx.clone(); - Box::pin(async move { + AcquiredConfig::Pending(Box::pin(async move { tokio::time::timeout(DEFAULT_READY_TIMEOUT, ready_rx.wait_for(|ready| *ready)) .await .map_err(|_| RoutingError::NotReady)? .map_err(|_| RoutingError::NotReady)?; - let rc = route_config_ref.load_full().ok_or(RoutingError::NotReady)?; - resolve_route(&rc, &authority, &headers) - }) + route_config_ref.load_full().ok_or(RoutingError::NotReady) + })) } - fn metadata(&self) -> Option { - self.route_config.load_full().map(|rc| rc.metadata.clone()) + fn route( + &self, + input: &RouteInput<'_>, + config: &RouteConfigResource, + ) -> Result { + resolve_route(config, input.authority, input.headers) } } @@ -1089,7 +1090,8 @@ mod tests { authority: "my-service", headers: &headers, }; - let decision = router.route(&input).await.unwrap(); + let config = router.snapshot().expect("config"); + let decision = router.route(&input, &config).unwrap(); assert_eq!(decision.cluster, "my-cluster"); } @@ -1107,13 +1109,15 @@ mod tests { headers: &headers, }; - let decision = router.route(&input).await.unwrap(); + let config = router.snapshot().expect("config"); + let decision = router.route(&input, &config).unwrap(); assert_eq!(decision.cluster, "cluster-a"); cache.update_route_config(make_route_config("cluster-b")); tokio::task::yield_now().await; - let decision = router.route(&input).await.unwrap(); + let config = router.snapshot().expect("config"); + let decision = router.route(&input, &config).unwrap(); assert_eq!(decision.cluster, "cluster-b"); } @@ -1122,15 +1126,15 @@ mod tests { let cache = XdsCache::new(); let router = XdsRouter::new(&cache); - let headers = http::HeaderMap::new(); - let input = RouteInput { - authority: "svc", - headers: &headers, - }; - // The router now blocks waiting for config; verify it returns - // NotReady after the timeout elapses. - let result = - tokio::time::timeout(std::time::Duration::from_millis(100), router.route(&input)).await; + // With no config yet, `acquire` yields the waiting variant, which + // blocks and then reports NotReady. + assert!(router.snapshot().is_none()); + assert!(matches!(router.acquire(), AcquiredConfig::Pending(_))); + let result = tokio::time::timeout( + std::time::Duration::from_millis(100), + router.acquire().get(), + ) + .await; // Either the inner timeout fires (NotReady) or the outer timeout // fires (config never arrived) — both are correct. match result { From f819c82ea9975e040969ff2b86c617efc6f28aa4 Mon Sep 17 00:00:00 2001 From: YutaoMa Date: Tue, 11 Aug 2026 13:39:04 -0700 Subject: [PATCH 2/2] fix: remove the duplicate timeout --- tonic-xds/src/xds/routing.rs | 27 ++++++++++++++------------- 1 file changed, 14 insertions(+), 13 deletions(-) diff --git a/tonic-xds/src/xds/routing.rs b/tonic-xds/src/xds/routing.rs index eed0adf06..4740487fe 100644 --- a/tonic-xds/src/xds/routing.rs +++ b/tonic-xds/src/xds/routing.rs @@ -1121,7 +1121,7 @@ mod tests { assert_eq!(decision.cluster, "cluster-b"); } - #[tokio::test] + #[tokio::test(start_paused = true)] async fn xds_router_returns_not_ready_without_config() { let cache = XdsCache::new(); let router = XdsRouter::new(&cache); @@ -1130,17 +1130,18 @@ mod tests { // blocks and then reports NotReady. assert!(router.snapshot().is_none()); assert!(matches!(router.acquire(), AcquiredConfig::Pending(_))); - let result = tokio::time::timeout( - std::time::Duration::from_millis(100), - router.acquire().get(), - ) - .await; - // Either the inner timeout fires (NotReady) or the outer timeout - // fires (config never arrived) — both are correct. - match result { - Ok(Err(RoutingError::NotReady)) => {} - Err(_elapsed) => {} - other => panic!("expected NotReady or timeout, got {other:?}"), - } + + let start = tokio::time::Instant::now(); + let result = router.acquire().get().await; + + assert!( + matches!(result, Err(RoutingError::NotReady)), + "expected NotReady, got {result:?}", + ); + assert!( + start.elapsed() >= DEFAULT_READY_TIMEOUT, + "expected the wait to span the full ready timeout, took {:?}", + start.elapsed(), + ); } }