Remove implicit requestConnection() on IDLE from MultiChildLB

One LB no longer needs to extend ChildLbState and one has to start, so it is a bit of a wash. There are more LBs that need the auto-request logic, but if we have an API where subclasses override it without calling super then we can't change the implementation in the future. Adding behavior on top of a base class allows subclasses to call super, which lets the base class change over time.
grpc · Aug 12, 2024 · a6f8ebf · a6f8ebf
1 parent 4ab3422
commit a6f8ebf
Show file tree

Hide file tree

Showing 6 changed files with 48 additions and 52 deletions.
diff --git a/util/src/main/java/io/grpc/util/MultiChildLoadBalancer.java b/util/src/main/java/io/grpc/util/MultiChildLoadBalancer.java
@@ -456,8 +456,6 @@ protected class ChildLbStateHelper extends ForwardingLoadBalancerHelper {
       /**
        * Update current state and picker for this child and then use
        * {@link #updateOverallBalancingState()} for the parent LB.
-       *
-       * <p/>Override this if you don't want to automatically request a connection when in IDLE
        */
       @Override
       public void updateBalancingState(final ConnectivityState newState,
@@ -471,9 +469,6 @@ public void updateBalancingState(final ConnectivityState newState,
         // If we are already in the process of resolving addresses, the overall balancing state
         // will be updated at the end of it, and we don't need to trigger that update here.
         if (!resolvingAddresses) {
-          if (newState == IDLE) {
-            lb.requestConnection();
-          }
           updateOverallBalancingState();
         }
       }

diff --git a/util/src/main/java/io/grpc/util/RoundRobinLoadBalancer.java b/util/src/main/java/io/grpc/util/RoundRobinLoadBalancer.java
@@ -95,6 +95,25 @@ private SubchannelPicker createReadyPicker(Collection<ChildLbState> children) {
     return new ReadyPicker(pickerList, sequence);
   }
 
+  @Override
+  protected ChildLbState createChildLbState(Object key, Object policyConfig,
+      SubchannelPicker initialPicker, ResolvedAddresses resolvedAddresses) {
+    return new ChildLbState(key, pickFirstLbProvider, policyConfig, initialPicker) {
+      @Override
+      protected ChildLbStateHelper createChildHelper() {
+        return new ChildLbStateHelper() {
+          @Override
+          public void updateBalancingState(ConnectivityState newState, SubchannelPicker newPicker) {
+            super.updateBalancingState(newState, newPicker);
+            if (!resolvingAddresses && newState == IDLE) {
+              getLb().requestConnection();
+            }
+          }
+        };
+      }
+    };
+  }
+
   @VisibleForTesting
   static class ReadyPicker extends SubchannelPicker {
     private final List<SubchannelPicker> subchannelPickers; // non-empty

diff --git a/xds/src/main/java/io/grpc/xds/LeastRequestLoadBalancer.java b/xds/src/main/java/io/grpc/xds/LeastRequestLoadBalancer.java
@@ -328,5 +328,18 @@ public LeastRequestLbState(Object key, LoadBalancerProvider policyProvider,
     int getActiveRequests() {
       return activeRequests.get();
     }
+
+    @Override
+    protected ChildLbStateHelper createChildHelper() {
+      return new ChildLbStateHelper() {
+        @Override
+        public void updateBalancingState(ConnectivityState newState, SubchannelPicker newPicker) {
+          super.updateBalancingState(newState, newPicker);
+          if (!resolvingAddresses && newState == IDLE) {
+            getLb().requestConnection();
+          }
+        }
+      };
+    }
   }
 }
diff --git a/xds/src/main/java/io/grpc/xds/RingHashLoadBalancer.java b/xds/src/main/java/io/grpc/xds/RingHashLoadBalancer.java
@@ -229,7 +229,7 @@ protected void updateOverallBalancingState() {
   @Override
   protected ChildLbState createChildLbState(Object key, Object policyConfig,
       SubchannelPicker initialPicker, ResolvedAddresses resolvedAddresses) {
-    return new RingHashChildLbState((Endpoint)key);
+    return new ChildLbState(key, lazyLbFactory, null, EMPTY_PICKER);
   }
 
   private Status validateAddrList(List<EquivalentAddressGroup> addrList) {
@@ -358,7 +358,7 @@ private RingHashPicker(
       this.ring = ring;
       pickableSubchannels = new HashMap<>(subchannels.size());
       for (Map.Entry<Object, ChildLbState> entry : subchannels.entrySet()) {
-        RingHashChildLbState childLbState = (RingHashChildLbState) entry.getValue();
+        ChildLbState childLbState = entry.getValue();
         pickableSubchannels.put((Endpoint)entry.getKey(),
             new SubchannelView(childLbState, childLbState.getCurrentState()));
       }
@@ -405,7 +405,7 @@ public PickResult pickSubchannel(PickSubchannelArgs args) {
       for (int i = 0; i < ring.size(); i++) {
         int index = (targetIndex + i) % ring.size();
         SubchannelView subchannelView = pickableSubchannels.get(ring.get(index).addrKey);
-        RingHashChildLbState childLbState = subchannelView.childLbState;
+        ChildLbState childLbState = subchannelView.childLbState;
 
         if (subchannelView.connectivityState  == READY) {
           return childLbState.getCurrentPicker().pickSubchannel(args);
@@ -427,7 +427,7 @@ public PickResult pickSubchannel(PickSubchannelArgs args) {
       }
 
       // return the pick from the original subchannel hit by hash, which is probably an error
-      RingHashChildLbState originalSubchannel =
+      ChildLbState originalSubchannel =
           pickableSubchannels.get(ring.get(targetIndex).addrKey).childLbState;
       return originalSubchannel.getCurrentPicker().pickSubchannel(args);
     }
@@ -439,10 +439,10 @@ public PickResult pickSubchannel(PickSubchannelArgs args) {
    * state changes.
    */
   private static final class SubchannelView {
-    private final RingHashChildLbState childLbState;
+    private final ChildLbState childLbState;
     private final ConnectivityState connectivityState;
 
-    private SubchannelView(RingHashChildLbState childLbState, ConnectivityState state) {
+    private SubchannelView(ChildLbState childLbState, ConnectivityState state) {
       this.childLbState = childLbState;
       this.connectivityState = state;
     }
@@ -487,41 +487,4 @@ public String toString() {
           .toString();
     }
   }
-
-  class RingHashChildLbState extends MultiChildLoadBalancer.ChildLbState {
-
-    public RingHashChildLbState(Endpoint key) {
-      super(key, lazyLbFactory, null, EMPTY_PICKER);
-    }
-
-    @Override
-    protected ChildLbStateHelper createChildHelper() {
-      return new RingHashChildHelper();
-    }
-
-    // Need to expose this to the LB class
-    @Override
-    protected void shutdown() {
-      super.shutdown();
-    }
-
-    private class RingHashChildHelper extends ChildLbStateHelper {
-      @Override
-      public void updateBalancingState(final ConnectivityState newState,
-                                       final SubchannelPicker newPicker) {
-        setCurrentState(newState);
-        setCurrentPicker(newPicker);
-
-        if (getChildLbState(getKey()) == null) {
-          return;
-        }
-
-        // If we are already in the process of resolving addresses, the overall balancing state
-        // will be updated at the end of it, and we don't need to trigger that update here.
-        if (!resolvingAddresses) {
-          updateOverallBalancingState();
-        }
-      }
-    }
-  }
 }
diff --git a/xds/src/main/java/io/grpc/xds/WeightedRoundRobinLoadBalancer.java b/xds/src/main/java/io/grpc/xds/WeightedRoundRobinLoadBalancer.java
@@ -340,6 +340,14 @@ final class WrrChildLbStateHelper extends ChildLbStateHelper {
       public Subchannel createSubchannel(CreateSubchannelArgs args) {
         return new WrrSubchannel(super.createSubchannel(args), WeightedChildLbState.this);
       }
+
+      @Override
+      public void updateBalancingState(ConnectivityState newState, SubchannelPicker newPicker) {
+        super.updateBalancingState(newState, newPicker);
+        if (!resolvingAddresses && newState == ConnectivityState.IDLE) {
+          getLb().requestConnection();
+        }
+      }
     }
 
     final class OrcaReportListener implements OrcaPerRequestReportListener, OrcaOobReportListener {

diff --git a/xds/src/test/java/io/grpc/xds/RingHashLoadBalancerTest.java b/xds/src/test/java/io/grpc/xds/RingHashLoadBalancerTest.java
@@ -66,7 +66,6 @@
 import io.grpc.testing.TestMethodDescriptors;
 import io.grpc.util.AbstractTestHelper;
 import io.grpc.util.MultiChildLoadBalancer.ChildLbState;
-import io.grpc.xds.RingHashLoadBalancer.RingHashChildLbState;
 import io.grpc.xds.RingHashLoadBalancer.RingHashConfig;
 import java.lang.Thread.UncaughtExceptionHandler;
 import java.net.SocketAddress;
@@ -177,8 +176,7 @@ public void subchannelNotAutoReconnectAfterReenteringIdle() {
     assertThat(addressesAcceptanceStatus.isOk()).isTrue();
     verify(helper).updateBalancingState(eq(IDLE), pickerCaptor.capture());
 
-    RingHashChildLbState childLbState =
-        (RingHashChildLbState) loadBalancer.getChildLbStates().iterator().next();
+    ChildLbState childLbState = loadBalancer.getChildLbStates().iterator().next();
     assertThat(subchannels.get(Collections.singletonList(childLbState.getEag()))).isNull();
 
     // Picking subchannel triggers connection.
@@ -422,7 +420,7 @@ public void skipFailingHosts_pickNextNonFailingHost() {
     assertThat(addressesAcceptanceStatus.isOk()).isTrue();
 
     // Create subchannel for the first address
-    ((RingHashChildLbState) loadBalancer.getChildLbStateEag(servers.get(0))).getCurrentPicker()
+    loadBalancer.getChildLbStateEag(servers.get(0)).getCurrentPicker()
         .pickSubchannel(getDefaultPickSubchannelArgs(hashFunc.hashVoid()));
     verifyConnection(1);