Wire A32 circuit breaking into xDS channels

tonic-xds/src/client/channel.rs

@@ -22,6 +22,9 @@
  *
  */
 
+use crate::client::circuit_breaking::{
+    CircuitBreakingClusterDiscovery, ClusterCircuitBreakerRegistry,
+};
 use crate::client::cluster::ClusterClientRegistryGrpc;
 use crate::client::endpoint::{EndpointAddress, EndpointChannel};
 use crate::client::lb::{ClusterDiscovery, XdsLbService};
@@ -342,14 +345,18 @@ impl XdsChannelBuilder {
         resource_manager: XdsResourceManager,
     ) -> XdsChannelGrpc {
         let router: Arc<dyn Router> = Arc::new(XdsRouter::new(&cache));
+        let circuit_breakers = ClusterCircuitBreakerRegistry::default();
         #[cfg(feature = "_tls-any")]
         let discovery: Arc<
             dyn ClusterDiscovery<EndpointAddress, EndpointChannel<Channel>>,
-        > = Arc::new(XdsClusterDiscovery::new(cache, cert_provider_registry));
+        > = Arc::new(XdsClusterDiscovery::new(
+            cache.clone(),
+            cert_provider_registry,
+        ));
         #[cfg(not(feature = "_tls-any"))]
         let discovery: Arc<
             dyn ClusterDiscovery<EndpointAddress, EndpointChannel<Channel>>,
-        > = Arc::new(XdsClusterDiscovery::new(cache));
+        > = Arc::new(XdsClusterDiscovery::new(cache.clone()));
         let retry_policy = GrpcRetryPolicy::new(GrpcRetryPolicyConfig::default());
 
         let resources = Arc::new(XdsChannelResources {
@@ -358,6 +365,10 @@ impl XdsChannelBuilder {
         });
 
         let routing_layer = XdsRoutingLayer::new(router, self.authority());
+        let discovery = Arc::new(
+            CircuitBreakingClusterDiscovery::new(discovery, circuit_breakers)
+                .with_cluster_cache(cache),
+        );
         let retry_layer = RetryLayer::new(retry_policy);
         let cluster_registry = Arc::new(ClusterClientRegistryGrpc::new());
         let lb_service = XdsLbService::new(cluster_registry, discovery);
@@ -390,8 +401,28 @@ impl XdsChannelBuilder {
         router: Arc<dyn Router>,
         discovery: Arc<dyn ClusterDiscovery<EndpointAddress, EndpointChannel<Channel>>>,
         retry_policy: GrpcRetryPolicy,
+    ) -> XdsChannelGrpc {
+        self.build_grpc_channel_from_parts_with_circuit_breakers(
+            router,
+            discovery,
+            retry_policy,
+            ClusterCircuitBreakerRegistry::new_for_test(),
+        )
+    }
+
+    #[cfg(test)]
+    pub(crate) fn build_grpc_channel_from_parts_with_circuit_breakers(
+        &self,
+        router: Arc<dyn Router>,
+        discovery: Arc<dyn ClusterDiscovery<EndpointAddress, EndpointChannel<Channel>>>,
+        retry_policy: GrpcRetryPolicy,
+        circuit_breakers: ClusterCircuitBreakerRegistry,
     ) -> XdsChannelGrpc {
         let routing_layer = XdsRoutingLayer::new(router, self.authority());
+        let discovery = Arc::new(CircuitBreakingClusterDiscovery::new(
+            discovery,
+            circuit_breakers,
+        ));
         let retry_layer = RetryLayer::new(retry_policy);
         let cluster_registry = Arc::new(ClusterClientRegistryGrpc::new());
         let lb_service = XdsLbService::new(cluster_registry, discovery);
@@ -421,6 +452,7 @@ mod tests {
     use super::{XdsChannelBuilder, XdsChannelConfig};
     use crate::XdsUri;
     use crate::client::channel::XdsChannelGrpc;
+    use crate::client::circuit_breaking::ClusterCircuitBreakerRegistry;
     use crate::client::endpoint::EndpointAddress;
     use crate::client::endpoint::EndpointChannel;
 
@@ -438,6 +470,7 @@ mod tests {
     use crate::testutil::grpc::TestServer;
     use crate::xds::cache::XdsCache;
     use crate::xds::resource::EndpointsResource;
+    use crate::xds::resource::circuit_breaking::CircuitBreakingConfig;
     use crate::xds::resource::route_config::RouteConfigResource;
     use std::sync::Arc;
     use tokio::sync::mpsc;
@@ -669,16 +702,140 @@ mod tests {
         assert_eq!(response.into_inner().message, "retry-server: retry-test");
     }
 
+    #[tokio::test]
+    async fn test_xds_channel_enforces_injected_circuit_breaking_limit() {
+        use crate::client::retry::{GrpcRetryPolicy, GrpcRetryPolicyConfig};
+
+        let (_, servers) = setup_grpc_servers(1).await;
+        let xds_manager = Arc::new(MockXdsManager::from_test_servers(&servers));
+        let circuit_breakers = ClusterCircuitBreakerRegistry::new_for_test();
+        circuit_breakers.set_config("test-cluster", CircuitBreakingConfig { max_requests: 0 });
+
+        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_with_circuit_breakers(
+                xds_manager.clone(),
+                xds_manager.clone(),
+                retry_policy,
+                circuit_breakers,
+            );
+        let mut client = GreeterClient::new(xds_channel);
+
+        let error = client
+            .say_hello(HelloRequest {
+                name: "limited".to_string(),
+            })
+            .await
+            .unwrap_err();
+
+        assert_eq!(error.code(), tonic::Code::Unavailable);
+        assert!(error.message().contains("max_requests limit 0"));
+
+        for server in servers {
+            let _ = server.shutdown.send(());
+            let _ = server.handle.await;
+        }
+    }
+
+    #[tokio::test]
+    async fn test_xds_channel_uses_cds_circuit_breaking_config() {
+        let cluster_name = "test-cluster";
+        let (_, servers) = setup_grpc_servers(1).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_with_circuit_breaking(
+                cluster_name,
+                CircuitBreakingConfig { max_requests: 0 },
+            ),
+        );
+        cache.update_endpoints(cluster_name, make_test_endpoints(cluster_name, &servers));
+
+        let channel = build_xds_channel_from_cache(cache).await;
+        let mut client = GreeterClient::new(channel);
+        let error = client
+            .say_hello(HelloRequest {
+                name: "cds-limited".to_string(),
+            })
+            .await
+            .unwrap_err();
+
+        assert_eq!(error.code(), tonic::Code::Unavailable);
+        assert!(error.message().contains("max_requests limit 0"));
+
+        for server in servers {
+            let _ = server.shutdown.send(());
+            let _ = server.handle.await;
+        }
+    }
+
+    #[tokio::test]
+    async fn test_xds_channel_waits_for_cds_before_circuit_breaking() {
+        let cluster_name = "test-cluster";
+        let (_, servers) = setup_grpc_servers(1).await;
+
+        let cache = Arc::new(XdsCache::new());
+        cache.update_route_config(make_test_route_config(cluster_name));
+        cache.update_endpoints(cluster_name, make_test_endpoints(cluster_name, &servers));
+
+        let channel = build_xds_channel_from_cache(cache.clone()).await;
+        let mut client = GreeterClient::new(channel);
+        let mut request = Box::pin(client.say_hello(HelloRequest {
+            name: "wait-for-cds".to_string(),
+        }));
+
+        let early =
+            tokio::time::timeout(tokio::time::Duration::from_millis(20), request.as_mut()).await;
+        assert!(
+            early.is_err(),
+            "request should wait for CDS before acquiring a circuit-breaking permit",
+        );
+
+        cache.update_cluster(
+            cluster_name,
+            make_test_cluster_with_circuit_breaking(
+                cluster_name,
+                CircuitBreakingConfig { max_requests: 0 },
+            ),
+        );
+
+        let result = tokio::time::timeout(tokio::time::Duration::from_secs(2), request)
+            .await
+            .expect("request should complete after CDS update");
+        let error = result.unwrap_err();
+        assert_eq!(error.code(), tonic::Code::Unavailable);
+        assert!(error.message().contains("max_requests limit 0"));
+
+        for server in servers {
+            let _ = server.shutdown.send(());
+            let _ = server.handle.await;
+        }
+    }
+
     /// 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<crate::xds::resource::ClusterResource> {
+        make_test_cluster_with_circuit_breaking(cluster_name, CircuitBreakingConfig::default())
+    }
+
+    fn make_test_cluster_with_circuit_breaking(
+        cluster_name: &str,
+        circuit_breaking: CircuitBreakingConfig,
+    ) -> Arc<crate::xds::resource::ClusterResource> {
         use crate::xds::resource::cluster::{ClusterResource, LbPolicy};
         Arc::new(ClusterResource {
             name: cluster_name.to_string(),
             eds_service_name: None,
             lb_policy: LbPolicy::RoundRobin,
             security: None,
+            circuit_breaking,
         })
     }
 
@@ -729,29 +886,25 @@ mod tests {
     /// Builds an XdsChannelGrpc using real XdsRouter and XdsClusterDiscovery
     /// backed by the given cache.
     async fn build_xds_channel_from_cache(cache: Arc<XdsCache>) -> XdsChannelGrpc {
-        use crate::xds::cluster_discovery::XdsClusterDiscovery;
-        use crate::xds::routing::XdsRouter;
-
-        let router: Arc<dyn Router> = Arc::new(XdsRouter::new(&cache));
+        use crate::xds::resource_manager::XdsResourceManager;
 
+        let xds_client = xds_client::XdsClient::disconnected();
+        let resource_manager =
+            XdsResourceManager::new(xds_client.clone(), cache.clone(), "test-listener".into());
+        let builder = XdsChannelBuilder::new(test_config());
         #[cfg(feature = "_tls-any")]
-        let discovery: Arc<
-            dyn ClusterDiscovery<EndpointAddress, EndpointChannel<Channel>>,
-        > = {
+        {
             use crate::xds::cert_provider::CertProviderRegistry;
             let registry = Arc::new(
                 CertProviderRegistry::from_bootstrap(&Default::default(), Default::default())
                     .unwrap(),
             );
-            Arc::new(XdsClusterDiscovery::new(cache, registry))
-        };
+            builder.build_from_cache(cache, registry, xds_client, resource_manager)
+        }
         #[cfg(not(feature = "_tls-any"))]
-        let discovery: Arc<
-            dyn ClusterDiscovery<EndpointAddress, EndpointChannel<Channel>>,
-        > = Arc::new(XdsClusterDiscovery::new(cache));
-
-        let builder = XdsChannelBuilder::new(test_config());
-        builder.build_grpc_channel_from_parts(router, discovery, GrpcRetryPolicy::default())
+        {
+            builder.build_from_cache(cache, xds_client, resource_manager)
+        }
     }
 
     /// Tests the full xDS stack (XdsRouter + XdsClusterDiscovery) with a