xds: Stop extending RR in WRR

ejona86 · ejona86 · commit dc83446d982b · 2024-07-31T13:32:49.000-07:00
They share very little code, and we really don't want RoundRobinLb to be public and non-final. Originally, WRR was expected to share much more code with RR, and even delegated to RR at times. The delegation was removed in 111ff60. After dca89b2, most of the sharing has been moved out into general-purpose tools that can be used by any LB policy. FixedResultPicker now has equals to makes it as a EmptyPicker replacement. RoundRobinLb still uses EmptyPicker because fixing its tests is a larger change. OutlierDetectionLbTest was changed because FixedResultPicker is used by PickFirstLeafLb, and now RoundRobinLb can squelch some of its updates for ready pickers.
diff --git a/api/src/main/java/io/grpc/LoadBalancer.java b/api/src/main/java/io/grpc/LoadBalancer.java
@@ -1526,5 +1526,19 @@ public PickResult pickSubchannel(PickSubchannelArgs args) {
     public String toString() {
       return "FixedResultPicker(" + result + ")";
     }
+
+    @Override
+    public int hashCode() {
+      return result.hashCode();
+    }
+
+    @Override
+    public boolean equals(Object o) {
+      if (!(o instanceof FixedResultPicker)) {
+        return false;
+      }
+      FixedResultPicker that = (FixedResultPicker) o;
+      return this.result.equals(that.result);
+    }
   }
 }
diff --git a/util/src/main/java/io/grpc/util/RoundRobinLoadBalancer.java b/util/src/main/java/io/grpc/util/RoundRobinLoadBalancer.java
@@ -27,7 +27,6 @@
 import com.google.common.base.Preconditions;
 import io.grpc.ConnectivityState;
 import io.grpc.EquivalentAddressGroup;
-import io.grpc.Internal;
 import io.grpc.LoadBalancer;
 import io.grpc.NameResolver;
 import java.util.ArrayList;
@@ -41,10 +40,9 @@
  * A {@link LoadBalancer} that provides round-robin load-balancing over the {@link
  * EquivalentAddressGroup}s from the {@link NameResolver}.
  */
