util: Remove EAG conveniences from MultiChildLb

This is a step toward removing ResolvedAddresses from ChildLbState,
which isn't actually used by MultiChildLb. Most usages of the EAG usages
can be served more directly without peering into MultiChildLb's
internals or even accessing ChildLbStates, which make the tests less
sensitive to implementation changes. Some changes do leverage the new
behavior of MultiChildLb where it preserves the order of the entries.

This does fix an important bug in shutdown tests. The tests looped over
the ChildLbStates after shutdown, but shutdown deleted all the children
so it looped over an entry collection. Fixing that exposed that
deliverSubchannelState() didn't function after shutdown, as the listener
was removed from the map when the subchannel was shut down. Moving the
listener onto the TestSubchannel allowed having access to the listener
even after shutdown.

A few places in LeastRequestLb lines were just deleted, but that's
because an existing assertion already provided the same check but
without digging into MultiChildLb.

Author: ejona86

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

@@ -24,7 +24,6 @@
 import static io.grpc.ConnectivityState.TRANSIENT_FAILURE;
 
 import com.google.common.annotations.VisibleForTesting;
-import com.google.common.base.Objects;
 import com.google.common.collect.Maps;
 import io.grpc.Attributes;
 import io.grpc.ConnectivityState;
@@ -237,21 +236,6 @@ public final Collection<ChildLbState> getChildLbStates() {
     return childLbStates;
   }
 
-  @VisibleForTesting
-  public final ChildLbState getChildLbState(Object key) {
-    for (ChildLbState state : childLbStates) {
-      if (Objects.equal(state.getKey(), key)) {
-        return state;
-      }
-    }
-    return null;
-  }
-
-  @VisibleForTesting
-  public final ChildLbState getChildLbStateEag(EquivalentAddressGroup eag) {
-    return getChildLbState(new Endpoint(eag));
-  }
-
   /**
    * Filters out non-ready child load balancers (subchannels).
    */
@@ -337,13 +321,6 @@ protected final void setCurrentPicker(SubchannelPicker newPicker) {
       currentPicker = newPicker;
     }
 
