core: Handle new PF enabled

grpclb/src/test/java/io/grpc/grpclb/GrpclbLoadBalancerTest.java

@@ -82,6 +82,7 @@
 import io.grpc.inprocess.InProcessServerBuilder;
 import io.grpc.internal.BackoffPolicy;
 import io.grpc.internal.FakeClock;
+import io.grpc.internal.PickFirstLoadBalancerProvider;
 import io.grpc.lb.v1.ClientStats;
 import io.grpc.lb.v1.ClientStatsPerToken;
 import io.grpc.lb.v1.FallbackResponse;
@@ -1921,10 +1922,16 @@ public void grpclbWorking_pickFirstMode() throws Exception {
     // With delegation, the child pick_first creates the subchannel
     inOrder.verify(helper).createSubchannel(createSubchannelArgsCaptor.capture());
     CreateSubchannelArgs createSubchannelArgs = createSubchannelArgsCaptor.getValue();
-    assertThat(createSubchannelArgs.getAddresses())
-        .containsExactly(
-            new EquivalentAddressGroup(backends1.get(0).addr, eagAttrsWithToken("token0001")),
-            new EquivalentAddressGroup(backends1.get(1).addr, eagAttrsWithToken("token0002")));
+    if (PickFirstLoadBalancerProvider.isEnabledNewPickFirst()) {
+      assertThat(createSubchannelArgs.getAddresses())
+          .containsExactly(
+              new EquivalentAddressGroup(backends1.get(0).addr, eagAttrsWithToken("token0001")));
+    } else {
+      assertThat(createSubchannelArgs.getAddresses())
+          .containsExactly(
+              new EquivalentAddressGroup(backends1.get(0).addr, eagAttrsWithToken("token0001")),
+              new EquivalentAddressGroup(backends1.get(1).addr, eagAttrsWithToken("token0002")));
+    }
 
     // Child pick_first eagerly connects, so we start in CONNECTING
     inOrder.verify(helper, atLeast(1)).updateBalancingState(eq(CONNECTING), pickerCaptor.capture());
@@ -1939,18 +1946,20 @@ public void grpclbWorking_pickFirstMode() throws Exception {
     // Child pick_first eagerly calls requestConnection()
     verify(subchannel).requestConnection();
 
-    // TRANSIENT_FAILURE
-    Status error = Status.UNAVAILABLE.withDescription("Simulated connection error");
-    deliverSubchannelState(subchannel, ConnectivityStateInfo.forTransientFailure(error));
-    // The child LB will notify our helper, which updates grpclb state
-    inOrder.verify(helper, atLeast(1))
-        .updateBalancingState(eq(TRANSIENT_FAILURE), pickerCaptor.capture());
-    RoundRobinPicker picker1 = (RoundRobinPicker) pickerCaptor.getValue();
-    assertThat(picker1.dropList).containsExactly(null, null);
-    ChildLbPickerEntry failureEntry = (ChildLbPickerEntry) picker1.pickList.get(0);
-    PickResult failureResult =
-        failureEntry.getChildPicker().pickSubchannel(mock(PickSubchannelArgs.class));
-    assertThat(failureResult.getStatus()).isEqualTo(error);
+    if (!PickFirstLoadBalancerProvider.isEnabledNewPickFirst()) {
+      // TRANSIENT_FAILURE
+      Status error = Status.UNAVAILABLE.withDescription("Simulated connection error");
+      deliverSubchannelState(subchannel, ConnectivityStateInfo.forTransientFailure(error));
+      // The child LB will notify our helper, which updates grpclb state
+      inOrder.verify(helper, atLeast(1))
+          .updateBalancingState(eq(TRANSIENT_FAILURE), pickerCaptor.capture());
+      RoundRobinPicker picker1 = (RoundRobinPicker) pickerCaptor.getValue();
+      assertThat(picker1.dropList).containsExactly(null, null);
+      ChildLbPickerEntry failureEntry = (ChildLbPickerEntry) picker1.pickList.get(0);
+      PickResult failureResult =
+          failureEntry.getChildPicker().pickSubchannel(mock(PickSubchannelArgs.class));
+      assertThat(failureResult.getStatus()).isEqualTo(error);
+    }
 
     // READY
     deliverSubchannelState(subchannel, ConnectivityStateInfo.forNonError(READY));
@@ -1977,11 +1986,15 @@ public void grpclbWorking_pickFirstMode() throws Exception {
     assertThat(mockSubchannels).isEmpty();
 
     // The child LB policy internally calls updateAddresses on the subchannel
-    verify(subchannel).updateAddresses(
-        eq(Arrays.asList(
-            new EquivalentAddressGroup(backends2.get(0).addr, eagAttrsWithToken("token0001")),
-            new EquivalentAddressGroup(backends2.get(2).addr,
-                eagAttrsWithToken("token0004")))));
+    List<EquivalentAddressGroup> expectedAddresses =
+        PickFirstLoadBalancerProvider.isEnabledNewPickFirst()
+            ? Collections.singletonList(
+                new EquivalentAddressGroup(backends2.get(0).addr, eagAttrsWithToken("token0001")))
+            : Arrays.asList(
+                new EquivalentAddressGroup(backends2.get(0).addr, eagAttrsWithToken("token0001")),
+                new EquivalentAddressGroup(backends2.get(2).addr,
+                    eagAttrsWithToken("token0004")));
+    verify(subchannel).updateAddresses(eq(expectedAddresses));
     inOrder.verify(helper, atLeast(1)).updateBalancingState(eq(READY), pickerCaptor.capture());
     RoundRobinPicker picker3 = (RoundRobinPicker) pickerCaptor.getValue();
     assertThat(picker3.dropList).containsExactly(
@@ -2004,7 +2017,8 @@ public void grpclbWorking_pickFirstMode() throws Exception {
     assertThat(pick.getSubchannel()).isNull();
     verify(subchannel, times(2)).requestConnection();
     balancer.requestConnection();
-    verify(subchannel, times(3)).requestConnection();
+    int expectedRequests = PickFirstLoadBalancerProvider.isEnabledNewPickFirst() ? 2 : 3;
+    verify(subchannel, times(expectedRequests)).requestConnection();
 
     // PICK_FIRST doesn't use subchannelPool
     verify(subchannelPool, never())
@@ -2045,10 +2059,16 @@ public void grpclbWorking_pickFirstMode_lbSendsEmptyAddress() throws Exception {
     // The child pick_first creates the first subchannel
     inOrder.verify(helper).createSubchannel(createSubchannelArgsCaptor.capture());
     CreateSubchannelArgs createSubchannelArgs = createSubchannelArgsCaptor.getValue();
-    assertThat(createSubchannelArgs.getAddresses())
-        .containsExactly(
-            new EquivalentAddressGroup(backends1.get(0).addr, eagAttrsWithToken("token0001")),
-            new EquivalentAddressGroup(backends1.get(1).addr, eagAttrsWithToken("token0002")));
+    if (PickFirstLoadBalancerProvider.isEnabledNewPickFirst()) {
+      assertThat(createSubchannelArgs.getAddresses())
+          .containsExactly(
+              new EquivalentAddressGroup(backends1.get(0).addr, eagAttrsWithToken("token0001")));
+    } else {
+      assertThat(createSubchannelArgs.getAddresses())
+          .containsExactly(
+              new EquivalentAddressGroup(backends1.get(0).addr, eagAttrsWithToken("token0001")),
+              new EquivalentAddressGroup(backends1.get(1).addr, eagAttrsWithToken("token0002")));
+    }
 
     // Child pick_first eagerly connects, so initial state is CONNECTING
     inOrder.verify(helper, atLeast(1)).updateBalancingState(eq(CONNECTING), pickerCaptor.capture());
@@ -2104,7 +2124,18 @@ public void grpclbWorking_pickFirstMode_lbSendsEmptyAddress() throws Exception {
     lbResponseObserver.onNext(buildLbResponse(backends2));
 
     // A NEW child LB and NEW subchannel are created upon recovery
-    inOrder.verify(helper).createSubchannel(any(CreateSubchannelArgs.class));
+    inOrder.verify(helper).createSubchannel(createSubchannelArgsCaptor.capture());
+    CreateSubchannelArgs createSubchannelArgs2 = createSubchannelArgsCaptor.getValue();
+    if (PickFirstLoadBalancerProvider.isEnabledNewPickFirst()) {
+      assertThat(createSubchannelArgs2.getAddresses())
+          .containsExactly(
+              new EquivalentAddressGroup(backends2.get(0).addr, eagAttrsWithToken("token0001")));
+    } else {
+      assertThat(createSubchannelArgs2.getAddresses())
+          .containsExactly(
+              new EquivalentAddressGroup(backends2.get(0).addr, eagAttrsWithToken("token0001")),
+              new EquivalentAddressGroup(backends2.get(2).addr, eagAttrsWithToken("token0004")));
+    }
     assertThat(mockSubchannels).hasSize(1);
     Subchannel subchannel2 = mockSubchannels.poll();
     inOrder.verify(helper, atLeast(1)).updateBalancingState(eq(CONNECTING), pickerCaptor.capture());
@@ -2180,8 +2211,13 @@ private void pickFirstModeFallback(long timeout) throws Exception {
     // Entering fallback mode - child LB is created for fallback backends
     inOrder.verify(helper).createSubchannel(createSubchannelArgsCaptor.capture());
     CreateSubchannelArgs createSubchannelArgs = createSubchannelArgsCaptor.getValue();
-    assertThat(createSubchannelArgs.getAddresses())
-        .containsExactly(backendList.get(0), backendList.get(1));
+    if (PickFirstLoadBalancerProvider.isEnabledNewPickFirst()) {
+      assertThat(createSubchannelArgs.getAddresses())
+          .containsExactly(backendList.get(0));
+    } else {
+      assertThat(createSubchannelArgs.getAddresses())
+          .containsExactly(backendList.get(0), backendList.get(1));
+    }
 
     assertThat(mockSubchannels).hasSize(1);
     Subchannel subchannel = mockSubchannels.poll();
@@ -2212,16 +2248,44 @@ private void pickFirstModeFallback(long timeout) throws Exception {
     lbResponseObserver.onNext(buildInitialResponse());
     lbResponseObserver.onNext(buildLbResponse(backends1));
 
-    // subchannel should be updated, NOT recreated
-    inOrder.verify(helper, never()).createSubchannel(any(CreateSubchannelArgs.class));
-    assertThat(mockSubchannels).isEmpty();
-    // The child LB internally calls updateAddresses on the existing subchannel
-    verify(subchannel).updateAddresses(
-        eq(Arrays.asList(
-            new EquivalentAddressGroup(backends1.get(0).addr, eagAttrsWithToken("token0001")),
-            new EquivalentAddressGroup(backends1.get(1).addr,
-                eagAttrsWithToken("token0002")))));
-    inOrder.verify(helper, atLeast(1)).updateBalancingState(eq(READY), pickerCaptor.capture());
+    if (PickFirstLoadBalancerProvider.isEnabledNewPickFirst()) {
+      // In new PF, active fallback subchannel is shut down because its address is not in backends1
+      verify(subchannel).shutdown();
+      // And LB state transitions to IDLE with RequestConnectionPicker
+      inOrder.verify(helper, atLeast(1)).updateBalancingState(eq(IDLE), pickerCaptor.capture());
+      RoundRobinPicker pickerIdle = (RoundRobinPicker) pickerCaptor.getValue();
+      ChildLbPickerEntry idleEntry = (ChildLbPickerEntry) pickerIdle.pickList.get(0);
+
+      // Picking while IDLE triggers requestConnection(), creating a subchannel for backends1.get(0)
+      PickSubchannelArgs args = mock(PickSubchannelArgs.class);
+      PickResult pick = idleEntry.getChildPicker().pickSubchannel(args);
+      assertThat(pick.getSubchannel()).isNull(); // BUFFERing while IDLE
+
+      inOrder.verify(helper).createSubchannel(createSubchannelArgsCaptor.capture());
+      CreateSubchannelArgs createSubchannelArgs2 = createSubchannelArgsCaptor.getValue();
+      assertThat(createSubchannelArgs2.getAddresses())
+          .containsExactly(
+              new EquivalentAddressGroup(backends1.get(0).addr, eagAttrsWithToken("token0001")));
+      assertThat(mockSubchannels).hasSize(1);
+      subchannel = mockSubchannels.poll();
+      verify(subchannel).requestConnection();
+
+      // subchannel becomes READY
+      deliverSubchannelState(subchannel, ConnectivityStateInfo.forNonError(READY));
+      inOrder.verify(helper, atLeast(1)).updateBalancingState(eq(READY), pickerCaptor.capture());
+    } else {
+      // subchannel should be updated, NOT recreated
+      inOrder.verify(helper, never()).createSubchannel(any(CreateSubchannelArgs.class));
+      assertThat(mockSubchannels).isEmpty();
+      // The child LB internally calls updateAddresses on the existing subchannel
+      verify(subchannel).updateAddresses(
+          eq(Arrays.asList(
+              new EquivalentAddressGroup(backends1.get(0).addr, eagAttrsWithToken("token0001")),
+              new EquivalentAddressGroup(backends1.get(1).addr,
+                  eagAttrsWithToken("token0002")))));
+      inOrder.verify(helper, atLeast(1)).updateBalancingState(eq(READY), pickerCaptor.capture());
+    }
+
     RoundRobinPicker picker2 = (RoundRobinPicker) pickerCaptor.getValue();
     assertThat(picker2.dropList).containsExactly(null, null);
 
@@ -2247,7 +2311,8 @@ private void pickFirstModeFallback(long timeout) throws Exception {
     verify(subchannel, times(2)).requestConnection();
 
     balancer.requestConnection();
-    verify(subchannel, times(3)).requestConnection();
+    int expectedRequests = PickFirstLoadBalancerProvider.isEnabledNewPickFirst() ? 2 : 3;
+    verify(subchannel, times(expectedRequests)).requestConnection();
 
     // PICK_FIRST doesn't use subchannelPool
     verify(subchannelPool, never())
@@ -2334,10 +2399,16 @@ public void switchMode() throws Exception {
     // PICK_FIRST Subchannel: child LB creates it
     inOrder.verify(helper).createSubchannel(createSubchannelArgsCaptor.capture());
     CreateSubchannelArgs createSubchannelArgs = createSubchannelArgsCaptor.getValue();
-    assertThat(createSubchannelArgs.getAddresses())
-        .containsExactly(
-            new EquivalentAddressGroup(backends1.get(0).addr, eagAttrsWithToken("token0001")),
-            new EquivalentAddressGroup(backends1.get(1).addr, eagAttrsWithToken("token0002")));
+    if (PickFirstLoadBalancerProvider.isEnabledNewPickFirst()) {
+      assertThat(createSubchannelArgs.getAddresses())
+          .containsExactly(
+              new EquivalentAddressGroup(backends1.get(0).addr, eagAttrsWithToken("token0001")));
+    } else {
+      assertThat(createSubchannelArgs.getAddresses())
+          .containsExactly(
+              new EquivalentAddressGroup(backends1.get(0).addr, eagAttrsWithToken("token0001")),
+              new EquivalentAddressGroup(backends1.get(1).addr, eagAttrsWithToken("token0002")));
+    }
 
     // Child pick_first eagerly connects, so initial state is CONNECTING (not IDLE)
     inOrder.verify(helper, atLeast(1))
@@ -2424,10 +2495,16 @@ public void switchMode_nullLbPolicy() throws Exception {
     // PICK_FIRST Subchannel: with delegation, child LB creates the subchannel
     inOrder.verify(helper).createSubchannel(createSubchannelArgsCaptor.capture());
     CreateSubchannelArgs createSubchannelArgs = createSubchannelArgsCaptor.getValue();
-    assertThat(createSubchannelArgs.getAddresses())
-        .containsExactly(
-            new EquivalentAddressGroup(backends1.get(0).addr, eagAttrsWithToken("token0001")),
-            new EquivalentAddressGroup(backends1.get(1).addr, eagAttrsWithToken("token0002")));
+    if (PickFirstLoadBalancerProvider.isEnabledNewPickFirst()) {
+      assertThat(createSubchannelArgs.getAddresses())
+          .containsExactly(
+              new EquivalentAddressGroup(backends1.get(0).addr, eagAttrsWithToken("token0001")));
+    } else {
+      assertThat(createSubchannelArgs.getAddresses())
+          .containsExactly(
+              new EquivalentAddressGroup(backends1.get(0).addr, eagAttrsWithToken("token0001")),
+              new EquivalentAddressGroup(backends1.get(1).addr, eagAttrsWithToken("token0002")));
+    }
 
     // Child pick_first eagerly connects, so state is CONNECTING (not IDLE)
     inOrder.verify(helper, atLeast(1))

rls/src/test/java/io/grpc/rls/RlsLoadBalancerTest.java

@@ -73,6 +73,7 @@
 import io.grpc.inprocess.InProcessServerBuilder;
 import io.grpc.internal.FakeClock;
 import io.grpc.internal.JsonParser;
+import io.grpc.internal.PickFirstLoadBalancerProvider;
 import io.grpc.internal.PickSubchannelArgsImpl;
 import io.grpc.internal.testing.StreamRecorder;
 import io.grpc.lookup.v1.RouteLookupServiceGrpc;
@@ -257,6 +258,9 @@ public void lb_working_withDefaultTarget_rlsResponding() throws Exception {
     InOrder inOrder = inOrder(helper);
     inOrder.verify(helper)
         .updateBalancingState(eq(ConnectivityState.CONNECTING), pickerCaptor.capture());
+    if (PickFirstLoadBalancerProvider.isEnabledNewPickFirst()) {
+      inOrder.verify(helper).createSubchannel(any(CreateSubchannelArgs.class));
+    }
     SubchannelPicker picker = pickerCaptor.getValue();
     // Warm-up pick; will be queued
     PickResult res = picker.pickSubchannel(searchSubchannelArgs);

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

@@ -255,7 +255,11 @@ public void pickAfterStateChange() throws Exception {
     assertThat(pickerCaptor.getValue().pickSubchannel(mockArgs).getStatus()).isEqualTo(error);
 
     deliverSubchannelState(subchannel, ConnectivityStateInfo.forNonError(IDLE));
-    inOrder.verify(mockHelper).refreshNameResolution();
+    if (PickFirstLoadBalancerProvider.isEnabledNewPickFirst()) {
+      inOrder.verify(mockHelper, never()).refreshNameResolution();
+    } else {
+      inOrder.verify(mockHelper).refreshNameResolution();
+    }
     inOrder.verify(mockHelper, never())
         .updateBalancingState(eq(TRANSIENT_FAILURE), any(SubchannelPicker.class));
 

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

@@ -257,7 +257,9 @@ public void pickAfterStateChange() throws Exception {
         .isEqualTo(PickResult.withNoResult());
 
     deliverSubchannelState(subchannel, ConnectivityStateInfo.forNonError(IDLE));
-    inOrder.verify(helper).refreshNameResolution();
+    if (!PickFirstLoadBalancerProvider.isEnabledNewPickFirst()) {
+      inOrder.verify(helper).refreshNameResolution();
+    }
     assertThat(childLbState.getCurrentState()).isEqualTo(TRANSIENT_FAILURE);
     assertThat(childLbState.getCurrentPicker().toString()).contains(error.toString());