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.

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