xds: fix NPE in wrr in TF state (#10874)

grpc · Feb 1, 2024 · fa36269 · fa36269
1 parent c83e5f9
commit fa36269
Show file tree

Hide file tree

Showing 2 changed files with 20 additions and 0 deletions.
diff --git a/xds/src/main/java/io/grpc/xds/WeightedRoundRobinLoadBalancer.java b/xds/src/main/java/io/grpc/xds/WeightedRoundRobinLoadBalancer.java
@@ -365,6 +365,9 @@ public PickResult pickSubchannel(PickSubchannelArgs args) {
       WeightedChildLbState wChild = (WeightedChildLbState) childLbState;
       PickResult pickResult = childLbState.getCurrentPicker().pickSubchannel(args);
       Subchannel subchannel = pickResult.getSubchannel();
+      if (subchannel == null) {
+        return pickResult;
+      }
       if (!enableOobLoadReport) {
         return PickResult.withSubchannel(subchannel,
             OrcaPerRequestUtil.getInstance().newOrcaClientStreamTracerFactory(

diff --git a/xds/src/test/java/io/grpc/xds/WeightedRoundRobinLoadBalancerTest.java b/xds/src/test/java/io/grpc/xds/WeightedRoundRobinLoadBalancerTest.java
@@ -47,6 +47,7 @@
 import io.grpc.LoadBalancer.Subchannel;
 import io.grpc.LoadBalancer.SubchannelPicker;
 import io.grpc.LoadBalancer.SubchannelStateListener;
+import io.grpc.Status;
 import io.grpc.SynchronizationContext;
 import io.grpc.internal.FakeClock;
 import io.grpc.internal.TestUtils;
@@ -162,6 +163,22 @@ public ClientCall<OrcaLoadReportRequest, OrcaLoadReport> answer(
     verify(helper, times(3)).createSubchannel(any(CreateSubchannelArgs.class));
   }
 
+  @Test
+  public void pickChildLbTF() throws Exception {
+    syncContext.execute(() -> wrr.acceptResolvedAddresses(ResolvedAddresses.newBuilder()
+        .setAddresses(servers.subList(0, 1)).setLoadBalancingPolicyConfig(weightedConfig)
+        .setAttributes(affinity).build()));
+    Iterator<Subchannel> it = subchannels.values().iterator();
+    Subchannel readySubchannel1 = it.next();
+    getSubchannelStateListener(readySubchannel1).onSubchannelState(ConnectivityStateInfo
+        .forTransientFailure(Status.UNAVAILABLE));
+    verify(helper).updateBalancingState(
+        eq(ConnectivityState.TRANSIENT_FAILURE), pickerCaptor.capture());
+    final WeightedRoundRobinPicker weightedPicker =
+        (WeightedRoundRobinPicker) pickerCaptor.getValue();
+    weightedPicker.pickSubchannel(mockArgs);
+  }
+
   @Test
   public void wrrLifeCycle() {
     syncContext.execute(() -> wrr.acceptResolvedAddresses(ResolvedAddresses.newBuilder()