[wip]Lbv2 wiring

tonic-xds/Cargo.toml

@@ -70,8 +70,17 @@ tempfile = "3.27.0"
 rcgen = "0.14"
 
 [features]
+default = ["tower-lb"]
 testutil = ["dep:tonic-prost"]
 
+# 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` takes precedence, so enabling both (e.g.
+# `--all-features`) is safe and selects `tonic-xds-lb`.
+tower-lb = []
+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"]

tonic-xds/src/client/loadbalance/channel_state.rs

@@ -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<S> ReadyChannel<S> {
     /// 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<C: Connector<Service = S>>(
         self,
         connector: Arc<C>,

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)
     }

tonic-xds/src/client/loadbalance/mod.rs

@@ -1,7 +1,11 @@
+// `LbChannel` is not constructed by the current P2C wiring (the LoadBalancer
+// tracks load via `ReadyChannel` + `EndpointChannel`); retained for now.
+#[allow(dead_code)]
 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;

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()
     }

tonic-xds/src/client/loadbalance/pickers/mod.rs

@@ -1,4 +1,7 @@
 pub(crate) mod p2c;
+// Ring-hash (gRFC A42) is implemented but not wired into the first-cut LB,
+// which selects P2C. Wired in a follow-up.
+#[allow(dead_code)]
 pub(crate) mod ring_hash;
 
 use indexmap::{IndexMap, IndexSet};

tonic-xds/src/client/loadbalance/service.rs

@@ -0,0 +1,156 @@
+//! Routing → per-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.
+//!
+//! First-cut policy: power-of-two-choices ([`P2cPicker`]) with outlier
+//! detection disabled ([`OutlierDetectionConfig::default`]). Ring-hash
+//! selection and xDS-driven outlier config are 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 (after the routing/retry layers
+/// and the `SharedBody` → `TonicBody` remap).
+type LbRequest = Request<TonicBody>;
+/// The response type produced by an endpoint channel.
+type LbResponse = Response<TonicBody>;
+
+/// 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<LbRequest, LbFuture<LbResponse>>;
+
+/// 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<XdsCache>,
+    #[cfg(feature = "_tls-any")]
+    cert_provider_registry: Arc<CertProviderRegistry>,
+    clusters: DashMap<String, ClusterChannel>,
+}
+
+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<ReadyChannel<EndpointChannel<Channel>>, 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<ClusterLbRegistry>,
+}
+
+impl XdsLoadBalanceService {
+    #[cfg(feature = "_tls-any")]
+    pub(crate) fn new(
+        cache: Arc<XdsCache>,
+        cert_provider_registry: Arc<CertProviderRegistry>,
+    ) -> Self {
+        Self {
+            registry: Arc::new(ClusterLbRegistry {
+                cache,
+                cert_provider_registry,
+                clusters: DashMap::new(),
+            }),
+        }
+    }
+
+    #[cfg(not(feature = "_tls-any"))]
+    pub(crate) fn new(cache: Arc<XdsCache>) -> Self {
+        Self {
+            registry: Arc::new(ClusterLbRegistry {
+                cache,
+                clusters: DashMap::new(),
+            }),
+        }
+    }
+}
+
+impl Service<LbRequest> for XdsLoadBalanceService {
+    type Response = LbResponse;
+    type Error = BoxError;
+    type Future = BoxFuture<Result<Self::Response, Self::Error>>;
+
+    fn poll_ready(&mut self, _cx: &mut Context<'_>) -> Poll<Result<(), Self::Error>> {
+        // 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::<RouteDecision>().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
+        })
+    }
+}

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;
+}

tonic-xds/src/lib.rs

@@ -140,6 +140,34 @@
 //! [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. Backend code is gated on the *absence* of `tonic-xds-lb`
+// (not on a `tower-lb` feature) so every feature combination — including the
+// per-feature checks run by `cargo hack --each-feature` — compiles with exactly
+// one backend in production.
+//
+// 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;

tonic-xds/src/testutil/grpc.rs

@@ -60,7 +60,10 @@ impl Greeter for FailFirstNGreeter {
 
 /// A test server that runs a gRPC service and provides a channel for clients to connect.
 pub(crate) struct TestServer {
-    /// The gRPC channel for talking to the test server.
+    /// The gRPC channel for talking to the test server. Currently unused by the
+    /// cache-driven channel tests (which connect via discovered endpoints), but
+    /// retained as a convenience for direct-connection tests.
+    #[allow(dead_code)]
     pub channel: Channel,
     /// Signal the server to shutdown.
     pub shutdown: oneshot::Sender<()>,