util: Remove resolvedAddresses from MultiChildLb.ChildLbState

It isn't actually used by MultiChildLb, and using the health API gives
us more confidence that health is properly plumbed.

Author: ejona86

core/src/main/java/io/grpc/internal/PickFirstLoadBalancerProvider.java

@@ -36,7 +36,7 @@ public final class PickFirstLoadBalancerProvider extends LoadBalancerProvider {
   public static final String GRPC_PF_USE_HAPPY_EYEBALLS = "GRPC_PF_USE_HAPPY_EYEBALLS";
   private static final String SHUFFLE_ADDRESS_LIST_KEY = "shuffleAddressList";
 
-  private static boolean enableNewPickFirst =
+  static boolean enableNewPickFirst =
       GrpcUtil.getFlag("GRPC_EXPERIMENTAL_ENABLE_NEW_PICK_FIRST", false);
 
   public static boolean isEnabledHappyEyeballs() {

core/src/testFixtures/java/io/grpc/internal/PickFirstLoadBalancerProviderAccessor.java

@@ -0,0 +1,28 @@
+/*
+ * Copyright 2024 The gRPC Authors
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *     http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package io.grpc.internal;
+
+/**
+ * Accessor for PickFirstLoadBalancerProvider, allowing access only during tests.
+ */
+public final class PickFirstLoadBalancerProviderAccessor {
+  private PickFirstLoadBalancerProviderAccessor() {}
+
+  public static void setEnableNewPickFirst(boolean enableNewPickFirst) {
+    PickFirstLoadBalancerProvider.enableNewPickFirst = enableNewPickFirst;
+  }
+}

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

@@ -191,7 +191,6 @@ private List<ChildLbState> updateChildrenWithResolvedAddresses(
       }
       newChildLbStates.add(childLbState);
       if (entry.getValue() != null) {
-        childLbState.setResolvedAddresses(entry.getValue()); // update child
         childLbState.lb.handleResolvedAddresses(entry.getValue()); // update child LB
       }
     }
@@ -262,8 +261,6 @@ protected final List<ChildLbState> getReadyChildren() {
    */
   public class ChildLbState {
     private final Object key;
-    private ResolvedAddresses resolvedAddresses;
-
     private final LoadBalancer lb;
     private ConnectivityState currentState;
     private SubchannelPicker currentPicker = new FixedResultPicker(PickResult.withNoResult());
@@ -321,16 +318,6 @@ protected final void setCurrentPicker(SubchannelPicker newPicker) {
       currentPicker = newPicker;
     }
 
-    protected final void setResolvedAddresses(ResolvedAddresses newAddresses) {
-      checkNotNull(newAddresses, "Missing address list for child");
-      resolvedAddresses = newAddresses;
-    }
-
-    @VisibleForTesting
-    public final ResolvedAddresses getResolvedAddresses() {
-      return resolvedAddresses;
-    }
-
     /**
      * ChildLbStateHelper is the glue between ChildLbState and the helpers associated with the
      * petiole policy above and the PickFirstLoadBalancer's helper below.

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

@@ -22,7 +22,6 @@
 import static io.grpc.ConnectivityState.READY;
 import static io.grpc.ConnectivityState.SHUTDOWN;
 import static io.grpc.ConnectivityState.TRANSIENT_FAILURE;
-import static io.grpc.util.MultiChildLoadBalancer.IS_PETIOLE_POLICY;
 import static org.junit.Assert.assertEquals;
 import static org.junit.Assert.assertNull;
 import static org.junit.Assert.fail;
@@ -54,13 +53,15 @@
 import io.grpc.LoadBalancer.Subchannel;
 import io.grpc.LoadBalancer.SubchannelPicker;
 import io.grpc.Status;
+import io.grpc.internal.PickFirstLoadBalancerProvider;
+import io.grpc.internal.PickFirstLoadBalancerProviderAccessor;
 import io.grpc.internal.TestUtils;
-import io.grpc.util.MultiChildLoadBalancer.ChildLbState;
 import io.grpc.util.RoundRobinLoadBalancer.ReadyPicker;
 import java.net.SocketAddress;
 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;
@@ -104,6 +105,7 @@ public class RoundRobinLoadBalancerTest {
   private ArgumentCaptor<CreateSubchannelArgs> createArgsCaptor;
   private TestHelper testHelperInst = new TestHelper();
   private Helper mockHelper = mock(Helper.class, delegatesTo(testHelperInst));
+  private boolean defaultNewPickFirst = PickFirstLoadBalancerProvider.isEnabledNewPickFirst();
 
   @Mock // This LoadBalancer doesn't use any of the arg fields, as verified in tearDown().
   private PickSubchannelArgs mockArgs;
@@ -126,6 +128,7 @@ private Status acceptAddresses(List<EquivalentAddressGroup> eagList, Attributes
 
   @After
   public void tearDown() throws Exception {
+    PickFirstLoadBalancerProviderAccessor.setEnableNewPickFirst(defaultNewPickFirst);
     verifyNoMoreInteractions(mockArgs);
   }
 
@@ -201,12 +204,6 @@ public void pickAfterResolvedUpdatedHosts() throws Exception {
     verify(removedSubchannel, times(1)).requestConnection();
     verify(oldSubchannel, times(1)).requestConnection();
 
-    assertThat(loadBalancer.getChildLbStates().size()).isEqualTo(2);
-    for (ChildLbState childLbState : loadBalancer.getChildLbStates()) {
-      assertThat(childLbState.getResolvedAddresses().getAttributes().get(IS_PETIOLE_POLICY))
-          .isTrue();
-    }
-
     // This time with Attributes
     List<EquivalentAddressGroup> latestServers = Lists.newArrayList(oldEag2, newEag);
 
@@ -464,6 +461,60 @@ public void subchannelStateIsolation() throws Exception {
     assertThat(pickers.hasNext()).isFalse();
   }
 
+  @Test
+  public void subchannelHealthObserved() throws Exception {
+    // Only the new PF policy observes the new separate listener for health
+    PickFirstLoadBalancerProviderAccessor.setEnableNewPickFirst(true);
+    // PickFirst does most of this work. If the test fails, check IS_PETIOLE_POLICY
+    Map<Subchannel, LoadBalancer.SubchannelStateListener> healthListeners = new HashMap<>();
+    loadBalancer = new RoundRobinLoadBalancer(new ForwardingLoadBalancerHelper() {
+      @Override
+      public Subchannel createSubchannel(CreateSubchannelArgs args) {
+        Subchannel subchannel = super.createSubchannel(args.toBuilder()
+            .setAttributes(args.getAttributes().toBuilder()
+              .set(LoadBalancer.HAS_HEALTH_PRODUCER_LISTENER_KEY, true)
+              .build())
+            .build());
+        healthListeners.put(
+            subchannel, args.getOption(LoadBalancer.HEALTH_CONSUMER_LISTENER_ARG_KEY));
+        return subchannel;
+      }
+
+      @Override
+      protected Helper delegate() {
+        return mockHelper;
+      }
+    });
+
+    InOrder inOrder = inOrder(mockHelper);
+    Status addressesAcceptanceStatus = acceptAddresses(servers, Attributes.EMPTY);
+    assertThat(addressesAcceptanceStatus.isOk()).isTrue();
+    Subchannel subchannel0 = subchannels.get(Arrays.asList(servers.get(0)));
+    Subchannel subchannel1 = subchannels.get(Arrays.asList(servers.get(1)));
+    Subchannel subchannel2 = subchannels.get(Arrays.asList(servers.get(2)));
+
+    // Subchannels go READY, but the LB waits for health
+    for (Subchannel subchannel : subchannels.values()) {
+      deliverSubchannelState(subchannel, ConnectivityStateInfo.forNonError(READY));
+    }
+    inOrder.verify(mockHelper, times(0))
+        .updateBalancingState(eq(READY), any(SubchannelPicker.class));
+
+    // Health results lets subchannels go READY
+    healthListeners.get(subchannel0).onSubchannelState(
+        ConnectivityStateInfo.forTransientFailure(Status.UNAVAILABLE.withDescription("oh no")));
+    healthListeners.get(subchannel1).onSubchannelState(ConnectivityStateInfo.forNonError(READY));
+    healthListeners.get(subchannel2).onSubchannelState(ConnectivityStateInfo.forNonError(READY));
+    inOrder.verify(mockHelper, times(2)).updateBalancingState(eq(READY), pickerCaptor.capture());
+    SubchannelPicker picker = pickerCaptor.getValue();
+    List<Subchannel> picks = Arrays.asList(
+        picker.pickSubchannel(mockArgs).getSubchannel(),
+        picker.pickSubchannel(mockArgs).getSubchannel(),
+        picker.pickSubchannel(mockArgs).getSubchannel(),
+        picker.pickSubchannel(mockArgs).getSubchannel());
+    assertThat(picks).containsExactly(subchannel1, subchannel2, subchannel1, subchannel2);
+  }
+
   @Test
   public void readyPicker_emptyList() {
     // ready picker list must be non-empty

xds/src/test/java/io/grpc/xds/RingHashLoadBalancerTest.java

@@ -23,7 +23,6 @@
 import static io.grpc.ConnectivityState.READY;
 import static io.grpc.ConnectivityState.SHUTDOWN;
 import static io.grpc.ConnectivityState.TRANSIENT_FAILURE;
-import static io.grpc.util.MultiChildLoadBalancer.IS_PETIOLE_POLICY;
 import static io.grpc.xds.RingHashLoadBalancerTest.InitializationFlags.DO_NOT_RESET_HELPER;
 import static io.grpc.xds.RingHashLoadBalancerTest.InitializationFlags.DO_NOT_VERIFY;
 import static io.grpc.xds.RingHashLoadBalancerTest.InitializationFlags.RESET_SUBCHANNEL_MOCKS;
@@ -48,6 +47,7 @@
 import io.grpc.ConnectivityState;
 import io.grpc.ConnectivityStateInfo;
 import io.grpc.EquivalentAddressGroup;
+import io.grpc.LoadBalancer;
 import io.grpc.LoadBalancer.CreateSubchannelArgs;
 import io.grpc.LoadBalancer.Helper;
 import io.grpc.LoadBalancer.PickDetailsConsumer;
@@ -62,9 +62,11 @@
 import io.grpc.SynchronizationContext;
 import io.grpc.internal.FakeClock;
 import io.grpc.internal.PickFirstLoadBalancerProvider;
+import io.grpc.internal.PickFirstLoadBalancerProviderAccessor;
 import io.grpc.internal.PickSubchannelArgsImpl;
 import io.grpc.testing.TestMethodDescriptors;
 import io.grpc.util.AbstractTestHelper;
+import io.grpc.util.ForwardingLoadBalancerHelper;
 import io.grpc.util.MultiChildLoadBalancer.ChildLbState;
 import io.grpc.xds.RingHashLoadBalancer.RingHashConfig;
 import java.lang.Thread.UncaughtExceptionHandler;
@@ -74,8 +76,11 @@
 import java.util.Collections;
 import java.util.Deque;
 import java.util.HashMap;
+import java.util.HashSet;
 import java.util.List;
 import java.util.Map;
+import java.util.Random;
+import java.util.Set;
 import org.junit.After;
 import org.junit.Before;
 import org.junit.Rule;
@@ -115,6 +120,7 @@ public void uncaughtException(Thread t, Throwable e) {
   @Captor
   private ArgumentCaptor<SubchannelPicker> pickerCaptor;
   private RingHashLoadBalancer loadBalancer;
+  private boolean defaultNewPickFirst = PickFirstLoadBalancerProvider.isEnabledNewPickFirst();
 
   @Before
   public void setUp() {
@@ -126,6 +132,7 @@ public void setUp() {
 
   @After
   public void tearDown() {
+    PickFirstLoadBalancerProviderAccessor.setEnableNewPickFirst(defaultNewPickFirst);
     loadBalancer.shutdown();
     for (Subchannel subchannel : subchannels.values()) {
       verify(subchannel).shutdown();
@@ -906,6 +913,70 @@ public void duplicateAddresses() {
     assertThat(description).contains("Address: FakeSocketAddress-server2, count: 3");
   }
 
+  @Test
+  public void subchannelHealthObserved() throws Exception {
+    // Only the new PF policy observes the new separate listener for health
+    PickFirstLoadBalancerProviderAccessor.setEnableNewPickFirst(true);
+    // PickFirst does most of this work. If the test fails, check IS_PETIOLE_POLICY
+    Map<Subchannel, LoadBalancer.SubchannelStateListener> healthListeners = new HashMap<>();
+    loadBalancer = new RingHashLoadBalancer(new ForwardingLoadBalancerHelper() {
+      @Override
+      public Subchannel createSubchannel(CreateSubchannelArgs args) {
+        Subchannel subchannel = super.createSubchannel(args.toBuilder()
+            .setAttributes(args.getAttributes().toBuilder()
+              .set(LoadBalancer.HAS_HEALTH_PRODUCER_LISTENER_KEY, true)
+              .build())
+            .build());
+        healthListeners.put(
+            subchannel, args.getOption(LoadBalancer.HEALTH_CONSUMER_LISTENER_ARG_KEY));
+        return subchannel;
+      }
+
+      @Override
+      protected Helper delegate() {
+        return helper;
+      }
+    });
+
+    InOrder inOrder = Mockito.inOrder(helper);
+    List<EquivalentAddressGroup> servers = createWeightedServerAddrs(1, 1);
+    initializeLbSubchannels(new RingHashConfig(10, 100), servers);
+    Subchannel subchannel0 = subchannels.get(Collections.singletonList(servers.get(0)));
+    Subchannel subchannel1 = subchannels.get(Collections.singletonList(servers.get(1)));
+
+    // Subchannels go READY, but the LB waits for health
+    for (Subchannel subchannel : subchannels.values()) {
+      deliverSubchannelState(subchannel, ConnectivityStateInfo.forNonError(READY));
+    }
+    inOrder.verify(helper, times(0)).updateBalancingState(eq(READY), any(SubchannelPicker.class));
+
+    // Health results lets subchannels go READY
+    healthListeners.get(subchannel0).onSubchannelState(ConnectivityStateInfo.forNonError(READY));
+    healthListeners.get(subchannel1).onSubchannelState(ConnectivityStateInfo.forNonError(READY));
+    inOrder.verify(helper, times(2)).updateBalancingState(eq(READY), pickerCaptor.capture());
+    SubchannelPicker picker = pickerCaptor.getValue();
+    Random random = new Random(1);
+    Set<Subchannel> picks = new HashSet<>();
+    for (int i = 0; i < 10; i++) {
+      picks.add(
+          picker.pickSubchannel(getDefaultPickSubchannelArgs(random.nextLong())).getSubchannel());
+    }
+    assertThat(picks).containsExactly(subchannel0, subchannel1);
+
+    // Unhealthy subchannel skipped
+    healthListeners.get(subchannel0).onSubchannelState(
+        ConnectivityStateInfo.forTransientFailure(Status.UNAVAILABLE.withDescription("oh no")));
+    inOrder.verify(helper).updateBalancingState(eq(READY), pickerCaptor.capture());
+    picker = pickerCaptor.getValue();
+    random.setSeed(1);
+    picks.clear();
+    for (int i = 0; i < 10; i++) {
+      picks.add(
+          picker.pickSubchannel(getDefaultPickSubchannelArgs(random.nextLong())).getSubchannel());
+    }
+    assertThat(picks).containsExactly(subchannel1);
+  }
+
   private List<Subchannel> initializeLbSubchannels(RingHashConfig config,
       List<EquivalentAddressGroup> servers, InitializationFlags... initFlags) {
 
@@ -950,8 +1021,6 @@ private List<Subchannel> initializeLbSubchannels(RingHashConfig config,
     for (ChildLbState childLbState : loadBalancer.getChildLbStates()) {
       childLbState.getCurrentPicker()
           .pickSubchannel(getDefaultPickSubchannelArgs(hashFunc.hashVoid()));
-      assertThat(childLbState.getResolvedAddresses().getAttributes().get(IS_PETIOLE_POLICY))
-          .isTrue();
     }
 
     if (doVerifies) {