A75 Aggregate cluster fixes

xds/src/main/java/io/grpc/xds/CdsLoadBalancer2.java

@@ -18,10 +18,10 @@
 
 import static com.google.common.base.Preconditions.checkNotNull;
 import static io.grpc.ConnectivityState.TRANSIENT_FAILURE;
+import static io.grpc.xds.XdsLbPolicies.CDS_POLICY_NAME;
 import static io.grpc.xds.XdsLbPolicies.CLUSTER_RESOLVER_POLICY_NAME;
+import static io.grpc.xds.XdsLbPolicies.PRIORITY_POLICY_NAME;
 
-import com.google.common.annotations.VisibleForTesting;
-import com.google.common.collect.ImmutableList;
 import com.google.errorprone.annotations.CheckReturnValue;
 import io.grpc.InternalLogId;
 import io.grpc.LoadBalancer;
@@ -33,6 +33,7 @@
 import io.grpc.xds.CdsLoadBalancerProvider.CdsConfig;
 import io.grpc.xds.ClusterResolverLoadBalancerProvider.ClusterResolverConfig;
 import io.grpc.xds.ClusterResolverLoadBalancerProvider.ClusterResolverConfig.DiscoveryMechanism;
+import io.grpc.xds.PriorityLoadBalancerProvider.PriorityLbConfig.PriorityChildConfig;
 import io.grpc.xds.XdsClusterResource.CdsUpdate;
 import io.grpc.xds.XdsClusterResource.CdsUpdate.ClusterType;
 import io.grpc.xds.XdsConfig.Subscription;
@@ -41,10 +42,11 @@
 import io.grpc.xds.XdsConfig.XdsClusterConfig.EndpointConfig;
 import io.grpc.xds.client.XdsLogger;
 import io.grpc.xds.client.XdsLogger.XdsLogLevel;
-import java.util.ArrayList;
 import java.util.Arrays;
 import java.util.Collections;
+import java.util.HashMap;
 import java.util.List;
