From 9ab35a761bf4bc9d39ade5fd15b1e29fb31f061c Mon Sep 17 00:00:00 2001
From: Eric Anderson <ejona@google.com>
Date: Wed, 2 Oct 2024 11:03:44 -0700
Subject: [PATCH] util: Store only a list of children in MultiChildLB

A map of children is still needed, but is created temporarily on update.
The order of children is currently preserved, but we could use regular
HashMaps if that is not useful.
---
 .../io/grpc/util/MultiChildLoadBalancer.java  | 61 ++++++++++---------
 1 file changed, 32 insertions(+), 29 deletions(-)

diff --git a/util/src/main/java/io/grpc/util/MultiChildLoadBalancer.java b/util/src/main/java/io/grpc/util/MultiChildLoadBalancer.java
index 330ec9d5357..3e56f41d038 100644
--- a/util/src/main/java/io/grpc/util/MultiChildLoadBalancer.java
+++ b/util/src/main/java/io/grpc/util/MultiChildLoadBalancer.java
@@ -24,7 +24,8 @@
 import static io.grpc.ConnectivityState.TRANSIENT_FAILURE;
 
 import com.google.common.annotations.VisibleForTesting;
-import com.google.common.collect.ImmutableList;
+import com.google.common.base.Objects;
+import com.google.common.collect.Maps;
 import io.grpc.Attributes;
 import io.grpc.ConnectivityState;
 import io.grpc.EquivalentAddressGroup;
@@ -37,12 +38,9 @@
 import java.util.ArrayList;
 import java.util.Collection;
 import java.util.Collections;
-import java.util.HashMap;
 import java.util.HashSet;
-import java.util.LinkedHashMap;
 import java.util.List;
 import java.util.Map;
-import java.util.Set;
 import java.util.logging.Level;
 import java.util.logging.Logger;
 import javax.annotation.Nullable;
@@ -55,7 +53,8 @@
 public abstract class MultiChildLoadBalancer extends LoadBalancer {
 
   private static final Logger logger = Logger.getLogger(MultiChildLoadBalancer.class.getName());
-  private final Map<Object, ChildLbState> childLbStates = new LinkedHashMap<>();
+  // Modify by replacing the list to release memory when no longer used.
+  private List<ChildLbState> childLbStates = new ArrayList<>(0);
   private final Helper helper;
   // Set to true if currently in the process of handling resolved addresses.
   protected boolean resolvingAddresses;
@@ -84,7 +83,8 @@ protected MultiChildLoadBalancer(Helper helper) {
    */
   protected Map<Object, ResolvedAddresses> createChildAddressesMap(
       ResolvedAddresses resolvedAddresses) {
-    Map<Object, ResolvedAddresses> childAddresses = new HashMap<>();
+    Map<Object, ResolvedAddresses> childAddresses =
+        Maps.newLinkedHashMapWithExpectedSize(resolvedAddresses.getAddresses().size());
     for (EquivalentAddressGroup eag : resolvedAddresses.getAddresses()) {
       ResolvedAddresses addresses = resolvedAddresses.toBuilder()
           .setAddresses(Collections.singletonList(eag))
@@ -144,7 +144,7 @@ public void handleNameResolutionError(Status error) {
   @Override
   public void shutdown() {
     logger.log(Level.FINE, "Shutdown");
-    for (ChildLbState state : childLbStates.values()) {
+    for (ChildLbState state : childLbStates) {
       state.shutdown();
     }
     childLbStates.clear();
@@ -169,39 +169,37 @@ protected final AcceptResolvedAddrRetVal acceptResolvedAddressesInternal(
       return new AcceptResolvedAddrRetVal(unavailableStatus, null);
     }
 
-    updateChildrenWithResolvedAddresses(newChildAddresses);
-
-    return new AcceptResolvedAddrRetVal(Status.OK, getRemovedChildren(newChildAddresses.keySet()));
+    List<ChildLbState> removed = updateChildrenWithResolvedAddresses(newChildAddresses);
+    return new AcceptResolvedAddrRetVal(Status.OK, removed);
   }
 
-  private void updateChildrenWithResolvedAddresses(
+  /** Returns removed children. */
+  private List<ChildLbState> updateChildrenWithResolvedAddresses(
       Map<Object, ResolvedAddresses> newChildAddresses) {
+    // Create a map with the old values
+    Map<Object, ChildLbState> oldStatesMap =
+        Maps.newLinkedHashMapWithExpectedSize(childLbStates.size());
+    for (ChildLbState state : childLbStates) {
+      oldStatesMap.put(state.getKey(), state);
+    }
+
+    // Move ChildLbStates from the map to a new list (preserving the new map's order)
+    List<ChildLbState> newChildLbStates = new ArrayList<>(newChildAddresses.size());
     for (Map.Entry<Object, ResolvedAddresses> entry : newChildAddresses.entrySet()) {
-      ChildLbState childLbState = childLbStates.get(entry.getKey());
+      ChildLbState childLbState = oldStatesMap.remove(entry.getKey());
       if (childLbState == null) {
         childLbState = createChildLbState(entry.getKey());
-        childLbStates.put(entry.getKey(), childLbState);
       }
+      newChildLbStates.add(childLbState);
       if (entry.getValue() != null) {
         childLbState.setResolvedAddresses(entry.getValue()); // update child
         childLbState.lb.handleResolvedAddresses(entry.getValue()); // update child LB
       }
     }
-  }
 
-  /**
-   * Identifies which children have been removed (are not part of the newChildKeys).
-   */
-  private List<ChildLbState> getRemovedChildren(Set<Object> newChildKeys) {
-    List<ChildLbState> removedChildren = new ArrayList<>();
-    // Do removals
-    for (Object key : ImmutableList.copyOf(childLbStates.keySet())) {
-      if (!newChildKeys.contains(key)) {
-        ChildLbState childLbState = childLbStates.remove(key);
-        removedChildren.add(childLbState);
-      }
-    }
-    return removedChildren;
+    childLbStates = newChildLbStates;
+    // Remaining entries in map are orphaned
+    return new ArrayList<>(oldStatesMap.values());
   }
 
   protected final void shutdownRemoved(List<ChildLbState> removedChildren) {
@@ -236,12 +234,17 @@ protected final Helper getHelper() {
 
   @VisibleForTesting
   public final Collection<ChildLbState> getChildLbStates() {
-    return childLbStates.values();
+    return childLbStates;
   }
 
   @VisibleForTesting
   public final ChildLbState getChildLbState(Object key) {
-    return childLbStates.get(key);
+    for (ChildLbState state : childLbStates) {
+      if (Objects.equal(state.getKey(), key)) {
+        return state;
+      }
+    }
+    return null;
   }
 
   @VisibleForTesting