Focus MultiChildLB updates around ResolvedAddresses of children

This makes ClusterManagerLB more straight-forward, focusing on just the
things that are relevant to it, and it avoids specialized map key
handling in updateChildrenWithResolvedAddresses().

Author: ejona86

util/src/main/java/io/grpc/util/MultiChildLoadBalancer.java

@@ -16,7 +16,6 @@
 
 package io.grpc.util;
 
-import static com.google.common.base.Preconditions.checkArgument;
 import static com.google.common.base.Preconditions.checkNotNull;
 import static io.grpc.ConnectivityState.CONNECTING;
 import static io.grpc.ConnectivityState.IDLE;
@@ -80,20 +79,20 @@ protected MultiChildLoadBalancer(Helper helper) {
 
   /**
    * Override to utilize parsing of the policy configuration or alternative helper/lb generation.
+   * Override this if keys are not Endpoints or if child policies have configuration.
    */
-  protected Map<Object, ChildLbState> createChildLbMap(ResolvedAddresses resolvedAddresses) {
-    Map<Object, ChildLbState> childLbMap = new HashMap<>();
-    List<EquivalentAddressGroup> addresses = resolvedAddresses.getAddresses();
-    for (EquivalentAddressGroup eag : addresses) {
-      Endpoint endpoint = new Endpoint(eag); // keys need to be just addresses
-      ChildLbState existingChildLbState = childLbStates.get(endpoint);
-      if (existingChildLbState != null) {
-        childLbMap.put(endpoint, existingChildLbState);
-      } else {
-        childLbMap.put(endpoint, createChildLbState(endpoint));
-      }
-    }
-    return childLbMap;
+  protected Map<Object, ResolvedAddresses> createChildAddressesMap(
+      ResolvedAddresses resolvedAddresses) {
+    Map<Object, ResolvedAddresses> childAddresses = new HashMap<>();
+    for (EquivalentAddressGroup eag : resolvedAddresses.getAddresses()) {
+      ResolvedAddresses addresses = resolvedAddresses.toBuilder()
+          .setAddresses(Collections.singletonList(eag))
+          .setAttributes(Attributes.newBuilder().set(IS_PETIOLE_POLICY, true).build())
+          .setLoadBalancingPolicyConfig(null)
+          .build();
+      childAddresses.put(new Endpoint(eag), addresses);
+    }
+    return childAddresses;
   }
 
   /**
@@ -128,39 +127,6 @@ public Status acceptResolvedAddresses(ResolvedAddresses resolvedAddresses) {
     }
   }
 
-  /**
-   * Override this if your keys are not of type Endpoint.
-   * @param key Key to identify the ChildLbState
-   * @param resolvedAddresses list of addresses which include attributes
-   * @return a fully loaded ResolvedAddresses object for the specified key
-   */
-  protected ResolvedAddresses getChildAddresses(Object key, ResolvedAddresses resolvedAddresses) {
-    Endpoint endpointKey;
-    if (key instanceof EquivalentAddressGroup) {
-      endpointKey = new Endpoint((EquivalentAddressGroup) key);
-    } else {
-      checkArgument(key instanceof Endpoint, "key is wrong type");
-      endpointKey = (Endpoint) key;
-    }
-
-    // Retrieve the non-stripped version
-    EquivalentAddressGroup eagToUse = null;
-    for (EquivalentAddressGroup currEag : resolvedAddresses.getAddresses()) {
-      if (endpointKey.equals(new Endpoint(currEag))) {
-        eagToUse = currEag;
-        break;
-      }
-    }
-
-    checkNotNull(eagToUse, key + " no longer present in load balancer children");
-
-    return resolvedAddresses.toBuilder()
-        .setAddresses(Collections.singletonList(eagToUse))
-        .setAttributes(Attributes.newBuilder().set(IS_PETIOLE_POLICY, true).build())
-        .setLoadBalancingPolicyConfig(null)
-        .build();
-  }
-
   /**
    * Handle the name resolution error.
    *
@@ -192,41 +158,31 @@ protected final AcceptResolvedAddrRetVal acceptResolvedAddressesInternal(
       ResolvedAddresses resolvedAddresses) {
     logger.log(Level.FINE, "Received resolution result: {0}", resolvedAddresses);
 
-    // Subclass handles any special manipulation to create appropriate types of keyed ChildLbStates
-    Map<Object, ChildLbState> newChildren = createChildLbMap(resolvedAddresses);
+    Map<Object, ResolvedAddresses> newChildAddresses = createChildAddressesMap(resolvedAddresses);
 
     // Handle error case
-    if (newChildren.isEmpty()) {
+    if (newChildAddresses.isEmpty()) {
       Status unavailableStatus = Status.UNAVAILABLE.withDescription(
           "NameResolver returned no usable address. " + resolvedAddresses);
       handleNameResolutionError(unavailableStatus);
       return new AcceptResolvedAddrRetVal(unavailableStatus, null);
     }
 
-    addMissingChildren(newChildren);
+    updateChildrenWithResolvedAddresses(newChildAddresses);
 
-    updateChildrenWithResolvedAddresses(resolvedAddresses, newChildren);
-
-    return new AcceptResolvedAddrRetVal(Status.OK, getRemovedChildren(newChildren.keySet()));
+    return new AcceptResolvedAddrRetVal(Status.OK, getRemovedChildren(newChildAddresses.keySet()));
   }
 
-  private void addMissingChildren(Map<Object, ChildLbState> newChildren) {
-    // Do adds and identify reused children
-    for (Map.Entry<Object, ChildLbState> entry : newChildren.entrySet()) {
-      final Object key = entry.getKey();
-      if (!childLbStates.containsKey(key)) {
-        childLbStates.put(key, entry.getValue());
-      }
-    }
-  }
-
-  private void updateChildrenWithResolvedAddresses(ResolvedAddresses resolvedAddresses,
-                                                     Map<Object, ChildLbState> newChildren) {
-    for (Map.Entry<Object, ChildLbState> entry : newChildren.entrySet()) {
+  private void updateChildrenWithResolvedAddresses(
+      Map<Object, ResolvedAddresses> newChildAddresses) {
+    for (Map.Entry<Object, ResolvedAddresses> entry : newChildAddresses.entrySet()) {
       ChildLbState childLbState = childLbStates.get(entry.getKey());
-      ResolvedAddresses childAddresses = getChildAddresses(entry.getKey(), resolvedAddresses);
-      childLbState.setResolvedAddresses(childAddresses); // update child
-      childLbState.lb.handleResolvedAddresses(childAddresses); // update child LB
+      if (childLbState == null) {
+        childLbState = createChildLbState(entry.getKey());
+        childLbStates.put(entry.getKey(), childLbState);
+      }
+      childLbState.setResolvedAddresses(entry.getValue()); // update child
+      childLbState.lb.handleResolvedAddresses(entry.getValue()); // update child LB
     }
   }
 

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

@@ -70,31 +70,28 @@ class ClusterManagerLoadBalancer extends MultiChildLoadBalancer {
   }
 
   @Override
-  protected ResolvedAddresses getChildAddresses(Object key, ResolvedAddresses resolvedAddresses) {
-    ClusterManagerConfig config = (ClusterManagerConfig)
-        resolvedAddresses.getLoadBalancingPolicyConfig();
-    Object childConfig = config.childPolicies.get(key);
-    return resolvedAddresses.toBuilder().setLoadBalancingPolicyConfig(childConfig).build();
+  protected ChildLbState createChildLbState(Object key) {
+    return new ClusterManagerLbState(key, GracefulSwitchLoadBalancerFactory.INSTANCE);
   }
 
   @Override
-  protected Map<Object, ChildLbState> createChildLbMap(ResolvedAddresses resolvedAddresses) {
+  protected Map<Object, ResolvedAddresses> createChildAddressesMap(
+      ResolvedAddresses resolvedAddresses) {
     ClusterManagerConfig config = (ClusterManagerConfig)
         resolvedAddresses.getLoadBalancingPolicyConfig();
-    Map<Object, ChildLbState> newChildPolicies = new HashMap<>();
+    Map<Object, ResolvedAddresses> childAddresses = new HashMap<>();
     if (config != null) {
-      for (String key : config.childPolicies.keySet()) {
-        ChildLbState child = getChildLbState(key);
-        if (child == null) {
-          child = new ClusterManagerLbState(key, GracefulSwitchLoadBalancerFactory.INSTANCE);
-        }
-        newChildPolicies.put(key, child);
+      for (Map.Entry<String, Object> childPolicy : config.childPolicies.entrySet()) {
+        ResolvedAddresses addresses = resolvedAddresses.toBuilder()
+            .setLoadBalancingPolicyConfig(childPolicy.getValue())
+            .build();
+        childAddresses.put(childPolicy.getKey(), addresses);
       }
     }
     logger.log(
         XdsLogLevel.INFO,
-        "Received cluster_manager lb config: child names={0}", newChildPolicies.keySet());
-    return newChildPolicies;
+        "Received cluster_manager lb config: child names={0}", childAddresses.keySet());
+    return childAddresses;
   }
 
   /**