+import java.util.Map;
 
 /**
  * Load balancer for cds_experimental LB policy. One instance per top-level cluster.
@@ -55,19 +57,15 @@ final class CdsLoadBalancer2 extends LoadBalancer {
   private final XdsLogger logger;
   private final Helper helper;
   private final LoadBalancerRegistry lbRegistry;
+  private GracefulSwitchLoadBalancer delegate;
   // Following fields are effectively final.
   private String clusterName;
   private Subscription clusterSubscription;
-  private LoadBalancer childLb;
 
-  CdsLoadBalancer2(Helper helper) {
-    this(helper, LoadBalancerRegistry.getDefaultRegistry());
-  }
-
-  @VisibleForTesting
   CdsLoadBalancer2(Helper helper, LoadBalancerRegistry lbRegistry) {
     this.helper = checkNotNull(helper, "helper");
     this.lbRegistry = checkNotNull(lbRegistry, "lbRegistry");
+    this.delegate = new GracefulSwitchLoadBalancer(helper);
     logger = XdsLogger.withLogId(InternalLogId.allocate("cds-lb", helper.getAuthority()));
     logger.log(XdsLogLevel.INFO, "Created");
   }
@@ -91,7 +89,7 @@ public Status acceptResolvedAddresses(ResolvedAddresses resolvedAddresses) {
       if (clusterSubscription == null) {
         // Should be impossible, because XdsDependencyManager wouldn't have generated this
         return fail(Status.INTERNAL.withDescription(
-            errorPrefix() + "Unable to find non-dynamic root cluster"));
+            errorPrefix() + "Unable to find non-dynamic cluster"));
       }
       // The dynamic cluster must not have loaded yet
       return Status.OK;
@@ -100,42 +98,25 @@ public Status acceptResolvedAddresses(ResolvedAddresses resolvedAddresses) {
       return fail(clusterConfigOr.getStatus());
     }
     XdsClusterConfig clusterConfig = clusterConfigOr.getValue();
-    List<String> leafNames;
-    if (clusterConfig.getChildren() instanceof AggregateConfig) {
-      leafNames = ((AggregateConfig) clusterConfig.getChildren()).getLeafNames();
-    } else if (clusterConfig.getChildren() instanceof EndpointConfig) {
-      leafNames = ImmutableList.of(clusterName);
-    } else {
-      return fail(Status.INTERNAL.withDescription(
-          errorPrefix() + "Unexpected cluster children type: "
-          + clusterConfig.getChildren().getClass()));
-    }
-    if (leafNames.isEmpty()) {
-      // Should be impossible, because XdsClusterResource validated this
-      return fail(Status.UNAVAILABLE.withDescription(
-          errorPrefix() + "Zero leaf clusters for root cluster " + clusterName));
-    }
 
-    Status noneFoundError = Status.INTERNAL
-        .withDescription(errorPrefix() + "No leaves and no error; this is a bug");
-    List<DiscoveryMechanism> instances = new ArrayList<>();
-    for (String leafName : leafNames) {
-      StatusOr<XdsClusterConfig> leafConfigOr = xdsConfig.getClusters().get(leafName);
-      if (!leafConfigOr.hasValue()) {
-        noneFoundError = leafConfigOr.getStatus();
-        continue;
-      }
-      if (!(leafConfigOr.getValue().getChildren() instanceof EndpointConfig)) {
-        noneFoundError = Status.INTERNAL.withDescription(
-            errorPrefix() + "Unexpected child " + leafName + " cluster children type: "
-            + leafConfigOr.getValue().getChildren().getClass());
-        continue;
+    NameResolver.ConfigOrError configOrError;
+    Object gracefulConfig;
+    if (clusterConfig.getChildren() instanceof EndpointConfig) {
+      // The LB policy config is provided in service_config.proto/JSON format.
+      configOrError =
+              GracefulSwitchLoadBalancer.parseLoadBalancingPolicyConfig(
+                      Arrays.asList(clusterConfig.getClusterResource().lbPolicyConfig()),
+                      lbRegistry);
+      if (configOrError.getError() != null) {
+        // Should be impossible, because XdsClusterResource validated this
+        return fail(Status.INTERNAL.withDescription(
+                errorPrefix() + "Unable to parse the LB config: " + configOrError.getError()));
       }
-      CdsUpdate result = leafConfigOr.getValue().getClusterResource();
+      CdsUpdate result = clusterConfig.getClusterResource();
       DiscoveryMechanism instance;
       if (result.clusterType() == ClusterType.EDS) {
         instance = DiscoveryMechanism.forEds(
-            leafName,
+            clusterName,
             result.edsServiceName(),
             result.lrsServerInfo(),
             result.maxConcurrentRequests(),
@@ -144,45 +125,49 @@ public Status acceptResolvedAddresses(ResolvedAddresses resolvedAddresses) {
             result.outlierDetection());
       } else {
         instance = DiscoveryMechanism.forLogicalDns(
-            leafName,
+            clusterName,
             result.dnsHostName(),
             result.lrsServerInfo(),
             result.maxConcurrentRequests(),
             result.upstreamTlsContext(),
             result.filterMetadata());
       }
-      instances.add(instance);
-    }
-    if (instances.isEmpty()) {
-      return fail(noneFoundError);
-    }
-
-    // The LB policy config is provided in service_config.proto/JSON format.
-    NameResolver.ConfigOrError configOrError =
-        GracefulSwitchLoadBalancer.parseLoadBalancingPolicyConfig(
-            Arrays.asList(clusterConfig.getClusterResource().lbPolicyConfig()), lbRegistry);
-    if (configOrError.getError() != null) {
-      // Should be impossible, because XdsClusterResource validated this
+      gracefulConfig = GracefulSwitchLoadBalancer.createLoadBalancingPolicyConfig(
+          lbRegistry.getProvider(CLUSTER_RESOLVER_POLICY_NAME),
+          new ClusterResolverConfig(
+              instance,
+              configOrError.getConfig(),
+              clusterConfig.getClusterResource().isHttp11ProxyAvailable()));
+    } else if (clusterConfig.getChildren() instanceof AggregateConfig) {
+      Map<String, PriorityChildConfig> priorityChildConfigs = new HashMap<>();
+      List<String> leafClusters = ((AggregateConfig) clusterConfig.getChildren()).getLeafNames();
+      for (String childCluster: leafClusters) {
+        priorityChildConfigs.put(childCluster,
+                new PriorityChildConfig(
+                        GracefulSwitchLoadBalancer.createLoadBalancingPolicyConfig(
+                                lbRegistry.getProvider(CDS_POLICY_NAME),
+                                new CdsConfig(childCluster)),
+                        false));
+      }
+      gracefulConfig = GracefulSwitchLoadBalancer.createLoadBalancingPolicyConfig(
+          lbRegistry.getProvider(PRIORITY_POLICY_NAME),
+          new PriorityLoadBalancerProvider.PriorityLbConfig(
+              Collections.unmodifiableMap(priorityChildConfigs), leafClusters));
+    } else {
       return fail(Status.INTERNAL.withDescription(
-          errorPrefix() + "Unable to parse the LB config: " + configOrError.getError()));
+              errorPrefix() + "Unexpected cluster children type: "
+                      + clusterConfig.getChildren().getClass()));
     }
 
-    ClusterResolverConfig config = new ClusterResolverConfig(
-        Collections.unmodifiableList(instances),
-        configOrError.getConfig(),
-        clusterConfig.getClusterResource().isHttp11ProxyAvailable());
-    if (childLb == null) {
-      childLb = lbRegistry.getProvider(CLUSTER_RESOLVER_POLICY_NAME).newLoadBalancer(helper);
-    }
-    return childLb.acceptResolvedAddresses(
-        resolvedAddresses.toBuilder().setLoadBalancingPolicyConfig(config).build());
+    return delegate.acceptResolvedAddresses(
+        resolvedAddresses.toBuilder().setLoadBalancingPolicyConfig(gracefulConfig).build());
   }
 
   @Override
   public void handleNameResolutionError(Status error) {
     logger.log(XdsLogLevel.WARNING, "Received name resolution error: {0}", error);
-    if (childLb != null) {
-      childLb.handleNameResolutionError(error);
+    if (delegate != null) {
+      delegate.handleNameResolutionError(error);
     } else {
       helper.updateBalancingState(
           TRANSIENT_FAILURE, new FixedResultPicker(PickResult.withError(error)));
@@ -192,10 +177,8 @@ public void handleNameResolutionError(Status error) {
   @Override
   public void shutdown() {
     logger.log(XdsLogLevel.INFO, "Shutdown");
-    if (childLb != null) {
-      childLb.shutdown();
-      childLb = null;
-    }
+    delegate.shutdown();
+    delegate = new GracefulSwitchLoadBalancer(helper);
     if (clusterSubscription != null) {
       clusterSubscription.close();
       clusterSubscription = null;
@@ -204,10 +187,7 @@ public void shutdown() {
 
   @CheckReturnValue // don't forget to return up the stack after the fail call
   private Status fail(Status error) {
-    if (childLb != null) {
-      childLb.shutdown();
-      childLb = null;
-    }
+    delegate.shutdown();
     helper.updateBalancingState(
         TRANSIENT_FAILURE, new FixedResultPicker(PickResult.withError(error)));
     return Status.OK; // XdsNameResolver isn't a polling NR, so this value doesn't matter

xds/src/main/java/io/grpc/xds/CdsLoadBalancerProvider.java

@@ -23,6 +23,7 @@
 import io.grpc.LoadBalancer;
 import io.grpc.LoadBalancer.Helper;
 import io.grpc.LoadBalancerProvider;
+import io.grpc.LoadBalancerRegistry;
 import io.grpc.NameResolver.ConfigOrError;
 import io.grpc.Status;
 import io.grpc.internal.JsonUtil;
@@ -51,9 +52,24 @@ public String getPolicyName() {
     return XdsLbPolicies.CDS_POLICY_NAME;
   }
 
+  private final LoadBalancerRegistry loadBalancerRegistry;
+
+  public CdsLoadBalancerProvider() {
+    this.loadBalancerRegistry = null;
+  }
+
+  public CdsLoadBalancerProvider(LoadBalancerRegistry loadBalancerRegistry) {
+    this.loadBalancerRegistry = loadBalancerRegistry;
+  }
+
   @Override
   public LoadBalancer newLoadBalancer(Helper helper) {
-    return new CdsLoadBalancer2(helper);
+    LoadBalancerRegistry loadBalancerRegistry = this.loadBalancerRegistry;
+    if (loadBalancerRegistry == null) {
+      loadBalancerRegistry = LoadBalancerRegistry.getDefaultRegistry();
+    }
+
+    return new CdsLoadBalancer2(helper, loadBalancerRegistry);
   }
 
   @Override

xds/src/main/java/io/grpc/xds/ClusterResolverLoadBalancer.java

@@ -168,8 +168,8 @@ public LoadBalancer newLoadBalancer(Helper helper) {
    */
   private final class ClusterResolverLbState extends LoadBalancer {
     private final Helper helper;
-    private final List<String> clusters = new ArrayList<>();
-    private final Map<String, ClusterState> clusterStates = new HashMap<>();
+    private ClusterState clusterState;
+    private String cluster;
     private Object endpointLbConfig;
     private ResolvedAddresses resolvedAddresses;
     private LoadBalancer childLb;
@@ -185,21 +185,18 @@ public Status acceptResolvedAddresses(ResolvedAddresses resolvedAddresses) {
       ClusterResolverConfig config =
           (ClusterResolverConfig) resolvedAddresses.getLoadBalancingPolicyConfig();
       endpointLbConfig = config.lbConfig;
-      for (DiscoveryMechanism instance : config.discoveryMechanisms) {
-        clusters.add(instance.cluster);
-        ClusterState state;
-        if (instance.type == DiscoveryMechanism.Type.EDS) {
-          state = new EdsClusterState(instance.cluster, instance.edsServiceName,
-              instance.lrsServerInfo, instance.maxConcurrentRequests, instance.tlsContext,
-              instance.filterMetadata, instance.outlierDetection);
-        } else {  // logical DNS
-          state = new LogicalDnsClusterState(instance.cluster, instance.dnsHostName,
-              instance.lrsServerInfo, instance.maxConcurrentRequests, instance.tlsContext,
-              instance.filterMetadata);
-        }
-        clusterStates.put(instance.cluster, state);
-        state.start();
-      }
+      DiscoveryMechanism instance = config.discoveryMechanism;
+      cluster = instance.cluster;
+      if (instance.type == DiscoveryMechanism.Type.EDS) {
+        clusterState = new EdsClusterState(instance.cluster, instance.edsServiceName,
+            instance.lrsServerInfo, instance.maxConcurrentRequests, instance.tlsContext,
+            instance.filterMetadata, instance.outlierDetection);
+      } else {  // logical DNS
+        clusterState = new LogicalDnsClusterState(instance.cluster, instance.dnsHostName,
+            instance.lrsServerInfo, instance.maxConcurrentRequests, instance.tlsContext,
+            instance.filterMetadata);
+      }
+      clusterState.start();
       return Status.OK;
     }
 
@@ -215,9 +212,7 @@ public void handleNameResolutionError(Status error) {
 
     @Override
     public void shutdown() {
-      for (ClusterState state : clusterStates.values()) {
-        state.shutdown();
-      }
+      clusterState.shutdown();
       if (childLb != null) {
         childLb.shutdown();
       }
@@ -229,24 +224,21 @@ private void handleEndpointResourceUpdate() {
       List<String> priorities = new ArrayList<>();  // totally ordered priority list
 
       Status endpointNotFound = Status.OK;
-      for (String cluster : clusters) {
-        ClusterState state = clusterStates.get(cluster);
-        // Propagate endpoints to the child LB policy only after all clusters have been resolved.
-        if (!state.resolved && state.status.isOk()) {
-          return;
-        }
-        if (state.result != null) {
-          addresses.addAll(state.result.addresses);
-          priorityChildConfigs.putAll(state.result.priorityChildConfigs);
-          priorities.addAll(state.result.priorities);
-        } else {
-          endpointNotFound = state.status;
-        }
+      // Propagate endpoints to the child LB policy only after all clusters have been resolved.
+      if (!clusterState.resolved && clusterState.status.isOk()) {
+        return;
+      }
+      if (clusterState.result != null) {
+        addresses.addAll(clusterState.result.addresses);
+        priorityChildConfigs.putAll(clusterState.result.priorityChildConfigs);
+        priorities.addAll(clusterState.result.priorities);
+      } else {
+        endpointNotFound = clusterState.status;
       }
       if (addresses.isEmpty()) {
         if (endpointNotFound.isOk()) {
           endpointNotFound = Status.UNAVAILABLE.withDescription(
-              "No usable endpoint from cluster(s): " + clusters);
+              "No usable endpoint from cluster: " + cluster);
         } else {
           endpointNotFound =
               Status.UNAVAILABLE.withCause(endpointNotFound.getCause())
@@ -274,22 +266,12 @@ private void handleEndpointResourceUpdate() {
     }
 
     private void handleEndpointResolutionError() {
-      boolean allInError = true;
-      Status error = null;
-      for (String cluster : clusters) {
-        ClusterState state = clusterStates.get(cluster);
-        if (state.status.isOk()) {
-          allInError = false;
-        } else {
-          error = state.status;
-        }
-      }
-      if (allInError) {
+      if (!clusterState.status.isOk()) {
         if (childLb != null) {
-          childLb.handleNameResolutionError(error);
+          childLb.handleNameResolutionError(clusterState.status);
         } else {
           helper.updateBalancingState(
-              TRANSIENT_FAILURE, new FixedResultPicker(PickResult.withError(error)));
+              TRANSIENT_FAILURE, new FixedResultPicker(PickResult.withError(clusterState.status)));
         }
       }
     }
@@ -306,10 +288,8 @@ private RefreshableHelper(Helper delegate) {
 
       @Override
       public void refreshNameResolution() {
-        for (ClusterState state : clusterStates.values()) {
-          if (state instanceof LogicalDnsClusterState) {
-            ((LogicalDnsClusterState) state).refresh();
-          }
+        if (clusterState instanceof LogicalDnsClusterState) {
+          ((LogicalDnsClusterState) clusterState).refresh();
         }
       }