xds: Have ClusterManagerLB use child map for preserving children

Instead of doing a dance of supplementing config so the later createChildAddressesMap() won't delete children, just look at the existing children and don't delete any that shouldn't be deleted.
larry-safran · Sep 30, 2024 · 8c34969 · 8c34969
1 parent a908b5e
commit 8c34969
Show file tree

Hide file tree

Showing 2 changed files with 33 additions and 52 deletions.
diff --git a/util/src/main/java/io/grpc/util/MultiChildLoadBalancer.java b/util/src/main/java/io/grpc/util/MultiChildLoadBalancer.java
@@ -79,7 +79,8 @@ 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.
+   * Override this if keys are not Endpoints or if child policies have configuration. Null map
+   * values preserve the child without delivering the child an update.
    */
   protected Map<Object, ResolvedAddresses> createChildAddressesMap(
       ResolvedAddresses resolvedAddresses) {
@@ -181,8 +182,10 @@ private void updateChildrenWithResolvedAddresses(
         childLbState = createChildLbState(entry.getKey());
         childLbStates.put(entry.getKey(), childLbState);
       }
-      childLbState.setResolvedAddresses(entry.getValue()); // update child
-      childLbState.lb.handleResolvedAddresses(entry.getValue()); // update child LB
+      if (entry.getValue() != null) {
+        childLbState.setResolvedAddresses(entry.getValue()); // update child
+        childLbState.lb.handleResolvedAddresses(entry.getValue()); // update child LB
+      }
     }
   }
 

diff --git a/xds/src/main/java/io/grpc/xds/ClusterManagerLoadBalancer.java b/xds/src/main/java/io/grpc/xds/ClusterManagerLoadBalancer.java
@@ -77,57 +77,43 @@ protected ChildLbState createChildLbState(Object key) {
   @Override
   protected Map<Object, ResolvedAddresses> createChildAddressesMap(
       ResolvedAddresses resolvedAddresses) {
+    lastResolvedAddresses = resolvedAddresses;
+
     ClusterManagerConfig config = (ClusterManagerConfig)
         resolvedAddresses.getLoadBalancingPolicyConfig();
     Map<Object, ResolvedAddresses> childAddresses = new HashMap<>();
-    if (config != null) {
-      for (Map.Entry<String, Object> childPolicy : config.childPolicies.entrySet()) {
-        ResolvedAddresses addresses = resolvedAddresses.toBuilder()
-            .setLoadBalancingPolicyConfig(childPolicy.getValue())
-            .build();
-        childAddresses.put(childPolicy.getKey(), addresses);
+
+    // Reactivate children with config; deactivate children without config
+    for (ChildLbState rawState : getChildLbStates()) {
+      ClusterManagerLbState state = (ClusterManagerLbState) rawState;
+      if (config.childPolicies.containsKey(state.getKey())) {
+        // Active child
+        if (state.deletionTimer != null) {
+          state.reactivateChild();
+        }
+      } else {
+        // Inactive child
+        if (state.deletionTimer == null) {
+          state.deactivateChild();
+        }
+        if (state.deletionTimer.isPending()) {
+          childAddresses.put(state.getKey(), null); // Preserve child, without config update
+        }
       }
     }
+
+    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}", childAddresses.keySet());
+        "Received cluster_manager lb config: child names={0}", config.childPolicies.keySet());
     return childAddresses;
   }
 
-  /**
-   * This is like the parent except that it doesn't shutdown the removed children since we want that
-   * to be done by the timer.
-   */
-  @Override
-  public Status acceptResolvedAddresses(ResolvedAddresses resolvedAddresses) {
-    if (lastResolvedAddresses != null) {
-      // Handle deactivated children
-      ClusterManagerConfig config = (ClusterManagerConfig)
-          resolvedAddresses.getLoadBalancingPolicyConfig();
-      ClusterManagerConfig lastConfig = (ClusterManagerConfig)
-          lastResolvedAddresses.getLoadBalancingPolicyConfig();
-      Map<String, Object> adjChildPolicies = new HashMap<>(config.childPolicies);
-      for (Entry<String, Object> entry : lastConfig.childPolicies.entrySet()) {
-        ClusterManagerLbState state = (ClusterManagerLbState) getChildLbState(entry.getKey());
-        if (adjChildPolicies.containsKey(entry.getKey())) {
-          if (state.deletionTimer != null) {
-            state.reactivateChild();
-          }
-        } else if (state != null) {
-          adjChildPolicies.put(entry.getKey(), entry.getValue());
-          if (state.deletionTimer == null) {
-            state.deactivateChild();
-          }
-        }
-      }
-      config = new ClusterManagerConfig(adjChildPolicies);
-      resolvedAddresses =
-          resolvedAddresses.toBuilder().setLoadBalancingPolicyConfig(config).build();
-    }
-    lastResolvedAddresses = resolvedAddresses;
-    return super.acceptResolvedAddresses(resolvedAddresses);
-  }
-
   /**
    * Using the state of all children will calculate the current connectivity state,
    * update currentConnectivityState, generate a picker and then call
@@ -232,14 +218,6 @@ class DeletionTask implements Runnable {
 
         @Override
         public void run() {
-          ClusterManagerConfig config = (ClusterManagerConfig)
-              lastResolvedAddresses.getLoadBalancingPolicyConfig();
-          Map<String, Object> childPolicies = new HashMap<>(config.childPolicies);
-          Object removed = childPolicies.remove(getKey());
-          assert removed != null;
-          config = new ClusterManagerConfig(childPolicies);
-          lastResolvedAddresses =
-              lastResolvedAddresses.toBuilder().setLoadBalancingPolicyConfig(config).build();
           acceptResolvedAddresses(lastResolvedAddresses);
         }
       }