diff --git a/tonic-xds/Cargo.toml b/tonic-xds/Cargo.toml index 3502905cc..a68628199 100644 --- a/tonic-xds/Cargo.toml +++ b/tonic-xds/Cargo.toml @@ -73,6 +73,7 @@ rcgen = "0.14" google-cloud-auth = { version = "1.9", default-features = false } [features] +default = [] testutil = ["dep:tonic-prost"] # Bundled OpenTelemetry metrics recorder for the gRFC A78 xDS client metrics. @@ -80,6 +81,12 @@ testutil = ["dep:tonic-prost"] # OpenTelemetry version is decoupled from the core crates. otel = ["dep:opentelemetry", "dep:xds-client-opentelemetry"] +# Load-balancer implementation. The default `tower-lb` stack (tower p2c +# `Balance` + `Buffer`) is used unless `tonic-xds-lb` is enabled, which switches +# to the in-crate `loadbalance/` implementation (`LoadBalancer` + pickers + +# outlier detection). +tonic-xds-lb = [] + # TLS crypto backend — pick exactly one. _tls-any = ["dep:rustls", "dep:rustls-pemfile", "dep:x509-parser"] tls-ring = ["_tls-any", "tonic/tls-ring", "xds-client/tonic-tls-ring"] diff --git a/tonic-xds/src/client/channel.rs b/tonic-xds/src/client/channel.rs index 7a6884da1..e82ffe474 100644 --- a/tonic-xds/src/client/channel.rs +++ b/tonic-xds/src/client/channel.rs @@ -1,26 +1,34 @@ -use crate::client::cluster::ClusterClientRegistryGrpc; -use crate::client::endpoint::{EndpointAddress, EndpointChannel}; -use crate::client::lb::{ClusterDiscovery, XdsLbService}; +use crate::TonicCallCredentials; +use crate::XdsUri; +use crate::client::retry::{GrpcRetryPolicy, GrpcRetryPolicyConfig, RetryLayer}; use crate::client::route::{Router, XdsRoutingLayer}; use crate::xds::bootstrap::{BootstrapConfig, BootstrapError}; use crate::xds::cache::XdsCache; #[cfg(feature = "_tls-any")] use crate::xds::cert_provider::{CertProviderError, CertProviderRegistry}; -use crate::xds::cluster_discovery::XdsClusterDiscovery; use crate::xds::resource_manager::XdsResourceManager; use crate::xds::routing::XdsRouter; -use crate::{TonicCallCredentials, XdsUri}; use http::Request; use std::fmt::Debug; use std::sync::Arc; use std::task::{Context, Poll}; -use tonic::{body::Body as TonicBody, client::GrpcService, transport::channel::Channel}; +use tonic::{body::Body as TonicBody, client::GrpcService}; use tower::{BoxError, Service, ServiceBuilder, util::BoxCloneSyncService}; use xds_client::{ ClientConfig, MetricsRecorder, Node, ProstCodec, TokioRuntime, TonicTransportBuilder, XdsClient, }; -use crate::client::retry::{GrpcRetryPolicy, GrpcRetryPolicyConfig, RetryLayer}; +cfg_tower_lb! { + use crate::client::cluster::ClusterClientRegistryGrpc; + use crate::client::endpoint::{EndpointAddress, EndpointChannel}; + use crate::client::lb::{ClusterDiscovery, XdsLbService}; + use crate::xds::cluster_discovery::XdsClusterDiscovery; + use tonic::transport::channel::Channel; +} + +cfg_tonic_xds_lb! { + use crate::client::loadbalance::service::XdsLoadBalanceService; +} /// Configuration for building [`XdsChannel`] / [`XdsChannelGrpc`]. #[derive(Clone, Debug)] @@ -291,14 +299,28 @@ impl XdsChannelBuilder { resource_manager: XdsResourceManager, ) -> XdsChannelGrpc { let router: Arc = Arc::new(XdsRouter::new(&cache)); - #[cfg(feature = "_tls-any")] - let discovery: Arc< - dyn ClusterDiscovery>, - > = Arc::new(XdsClusterDiscovery::new(cache, cert_provider_registry)); - #[cfg(not(feature = "_tls-any"))] - let discovery: Arc< - dyn ClusterDiscovery>, - > = Arc::new(XdsClusterDiscovery::new(cache)); + + #[cfg(not(feature = "tonic-xds-lb"))] + let lb_service = { + #[cfg(feature = "_tls-any")] + let discovery: Arc< + dyn ClusterDiscovery>, + > = Arc::new(XdsClusterDiscovery::new(cache, cert_provider_registry)); + #[cfg(not(feature = "_tls-any"))] + let discovery: Arc< + dyn ClusterDiscovery>, + > = Arc::new(XdsClusterDiscovery::new(cache)); + let cluster_registry = Arc::new(ClusterClientRegistryGrpc::new()); + XdsLbService::new(cluster_registry, discovery) + }; + + #[cfg(feature = "tonic-xds-lb")] + let lb_service = XdsLoadBalanceService::new( + cache, + #[cfg(feature = "_tls-any")] + cert_provider_registry, + ); + let retry_policy = GrpcRetryPolicy::new(GrpcRetryPolicyConfig::default()); let resources = Arc::new(XdsChannelResources { @@ -308,8 +330,6 @@ impl XdsChannelBuilder { let routing_layer = XdsRoutingLayer::new(router, self.authority()); let retry_layer = RetryLayer::new(retry_policy); - let cluster_registry = Arc::new(ClusterClientRegistryGrpc::new()); - let lb_service = XdsLbService::new(cluster_registry, discovery); let inner = ServiceBuilder::new() .layer(routing_layer) .layer(retry_layer) @@ -332,18 +352,65 @@ impl XdsChannelBuilder { self.build_tonic_grpc_channel() } - /// Builds an `XdsChannelGrpc` from the given router, cluster discovery, and retry policy. + /// Test-only: builds an `XdsChannelGrpc` for the `tower-lb` backend + /// (router + `XdsLbService`, no resource manager) from a pre-populated + /// cache. Both backend constructors are compiled in test builds (see the + /// `any(test, …)` module gates), so the channel tests run against each. #[cfg(test)] - pub(crate) fn build_grpc_channel_from_parts( + pub(crate) fn build_grpc_channel_from_cache_tower( &self, - router: Arc, - discovery: Arc>>, + cache: Arc, retry_policy: GrpcRetryPolicy, ) -> XdsChannelGrpc { - let routing_layer = XdsRoutingLayer::new(router, self.authority()); - let retry_layer = RetryLayer::new(retry_policy); + let router: Arc = Arc::new(XdsRouter::new(&cache)); + #[cfg(feature = "_tls-any")] + let discovery: Arc< + dyn ClusterDiscovery>, + > = Arc::new(XdsClusterDiscovery::new( + cache, + Arc::new(CertProviderRegistry::from_bootstrap(&Default::default()).unwrap()), + )); + #[cfg(not(feature = "_tls-any"))] + let discovery: Arc< + dyn ClusterDiscovery>, + > = Arc::new(XdsClusterDiscovery::new(cache)); let cluster_registry = Arc::new(ClusterClientRegistryGrpc::new()); let lb_service = XdsLbService::new(cluster_registry, discovery); + + let routing_layer = XdsRoutingLayer::new(router, self.authority()); + let retry_layer = RetryLayer::new(retry_policy); + let inner = ServiceBuilder::new() + .layer(routing_layer) + .layer(retry_layer) + .map_request(|req: Request>| { + req.map(TonicBody::new) + }) + .service(lb_service); + BoxCloneSyncService::new(XdsChannel { + config: self.config.clone(), + inner, + _resources: None, + }) + } + + /// Test-only: builds an `XdsChannelGrpc` for the `tonic-xds-lb` backend + /// (router + `XdsLoadBalanceService`, no resource manager) from a + /// pre-populated cache. + #[cfg(test)] + pub(crate) fn build_grpc_channel_from_cache_xds( + &self, + cache: Arc, + retry_policy: GrpcRetryPolicy, + ) -> XdsChannelGrpc { + let router: Arc = Arc::new(XdsRouter::new(&cache)); + let lb_service = XdsLoadBalanceService::new( + cache, + #[cfg(feature = "_tls-any")] + Arc::new(CertProviderRegistry::from_bootstrap(&Default::default()).unwrap()), + ); + + let routing_layer = XdsRoutingLayer::new(router, self.authority()); + let retry_layer = RetryLayer::new(retry_policy); let inner = ServiceBuilder::new() .layer(routing_layer) .layer(retry_layer) @@ -365,263 +432,38 @@ impl XdsChannelBuilder { } } +/// Feature-agnostic test helpers shared by both LB backends' channel tests. #[cfg(test)] -mod tests { - use super::{XdsChannelBuilder, XdsChannelConfig}; +mod test_support { + use super::{XdsChannelConfig, XdsChannelGrpc}; use crate::XdsUri; - use crate::client::channel::XdsChannelGrpc; use crate::client::endpoint::EndpointAddress; - use crate::client::endpoint::EndpointChannel; - - fn test_config() -> XdsChannelConfig { - XdsChannelConfig::new(XdsUri::parse("xds:///test-service").unwrap()) - } - use crate::client::lb::{BoxDiscover, ClusterDiscovery}; - use crate::client::retry::GrpcRetryPolicy; - 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; - use crate::xds::cache::XdsCache; + use crate::testutil::grpc::{GreeterClient, HelloRequest, TestServer, spawn_greeter_server}; use crate::xds::resource::EndpointsResource; use crate::xds::resource::route_config::RouteConfigResource; + use std::collections::HashMap; use std::sync::Arc; - use tokio::sync::mpsc; - use tonic::transport::Channel; - use tower::discover::Change; - /// Sets up multiple gRPC test servers and returns their addresses, clients and shutdown handles. - async fn setup_grpc_servers( - count: usize, - ) -> (Vec, Vec) { - use crate::testutil::grpc::spawn_greeter_server; + pub(super) fn test_config() -> XdsChannelConfig { + XdsChannelConfig::new(XdsUri::parse("xds:///test-service").unwrap()) + } + /// Spawns `count` greeter servers named `server-0` .. `server-{count-1}`. + pub(super) async fn setup_grpc_servers(count: usize) -> Vec { let mut servers = Vec::new(); - let mut server_addrs = Vec::new(); - for i in 0..count { - let server_name = format!("server-{i}"); - let server = spawn_greeter_server(&server_name, None, None) + let server = spawn_greeter_server(&format!("server-{i}"), None) .await .expect("Failed to spawn gRPC server"); - - server_addrs.push(server.addr.to_string()); servers.push(server); } - - (server_addrs, servers) - } - - /// A mock XdsManager that provides pre-configured endpoints for testing. - struct MockXdsManager { - endpoints: Vec<(EndpointAddress, Channel)>, - } - - impl MockXdsManager { - /// Creates a new MockXdsManager from test servers. - fn from_test_servers(servers: &[TestServer]) -> Self { - let endpoints = servers - .iter() - .map(|s| { - let addr = EndpointAddress::from(s.addr); - (addr, s.channel.clone()) - }) - .collect(); - Self { endpoints } - } - } - - impl Router for MockXdsManager { - fn route( - &self, - _input: &RouteInput<'_>, - ) -> BoxFuture> { - Box::pin(async move { - Ok(RouteDecision { - cluster: "test-cluster".to_string(), - request_hash: None, - }) - }) - } - } - - impl ClusterDiscovery> for MockXdsManager { - fn discover_cluster( - &self, - _cluster_name: &str, - ) -> BoxDiscover> { - let endpoints = self.endpoints.clone(); - let (tx, rx) = mpsc::channel(16); - - tokio::spawn(async move { - for (addr, channel) in endpoints { - let endpoint_channel = EndpointChannel::new(channel); - let change = Change::Insert(addr, endpoint_channel); - tx.send(Ok(change)).await.expect("Failed to send SD change"); - } - }); - - Box::pin(tokio_stream::wrappers::ReceiverStream::new(rx)) - } - } - - /// Sends multiple gRPC requests using the provided client and returns statistics about the requests. - async fn send_grpc_requests( - mut grpc_client: crate::testutil::grpc::GreeterClient, - num_requests: usize, - ) -> ( - usize, - std::collections::HashMap, - std::collections::HashMap, - ) { - let mut successful_requests = 0; - let mut error_types = std::collections::HashMap::new(); - let mut server_counts = std::collections::HashMap::new(); - - for i in 0..num_requests { - let request_timeout = tokio::time::Duration::from_secs(3); - let request_future = grpc_client.say_hello(HelloRequest { - name: format!("test-request-{i}"), - }); - - match tokio::time::timeout(request_timeout, request_future).await { - Ok(Ok(response)) => { - successful_requests += 1; - // Extract server name from response message (format: "server-X: test-request-Y") - let message = response.into_inner().message; - if let Some(server_name) = message.split(':').next() { - *server_counts.entry(server_name.to_string()).or_insert(0) += 1; - } - } - Ok(Err(e)) => { - let error_type = format!("{e:?}").chars().take(80).collect::(); - *error_types.entry(error_type).or_insert(0) += 1; - } - Err(_) => { - *error_types.entry("Timeout".to_string()).or_insert(0) += 1; - if error_types.get("Timeout").unwrap_or(&0) > &2 { - break; - } - } - } - } - - (successful_requests, error_types, server_counts) - } - - #[tokio::test] - /// Tests the `XdsChannelGrpc` with a power-of-two-choices load balancer. - async fn test_xds_channel_grpc_with_p2c_lb() { - let num_requests = 1000; - let num_servers = 5; - let (_, servers) = setup_grpc_servers(num_servers).await; - - // Create a mock XdsManager with the test servers - let xds_manager = Arc::new(MockXdsManager::from_test_servers(&servers)); - - let xds_channel_builder = XdsChannelBuilder::new(test_config()); - let xds_channel = xds_channel_builder.build_grpc_channel_from_parts( - xds_manager.clone(), - xds_manager.clone(), - GrpcRetryPolicy::default(), - ); - - let client = GreeterClient::new(xds_channel); - - let (successful_requests, error_types, server_counts) = - send_grpc_requests(client, num_requests).await; - - println!("Successful requests: {successful_requests}"); - println!("Error types: {error_types:?}"); - println!("Per-server call counts: {server_counts:?}"); - - assert_eq!( - successful_requests, num_requests, - "Expected 100% success rate. Got {successful_requests} successful out of {num_requests} requests. Errors: {error_types:?}", - ); - - assert!( - error_types.is_empty(), - "Expected no errors but got: {error_types:?}", - ); - - let actual_server_count = server_counts.len(); - assert_eq!( - actual_server_count, num_servers, - "Expected all {num_servers} servers to receive requests, but only {actual_server_count} servers received traffic. Server counts: {server_counts:?}", - ); - - let expected_per_server = num_requests / num_servers; - let min_requests_per_server = (expected_per_server as f64 / 1.5) as usize; - let max_requests_per_server = (expected_per_server as f64 * 1.5) as usize; - - for (server_name, count) in &server_counts { - assert!( - *count >= min_requests_per_server, - "Server {server_name} received only {count} requests, expected at least {min_requests_per_server} (expected ~{expected_per_server} per server with 1.5x variance)", - ); - assert!( - *count <= max_requests_per_server, - "Server {server_name} received {count} requests, expected at most {max_requests_per_server} (expected ~{expected_per_server} per server with 1.5x variance)", - ); - } - - let total_server_requests: usize = server_counts.values().sum(); - assert_eq!( - total_server_requests, successful_requests, - "Total server requests ({total_server_requests}) should equal successful requests ({successful_requests}). Server counts: {server_counts:?}", - ); - - for server in servers { - let _ = server.shutdown.send(()); - let _ = server.handle.await; - } + servers } - #[tokio::test] - async fn test_retry_once_on_unavailable() { - use crate::client::retry::{GrpcRetryPolicy, GrpcRetryPolicyConfig}; - use crate::testutil::grpc::spawn_fail_first_n_server; - - // Server fails the first request with UNAVAILABLE, succeeds on retry. - let server = spawn_fail_first_n_server("retry-server", 1) - .await - .expect("Failed to spawn server"); - - let servers = vec![server]; - let xds_manager = Arc::new(MockXdsManager::from_test_servers(&servers)); - - let retry_policy = GrpcRetryPolicy::new( - GrpcRetryPolicyConfig::new() - .retry_on(vec![tonic::Code::Unavailable]) - .num_retries(1), - ); - - let xds_channel = XdsChannelBuilder::new(test_config()).build_grpc_channel_from_parts( - xds_manager.clone(), - xds_manager.clone(), - retry_policy, - ); - - let mut client = GreeterClient::new(xds_channel); - - let response = client - .say_hello(HelloRequest { - name: "retry-test".to_string(), - }) - .await - .expect("request should succeed after retry"); - - assert_eq!(response.into_inner().message, "retry-server: retry-test"); - } - - /// Helper: creates a minimal plaintext `ClusterResource` for tests that - /// drive `XdsClusterDiscovery`. The cluster watch in `discover_cluster` - /// blocks until a cluster is in the cache. - fn make_test_cluster(cluster_name: &str) -> Arc { + /// A minimal plaintext `ClusterResource`. + pub(super) fn make_test_cluster( + cluster_name: &str, + ) -> Arc { use crate::xds::resource::cluster::{ClusterResource, LbPolicy}; Arc::new(ClusterResource { name: cluster_name.to_string(), @@ -631,10 +473,12 @@ mod tests { }) } - /// Helper: creates a `RouteConfigResource` that routes all traffic to the given cluster. - fn make_test_route_config(cluster_name: &str) -> Arc { - use crate::xds::resource::route_config::*; - + /// A `RouteConfigResource` that routes all traffic to `cluster_name`. + pub(super) fn make_test_route_config(cluster_name: &str) -> Arc { + use crate::xds::resource::route_config::{ + PathSpecifierConfig, RouteConfig, RouteConfigAction, RouteConfigMatch, + VirtualHostConfig, + }; Arc::new(RouteConfigResource { name: "test-route".to_string(), virtual_hosts: vec![VirtualHostConfig { @@ -653,10 +497,12 @@ mod tests { }) } - /// Helper: creates an `EndpointsResource` from test server addresses. - fn make_test_endpoints(cluster_name: &str, servers: &[TestServer]) -> Arc { + /// An `EndpointsResource` built from test server addresses. + pub(super) fn make_test_endpoints( + cluster_name: &str, + servers: &[TestServer], + ) -> Arc { use crate::xds::resource::endpoints::{HealthStatus, LocalityEndpoints, ResolvedEndpoint}; - Arc::new(EndpointsResource { cluster_name: cluster_name.to_string(), localities: vec![LocalityEndpoints { @@ -675,115 +521,235 @@ mod tests { }) } - /// Builds an XdsChannelGrpc using real XdsRouter and XdsClusterDiscovery - /// backed by the given cache. - async fn build_xds_channel_from_cache(cache: Arc) -> XdsChannelGrpc { - use crate::xds::cluster_discovery::XdsClusterDiscovery; - use crate::xds::routing::XdsRouter; + /// Sends `num_requests` gRPC requests and returns + /// `(successful, error_types, per_server_counts)`. The message format is + /// `"{server-name}: {request}"`, so `per_server_counts` reflects the LB + /// distribution. + pub(super) async fn send_grpc_requests( + mut grpc_client: GreeterClient, + num_requests: usize, + ) -> (usize, HashMap, HashMap) { + let mut successful_requests = 0; + let mut error_types = HashMap::new(); + let mut server_counts = HashMap::new(); - let router: Arc = Arc::new(XdsRouter::new(&cache)); + for i in 0..num_requests { + let request_timeout = tokio::time::Duration::from_secs(3); + let request_future = grpc_client.say_hello(HelloRequest { + name: format!("test-request-{i}"), + }); - #[cfg(feature = "_tls-any")] - let discovery: Arc< - dyn ClusterDiscovery>, - > = { - use crate::xds::cert_provider::CertProviderRegistry; - let registry = - Arc::new(CertProviderRegistry::from_bootstrap(&Default::default()).unwrap()); - Arc::new(XdsClusterDiscovery::new(cache, registry)) - }; - #[cfg(not(feature = "_tls-any"))] - let discovery: Arc< - dyn ClusterDiscovery>, - > = Arc::new(XdsClusterDiscovery::new(cache)); + match tokio::time::timeout(request_timeout, request_future).await { + Ok(Ok(response)) => { + successful_requests += 1; + let message = response.into_inner().message; + if let Some(server_name) = message.split(':').next() { + *server_counts.entry(server_name.to_string()).or_insert(0) += 1; + } + } + Ok(Err(e)) => { + let error_type = format!("{e:?}").chars().take(80).collect::(); + *error_types.entry(error_type).or_insert(0) += 1; + } + Err(_) => { + *error_types.entry("Timeout".to_string()).or_insert(0) += 1; + if error_types.get("Timeout").unwrap_or(&0) > &2 { + break; + } + } + } + } - let builder = XdsChannelBuilder::new(test_config()); - builder.build_grpc_channel_from_parts(router, discovery, GrpcRetryPolicy::default()) + (successful_requests, error_types, server_counts) } +} - /// Tests the full xDS stack (XdsRouter + XdsClusterDiscovery) with a - /// pre-populated cache, validating that requests are routed and - /// load-balanced across real backend servers. - #[tokio::test] - async fn test_xds_channel_with_real_router_and_discovery() { - let num_servers = 3; - let num_requests = 300; - let cluster_name = "test-cluster"; - let (_, servers) = setup_grpc_servers(num_servers).await; +#[cfg(test)] +mod tests { + use super::XdsChannelBuilder; + use super::XdsChannelGrpc; + use super::test_support::{ + make_test_cluster, make_test_endpoints, make_test_route_config, send_grpc_requests, + setup_grpc_servers, test_config, + }; + use crate::client::retry::{GrpcRetryPolicy, GrpcRetryPolicyConfig}; + use crate::testutil::grpc::{GreeterClient, HelloRequest}; + use crate::xds::cache::XdsCache; + use crate::{XdsChannelConfig, XdsUri}; + use std::sync::Arc; - let cache = Arc::new(XdsCache::new()); - cache.update_route_config(make_test_route_config(cluster_name)); - cache.update_cluster(cluster_name, make_test_cluster(cluster_name)); - cache.update_endpoints(cluster_name, make_test_endpoints(cluster_name, &servers)); + /// A cache-backed channel constructor for a specific LB backend. + type BackendCtor = fn(&XdsChannelBuilder, Arc, GrpcRetryPolicy) -> XdsChannelGrpc; + + /// Both LB backends are compiled in test builds (see the `any(test, …)` + /// module gates), so every channel test below runs against each. + fn backends() -> [(&'static str, BackendCtor); 2] { + [ + ( + "tower-lb", + XdsChannelBuilder::build_grpc_channel_from_cache_tower, + ), + ( + "tonic-xds-lb", + XdsChannelBuilder::build_grpc_channel_from_cache_xds, + ), + ] + } + + /// Power-of-two-choices distribution: with a pre-populated cache, requests + /// are routed and load-balanced roughly evenly across all backends. + #[tokio::test] + async fn test_xds_channel_grpc_with_p2c_lb() { + for (backend, build) in backends() { + let cluster_name = "test-cluster"; + let num_requests = 1000; + let num_servers = 5; + let servers = setup_grpc_servers(num_servers).await; + + let cache = Arc::new(XdsCache::new()); + cache.update_route_config(make_test_route_config(cluster_name)); + cache.update_cluster(cluster_name, make_test_cluster(cluster_name)); + cache.update_endpoints(cluster_name, make_test_endpoints(cluster_name, &servers)); + + let channel = build( + &XdsChannelBuilder::new(test_config()), + cache, + GrpcRetryPolicy::default(), + ); + let client = GreeterClient::new(channel); - let channel = build_xds_channel_from_cache(cache).await; - let client = GreeterClient::new(channel); + let (successful, error_types, server_counts) = + send_grpc_requests(client, num_requests).await; - let (successful, error_types, server_counts) = - send_grpc_requests(client, num_requests).await; + assert_eq!( + successful, num_requests, + "[{backend}] expected 100% success. Errors: {error_types:?}", + ); + assert!( + error_types.is_empty(), + "[{backend}] expected no errors: {error_types:?}", + ); + assert_eq!( + server_counts.len(), + num_servers, + "[{backend}] all {num_servers} servers should receive traffic: {server_counts:?}", + ); - assert_eq!( - successful, num_requests, - "Expected 100% success rate. Errors: {error_types:?}", - ); - assert_eq!( - server_counts.len(), - num_servers, - "Expected all {num_servers} servers to receive traffic. Counts: {server_counts:?}", - ); + let expected = num_requests / num_servers; + let min = (expected as f64 / 1.5) as usize; + let max = (expected as f64 * 1.5) as usize; + for (server, count) in &server_counts { + assert!( + *count >= min && *count <= max, + "[{backend}] server {server} received {count}, expected ~{expected} (±1.5x): {server_counts:?}", + ); + } - for server in servers { - let _ = server.shutdown.send(()); - let _ = server.handle.await; + for server in servers { + let _ = server.shutdown.send(()); + let _ = server.handle.await; + } } } - /// Tests that endpoint changes are picked up dynamically by the - /// XdsClusterDiscovery while the channel is serving requests. + /// The retry layer retries UNAVAILABLE and succeeds on the second attempt. #[tokio::test] - async fn test_xds_channel_handles_dynamic_endpoint_updates() { - let cluster_name = "test-cluster"; - let (_, servers) = setup_grpc_servers(2).await; + async fn test_retry_once_on_unavailable() { + use crate::testutil::grpc::spawn_fail_first_n_server; - let cache = Arc::new(XdsCache::new()); - cache.update_route_config(make_test_route_config(cluster_name)); - cache.update_cluster(cluster_name, make_test_cluster(cluster_name)); - // Start with only the first server. - cache.update_endpoints( - cluster_name, - make_test_endpoints(cluster_name, &servers[..1]), - ); + for (backend, build) in backends() { + let cluster_name = "test-cluster"; + // Server fails the first request with UNAVAILABLE, succeeds on retry. + let server = spawn_fail_first_n_server("retry-server", 1) + .await + .expect("Failed to spawn server"); + let servers = vec![server]; + + let cache = Arc::new(XdsCache::new()); + cache.update_route_config(make_test_route_config(cluster_name)); + cache.update_cluster(cluster_name, make_test_cluster(cluster_name)); + cache.update_endpoints(cluster_name, make_test_endpoints(cluster_name, &servers)); + + let retry_policy = GrpcRetryPolicy::new( + GrpcRetryPolicyConfig::new() + .retry_on(vec![tonic::Code::Unavailable]) + .num_retries(1), + ); + let channel = build(&XdsChannelBuilder::new(test_config()), cache, retry_policy); + let mut client = GreeterClient::new(channel); - let channel = build_xds_channel_from_cache(cache.clone()).await; - let client = GreeterClient::new(channel.clone()); + let response = client + .say_hello(HelloRequest { + name: "retry-test".to_string(), + }) + .await + .unwrap_or_else(|e| { + panic!("[{backend}] request should succeed after retry: {e:?}") + }); + assert_eq!( + response.into_inner().message, + "retry-server: retry-test", + "[{backend}]", + ); - // Phase 1: all traffic goes to server-0. - let (successful, _, server_counts) = send_grpc_requests(client, 50).await; - assert_eq!(successful, 50); - assert_eq!( - server_counts.len(), - 1, - "Only 1 server should receive traffic before update. Counts: {server_counts:?}", - ); + for server in servers { + let _ = server.shutdown.send(()); + let _ = server.handle.await; + } + } + } - // Add second server. - cache.update_endpoints(cluster_name, make_test_endpoints(cluster_name, &servers)); - // Give the endpoint manager diff loop time to process the update. - tokio::time::sleep(tokio::time::Duration::from_millis(50)).await; - - // Phase 2: traffic should go to both servers. - let client2 = GreeterClient::new(channel); - let (successful, _, server_counts) = send_grpc_requests(client2, 200).await; - assert_eq!(successful, 200); - assert_eq!( - server_counts.len(), - 2, - "Both servers should receive traffic after update. Counts: {server_counts:?}", - ); + /// Endpoint changes in the cache are picked up dynamically while serving. + #[tokio::test] + async fn test_xds_channel_handles_dynamic_endpoint_updates() { + for (backend, build) in backends() { + let cluster_name = "test-cluster"; + let servers = setup_grpc_servers(2).await; + + let cache = Arc::new(XdsCache::new()); + cache.update_route_config(make_test_route_config(cluster_name)); + cache.update_cluster(cluster_name, make_test_cluster(cluster_name)); + // Start with only the first server. + cache.update_endpoints( + cluster_name, + make_test_endpoints(cluster_name, &servers[..1]), + ); + + let channel = build( + &XdsChannelBuilder::new(test_config()), + cache.clone(), + GrpcRetryPolicy::default(), + ); + let client = GreeterClient::new(channel.clone()); + + // Phase 1: all traffic goes to server-0. + let (successful, _, server_counts) = send_grpc_requests(client, 50).await; + assert_eq!(successful, 50, "[{backend}]"); + assert_eq!( + server_counts.len(), + 1, + "[{backend}] only 1 server should receive traffic before update: {server_counts:?}", + ); - for server in servers { - let _ = server.shutdown.send(()); - let _ = server.handle.await; + // Add second server. + cache.update_endpoints(cluster_name, make_test_endpoints(cluster_name, &servers)); + // Give the endpoint diff loop time to process the update. + tokio::time::sleep(tokio::time::Duration::from_millis(50)).await; + + // Phase 2: traffic should go to both servers. + let client2 = GreeterClient::new(channel); + let (successful, _, server_counts) = send_grpc_requests(client2, 200).await; + assert_eq!(successful, 200, "[{backend}]"); + assert_eq!( + server_counts.len(), + 2, + "[{backend}] both servers should receive traffic after update: {server_counts:?}", + ); + + for server in servers { + let _ = server.shutdown.send(()); + let _ = server.handle.await; + } } } @@ -805,8 +771,9 @@ mod tests { assert!(config.call_creds.is_some()); } - /// Smoke test: verifies builder wiring with a disconnected XdsClient - /// doesn't panic during construction. + /// Smoke test: building the full stack (with a disconnected client) does + /// not panic. Uses the production `build_from_cache`, which selects the LB + /// backend by feature. #[tokio::test] async fn test_build_from_cache_smoke() { use crate::xds::resource_manager::XdsResourceManager; diff --git a/tonic-xds/src/client/loadbalance/channel.rs b/tonic-xds/src/client/loadbalance/channel.rs deleted file mode 100644 index 85368d28f..000000000 --- a/tonic-xds/src/client/loadbalance/channel.rs +++ /dev/null @@ -1,199 +0,0 @@ -//! LbChannel: an instrumented channel wrapper with in-flight request tracking. - -use std::future::Future; -use std::pin::Pin; -use std::sync::Arc; -use std::sync::atomic::{AtomicU64, Ordering}; -use std::task::{Context, Poll}; - -use pin_project_lite::pin_project; -use tower::Service; -use tower::load::Load; - -use crate::client::endpoint::EndpointAddress; - -/// RAII guard that increments an in-flight counter on creation and decrements on drop. -/// Ensures accurate tracking even when futures are cancelled. -struct InFlightGuard { - counter: Arc, -} - -impl InFlightGuard { - fn acquire(counter: Arc) -> Self { - counter.fetch_add(1, Ordering::Relaxed); - Self { counter } - } -} - -impl Drop for InFlightGuard { - fn drop(&mut self) { - self.counter.fetch_sub(1, Ordering::Relaxed); - } -} - -pin_project! { - /// A future that holds an [`InFlightGuard`] for the duration of a request. - /// - /// Preserves the inner future's output type — no boxing or error mapping. - pub(crate) struct InFlightFuture { - #[pin] - inner: F, - _guard: InFlightGuard, - } -} - -impl Future for InFlightFuture { - type Output = F::Output; - - fn poll(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll { - self.project().inner.poll(cx) - } -} - -/// A channel wrapper that tracks in-flight requests for load balancing. -/// -/// `LbChannel` wraps an inner service `S` and maintains an atomic counter of -/// in-flight requests. This counter is used by P2C load balancers (via the -/// [`Load`] trait) to prefer endpoints with fewer active requests. -/// -/// All clones of an `LbChannel` share the same in-flight counter. -pub(crate) struct LbChannel { - addr: EndpointAddress, - inner: S, - in_flight: Arc, -} - -impl LbChannel { - /// Create a new `LbChannel` wrapping the given service. - pub(crate) fn new(addr: EndpointAddress, inner: S) -> Self { - Self { - addr, - inner, - in_flight: Arc::new(AtomicU64::new(0)), - } - } - - /// Returns the endpoint address. - pub(crate) fn addr(&self) -> &EndpointAddress { - &self.addr - } - - /// Returns the current number of in-flight requests. - #[cfg(test)] - pub(crate) fn in_flight(&self) -> u64 { - self.in_flight.load(Ordering::Relaxed) - } -} - -impl Clone for LbChannel { - fn clone(&self) -> Self { - Self { - addr: self.addr.clone(), - inner: self.inner.clone(), - in_flight: self.in_flight.clone(), - } - } -} - -impl Service for LbChannel -where - S: Service, -{ - type Response = S::Response; - type Error = S::Error; - type Future = InFlightFuture; - - fn poll_ready(&mut self, cx: &mut Context<'_>) -> Poll> { - self.inner.poll_ready(cx) - } - - fn call(&mut self, req: Req) -> Self::Future { - let guard = InFlightGuard::acquire(self.in_flight.clone()); - InFlightFuture { - inner: self.inner.call(req), - _guard: guard, - } - } -} - -impl Load for LbChannel { - type Metric = u64; - - fn load(&self) -> Self::Metric { - self.in_flight.load(Ordering::Relaxed) - } -} - -#[cfg(test)] -mod tests { - use super::*; - use std::future; - use std::task::Poll; - - fn test_addr() -> EndpointAddress { - EndpointAddress::new("127.0.0.1", 8080) - } - - #[derive(Clone)] - struct MockService; - - impl Service<&'static str> for MockService { - type Response = &'static str; - type Error = &'static str; - type Future = future::Ready>; - - fn poll_ready(&mut self, _cx: &mut Context<'_>) -> Poll> { - Poll::Ready(Ok(())) - } - - fn call(&mut self, _req: &'static str) -> Self::Future { - future::ready(Ok("ok")) - } - } - - #[tokio::test] - async fn test_in_flight_increments_and_decrements() { - let mut ch = LbChannel::new(test_addr(), MockService); - assert_eq!(ch.in_flight(), 0); - - let fut = ch.call("hello"); - assert_eq!(ch.in_flight(), 1); - - let resp = fut.await.unwrap(); - assert_eq!(resp, "ok"); - assert_eq!(ch.in_flight(), 0); - } - - #[tokio::test] - async fn test_in_flight_on_future_drop() { - let mut ch = LbChannel::new(test_addr(), MockService); - let fut = ch.call("hello"); - assert_eq!(ch.in_flight(), 1); - - drop(fut); - assert_eq!(ch.in_flight(), 0); - } - - #[tokio::test] - async fn test_clone_shares_in_flight() { - let mut ch1 = LbChannel::new(test_addr(), MockService); - let ch2 = ch1.clone(); - - let fut = ch1.call("hello"); - assert_eq!(ch1.in_flight(), 1); - assert_eq!(ch2.in_flight(), 1); - - let _ = fut.await; - assert_eq!(ch1.in_flight(), 0); - assert_eq!(ch2.in_flight(), 0); - } - - #[test] - fn test_load_returns_in_flight() { - let ch = LbChannel::new(test_addr(), MockService); - assert_eq!(Load::load(&ch), 0); - - ch.in_flight.fetch_add(3, Ordering::Relaxed); - assert_eq!(Load::load(&ch), 3); - } -} diff --git a/tonic-xds/src/client/loadbalance/channel_state.rs b/tonic-xds/src/client/loadbalance/channel_state.rs index e9e7f264d..8860b6d8b 100644 --- a/tonic-xds/src/client/loadbalance/channel_state.rs +++ b/tonic-xds/src/client/loadbalance/channel_state.rs @@ -18,10 +18,10 @@ //! manage multiple in-flight state changes and handle cancellation by key. //! //! The state types hold the raw service `S` directly. In-flight tracking and -//! load reporting are handled separately by [`LbChannel`] at the pool level. +//! load reporting are handled separately by [`EndpointChannel`] at the pool level. //! //! [`KeyedFutures`]: crate::client::loadbalance::keyed_futures::KeyedFutures -//! [`LbChannel`]: crate::client::loadbalance::channel::LbChannel +//! [`EndpointChannel`]: crate::client::endpoint::EndpointChannel use std::future::Future; use std::pin::Pin; @@ -80,6 +80,7 @@ impl EndpointCounters { /// can update registry-level counters exactly once per transition. #[derive(Debug)] pub(crate) struct OutlierChannelState { + #[allow(dead_code)] addr: EndpointAddress, counters: EndpointCounters, /// Bumped on each ejection; decremented (saturating) on each @@ -104,6 +105,7 @@ impl OutlierChannelState { } /// Endpoint address this state belongs to. + #[allow(dead_code)] pub(crate) fn addr(&self) -> &EndpointAddress { &self.addr } @@ -363,6 +365,7 @@ impl ReadyChannel { /// Drop the connection and start a fresh connect for the same /// address. The outlier state is re-attached from `registry` /// when the new connect resolves. + #[allow(dead_code)] pub(crate) fn reconnect>( self, connector: Arc, diff --git a/tonic-xds/src/client/loadbalance/keyed_futures.rs b/tonic-xds/src/client/loadbalance/keyed_futures.rs index da7c6c3ce..fdba04243 100644 --- a/tonic-xds/src/client/loadbalance/keyed_futures.rs +++ b/tonic-xds/src/client/loadbalance/keyed_futures.rs @@ -90,6 +90,7 @@ where } /// True if a live (non-cancelled) future is tracked for `key`. + #[allow(dead_code)] pub(crate) fn contains_key(&self, key: &K) -> bool { self.cancellations.contains_key(key) } diff --git a/tonic-xds/src/client/loadbalance/mod.rs b/tonic-xds/src/client/loadbalance/mod.rs index 1c4ffa395..4da8c8ef7 100644 --- a/tonic-xds/src/client/loadbalance/mod.rs +++ b/tonic-xds/src/client/loadbalance/mod.rs @@ -1,7 +1,7 @@ -pub(crate) mod channel; pub(crate) mod channel_state; pub(crate) mod errors; pub(crate) mod keyed_futures; pub(crate) mod loadbalancer; pub(crate) mod outlier_detection; pub(crate) mod pickers; +pub(crate) mod service; diff --git a/tonic-xds/src/client/loadbalance/outlier_detection.rs b/tonic-xds/src/client/loadbalance/outlier_detection.rs index 0b00ef6dd..dc3fde326 100644 --- a/tonic-xds/src/client/loadbalance/outlier_detection.rs +++ b/tonic-xds/src/client/loadbalance/outlier_detection.rs @@ -117,6 +117,7 @@ impl OutlierStatsRegistry { } /// Number of registered channels. + #[allow(dead_code)] pub(crate) fn len(&self) -> usize { self.channels.len() } diff --git a/tonic-xds/src/client/loadbalance/pickers/mod.rs b/tonic-xds/src/client/loadbalance/pickers/mod.rs index 08a313f25..2d4daf3be 100644 --- a/tonic-xds/src/client/loadbalance/pickers/mod.rs +++ b/tonic-xds/src/client/loadbalance/pickers/mod.rs @@ -1,4 +1,8 @@ pub(crate) mod p2c; +// Ring-hash (gRFC A42) is implemented but not wired into the first-cut LB, +// which selects P2C. +// TODO: wire ring-hash into the LB. +#[allow(dead_code)] pub(crate) mod ring_hash; use indexmap::{IndexMap, IndexSet}; diff --git a/tonic-xds/src/client/loadbalance/service.rs b/tonic-xds/src/client/loadbalance/service.rs new file mode 100644 index 000000000..88ea4436a --- /dev/null +++ b/tonic-xds/src/client/loadbalance/service.rs @@ -0,0 +1,155 @@ +//! XDS cluster load-balancing service for the `tonic-xds-lb` path. +//! +//! [`XdsLoadBalanceService`] is the analogue of the `tower-lb` +//! `XdsLbService`: it reads the [`RouteDecision`] attached by the routing +//! layer, looks up (or lazily builds) the target cluster's [`LoadBalancer`], +//! and dispatches the request to it. +//! +//! Each cluster's `LoadBalancer` is not `Clone` and needs `&mut self` to +//! serve, so it is wrapped in a [`tower::buffer::Buffer`] — giving a cheap, +//! cloneable handle shared across concurrent requests. The buffer maps the +//! LB's [`LbError`](crate::client::loadbalance::errors::LbError) to +//! [`BoxError`] automatically. +//! +//! TODO: only the power-of-two-choices picker ([`P2cPicker`]) is enabled; +//! outlier detection ([`OutlierDetectionConfig::default`]) and other pickers (e.g. ring-hash) +//! will be wired in by follow-ups. + +use std::sync::Arc; +use std::task::{Context, Poll}; + +use arc_swap::ArcSwap; +use dashmap::DashMap; +use http::{Request, Response}; +use tonic::body::Body as TonicBody; +use tonic::transport::Channel; +use tower::buffer::Buffer; +use tower::{BoxError, Service, ServiceExt}; + +use crate::client::endpoint::EndpointChannel; +use crate::client::loadbalance::channel_state::ReadyChannel; +use crate::client::loadbalance::loadbalancer::{LbFuture, LoadBalancer}; +use crate::client::loadbalance::pickers::ChannelPicker; +use crate::client::loadbalance::pickers::p2c::P2cPicker; +use crate::client::route::RouteDecision; +use crate::common::async_util::BoxFuture; +use crate::xds::cache::XdsCache; +#[cfg(feature = "_tls-any")] +use crate::xds::cert_provider::CertProviderRegistry; +use crate::xds::lb_discovery::{XdsLbConnector, discover_endpoints}; +use crate::xds::resource::outlier_detection::OutlierDetectionConfig; + +/// Buffer capacity between callers and a cluster's `LoadBalancer` worker. +const DEFAULT_BUFFER_CAPACITY: usize = 1024; + +/// The request type flowing into the LB layer. +type LbRequest = Request; +/// The response type produced by an endpoint channel. +type LbResponse = Response; + +/// A cloneable handle to one cluster's `LoadBalancer`, buffered so concurrent +/// callers share a single balancer. `Buffer`'s second type parameter is the +/// wrapped service's future — here the `LoadBalancer`'s [`LbFuture`]. +type ClusterChannel = Buffer>; + +/// Error returned when a request reaches the LB layer without a routing +/// decision (i.e. the routing layer did not run or produced nothing). +#[derive(Debug, thiserror::Error)] +#[error("no routing decision extension from the routing layer available")] +struct NoRoutingDecision; + +/// Registry of per-cluster [`LoadBalancer`]s, built lazily on first use. +struct ClusterLbRegistry { + cache: Arc, + #[cfg(feature = "_tls-any")] + cert_provider_registry: Arc, + clusters: DashMap, +} + +impl ClusterLbRegistry { + /// Returns a cloneable channel to the cluster's balancer, building it on + /// first access. + fn cluster_channel(&self, cluster_name: &str) -> ClusterChannel { + self.clusters + .entry(cluster_name.to_string()) + .or_insert_with(|| self.build_cluster_channel(cluster_name)) + .clone() + } + + /// Constructs a fresh `LoadBalancer` for `cluster_name` and wraps it in a + /// buffer. Discovery yields idle endpoints; [`XdsLbConnector`] resolves the + /// cluster's CDS security config lazily inside `connect`. + fn build_cluster_channel(&self, cluster_name: &str) -> ClusterChannel { + let discover = discover_endpoints(&self.cache, cluster_name); + let connector = Arc::new(XdsLbConnector::new( + self.cache.clone(), + cluster_name.to_string(), + #[cfg(feature = "_tls-any")] + self.cert_provider_registry.clone(), + )); + let picker: Arc< + dyn ChannelPicker>, LbRequest> + Send + Sync, + > = Arc::new(P2cPicker); + let config = Arc::new(ArcSwap::from_pointee(OutlierDetectionConfig::default())); + let lb = LoadBalancer::new(discover, connector, picker, config); + Buffer::new(lb, DEFAULT_BUFFER_CAPACITY) + } +} + +/// Tower service that routes each request to its cluster's `LoadBalancer`. +#[derive(Clone)] +pub(crate) struct XdsLoadBalanceService { + registry: Arc, +} + +impl XdsLoadBalanceService { + #[cfg(feature = "_tls-any")] + pub(crate) fn new( + cache: Arc, + cert_provider_registry: Arc, + ) -> Self { + Self { + registry: Arc::new(ClusterLbRegistry { + cache, + cert_provider_registry, + clusters: DashMap::new(), + }), + } + } + + #[cfg(not(feature = "_tls-any"))] + pub(crate) fn new(cache: Arc) -> Self { + Self { + registry: Arc::new(ClusterLbRegistry { + cache, + clusters: DashMap::new(), + }), + } + } +} + +impl Service for XdsLoadBalanceService { + type Response = LbResponse; + type Error = BoxError; + type Future = BoxFuture>; + + fn poll_ready(&mut self, _cx: &mut Context<'_>) -> Poll> { + // The target cluster is decided per-request by the routing layer, so + // readiness cannot be determined without the request. + Poll::Ready(Ok(())) + } + + fn call(&mut self, request: LbRequest) -> Self::Future { + let Some(decision) = request.extensions().get::().cloned() else { + return Box::pin(async move { Err(BoxError::from(NoRoutingDecision)) }); + }; + + let mut channel = self.registry.cluster_channel(&decision.cluster); + + Box::pin(async move { + // Blocks until the balancer has a ready endpoint. + channel.ready().await?; + channel.call(request).await + }) + } +} diff --git a/tonic-xds/src/client/mod.rs b/tonic-xds/src/client/mod.rs index 3e02c9b29..2835a9b93 100644 --- a/tonic-xds/src/client/mod.rs +++ b/tonic-xds/src/client/mod.rs @@ -1,9 +1,14 @@ pub(crate) mod channel; -pub(crate) mod cluster; pub(crate) mod endpoint; -pub(crate) mod lb; -#[allow(dead_code)] -pub(crate) mod loadbalance; #[allow(dead_code)] pub(crate) mod retry; pub(crate) mod route; + +cfg_tower_lb! { + pub(crate) mod cluster; + pub(crate) mod lb; +} + +cfg_tonic_xds_lb! { + pub(crate) mod loadbalance; +} diff --git a/tonic-xds/src/lib.rs b/tonic-xds/src/lib.rs index 3392f2f97..d50437c4e 100644 --- a/tonic-xds/src/lib.rs +++ b/tonic-xds/src/lib.rs @@ -140,6 +140,31 @@ //! [A48]: https://github.com/grpc/proposal/blob/master/A48-xds-least-request-lb-policy.md //! [A63]: https://github.com/grpc/proposal/blob/master/A63-xds-string-matcher-ignore-case.md +//! Load-balancer selection. The default `tower-lb` stack is used unless the +//! `tonic-xds-lb` feature is enabled, which switches to the in-crate +//! implementation. +//! +//! In test builds both backends are compiled (the `any(test, …)` arm) so the +//! channel tests can exercise each in a single run. These two macros stamp that +//! cfg onto a group of items, keeping the selection logic in one place. + +/// Compiles each item for the `tower-lb` backend: selected in production when +/// `tonic-xds-lb` is disabled (the default), and always compiled in test builds +/// so both backends can be exercised in one test run. +macro_rules! cfg_tower_lb { + ($($item:item)*) => { + $( #[cfg(any(test, not(feature = "tonic-xds-lb")))] $item )* + }; +} + +/// Compiles each item for the `tonic-xds-lb` backend: selected in production +/// when its feature is enabled, and always compiled in test builds. +macro_rules! cfg_tonic_xds_lb { + ($($item:item)*) => { + $( #[cfg(any(test, feature = "tonic-xds-lb"))] $item )* + }; +} + pub(crate) mod client; pub(crate) mod common; pub(crate) mod xds; diff --git a/tonic-xds/src/testutil/grpc.rs b/tonic-xds/src/testutil/grpc.rs index fe9a69232..a80e1cb33 100644 --- a/tonic-xds/src/testutil/grpc.rs +++ b/tonic-xds/src/testutil/grpc.rs @@ -3,7 +3,7 @@ use std::error::Error; use std::net::SocketAddr; use tokio::{net::TcpListener, sync::oneshot}; use tonic::server::NamedService; -use tonic::transport::{Channel, ClientTlsConfig, Endpoint, Server, ServerTlsConfig}; +use tonic::transport::{Server, ServerTlsConfig}; use tonic::{Request, Response, Status}; pub(crate) use crate::testutil::proto::helloworld::{ @@ -58,10 +58,9 @@ impl Greeter for FailFirstNGreeter { } } -/// A test server that runs a gRPC service and provides a channel for clients to connect. +/// A test server that runs a gRPC service. Tests reach it via its +/// [`addr`](Self::addr) through xDS-discovered endpoints. pub(crate) struct TestServer { - /// The gRPC channel for talking to the test server. - pub channel: Channel, /// Signal the server to shutdown. pub shutdown: oneshot::Sender<()>, /// Handle to wait for server to exit. @@ -78,7 +77,6 @@ impl NamedService for TestServer { pub(crate) async fn spawn_greeter_server( msg: &str, server_tls: Option, - client_tls: Option, ) -> Result> { // Bind to an ephemeral port (random free port assigned by OS) let listener = TcpListener::bind("127.0.0.1:0").await?; @@ -111,19 +109,7 @@ pub(crate) async fn spawn_greeter_server( Ok(()) }); - let channel = if let Some(client_tls) = client_tls { - let endpoint_str = format!("https://{addr}"); - Endpoint::from_shared(endpoint_str)? - .tls_config(client_tls)? - .connect() - .await? - } else { - let endpoint_str = format!("http://{addr}"); - Endpoint::from_shared(endpoint_str)?.connect().await? - }; - Ok(TestServer { - channel, shutdown: tx, handle, addr, @@ -151,12 +137,7 @@ pub(crate) async fn spawn_fail_first_n_server( .await }); - let channel = Endpoint::from_shared(format!("http://{addr}"))? - .connect() - .await?; - Ok(TestServer { - channel, shutdown: tx, handle, addr, diff --git a/tonic-xds/src/xds/cluster_discovery.rs b/tonic-xds/src/xds/cluster_discovery.rs index 61251b9a9..3f69d1381 100644 --- a/tonic-xds/src/xds/cluster_discovery.rs +++ b/tonic-xds/src/xds/cluster_discovery.rs @@ -3,10 +3,11 @@ //! Per cluster, [`XdsClusterDiscovery::discover_cluster`] spawns a task that //! drives two concurrent watches: //! -//! 1. The cluster resource watch — produces a fresh [`Connector`] on each -//! CDS update (e.g. when `transport_socket` changes). The connector is -//! held inside a [`ConnectorSwap`] so the diff loop reads the latest -//! snapshot per endpoint connection. +//! 1. The cluster resource watch — produces a fresh +//! [`Connector`](crate::client::endpoint::Connector) on each CDS update +//! (e.g. when `transport_socket` changes). The connector is held inside a +//! [`ConnectorSwap`] so the diff loop reads the latest snapshot per +//! endpoint connection. //! 2. The endpoint watch — produces `Change::Insert` / `Change::Remove` //! events forwarded to the LB layer. //! @@ -19,20 +20,15 @@ use arc_swap::ArcSwap; use tokio::sync::mpsc; use tokio_stream::StreamExt as _; use tokio_stream::wrappers::ReceiverStream; -use tonic::transport::{Channel, Endpoint}; +use tonic::transport::Channel; -use crate::client::endpoint::{Connector, EndpointAddress, EndpointChannel}; +use crate::client::endpoint::{EndpointAddress, EndpointChannel}; use crate::client::lb::{BoxDiscover, ClusterDiscovery}; -use crate::common::async_util::BoxFuture; use crate::xds::cache::XdsCache; #[cfg(feature = "_tls-any")] -use crate::xds::cert_provider::verifier::XdsServerCertVerifier; -#[cfg(feature = "_tls-any")] -use crate::xds::cert_provider::{CertProviderRegistry, CertificateProvider}; +use crate::xds::cert_provider::CertProviderRegistry; +use crate::xds::connector::build_connector; use crate::xds::endpoint_manager::{ConnectorSwap, EndpointManager}; -use crate::xds::resource::ClusterResource; -#[cfg(feature = "_tls-any")] -use crate::xds::resource::security::ClusterSecurityConfig; /// Buffer capacity for the discovery channel between the spawned task and /// Tower's LB layer. @@ -132,299 +128,3 @@ impl ClusterDiscovery> for XdsClusterD Box::pin(ReceiverStream::new(rx)) } } - -/// Build a [`Connector`] for the given cluster. -/// -/// - `cluster.security == None` → [`PlaintextConnector`]. -/// - `cluster.security == Some(_)` under a TLS feature → [`TlsConnector`]. -/// - `cluster.security == Some(_)` without a TLS feature → error. -fn build_connector( - cluster: &ClusterResource, - #[cfg(feature = "_tls-any")] registry: &CertProviderRegistry, -) -> Result> + Send + Sync>, ConnectorBuildError> -{ - match &cluster.security { - None => Ok(Arc::new(PlaintextConnector)), - #[cfg(feature = "_tls-any")] - Some(sec) => Ok(Arc::new(TlsConnector::new(registry, sec)?)), - #[cfg(not(feature = "_tls-any"))] - Some(_) => Err(ConnectorBuildError::TlsFeatureMissing), - } -} - -/// Errors building a per-cluster [`Connector`] from a [`ClusterResource`]. -#[derive(Debug, thiserror::Error)] -pub(crate) enum ConnectorBuildError { - /// TLS connector build failed (unknown provider instance, etc.). - #[cfg(feature = "_tls-any")] - #[error("build TLS connector: {0}")] - Tls(#[from] TlsConnectorBuildError), - /// The cluster requires TLS but the binary was built without a TLS - /// crypto backend feature. - #[cfg(not(feature = "_tls-any"))] - #[error("cluster requires TLS but no TLS feature enabled (build with tls-ring or tls-aws-lc)")] - TlsFeatureMissing, -} - -/// Errors constructing a [`TlsConnector`]. -#[cfg(feature = "_tls-any")] -#[derive(Debug, thiserror::Error)] -pub(crate) enum TlsConnectorBuildError { - #[error("CA provider instance '{0}' is not configured in bootstrap.certificate_providers")] - UnknownCaInstance(String), - #[error( - "identity provider instance '{0}' is not configured in bootstrap.certificate_providers" - )] - UnknownIdentityInstance(String), -} - -/// Plaintext (non-TLS) [`Connector`] that produces a lazily-connected -/// `tonic::Channel` for each endpoint. -pub(crate) struct PlaintextConnector; - -impl Connector for PlaintextConnector { - type Service = EndpointChannel; - - fn connect(&self, addr: &EndpointAddress) -> BoxFuture { - // EndpointAddress only holds validated Ipv4/Ipv6/Hostname + u16 port, - // and its Display impl produces "ip:port" or "hostname:port". Prefixing - // with "http://" always yields a valid URI, so from_shared cannot fail. - let channel = Endpoint::from_shared(format!("http://{addr}")) - .expect("EndpointAddress Display guarantees valid URI") - .connect_lazy(); - let svc = EndpointChannel::new(channel); - Box::pin(async move { svc }) - } -} - -/// TLS [`Connector`] for clusters whose CDS resource carries an -/// `UpstreamTlsContext`. Holds: -/// -/// - a verifier that reads CA roots from its [`CertificateProvider`] on -/// each handshake (so `file_watcher`-driven CA rotation is picked up -/// automatically), and -/// - an optional identity provider for mTLS — fetched per `connect` call -/// so identity rotation is picked up on each new connection. -/// -/// The connector is rebuilt by [`build_connector`] on every CDS update, so -/// changes to `ca_instance_name` / `identity_instance_name` / SAN matchers -/// also propagate as the cluster watch swaps the connector. -#[cfg(feature = "_tls-any")] -pub(crate) struct TlsConnector { - verifier: Arc, - identity_provider: Option>, -} - -#[cfg(feature = "_tls-any")] -impl TlsConnector { - pub(crate) fn new( - registry: &CertProviderRegistry, - security: &ClusterSecurityConfig, - ) -> Result { - let ca_provider = registry - .get(&security.ca_instance_name) - .ok_or_else(|| { - TlsConnectorBuildError::UnknownCaInstance(security.ca_instance_name.clone()) - })? - .clone(); - let verifier = Arc::new(XdsServerCertVerifier::new( - ca_provider, - security.san_matchers.clone(), - )); - - let identity_provider = security - .identity_instance_name - .as_ref() - .map(|name| { - registry - .get(name) - .cloned() - .ok_or_else(|| TlsConnectorBuildError::UnknownIdentityInstance(name.clone())) - }) - .transpose()?; - - Ok(Self { - verifier, - identity_provider, - }) - } -} - -#[cfg(feature = "_tls-any")] -impl Connector for TlsConnector { - type Service = EndpointChannel; - - fn connect(&self, addr: &EndpointAddress) -> BoxFuture { - use rustls::client::danger::ServerCertVerifier; - - let verifier: Arc = self.verifier.clone(); - - // Identity is fetched per `connect` so file_watcher-driven identity - // rotation reaches each new connection. `Identity::from_pem` is - // bytes-only; the rustls parse happens inside `tls_config_with_verifier`. - let identity = self - .identity_provider - .as_ref() - .and_then(|p| match p.fetch() { - Ok(data) => data - .identity() - .map(|id| tonic::transport::Identity::from_pem(&id.cert_chain, &id.key)), - Err(e) => { - tracing::error!( - error = %e, - "identity provider fetch failed; falling back to TLS-only", - ); - None - } - }); - - let mut tls_config = tonic::transport::ClientTlsConfig::new(); - if let Some(id) = identity { - tls_config = tls_config.identity(id); - } - - let uri = format!("https://{addr}"); - let endpoint = Endpoint::from_shared(uri.clone()) - .expect("EndpointAddress Display guarantees valid URI"); - - let channel = match endpoint.tls_config_with_verifier(tls_config, verifier) { - Ok(ep) => ep.connect_lazy(), - Err(e) => { - // tls_config_with_verifier only errors on UDS endpoints - // (see tonic's endpoint.rs), which we never construct. The - // defensive fallback returns a non-TLS lazy channel — the - // request will fail at the wire, surfacing the misconfig. - tracing::error!( - error = %e, address = %addr, - "tls_config_with_verifier failed; non-TLS lazy fallback", - ); - Endpoint::from_shared(uri) - .expect("EndpointAddress Display guarantees valid URI") - .connect_lazy() - } - }; - let svc = EndpointChannel::new(channel); - Box::pin(async move { svc }) - } -} - -#[cfg(test)] -mod tests { - use super::*; - use crate::xds::resource::cluster::{ClusterResource, LbPolicy}; - - fn plaintext_cluster() -> ClusterResource { - ClusterResource { - name: "c".into(), - eds_service_name: None, - lb_policy: LbPolicy::RoundRobin, - security: None, - } - } - - #[cfg(feature = "_tls-any")] - fn empty_registry() -> CertProviderRegistry { - use std::collections::HashMap; - CertProviderRegistry::from_bootstrap(&HashMap::new()).unwrap() - } - - /// Plaintext dispatch under TLS feature. - #[cfg(feature = "_tls-any")] - #[test] - fn build_connector_plaintext_tls_feature_on() { - assert!(build_connector(&plaintext_cluster(), &empty_registry()).is_ok()); - } - - /// Plaintext dispatch without any TLS feature. - #[cfg(not(feature = "_tls-any"))] - #[test] - fn build_connector_plaintext_no_tls() { - assert!(build_connector(&plaintext_cluster()).is_ok()); - } - - /// Cluster with TLS pointing at an instance not in the registry surfaces - /// a clear error — useful for misconfig diagnostics. - #[cfg(feature = "_tls-any")] - #[test] - fn build_connector_tls_unknown_ca() { - use crate::xds::resource::security::ClusterSecurityConfig; - - let mut cluster = plaintext_cluster(); - cluster.security = Some(ClusterSecurityConfig { - ca_instance_name: "missing-ca".into(), - identity_instance_name: None, - san_matchers: vec![], - }); - let Err(err) = build_connector(&cluster, &empty_registry()) else { - panic!("expected UnknownCaInstance error"); - }; - assert!(matches!( - err, - ConnectorBuildError::Tls(TlsConnectorBuildError::UnknownCaInstance(ref name)) - if name == "missing-ca" - )); - } - - /// `TlsConnector::connect` fetches the identity provider on every call, - /// which is what gives us identity rotation between CDS updates. Counter - /// shim verifies the call count without standing up a TLS handshake. - #[cfg(feature = "_tls-any")] - #[tokio::test] - async fn tls_connector_fetches_identity_per_connect() { - use crate::xds::cert_provider::{ - CertProviderError, CertificateData, CertificateProvider, Identity, - }; - use rustls::RootCertStore; - use std::sync::atomic::{AtomicUsize, Ordering}; - - struct CountingIdentity { - count: AtomicUsize, - data: Arc, - } - impl CertificateProvider for CountingIdentity { - fn fetch(&self) -> Result, CertProviderError> { - self.count.fetch_add(1, Ordering::Relaxed); - Ok(self.data.clone()) - } - } - - struct StaticCa(Arc); - impl CertificateProvider for StaticCa { - fn fetch(&self) -> Result, CertProviderError> { - Ok(self.0.clone()) - } - } - - let ca_provider: Arc = - Arc::new(StaticCa(Arc::new(CertificateData::RootsOnly { - roots: Arc::new(RootCertStore::empty()), - }))); - let verifier = Arc::new(XdsServerCertVerifier::new(ca_provider, vec![])); - - let identity_data = Arc::new(CertificateData::IdentityOnly { - identity: Identity { - cert_chain: b"cert".to_vec(), - key: b"key".to_vec(), - }, - }); - let counter = Arc::new(CountingIdentity { - count: AtomicUsize::new(0), - data: identity_data, - }); - let identity_provider: Arc = counter.clone(); - let connector = TlsConnector { - verifier, - identity_provider: Some(identity_provider), - }; - - let addr = EndpointAddress::from("1.2.3.4:443".parse::().unwrap()); - let _ = connector.connect(&addr).await; - let _ = connector.connect(&addr).await; - - assert_eq!( - counter.count.load(Ordering::Relaxed), - 2, - "TlsConnector should fetch identity provider on every connect call", - ); - } -} diff --git a/tonic-xds/src/xds/connector.rs b/tonic-xds/src/xds/connector.rs new file mode 100644 index 000000000..8a728e8c6 --- /dev/null +++ b/tonic-xds/src/xds/connector.rs @@ -0,0 +1,309 @@ +//! Per-cluster [`Connector`] construction from validated CDS resources. +//! +//! `build_connector` maps a cluster's security config to a concrete +//! connector, which creates a [`EndpointChannel`] for sending requests to an [`Endpoint`]. + +use crate::client::endpoint::{Connector, EndpointAddress, EndpointChannel}; +use crate::common::async_util::BoxFuture; +use crate::xds::resource::ClusterResource; +#[cfg(feature = "_tls-any")] +use crate::xds::{ + cert_provider::CertProviderRegistry, cert_provider::CertificateProvider, + cert_provider::verifier::XdsServerCertVerifier, resource::security::ClusterSecurityConfig, +}; +use std::sync::Arc; +use tonic::transport::{Channel, Endpoint}; + +/// Build a [`Connector`] for the given cluster. +pub(crate) fn build_connector( + cluster: &ClusterResource, + #[cfg(feature = "_tls-any")] registry: &CertProviderRegistry, +) -> Result> + Send + Sync>, ConnectorBuildError> +{ + match &cluster.security { + None => Ok(Arc::new(PlaintextConnector)), + #[cfg(feature = "_tls-any")] + Some(sec) => Ok(Arc::new(TlsConnector::new(registry, sec)?)), + #[cfg(not(feature = "_tls-any"))] + Some(_) => Err(ConnectorBuildError::TlsFeatureMissing), + } +} + +/// Errors building a per-cluster [`Connector`] from a [`ClusterResource`]. +#[derive(Debug, thiserror::Error)] +pub(crate) enum ConnectorBuildError { + /// TLS connector build failed (unknown provider instance, etc.). + #[cfg(feature = "_tls-any")] + #[error("build TLS connector: {0}")] + Tls(#[from] TlsConnectorBuildError), + /// The cluster requires TLS but the binary was built without a TLS + /// crypto backend feature. + #[cfg(not(feature = "_tls-any"))] + #[error("cluster requires TLS but no TLS feature enabled (build with tls-ring or tls-aws-lc)")] + TlsFeatureMissing, +} + +/// Errors constructing a [`TlsConnector`]. +#[cfg(feature = "_tls-any")] +#[derive(Debug, thiserror::Error)] +pub(crate) enum TlsConnectorBuildError { + #[error("CA provider instance '{0}' is not configured in bootstrap.certificate_providers")] + UnknownCaInstance(String), + #[error( + "identity provider instance '{0}' is not configured in bootstrap.certificate_providers" + )] + UnknownIdentityInstance(String), +} + +/// Plaintext (non-TLS) [`Connector`] that produces a lazily-connected +/// `tonic::Channel` for each endpoint. +pub(crate) struct PlaintextConnector; + +impl Connector for PlaintextConnector { + type Service = EndpointChannel; + + fn connect(&self, addr: &EndpointAddress) -> BoxFuture { + // EndpointAddress only holds validated Ipv4/Ipv6/Hostname + u16 port, + // and its Display impl produces "ip:port" or "hostname:port". Prefixing + // with "http://" always yields a valid URI, so from_shared cannot fail. + let channel = Endpoint::from_shared(format!("http://{addr}")) + .expect("EndpointAddress Display guarantees valid URI") + .connect_lazy(); + let svc = EndpointChannel::new(channel); + Box::pin(async move { svc }) + } +} + +/// TLS [`Connector`] for clusters whose CDS resource carries an +/// `UpstreamTlsContext`. Holds: +/// +/// - a verifier that reads CA roots from its [`CertificateProvider`] on +/// each handshake (so `file_watcher`-driven CA rotation is picked up +/// automatically), and +/// - an optional identity provider for mTLS — fetched per [`connect`] call +/// so identity rotation is picked up on each new connection. +/// +/// The connector is rebuilt by [`build_connector`] on every CDS update, so +/// changes to `ca_instance_name` / `identity_instance_name` / SAN matchers +/// also propagate as the cluster watch swaps the connector. +/// +/// [`connect`]: Connector::connect +#[cfg(feature = "_tls-any")] +pub(crate) struct TlsConnector { + verifier: Arc, + identity_provider: Option>, +} + +#[cfg(feature = "_tls-any")] +impl TlsConnector { + pub(crate) fn new( + registry: &CertProviderRegistry, + security: &ClusterSecurityConfig, + ) -> Result { + let ca_provider = registry + .get(&security.ca_instance_name) + .ok_or_else(|| { + TlsConnectorBuildError::UnknownCaInstance(security.ca_instance_name.clone()) + })? + .clone(); + let verifier = Arc::new(XdsServerCertVerifier::new( + ca_provider, + security.san_matchers.clone(), + )); + + let identity_provider = security + .identity_instance_name + .as_ref() + .map(|name| { + registry + .get(name) + .cloned() + .ok_or_else(|| TlsConnectorBuildError::UnknownIdentityInstance(name.clone())) + }) + .transpose()?; + + Ok(Self { + verifier, + identity_provider, + }) + } +} + +#[cfg(feature = "_tls-any")] +impl Connector for TlsConnector { + type Service = EndpointChannel; + + fn connect(&self, addr: &EndpointAddress) -> BoxFuture { + use rustls::client::danger::ServerCertVerifier; + + let verifier: Arc = self.verifier.clone(); + + // Identity is fetched per `connect` so file_watcher-driven identity + // rotation reaches each new connection. `Identity::from_pem` is + // bytes-only; the rustls parse happens inside `tls_config_with_verifier`. + let identity = self + .identity_provider + .as_ref() + .and_then(|p| match p.fetch() { + Ok(data) => data + .identity() + .map(|id| tonic::transport::Identity::from_pem(&id.cert_chain, &id.key)), + Err(e) => { + tracing::error!( + error = %e, + "identity provider fetch failed; falling back to TLS-only", + ); + None + } + }); + + let mut tls_config = tonic::transport::ClientTlsConfig::new(); + if let Some(id) = identity { + tls_config = tls_config.identity(id); + } + + let uri = format!("https://{addr}"); + let endpoint = Endpoint::from_shared(uri.clone()) + .expect("EndpointAddress Display guarantees valid URI"); + + let channel = match endpoint.tls_config_with_verifier(tls_config, verifier) { + Ok(ep) => ep.connect_lazy(), + Err(e) => { + // tls_config_with_verifier only errors on UDS endpoints + // (see tonic's endpoint.rs), which we never construct. The + // defensive fallback returns a non-TLS lazy channel — the + // request will fail at the wire, surfacing the misconfig. + tracing::error!( + error = %e, address = %addr, + "tls_config_with_verifier failed; non-TLS lazy fallback", + ); + Endpoint::from_shared(uri) + .expect("EndpointAddress Display guarantees valid URI") + .connect_lazy() + } + }; + let svc = EndpointChannel::new(channel); + Box::pin(async move { svc }) + } +} + +#[cfg(test)] +mod tests { + use super::*; + use crate::xds::resource::cluster::{ClusterResource, LbPolicy}; + + fn plaintext_cluster() -> ClusterResource { + ClusterResource { + name: "c".into(), + eds_service_name: None, + lb_policy: LbPolicy::RoundRobin, + security: None, + } + } + + #[cfg(feature = "_tls-any")] + fn empty_registry() -> CertProviderRegistry { + use std::collections::HashMap; + CertProviderRegistry::from_bootstrap(&HashMap::new()).unwrap() + } + + /// Plaintext dispatch under TLS feature. + #[cfg(feature = "_tls-any")] + #[test] + fn build_connector_plaintext_tls_feature_on() { + assert!(build_connector(&plaintext_cluster(), &empty_registry()).is_ok()); + } + + /// Plaintext dispatch without any TLS feature. + #[cfg(not(feature = "_tls-any"))] + #[test] + fn build_connector_plaintext_no_tls() { + assert!(build_connector(&plaintext_cluster()).is_ok()); + } + + /// Cluster with TLS pointing at an instance not in the registry surfaces + /// a clear error — useful for misconfig diagnostics. + #[cfg(feature = "_tls-any")] + #[test] + fn build_connector_tls_unknown_ca() { + use crate::xds::resource::security::ClusterSecurityConfig; + + let mut cluster = plaintext_cluster(); + cluster.security = Some(ClusterSecurityConfig { + ca_instance_name: "missing-ca".into(), + identity_instance_name: None, + san_matchers: vec![], + }); + let Err(err) = build_connector(&cluster, &empty_registry()) else { + panic!("expected UnknownCaInstance error"); + }; + assert!(matches!( + err, + ConnectorBuildError::Tls(TlsConnectorBuildError::UnknownCaInstance(ref name)) + if name == "missing-ca" + )); + } + + /// `TlsConnector::connect` fetches the identity provider on every call, + /// which is what gives us identity rotation between CDS updates. Counter + /// shim verifies the call count without standing up a TLS handshake. + #[cfg(feature = "_tls-any")] + #[tokio::test] + async fn tls_connector_fetches_identity_per_connect() { + use crate::xds::cert_provider::{ + CertProviderError, CertificateData, CertificateProvider, Identity, + }; + use rustls::RootCertStore; + use std::sync::atomic::{AtomicUsize, Ordering}; + + struct CountingIdentity { + count: AtomicUsize, + data: Arc, + } + impl CertificateProvider for CountingIdentity { + fn fetch(&self) -> Result, CertProviderError> { + self.count.fetch_add(1, Ordering::Relaxed); + Ok(self.data.clone()) + } + } + + struct StaticCa(Arc); + impl CertificateProvider for StaticCa { + fn fetch(&self) -> Result, CertProviderError> { + Ok(self.0.clone()) + } + } + + let ca_provider: Arc = + Arc::new(StaticCa(Arc::new(CertificateData::RootsOnly { + roots: Arc::new(RootCertStore::empty()), + }))); + let verifier = Arc::new(XdsServerCertVerifier::new(ca_provider, vec![])); + + let identity_data = Arc::new(CertificateData::IdentityOnly { + identity: Identity { + cert_chain: b"cert".to_vec(), + key: b"key".to_vec(), + }, + }); + let counter = Arc::new(CountingIdentity { + count: AtomicUsize::new(0), + data: identity_data, + }); + let identity_provider: Arc = counter.clone(); + let connector = TlsConnector { + verifier, + identity_provider: Some(identity_provider), + }; + + let addr = EndpointAddress::from("1.2.3.4:443".parse::().unwrap()); + let _ = connector.connect(&addr).await; + let _ = connector.connect(&addr).await; + + assert_eq!( + counter.count.load(Ordering::Relaxed), + 2, + "TlsConnector should fetch identity provider on every connect call", + ); + } +} diff --git a/tonic-xds/src/xds/lb_discovery.rs b/tonic-xds/src/xds/lb_discovery.rs new file mode 100644 index 000000000..b712e083c --- /dev/null +++ b/tonic-xds/src/xds/lb_discovery.rs @@ -0,0 +1,176 @@ +//! xDS-backed discovery for the `tonic-xds-lb` load balancer. +//! +//! Unlike the `tower-lb` path (which connects inside discovery and yields +//! ready [`EndpointChannel`]s), the in-crate `LoadBalancer` manages the +//! connection lifecycle itself. So this module provides two decoupled pieces: +//! +//! 1. [`discover_endpoints`] — diffs EDS snapshots into +//! `Change` events (addresses only; no +//! connection is established here). +//! 2. [`XdsLbConnector`] — the [`Connector`] the `LoadBalancer` calls to turn +//! an address into a live channel. Because [`Connector::connect`] returns a +//! future, the connector is handed to the LB synchronously and resolves the +//! cluster's CDS security config *inside* that future via +//! [`build_connector`]. If CDS has not arrived yet (or a CDS update fails +//! validation), the connect future parks until a valid config is available. +//! +//! TODO: handle the case where the cluster is removed from the cache while a connect future is pending. +//! Currently, the future will park forever, it should instead try to re-establish the cluster watch and +//! only fail-fast when the client is closed. + +use std::collections::HashSet; +use std::pin::Pin; +use std::sync::Arc; + +use futures_core::Stream; +use tokio::sync::mpsc; +use tokio_stream::wrappers::ReceiverStream; +use tonic::transport::Channel; +use tower::BoxError; +use tower::discover::Change; + +use crate::client::endpoint::{Connector, EndpointAddress, EndpointChannel}; +use crate::client::loadbalance::channel_state::IdleChannel; +use crate::common::async_util::BoxFuture; +use crate::xds::cache::{CacheWatch, XdsCache}; +#[cfg(feature = "_tls-any")] +use crate::xds::cert_provider::CertProviderRegistry; +use crate::xds::connector::build_connector; +use crate::xds::resource::EndpointsResource; + +/// Buffer capacity for the endpoint change channel between the diff loop and +/// the load balancer's `Discover`. +const ENDPOINT_CHANNEL_CAPACITY: usize = 64; + +/// A pinned, boxed stream of endpoint changes, consumed by the `LoadBalancer` +/// through Tower's `Discover` blanket impl. Each inserted endpoint is carried +/// as an [`IdleChannel`] — the initial state of the channel state machine. +pub(crate) type EndpointDiscover = + Pin, BoxError>> + Send>>; + +// Compile-time assertion that `EndpointDiscover` satisfies the bounds the +// `LoadBalancer` requires from its discovery. +const _: fn() = || { + fn assert_discover() + where + D: tower::discover::Discover + Unpin, + { + } + assert_discover::(); +}; + +/// Returns a stream of endpoint changes for a cluster. +/// +/// Diffs each EDS snapshot against the previous set of healthy endpoints, +/// emitting `Change::Insert` (as a not-yet-connected [`IdleChannel`]) for new +/// endpoints and `Change::Remove` for gone ones. +pub(crate) fn discover_endpoints(cache: &Arc, cluster_name: &str) -> EndpointDiscover { + let (tx, rx) = mpsc::channel(ENDPOINT_CHANNEL_CAPACITY); + tokio::spawn(diff_loop(cache.watch_endpoints(cluster_name), tx)); + Box::pin(ReceiverStream::new(rx)) +} + +/// Background task: watches EDS snapshots and emits incremental endpoint +/// changes. Exits when the cache watch closes or the receiver is dropped. +async fn diff_loop( + mut watch: CacheWatch, + tx: mpsc::Sender, BoxError>>, +) { + let mut active: HashSet = HashSet::new(); + + while let Some(endpoints) = watch.next().await { + let new_set: HashSet = endpoints + .healthy_endpoints() + .map(|ep| ep.address.clone()) + .collect(); + + for added in new_set.difference(&active) { + let change = Change::Insert(added.clone(), IdleChannel::new(added.clone())); + if tx.send(Ok(change)).await.is_err() { + return; + } + } + + for removed in active.difference(&new_set) { + if tx.send(Ok(Change::Remove(removed.clone()))).await.is_err() { + return; + } + } + + active = new_set; + } +} + +/// The [`Connector`] used by the `tonic-xds-lb` `LoadBalancer`. +/// +/// Resolves the cluster's CDS security config lazily, inside the connect +/// future: on each `connect`, it reads the current +/// [`ClusterResource`](crate::xds::resource::ClusterResource) from +/// the cache and builds the appropriate plaintext/TLS connector via +/// [`build_connector`]. This lets the connector be constructed synchronously +/// (before CDS/EDS resolve) while the returned future waits for a valid +/// config. +pub(crate) struct XdsLbConnector { + cache: Arc, + cluster_name: String, + #[cfg(feature = "_tls-any")] + cert_provider_registry: Arc, +} + +impl XdsLbConnector { + #[cfg(feature = "_tls-any")] + pub(crate) fn new( + cache: Arc, + cluster_name: String, + registry: Arc, + ) -> Self { + Self { + cache, + cluster_name, + cert_provider_registry: registry, + } + } + + #[cfg(not(feature = "_tls-any"))] + pub(crate) fn new(cache: Arc, cluster_name: String) -> Self { + Self { + cache, + cluster_name, + } + } +} + +impl Connector for XdsLbConnector { + type Service = EndpointChannel; + + fn connect(&self, addr: &EndpointAddress) -> BoxFuture { + let mut cluster_watch = self.cache.watch_cluster(&self.cluster_name); + let cluster_name = self.cluster_name.clone(); + let addr = addr.clone(); + #[cfg(feature = "_tls-any")] + let registry = self.cert_provider_registry.clone(); + + Box::pin(async move { + loop { + let Some(cluster) = cluster_watch.next().await else { + // Cluster removed from cache: no connector can be built. + // The LB drops this future when the endpoint leaves + // discovery, so park until then. + return std::future::pending().await; + }; + match build_connector( + &cluster, + #[cfg(feature = "_tls-any")] + ®istry, + ) { + Ok(connector) => return connector.connect(&addr).await, + Err(e) => tracing::warn!( + cluster = %cluster_name, + error = %e, + "CDS connector build failed; awaiting next CDS update", + ), + } + } + }) + } +} diff --git a/tonic-xds/src/xds/mod.rs b/tonic-xds/src/xds/mod.rs index 479258a30..69a833c7c 100644 --- a/tonic-xds/src/xds/mod.rs +++ b/tonic-xds/src/xds/mod.rs @@ -2,9 +2,15 @@ pub(crate) mod bootstrap; pub(crate) mod cache; #[cfg(feature = "_tls-any")] pub(crate) mod cert_provider; -pub(crate) mod cluster_discovery; -pub(crate) mod endpoint_manager; +pub(crate) mod connector; pub(crate) mod resource; pub(crate) mod resource_manager; pub(crate) mod routing; pub(crate) mod uri; +cfg_tower_lb! { + pub(crate) mod cluster_discovery; + pub(crate) mod endpoint_manager; +} +cfg_tonic_xds_lb! { + pub(crate) mod lb_discovery; +}