-@Internal
-public class RoundRobinLoadBalancer extends MultiChildLoadBalancer {
+final class RoundRobinLoadBalancer extends MultiChildLoadBalancer {
   private final AtomicInteger sequence = new AtomicInteger(new Random().nextInt());
-  protected SubchannelPicker currentPicker = new EmptyPicker();
+  private SubchannelPicker currentPicker = new EmptyPicker();
 
   public RoundRobinLoadBalancer(Helper helper) {
     super(helper);
@@ -87,7 +85,7 @@ private void updateBalancingState(ConnectivityState state, SubchannelPicker pick
     }
   }
 
-  protected SubchannelPicker createReadyPicker(Collection<ChildLbState> children) {
+  private SubchannelPicker createReadyPicker(Collection<ChildLbState> children) {
     List<SubchannelPicker> pickerList = new ArrayList<>();
     for (ChildLbState child : children) {
       SubchannelPicker picker = child.getCurrentPicker();
diff --git a/util/src/test/java/io/grpc/util/OutlierDetectionLoadBalancerTest.java b/util/src/test/java/io/grpc/util/OutlierDetectionLoadBalancerTest.java
@@ -569,7 +569,7 @@ public void successRateOneOutlier_configChange() {
     loadBalancer.acceptResolvedAddresses(buildResolvedAddress(config, servers));
 
     // The PickFirstLeafLB has an extra level of indirection because of health
-    int expectedStateChanges = PickFirstLoadBalancerProvider.isEnabledNewPickFirst() ? 16 : 12;
+    int expectedStateChanges = PickFirstLoadBalancerProvider.isEnabledNewPickFirst() ? 8 : 12;
     generateLoad(ImmutableMap.of(subchannel2, Status.DEADLINE_EXCEEDED), expectedStateChanges);
 
     // Move forward in time to a point where the detection timer has fired.
@@ -604,7 +604,7 @@ public void successRateOneOutlier_unejected() {
     assertEjectedSubchannels(ImmutableSet.of(ImmutableSet.copyOf(servers.get(0).getAddresses())));
 
     // Now we produce more load, but the subchannel has started working and is no longer an outlier.
-    int expectedStateChanges = PickFirstLoadBalancerProvider.isEnabledNewPickFirst() ? 16 : 12;
+    int expectedStateChanges = PickFirstLoadBalancerProvider.isEnabledNewPickFirst() ? 8 : 12;
     generateLoad(ImmutableMap.of(), expectedStateChanges);
 
     // Move forward in time to a point where the detection timer has fired.
diff --git a/xds/src/main/java/io/grpc/xds/WeightedRoundRobinLoadBalancer.java b/xds/src/main/java/io/grpc/xds/WeightedRoundRobinLoadBalancer.java
@@ -42,7 +42,7 @@
 import io.grpc.services.MetricReport;
 import io.grpc.util.ForwardingLoadBalancerHelper;
 import io.grpc.util.ForwardingSubchannel;
-import io.grpc.util.RoundRobinLoadBalancer;
+import io.grpc.util.MultiChildLoadBalancer;
 import io.grpc.xds.orca.OrcaOobUtil;
 import io.grpc.xds.orca.OrcaOobUtil.OrcaOobReportListener;
 import io.grpc.xds.orca.OrcaPerRequestUtil;
@@ -90,7 +90,7 @@
  *  See related documentation: https://cloud.google.com/service-mesh/legacy/load-balancing-apis/proxyless-configure-advanced-traffic-management#custom-lb-config
  */
 @ExperimentalApi("https://github.com/grpc/grpc-java/issues/9885")
-final class WeightedRoundRobinLoadBalancer extends RoundRobinLoadBalancer {
+final class WeightedRoundRobinLoadBalancer extends MultiChildLoadBalancer {
 
   private static final LongCounterMetricInstrument RR_FALLBACK_COUNTER;
   private static final LongCounterMetricInstrument ENDPOINT_WEIGHT_NOT_YET_USEABLE_COUNTER;
@@ -107,6 +107,7 @@ final class WeightedRoundRobinLoadBalancer extends RoundRobinLoadBalancer {
   private final long infTime;
   private final Ticker ticker;
   private String locality = "";
+  private SubchannelPicker currentPicker = new FixedResultPicker(PickResult.withNoResult());
 
   // The metric instruments are only registered once and shared by all instances of this LB.
   static {
@@ -209,13 +210,51 @@ public Status acceptResolvedAddresses(ResolvedAddresses resolvedAddresses) {
     return acceptRetVal.status;
   }
 
+  /**
+   * Updates picker with the list of active subchannels (state == READY).
+   */
   @Override
-  public SubchannelPicker createReadyPicker(Collection<ChildLbState> activeList) {
+  protected void updateOverallBalancingState() {
+    List<ChildLbState> activeList = getReadyChildren();
+    if (activeList.isEmpty()) {
+      // No READY subchannels
+
+      // MultiChildLB will request connection immediately on subchannel IDLE.
+      boolean isConnecting = false;
+      for (ChildLbState childLbState : getChildLbStates()) {
+        ConnectivityState state = childLbState.getCurrentState();
+        if (state == ConnectivityState.CONNECTING || state == ConnectivityState.IDLE) {
+          isConnecting = true;
+          break;
+        }
+      }
+
+      if (isConnecting) {
+        updateBalancingState(
+            ConnectivityState.CONNECTING, new FixedResultPicker(PickResult.withNoResult()));
+      } else {
+        updateBalancingState(
+            ConnectivityState.TRANSIENT_FAILURE, createReadyPicker(getChildLbStates()));
+      }
+    } else {
+      updateBalancingState(ConnectivityState.READY, createReadyPicker(activeList));
+    }
+  }
+
+  private SubchannelPicker createReadyPicker(Collection<ChildLbState> activeList) {
     return new WeightedRoundRobinPicker(ImmutableList.copyOf(activeList),
         config.enableOobLoadReport, config.errorUtilizationPenalty, sequence, getHelper(),
         locality);
   }
 
+  private void updateBalancingState(ConnectivityState state, SubchannelPicker picker) {
+    if (state != currentConnectivityState || !picker.equals(currentPicker)) {
+      getHelper().updateBalancingState(state, picker);
+      currentConnectivityState = state;
+      currentPicker = picker;
+    }
+  }
+
   @VisibleForTesting
   final class WeightedChildLbState extends ChildLbState {
 
diff --git a/xds/src/test/java/io/grpc/xds/WeightedRoundRobinLoadBalancerTest.java b/xds/src/test/java/io/grpc/xds/WeightedRoundRobinLoadBalancerTest.java
@@ -536,8 +536,8 @@ public void emptyConfig() {
     verify(helper, times(3)).createSubchannel(
             any(CreateSubchannelArgs.class));
     verify(helper).updateBalancingState(eq(CONNECTING), pickerCaptor.capture());
-    assertThat(pickerCaptor.getValue().getClass().getName())
-        .isEqualTo("io.grpc.util.RoundRobinLoadBalancer$EmptyPicker");
+    assertThat(pickerCaptor.getValue().pickSubchannel(mockArgs))
+        .isEqualTo(PickResult.withNoResult());
     int expectedCount = isEnabledHappyEyeballs() ? servers.size() + 1 : 1;
     assertThat(fakeClock.forwardTime(11, TimeUnit.SECONDS)).isEqualTo( expectedCount);
   }