-    public final EquivalentAddressGroup getEag() {
-      if (resolvedAddresses == null || resolvedAddresses.getAddresses().isEmpty()) {
-        return null;
-      }
-      return resolvedAddresses.getAddresses().get(0);
-    }
-
     protected final void setResolvedAddresses(ResolvedAddresses newAddresses) {
       checkNotNull(newAddresses, "Missing address list for child");
       resolvedAddresses = newAddresses;

util/src/test/java/io/grpc/util/MultiChildLoadBalancerTest.java

@@ -52,7 +52,6 @@
 import java.util.List;
 import java.util.Map;
 import java.util.concurrent.ConcurrentHashMap;
-import java.util.stream.Collectors;
 import org.junit.Before;
 import org.junit.Rule;
 import org.junit.Test;
@@ -153,8 +152,7 @@ public void pickAfterResolvedUpdatedHosts() {
     LoadBalancer.Subchannel removedSubchannel = getSubchannel(removedEag);
     LoadBalancer.Subchannel oldSubchannel = getSubchannel(oldEag1);
     LoadBalancer.SubchannelStateListener removedListener =
-        testHelperInst.getSubchannelStateListeners()
-            .get(testHelperInst.getRealForMockSubChannel(removedSubchannel));
+        testHelperInst.getSubchannelStateListener(removedSubchannel);
 
     inOrder.verify(mockHelper).updateBalancingState(eq(CONNECTING), pickerCaptor.capture());
 
@@ -168,8 +166,6 @@ public void pickAfterResolvedUpdatedHosts() {
     verify(removedSubchannel, times(1)).requestConnection();
     verify(oldSubchannel, times(1)).requestConnection();
 
-    assertThat(getChildEags(loadBalancer)).containsExactly(removedEag, oldEag1);
-
     // This time with Attributes
     List<EquivalentAddressGroup> latestServers = Lists.newArrayList(oldEag2, newEag);
 
@@ -186,10 +182,10 @@ public void pickAfterResolvedUpdatedHosts() {
     removedListener.onSubchannelState(ConnectivityStateInfo.forNonError(SHUTDOWN));
     deliverSubchannelState(newSubchannel, ConnectivityStateInfo.forNonError(READY));
 
-    assertThat(getChildEags(loadBalancer)).containsExactly(oldEag2, newEag);
-
     verify(mockHelper, times(3)).createSubchannel(any(LoadBalancer.CreateSubchannelArgs.class));
     inOrder.verify(mockHelper, times(2)).updateBalancingState(eq(READY), pickerCaptor.capture());
+    picker = pickerCaptor.getValue();
+    assertThat(getList(picker)).containsExactly(oldSubchannel, newSubchannel);
 
     AbstractTestHelper.verifyNoMoreMeaningfulInteractions(mockHelper);
   }
@@ -328,12 +324,6 @@ private LoadBalancer.Subchannel getSubchannel(EquivalentAddressGroup eag) {
     return null;
   }
 
-  private static List<Object> getChildEags(MultiChildLoadBalancer loadBalancer) {
-    return loadBalancer.getChildLbStates().stream()
-        .map(ChildLbState::getEag)
-        .collect(Collectors.toList());
-  }
-
   private void deliverSubchannelState(LoadBalancer.Subchannel subchannel,
                                       ConnectivityStateInfo newState) {
     testHelperInst.deliverSubchannelState(subchannel, newState);
@@ -348,13 +338,16 @@ protected TestLb(Helper mockHelper) {
     protected void updateOverallBalancingState() {
       ConnectivityState overallState = null;
       final Map<Object, SubchannelPicker> childPickers = new HashMap<>();
+      final Map<Object, ConnectivityState> childConnStates = new HashMap<>();
       for (ChildLbState childLbState : getChildLbStates()) {
         childPickers.put(childLbState.getKey(), childLbState.getCurrentPicker());
+        childConnStates.put(childLbState.getKey(), childLbState.getCurrentState());
         overallState = aggregateState(overallState, childLbState.getCurrentState());
       }
 
       if (overallState != null) {
-        getHelper().updateBalancingState(overallState, new TestSubchannelPicker(childPickers));
+        getHelper().updateBalancingState(
+            overallState, new TestSubchannelPicker(childPickers, childConnStates));
         currentConnectivityState = overallState;
       }
 
@@ -364,18 +357,17 @@ private class TestSubchannelPicker extends SubchannelPicker {
       Map<Object, SubchannelPicker> childPickerMap;
       Map<Object, ConnectivityState> childStates = new HashMap<>();
 
-      TestSubchannelPicker(Map<Object, SubchannelPicker> childPickers) {
-        childPickerMap = childPickers;
-        for (Object key : childPickerMap.keySet()) {
-          childStates.put(key, getChildLbState(key).getCurrentState());
-        }
+      TestSubchannelPicker(
+          Map<Object, SubchannelPicker> childPickers, Map<Object, ConnectivityState> childStates) {
+        this.childPickerMap = childPickers;
+        this.childStates = childStates;
       }
 
       List<Subchannel>  getReadySubchannels() {
         List<Subchannel> readySubchannels = new ArrayList<>();
         for ( Map.Entry<Object, ConnectivityState> cur : childStates.entrySet()) {
           if (cur.getValue() == READY) {
-            Subchannel s = subchannels.get(Arrays.asList(getChildLbState(cur.getKey()).getEag()));
+            Subchannel s = childPickerMap.get(cur.getKey()).pickSubchannel(null).getSubchannel();
             readySubchannels.add(s);
           }
         }

util/src/test/java/io/grpc/util/RoundRobinLoadBalancerTest.java

@@ -61,7 +61,6 @@
 import java.util.ArrayList;
 import java.util.Arrays;
 import java.util.Collections;
-import java.util.HashMap;
 import java.util.Iterator;
 import java.util.List;
 import java.util.Map;
@@ -207,10 +206,6 @@ public void pickAfterResolvedUpdatedHosts() throws Exception {
       assertThat(childLbState.getResolvedAddresses().getAttributes().get(IS_PETIOLE_POLICY))
           .isTrue();
     }
-    assertThat(loadBalancer.getChildLbStateEag(removedEag).getCurrentPicker().pickSubchannel(null)
-        .getSubchannel()).isEqualTo(removedSubchannel);
-    assertThat(loadBalancer.getChildLbStateEag(oldEag1).getCurrentPicker().pickSubchannel(null)
-        .getSubchannel()).isEqualTo(oldSubchannel);
 
     // This time with Attributes
     List<EquivalentAddressGroup> latestServers = Lists.newArrayList(oldEag2, newEag);
@@ -225,12 +220,6 @@ public void pickAfterResolvedUpdatedHosts() throws Exception {
 
     deliverSubchannelState(newSubchannel, ConnectivityStateInfo.forNonError(READY));
 
-    assertThat(loadBalancer.getChildLbStates().size()).isEqualTo(2);
-    assertThat(loadBalancer.getChildLbStateEag(newEag).getCurrentPicker()
-        .pickSubchannel(null).getSubchannel()).isEqualTo(newSubchannel);
-    assertThat(loadBalancer.getChildLbStateEag(oldEag2).getCurrentPicker()
-        .pickSubchannel(null).getSubchannel()).isEqualTo(oldSubchannel);
-
     verify(mockHelper, times(6)).createSubchannel(any(CreateSubchannelArgs.class));
     inOrder.verify(mockHelper, times(2)).updateBalancingState(eq(READY), pickerCaptor.capture());
 
@@ -243,35 +232,35 @@ public void pickAfterResolvedUpdatedHosts() throws Exception {
   @Test
   public void pickAfterStateChange() throws Exception {
     InOrder inOrder = inOrder(mockHelper);
-    Status addressesAcceptanceStatus = acceptAddresses(servers, Attributes.EMPTY);
+    Status addressesAcceptanceStatus =
+        acceptAddresses(Arrays.asList(servers.get(0)), Attributes.EMPTY);
     assertThat(addressesAcceptanceStatus.isOk()).isTrue();
+    inOrder.verify(mockHelper).createSubchannel(any(CreateSubchannelArgs.class));
 
     // TODO figure out if this method testing the right things
 
-    ChildLbState childLbState = loadBalancer.getChildLbStates().iterator().next();
-    Subchannel subchannel = subchannels.get(Arrays.asList(childLbState.getEag()));
+    assertThat(subchannels).hasSize(1);
+    Subchannel subchannel = subchannels.values().iterator().next();
 
     inOrder.verify(mockHelper).updateBalancingState(eq(CONNECTING), eq(EMPTY_PICKER));
-    assertThat(childLbState.getCurrentState()).isEqualTo(CONNECTING);
 
     deliverSubchannelState(subchannel, ConnectivityStateInfo.forNonError(READY));
     inOrder.verify(mockHelper).updateBalancingState(eq(READY), pickerCaptor.capture());
     assertThat(pickerCaptor.getValue()).isInstanceOf(ReadyPicker.class);
-    assertThat(childLbState.getCurrentState()).isEqualTo(READY);
 
     Status error = Status.UNKNOWN.withDescription("¯\\_(ツ)_//¯");
     deliverSubchannelState(subchannel,
         ConnectivityStateInfo.forTransientFailure(error));
-    assertThat(childLbState.getCurrentState()).isEqualTo(TRANSIENT_FAILURE);
-    AbstractTestHelper.refreshInvokedAndUpdateBS(inOrder, CONNECTING, mockHelper, pickerCaptor);
-    assertThat(pickerCaptor.getValue()).isEqualTo(EMPTY_PICKER);
+    AbstractTestHelper.refreshInvokedAndUpdateBS(
+        inOrder, TRANSIENT_FAILURE, mockHelper, pickerCaptor);
+    assertThat(pickerCaptor.getValue().pickSubchannel(mockArgs).getStatus()).isEqualTo(error);
 
     deliverSubchannelState(subchannel, ConnectivityStateInfo.forNonError(IDLE));
     inOrder.verify(mockHelper).refreshNameResolution();
-    assertThat(childLbState.getCurrentState()).isEqualTo(TRANSIENT_FAILURE);
+    inOrder.verify(mockHelper, never())
+        .updateBalancingState(eq(TRANSIENT_FAILURE), any(SubchannelPicker.class));
 
     verify(subchannel, atLeastOnce()).requestConnection();
-    verify(mockHelper, times(3)).createSubchannel(any(CreateSubchannelArgs.class));
     AbstractTestHelper.verifyNoMoreMeaningfulInteractions(mockHelper);
   }
 
@@ -282,10 +271,10 @@ public void ignoreShutdownSubchannelStateChange() {
     assertThat(addressesAcceptanceStatus.isOk()).isTrue();
     inOrder.verify(mockHelper).updateBalancingState(eq(CONNECTING), eq(EMPTY_PICKER));
 
+    List<Subchannel> savedSubchannels = new ArrayList<>(subchannels.values());
     loadBalancer.shutdown();
-    for (ChildLbState child : loadBalancer.getChildLbStates()) {
-      Subchannel sc = child.getCurrentPicker().pickSubchannel(null).getSubchannel();
-      verify(child).shutdown();
+    for (Subchannel sc : savedSubchannels) {
+      verify(sc).shutdown();
       // When the subchannel is being shut down, a SHUTDOWN connectivity state is delivered
       // back to the subchannel state listener.
       deliverSubchannelState(sc, ConnectivityStateInfo.forNonError(SHUTDOWN));
@@ -300,34 +289,27 @@ public void stayTransientFailureUntilReady() {
     Status addressesAcceptanceStatus = acceptAddresses(servers, Attributes.EMPTY);
     assertThat(addressesAcceptanceStatus.isOk()).isTrue();
 
+    inOrder.verify(mockHelper, times(3)).createSubchannel(any(CreateSubchannelArgs.class));
     inOrder.verify(mockHelper).updateBalancingState(eq(CONNECTING), eq(EMPTY_PICKER));
 
-    Map<ChildLbState, Subchannel> childToSubChannelMap = new HashMap<>();
     // Simulate state transitions for each subchannel individually.
-    for ( ChildLbState child : loadBalancer.getChildLbStates()) {
-      Subchannel sc = subchannels.get(Arrays.asList(child.getEag()));
-      childToSubChannelMap.put(child, sc);
+    for (Subchannel sc : subchannels.values()) {
       Status error = Status.UNKNOWN.withDescription("connection broken");
       deliverSubchannelState(
           sc,
           ConnectivityStateInfo.forTransientFailure(error));
-      assertEquals(TRANSIENT_FAILURE, child.getCurrentState());
       deliverSubchannelState(
           sc,
           ConnectivityStateInfo.forNonError(CONNECTING));
-      assertEquals(TRANSIENT_FAILURE, child.getCurrentState());
     }
     inOrder.verify(mockHelper).updateBalancingState(eq(TRANSIENT_FAILURE), isA(ReadyPicker.class));
     inOrder.verify(mockHelper, atLeast(0)).refreshNameResolution();
     inOrder.verifyNoMoreInteractions();
 
-    ChildLbState child = loadBalancer.getChildLbStates().iterator().next();
-    Subchannel subchannel = childToSubChannelMap.get(child);
+    Subchannel subchannel = subchannels.values().iterator().next();
     deliverSubchannelState(subchannel, ConnectivityStateInfo.forNonError(READY));
-    assertThat(child.getCurrentState()).isEqualTo(READY);
     inOrder.verify(mockHelper).updateBalancingState(eq(READY), isA(ReadyPicker.class));
 
-    verify(mockHelper, times(3)).createSubchannel(any(CreateSubchannelArgs.class));
     inOrder.verify(mockHelper, atLeast(0)).refreshNameResolution();
     inOrder.verifyNoMoreInteractions();
   }
@@ -342,8 +324,7 @@ public void refreshNameResolutionWhenSubchannelConnectionBroken() {
     inOrder.verify(mockHelper).updateBalancingState(eq(CONNECTING), eq(EMPTY_PICKER));
 
     // Simulate state transitions for each subchannel individually.
-    for (ChildLbState child : loadBalancer.getChildLbStates()) {
-      Subchannel sc = subchannels.get(Arrays.asList(child.getEag()));
+    for (Subchannel sc : subchannels.values()) {
       verify(sc).requestConnection();
       deliverSubchannelState(sc, ConnectivityStateInfo.forNonError(CONNECTING));
       Status error = Status.UNKNOWN.withDescription("connection broken");

util/src/testFixtures/java/io/grpc/util/AbstractTestHelper.java

@@ -16,14 +16,14 @@
 
 package io.grpc.util;
 
+import static com.google.common.base.Preconditions.checkNotNull;
 import static org.mockito.AdditionalAnswers.delegatesTo;
 import static org.mockito.ArgumentMatchers.eq;
 import static org.mockito.Mockito.atLeast;
 import static org.mockito.Mockito.mock;
 import static org.mockito.Mockito.verify;
 import static org.mockito.Mockito.verifyNoMoreInteractions;
 
-import com.google.common.collect.Maps;
 import io.grpc.Attributes;
 import io.grpc.Channel;
 import io.grpc.ChannelLogger;
@@ -63,10 +63,8 @@
  */
 public abstract class AbstractTestHelper extends ForwardingLoadBalancerHelper {
 
-  private final Map<Subchannel, Subchannel> mockToRealSubChannelMap = new HashMap<>();
+  private final Map<Subchannel, TestSubchannel> mockToRealSubChannelMap = new HashMap<>();
   protected final Map<Subchannel, Subchannel> realToMockSubChannelMap = new HashMap<>();
-  private final Map<Subchannel, SubchannelStateListener> subchannelStateListeners =
-      Maps.newLinkedHashMap();
   private final FakeClock fakeClock;
   private final SynchronizationContext syncContext;
 
@@ -87,22 +85,14 @@ public AbstractTestHelper(FakeClock fakeClock, SynchronizationContext syncContex
     this.syncContext = syncContext;
   }
 
-  public Map<Subchannel, Subchannel> getMockToRealSubChannelMap() {
-    return mockToRealSubChannelMap;
-  }
-
-  public Subchannel getRealForMockSubChannel(Subchannel mock) {
-    Subchannel realSc = getMockToRealSubChannelMap().get(mock);
+  private TestSubchannel getRealForMockSubChannel(Subchannel mock) {
+    TestSubchannel realSc = mockToRealSubChannelMap.get(mock);
     if (realSc == null) {
-      realSc = mock;
+      realSc = (TestSubchannel) mock;
     }
     return realSc;
   }
 
-  public Map<Subchannel, SubchannelStateListener> getSubchannelStateListeners() {
-    return subchannelStateListeners;
-  }
-
   public static final FakeClock.TaskFilter NOT_START_NEXT_CONNECTION =
       new FakeClock.TaskFilter() {
         @Override
@@ -116,15 +106,15 @@ public static int getNumFilteredPendingTasks(FakeClock fakeClock) {
   }
 
   public void deliverSubchannelState(Subchannel subchannel, ConnectivityStateInfo newState) {
-    Subchannel realSc = getMockToRealSubChannelMap().get(subchannel);
-    if (realSc == null) {
-      realSc = subchannel;
-    }
-    SubchannelStateListener listener = getSubchannelStateListeners().get(realSc);
+    getSubchannelStateListener(subchannel).onSubchannelState(newState);
+  }
+
+  public SubchannelStateListener getSubchannelStateListener(Subchannel subchannel) {
+    SubchannelStateListener listener = getRealForMockSubChannel(subchannel).listener;
     if (listener == null) {
-      throw new IllegalArgumentException("subchannel does not have a matching listener");
+      throw new IllegalArgumentException("subchannel has not been started");
     }
-    listener.onSubchannelState(newState);
+    return listener;
   }
 
   @Override
@@ -144,7 +134,7 @@ public Subchannel createSubchannel(CreateSubchannelArgs args) {
       TestSubchannel delegate = createRealSubchannel(args);
       subchannel = mock(Subchannel.class, delegatesTo(delegate));
       getSubchannelMap().put(args.getAddresses(), subchannel);
-      getMockToRealSubChannelMap().put(subchannel, delegate);
+      mockToRealSubChannelMap.put(subchannel, delegate);
       realToMockSubChannelMap.put(delegate, subchannel);
     }
 
@@ -161,7 +151,7 @@ public void refreshNameResolution() {
   }
 
   public void setChannel(Subchannel subchannel, Channel channel) {
-    ((TestSubchannel)subchannel).channel = channel;
+    getRealForMockSubChannel(subchannel).channel = channel;
   }
 
   @Override
@@ -208,6 +198,7 @@ public static void verifyNoMoreMeaningfulInteractions(Helper helper, InOrder inO
 
   protected class TestSubchannel extends ForwardingSubchannel {
     CreateSubchannelArgs args;
+    SubchannelStateListener listener;
     Channel channel;
 
     public TestSubchannel(CreateSubchannelArgs args) {
@@ -250,12 +241,11 @@ public void updateAddresses(List<EquivalentAddressGroup> addrs) {
 
     @Override
     public void start(SubchannelStateListener listener) {
-      getSubchannelStateListeners().put(this, listener);
+      this.listener = checkNotNull(listener, "listener");
     }
 
     @Override
     public void shutdown() {
-      getSubchannelStateListeners().remove(this);
       for (EquivalentAddressGroup eag : getAllAddresses()) {
         getSubchannelMap().remove(Collections.singletonList(eag));
       }