From 85758e305d8b1f8569d5667140180ce2b7ed47d8 Mon Sep 17 00:00:00 2001 From: Kannan J Date: Mon, 31 Aug 2026 12:31:37 +0000 Subject: [PATCH 1/7] xds: Make RawMessageClientInterceptor conditional on ext_proc flags Only add RawMessageClientInterceptor to the interceptor chain in XdsNameResolver when either GRPC_EXPERIMENTAL_XDS_EXT_PROC_ON_CLIENT or GRPC_EXPERIMENTAL_XDS_EXT_PROC_ON_SERVER is true. This prevents RawMessageClientInterceptor from corrupting method descriptors and causing empty payloads on retry attempts for regular xDS channels. Fixes #13010 TAG=agy CONV=56c87717-243d-4f12-af9e-c532e509892a --- .../java/io/grpc/xds/XdsNameResolver.java | 5 +- .../java/io/grpc/xds/XdsNameResolverTest.java | 83 ++++++++++++++++++- 2 files changed, 85 insertions(+), 3 deletions(-) diff --git a/xds/src/main/java/io/grpc/xds/XdsNameResolver.java b/xds/src/main/java/io/grpc/xds/XdsNameResolver.java index 6221aa8fcae..70bdc25e334 100644 --- a/xds/src/main/java/io/grpc/xds/XdsNameResolver.java +++ b/xds/src/main/java/io/grpc/xds/XdsNameResolver.java @@ -907,7 +907,10 @@ private ClientInterceptor createFilters( } ImmutableList.Builder withRawMessage = ImmutableList.builder(); - withRawMessage.add(new RawMessageClientInterceptor()); + if (GrpcUtil.getFlag("GRPC_EXPERIMENTAL_XDS_EXT_PROC_ON_CLIENT", false) + || GrpcUtil.getFlag("GRPC_EXPERIMENTAL_XDS_EXT_PROC_ON_SERVER", false)) { + withRawMessage.add(new RawMessageClientInterceptor()); + } withRawMessage.addAll(filterInterceptors.build()); return combineInterceptors(withRawMessage.build()); } diff --git a/xds/src/test/java/io/grpc/xds/XdsNameResolverTest.java b/xds/src/test/java/io/grpc/xds/XdsNameResolverTest.java index 1a92598aaf3..7f4035db5d6 100644 --- a/xds/src/test/java/io/grpc/xds/XdsNameResolverTest.java +++ b/xds/src/test/java/io/grpc/xds/XdsNameResolverTest.java @@ -2957,7 +2957,7 @@ private final class TestChannel extends Channel { @Override public ClientCall newCall( MethodDescriptor methodDescriptor, CallOptions callOptions) { - TestCall call = new TestCall<>(callOptions); + TestCall call = new TestCall<>(methodDescriptor, callOptions); testCall = call; return call; } @@ -2969,11 +2969,13 @@ public String authority() { } private static final class TestCall extends NoopClientCall { + final MethodDescriptor methodDescriptor; // CallOptions actually received from the channel when the call is created. final CallOptions callOptions; ClientCall.Listener listener; - TestCall(CallOptions callOptions) { + TestCall(MethodDescriptor methodDescriptor, CallOptions callOptions) { + this.methodDescriptor = methodDescriptor; this.callOptions = callOptions; } @@ -3098,4 +3100,81 @@ public void onMessage(String message) { channel, METHOD_SAY_HELLO, CallOptions.DEFAULT, "World"); assertThat(response).isEqualTo("Hello World"); } + + @Test + public void rawMessageClientInterceptor_flagFalse() { + String origClientProp = System.getProperty("GRPC_EXPERIMENTAL_XDS_EXT_PROC_ON_CLIENT"); + String origServerProp = System.getProperty("GRPC_EXPERIMENTAL_XDS_EXT_PROC_ON_SERVER"); + System.setProperty("GRPC_EXPERIMENTAL_XDS_EXT_PROC_ON_CLIENT", "false"); + System.setProperty("GRPC_EXPERIMENTAL_XDS_EXT_PROC_ON_SERVER", "false"); + try { + filterStateTestSetupResolver(); + FakeXdsClient xdsClient = (FakeXdsClient) resolver.getXdsClient(); + VirtualHost vhost = filterStateTestVhost(); + + xdsClient.deliverLdsUpdateWithFilters(vhost, filterStateTestConfigs(STATEFUL_1)); + createAndDeliverClusterUpdates(xdsClient, cluster1); + + // When flags are false, RawMessageClientInterceptor is not added. + assertClusterResolutionResult(call1, cluster1); + assertThat(testCall.methodDescriptor).isSameInstanceAs(call1.methodDescriptor); + } finally { + restoreProperty("GRPC_EXPERIMENTAL_XDS_EXT_PROC_ON_CLIENT", origClientProp); + restoreProperty("GRPC_EXPERIMENTAL_XDS_EXT_PROC_ON_SERVER", origServerProp); + } + } + + @Test + public void rawMessageClientInterceptor_flagTrue() { + String origClientProp = System.getProperty("GRPC_EXPERIMENTAL_XDS_EXT_PROC_ON_CLIENT"); + String origServerProp = System.getProperty("GRPC_EXPERIMENTAL_XDS_EXT_PROC_ON_SERVER"); + + // When GRPC_EXPERIMENTAL_XDS_EXT_PROC_ON_CLIENT is true, RawMessageClientInterceptor is added. + System.setProperty("GRPC_EXPERIMENTAL_XDS_EXT_PROC_ON_CLIENT", "true"); + System.setProperty("GRPC_EXPERIMENTAL_XDS_EXT_PROC_ON_SERVER", "false"); + try { + filterStateTestSetupResolver(); + FakeXdsClient xdsClient = (FakeXdsClient) resolver.getXdsClient(); + VirtualHost vhost = filterStateTestVhost(); + + xdsClient.deliverLdsUpdateWithFilters(vhost, filterStateTestConfigs(STATEFUL_1)); + createAndDeliverClusterUpdates(xdsClient, cluster1); + + assertClusterResolutionResult(call1, cluster1); + assertThat(testCall.methodDescriptor).isNotSameInstanceAs(call1.methodDescriptor); + } finally { + restoreProperty("GRPC_EXPERIMENTAL_XDS_EXT_PROC_ON_CLIENT", origClientProp); + restoreProperty("GRPC_EXPERIMENTAL_XDS_EXT_PROC_ON_SERVER", origServerProp); + } + + resolver.shutdown(); + reset(mockListener); + when(mockListener.onResult2(any())).thenReturn(Status.OK); + + // When GRPC_EXPERIMENTAL_XDS_EXT_PROC_ON_SERVER is true, RawMessageClientInterceptor is added. + System.setProperty("GRPC_EXPERIMENTAL_XDS_EXT_PROC_ON_CLIENT", "false"); + System.setProperty("GRPC_EXPERIMENTAL_XDS_EXT_PROC_ON_SERVER", "true"); + try { + filterStateTestSetupResolver(); + FakeXdsClient xdsClient = (FakeXdsClient) resolver.getXdsClient(); + VirtualHost vhost = filterStateTestVhost(); + + xdsClient.deliverLdsUpdateWithFilters(vhost, filterStateTestConfigs(STATEFUL_1)); + createAndDeliverClusterUpdates(xdsClient, cluster1); + + assertClusterResolutionResult(call1, cluster1); + assertThat(testCall.methodDescriptor).isNotSameInstanceAs(call1.methodDescriptor); + } finally { + restoreProperty("GRPC_EXPERIMENTAL_XDS_EXT_PROC_ON_CLIENT", origClientProp); + restoreProperty("GRPC_EXPERIMENTAL_XDS_EXT_PROC_ON_SERVER", origServerProp); + } + } + + private static void restoreProperty(String key, @Nullable String value) { + if (value == null) { + System.clearProperty(key); + } else { + System.setProperty(key, value); + } + } } From d6f80bbe7f2c065edba8949c03ea20fa37521b62 Mon Sep 17 00:00:00 2001 From: Kannan J Date: Wed, 7 Oct 2026 06:08:42 +0000 Subject: [PATCH 2/7] core, grpclb: Enable new pick first LB by default and update grpclb tests Enable GRPC_EXPERIMENTAL_ENABLE_NEW_PICK_FIRST by default in PickFirstLoadBalancerProvider. Duplicate failing tests in GrpclbLoadBalancerTest to verify both legacy PickFirstLoadBalancer and modern PickFirstLeafLoadBalancer behaviors. TAG=agy CONV=2b34d6af-c0df-46b4-8326-2dea3641f8cc --- .../PickFirstLoadBalancerProvider.java | 2 +- .../grpc/grpclb/GrpclbLoadBalancerTest.java | 544 ++++++++++++++++++ 2 files changed, 545 insertions(+), 1 deletion(-) diff --git a/core/src/main/java/io/grpc/internal/PickFirstLoadBalancerProvider.java b/core/src/main/java/io/grpc/internal/PickFirstLoadBalancerProvider.java index 83b3fb7d8e6..9a204a6d740 100644 --- a/core/src/main/java/io/grpc/internal/PickFirstLoadBalancerProvider.java +++ b/core/src/main/java/io/grpc/internal/PickFirstLoadBalancerProvider.java @@ -37,7 +37,7 @@ public final class PickFirstLoadBalancerProvider extends LoadBalancerProvider { private static final String SHUFFLE_ADDRESS_LIST_KEY = "shuffleAddressList"; static boolean enableNewPickFirst = - GrpcUtil.getFlag("GRPC_EXPERIMENTAL_ENABLE_NEW_PICK_FIRST", false); + GrpcUtil.getFlag("GRPC_EXPERIMENTAL_ENABLE_NEW_PICK_FIRST", true); public static boolean isEnabledHappyEyeballs() { diff --git a/grpclb/src/test/java/io/grpc/grpclb/GrpclbLoadBalancerTest.java b/grpclb/src/test/java/io/grpc/grpclb/GrpclbLoadBalancerTest.java index ef31b318cb5..9ba8b4d3bf5 100644 --- a/grpclb/src/test/java/io/grpc/grpclb/GrpclbLoadBalancerTest.java +++ b/grpclb/src/test/java/io/grpc/grpclb/GrpclbLoadBalancerTest.java @@ -82,6 +82,8 @@ import io.grpc.inprocess.InProcessServerBuilder; import io.grpc.internal.BackoffPolicy; import io.grpc.internal.FakeClock; +import io.grpc.internal.PickFirstLoadBalancerProvider; +import io.grpc.internal.PickFirstLoadBalancerProviderAccessor; import io.grpc.lb.v1.ClientStats; import io.grpc.lb.v1.ClientStatsPerToken; import io.grpc.lb.v1.FallbackResponse; @@ -205,6 +207,7 @@ public void uncaughtException(Thread t, Throwable e) { private GrpclbLoadBalancer balancer; private final ArgumentCaptor createSubchannelArgsCaptor = ArgumentCaptor.forClass(CreateSubchannelArgs.class); + private final boolean defaultNewPickFirst = PickFirstLoadBalancerProvider.isEnabledNewPickFirst(); @Before public void setUp() throws Exception { @@ -264,6 +267,7 @@ public void run() { // No timer should linger after shutdown assertThat(fakeClock.getPendingTasks()).isEmpty(); } finally { + PickFirstLoadBalancerProviderAccessor.setEnableNewPickFirst(defaultNewPickFirst); if (fakeLbServer != null) { fakeLbServer.shutdownNow(); } @@ -1890,6 +1894,7 @@ public void grpclbBalancerStreamClosedAndRetried() throws Exception { @Test public void grpclbWorking_pickFirstMode() throws Exception { + PickFirstLoadBalancerProviderAccessor.setEnableNewPickFirst(false); InOrder inOrder = inOrder(helper); List grpclbBalancerList = createResolvedBalancerAddresses(1); @@ -2013,8 +2018,119 @@ public void grpclbWorking_pickFirstMode() throws Exception { .returnSubchannel(any(Subchannel.class), any(ConnectivityStateInfo.class)); } + @Test + public void grpclbWorking_pickFirstMode_newPickFirst() throws Exception { + PickFirstLoadBalancerProviderAccessor.setEnableNewPickFirst(true); + InOrder inOrder = inOrder(helper); + + List grpclbBalancerList = createResolvedBalancerAddresses(1); + + deliverResolvedAddresses( + Collections.emptyList(), + grpclbBalancerList, + GrpclbConfig.create(Mode.PICK_FIRST)); + + assertEquals(1, fakeOobChannels.size()); + verify(mockLbService).balanceLoad(lbResponseObserverCaptor.capture()); + StreamObserver lbResponseObserver = lbResponseObserverCaptor.getValue(); + assertEquals(1, lbRequestObservers.size()); + StreamObserver lbRequestObserver = lbRequestObservers.poll(); + verify(lbRequestObserver).onNext( + eq(LoadBalanceRequest.newBuilder().setInitialRequest( + InitialLoadBalanceRequest.newBuilder().setName(SERVICE_AUTHORITY).build()) + .build())); + + // Simulate receiving LB response + List backends1 = Arrays.asList( + new ServerEntry("127.0.0.1", 2000, "token0001"), + new ServerEntry("127.0.0.1", 2010, "token0002")); + inOrder.verify(helper, never()) + .updateBalancingState(any(ConnectivityState.class), any(SubchannelPicker.class)); + lbResponseObserver.onNext(buildInitialResponse()); + lbResponseObserver.onNext(buildLbResponse(backends1)); + + // With delegation, the child pick_first creates the subchannel for the 1st address + inOrder.verify(helper).createSubchannel(createSubchannelArgsCaptor.capture()); + CreateSubchannelArgs createSubchannelArgs0 = createSubchannelArgsCaptor.getValue(); + assertThat(createSubchannelArgs0.getAddresses()) + .containsExactly( + new EquivalentAddressGroup(backends1.get(0).addr, eagAttrsWithToken("token0001"))); + + // Child pick_first eagerly connects, so we start in CONNECTING + inOrder.verify(helper, atLeast(1)).updateBalancingState(eq(CONNECTING), pickerCaptor.capture()); + RoundRobinPicker picker0 = (RoundRobinPicker) pickerCaptor.getValue(); + // Only one subchannel is created by the child LB initially + assertThat(mockSubchannels).hasSize(1); + Subchannel subchannel0 = mockSubchannels.poll(); + assertThat(picker0.dropList).containsExactly(null, null); + assertThat(picker0.pickList).hasSize(1); + assertThat(picker0.pickList.get(0)).isInstanceOf(ChildLbPickerEntry.class); + + // Child pick_first eagerly calls requestConnection() + verify(subchannel0).requestConnection(); + + // READY on subchannel0 + deliverSubchannelState(subchannel0, ConnectivityStateInfo.forNonError(READY)); + inOrder.verify(helper, atLeast(1)).updateBalancingState(eq(READY), pickerCaptor.capture()); + RoundRobinPicker picker1 = (RoundRobinPicker) pickerCaptor.getValue(); + assertThat(picker1.dropList).containsExactly(null, null); + ChildLbPickerEntry readyEntry = (ChildLbPickerEntry) picker1.pickList.get(0); + PickResult readyResult = + readyEntry.getChildPicker().pickSubchannel(mock(PickSubchannelArgs.class)); + assertThat(readyResult.getSubchannel()).isEqualTo(subchannel0); + + // New server list with drops + List backends2 = Arrays.asList( + new ServerEntry("127.0.0.1", 2000, "token0001"), + new ServerEntry("token0003"), // drop + new ServerEntry("127.0.0.1", 2020, "token0004")); + inOrder.verify(helper, never()) + .updateBalancingState(any(ConnectivityState.class), any(SubchannelPicker.class)); + lbResponseObserver.onNext(buildLbResponse(backends2)); + + // Verify child LB updates subchannel0 in-place with backends2.get(0) attributes + inOrder.verify(helper, never()).createSubchannel(any(CreateSubchannelArgs.class)); + verify(helper, times(1)).createSubchannel(any(CreateSubchannelArgs.class)); + assertThat(mockSubchannels).isEmpty(); + + // The child LB policy calls updateAddresses on subchannel0 with ONLY the current address + verify(subchannel0).updateAddresses( + eq(Collections.singletonList( + new EquivalentAddressGroup(backends2.get(0).addr, eagAttrsWithToken("token0001"))))); + inOrder.verify(helper, atLeast(1)).updateBalancingState(eq(READY), pickerCaptor.capture()); + RoundRobinPicker picker2 = (RoundRobinPicker) pickerCaptor.getValue(); + assertThat(picker2.dropList).containsExactly( + null, new DropEntry(getLoadRecorder(), "token0003"), null); + ChildLbPickerEntry updatedEntry = (ChildLbPickerEntry) picker2.pickList.get(0); + PickResult updatedResult = + updatedEntry.getChildPicker().pickSubchannel(mock(PickSubchannelArgs.class)); + assertThat(updatedResult.getSubchannel()).isEqualTo(subchannel0); + + // Subchannel goes IDLE, grpclb state should follow + deliverSubchannelState(subchannel0, ConnectivityStateInfo.forNonError(IDLE)); + inOrder.verify(helper).updateBalancingState(eq(IDLE), pickerCaptor.capture()); + RoundRobinPicker picker3 = (RoundRobinPicker) pickerCaptor.getValue(); + + // No new connection request should have happened yet (beyond the first eager one) + verify(subchannel0, times(1)).requestConnection(); + PickSubchannelArgs args = mock(PickSubchannelArgs.class); + PickResult pick = picker3.pickSubchannel(args); + // Child pick_first picker returns withNoResult() when IDLE and requests connection + assertThat(pick.getSubchannel()).isNull(); + verify(subchannel0, times(2)).requestConnection(); + balancer.requestConnection(); + verify(subchannel0, times(3)).requestConnection(); + + // PICK_FIRST doesn't use subchannelPool + verify(subchannelPool, never()) + .takeOrCreateSubchannel(any(EquivalentAddressGroup.class), any(Attributes.class)); + verify(subchannelPool, never()) + .returnSubchannel(any(Subchannel.class), any(ConnectivityStateInfo.class)); + } + @Test public void grpclbWorking_pickFirstMode_lbSendsEmptyAddress() throws Exception { + PickFirstLoadBalancerProviderAccessor.setEnableNewPickFirst(false); InOrder inOrder = inOrder(helper); List grpclbBalancerList = createResolvedBalancerAddresses(1); @@ -2121,6 +2237,118 @@ public void grpclbWorking_pickFirstMode_lbSendsEmptyAddress() throws Exception { .isEqualTo(subchannel2); } + @Test + public void grpclbWorking_pickFirstMode_lbSendsEmptyAddress_newPickFirst() throws Exception { + PickFirstLoadBalancerProviderAccessor.setEnableNewPickFirst(true); + InOrder inOrder = inOrder(helper); + + List grpclbBalancerList = createResolvedBalancerAddresses(1); + deliverResolvedAddresses( + Collections.emptyList(), + grpclbBalancerList, + GrpclbConfig.create(Mode.PICK_FIRST)); + + assertEquals(1, fakeOobChannels.size()); + verify(mockLbService).balanceLoad(lbResponseObserverCaptor.capture()); + StreamObserver lbResponseObserver = lbResponseObserverCaptor.getValue(); + assertEquals(1, lbRequestObservers.size()); + StreamObserver lbRequestObserver = lbRequestObservers.poll(); + verify(lbRequestObserver).onNext( + eq(LoadBalanceRequest.newBuilder().setInitialRequest( + InitialLoadBalanceRequest.newBuilder().setName(SERVICE_AUTHORITY).build()) + .build())); + + // Simulate receiving LB response + List backends1 = Arrays.asList( + new ServerEntry("127.0.0.1", 2000, "token0001"), + new ServerEntry("127.0.0.1", 2010, "token0002")); + inOrder.verify(helper, never()) + .updateBalancingState(any(ConnectivityState.class), any(SubchannelPicker.class)); + lbResponseObserver.onNext(buildInitialResponse()); + lbResponseObserver.onNext(buildLbResponse(backends1)); + + // The child pick_first creates the first subchannel with only the first address + inOrder.verify(helper).createSubchannel(createSubchannelArgsCaptor.capture()); + CreateSubchannelArgs createSubchannelArgs = createSubchannelArgsCaptor.getValue(); + assertThat(createSubchannelArgs.getAddresses()) + .containsExactly( + new EquivalentAddressGroup(backends1.get(0).addr, eagAttrsWithToken("token0001"))); + + // Child pick_first eagerly connects, so initial state is CONNECTING + inOrder.verify(helper, atLeast(1)).updateBalancingState(eq(CONNECTING), pickerCaptor.capture()); + RoundRobinPicker picker0 = (RoundRobinPicker) pickerCaptor.getValue(); + // Verify subchannel creation by child LB + assertThat(mockSubchannels).hasSize(1); + Subchannel subchannel = mockSubchannels.poll(); + assertThat(picker0.dropList).containsExactly(null, null); + assertThat(picker0.pickList).hasSize(1); + assertThat(picker0.pickList.get(0)).isInstanceOf(ChildLbPickerEntry.class); + + // Child pick_first eagerly calls requestConnection() + verify(subchannel).requestConnection(); + + // READY + deliverSubchannelState(subchannel, ConnectivityStateInfo.forNonError(READY)); + inOrder.verify(helper, atLeast(1)).updateBalancingState(eq(READY), pickerCaptor.capture()); + RoundRobinPicker pickerReady = (RoundRobinPicker) pickerCaptor.getValue(); + // Verify the subchannel in the delegated picker + ChildLbPickerEntry readyEntry = (ChildLbPickerEntry) pickerReady.pickList.get(0); + assertThat( + readyEntry.getChildPicker().pickSubchannel(mock(PickSubchannelArgs.class)).getSubchannel()) + .isEqualTo(subchannel); + inOrder.verify(helper, never()) + .updateBalancingState(any(ConnectivityState.class), any(SubchannelPicker.class)); + + // Empty addresses from LB - child LB is shutdown + lbResponseObserver.onNext(buildLbResponse(Collections.emptyList())); + + // Child LB is shutdown (which shuts down its subchannel) + // createSubchannel() has ever been called only once + inOrder.verify(helper, never()).createSubchannel(any(CreateSubchannelArgs.class)); + assertThat(mockSubchannels).isEmpty(); + verify(subchannel).shutdown(); + + // RPC error status includes message of no backends provided by balancer + inOrder.verify(helper, atLeast(1)) + .updateBalancingState(eq(TRANSIENT_FAILURE), pickerCaptor.capture()); + RoundRobinPicker errorPicker = (RoundRobinPicker) pickerCaptor.getValue(); + assertThat(errorPicker.pickList) + .containsExactly(new ErrorEntry(GrpclbState.NO_AVAILABLE_BACKENDS_STATUS)); + + lbResponseObserver.onNext(buildLbResponse(Collections.emptyList())); + + // Test recover from new LB response with addresses + // New server list with drops + List backends2 = Arrays.asList( + new ServerEntry("127.0.0.1", 2000, "token0001"), + new ServerEntry("token0003"), // drop + new ServerEntry("127.0.0.1", 2020, "token0004")); + inOrder.verify(helper, never()) + .updateBalancingState(any(ConnectivityState.class), any(SubchannelPicker.class)); + lbResponseObserver.onNext(buildLbResponse(backends2)); + + // A NEW child LB and NEW subchannel are created upon recovery with only backends2.get(0) + inOrder.verify(helper).createSubchannel(createSubchannelArgsCaptor.capture()); + CreateSubchannelArgs createSubchannelArgs2 = createSubchannelArgsCaptor.getValue(); + assertThat(createSubchannelArgs2.getAddresses()) + .containsExactly( + new EquivalentAddressGroup(backends2.get(0).addr, eagAttrsWithToken("token0001"))); + assertThat(mockSubchannels).hasSize(1); + Subchannel subchannel2 = mockSubchannels.poll(); + inOrder.verify(helper, atLeast(1)).updateBalancingState(eq(CONNECTING), pickerCaptor.capture()); + + // Subchannel became READY + deliverSubchannelState(subchannel2, ConnectivityStateInfo.forNonError(READY)); + inOrder.verify(helper, atLeast(1)).updateBalancingState(eq(READY), pickerCaptor.capture()); + RoundRobinPicker pickerFinal = (RoundRobinPicker) pickerCaptor.getValue(); + assertThat(pickerFinal.dropList).containsExactly( + null, new DropEntry(getLoadRecorder(), "token0003"), null); + ChildLbPickerEntry finalEntry = (ChildLbPickerEntry) pickerFinal.pickList.get(0); + assertThat( + finalEntry.getChildPicker().pickSubchannel(mock(PickSubchannelArgs.class)).getSubchannel()) + .isEqualTo(subchannel2); + } + @Test public void shutdownWithoutSubchannel_roundRobin() throws Exception { subtestShutdownWithoutSubchannel(GrpclbConfig.create(Mode.ROUND_ROBIN)); @@ -2151,15 +2379,30 @@ private void subtestShutdownWithoutSubchannel(GrpclbConfig grpclbConfig) { @Test public void pickFirstMode_defaultTimeout_fallback() throws Exception { + PickFirstLoadBalancerProviderAccessor.setEnableNewPickFirst(false); pickFirstModeFallback(GrpclbState.FALLBACK_TIMEOUT_MS); } + @Test + public void pickFirstMode_defaultTimeout_fallback_newPickFirst() throws Exception { + PickFirstLoadBalancerProviderAccessor.setEnableNewPickFirst(true); + pickFirstModeFallback_newPickFirst(GrpclbState.FALLBACK_TIMEOUT_MS); + } + @Test public void pickFirstMode_serviceConfigTimeout_fallback() throws Exception { + PickFirstLoadBalancerProviderAccessor.setEnableNewPickFirst(false); pickFirstModeFallback(12345); } + @Test + public void pickFirstMode_serviceConfigTimeout_fallback_newPickFirst() throws Exception { + PickFirstLoadBalancerProviderAccessor.setEnableNewPickFirst(true); + pickFirstModeFallback_newPickFirst(12345); + } + private void pickFirstModeFallback(long timeout) throws Exception { + PickFirstLoadBalancerProviderAccessor.setEnableNewPickFirst(false); InOrder inOrder = inOrder(helper); // Name resolver returns balancer and backend addresses @@ -2256,8 +2499,134 @@ private void pickFirstModeFallback(long timeout) throws Exception { .returnSubchannel(any(Subchannel.class), any(ConnectivityStateInfo.class)); } + private void pickFirstModeFallback_newPickFirst(long timeout) throws Exception { + InOrder inOrder = inOrder(helper); + + // Name resolver returns balancer and backend addresses + List backendList = createResolvedBackendAddresses(2); + List grpclbBalancerList = createResolvedBalancerAddresses(1); + deliverResolvedAddresses( + backendList, grpclbBalancerList, GrpclbConfig.create(Mode.PICK_FIRST, null, timeout)); + + // Attempted to connect to balancer + assertEquals(1, fakeOobChannels.size()); + verify(mockLbService).balanceLoad(lbResponseObserverCaptor.capture()); + StreamObserver lbResponseObserver = lbResponseObserverCaptor.getValue(); + assertEquals(1, lbRequestObservers.size()); + + // Fallback timer expires with no response + fakeClock.forwardTime(timeout, TimeUnit.MILLISECONDS); + + // Entering fallback mode - child LB creates 1st subchannel for 1st backend address + inOrder.verify(helper).createSubchannel(createSubchannelArgsCaptor.capture()); + CreateSubchannelArgs createSubchannelArgs0 = createSubchannelArgsCaptor.getValue(); + assertThat(createSubchannelArgs0.getAddresses()) + .containsExactly(backendList.get(0)); + + assertThat(mockSubchannels).hasSize(1); + Subchannel subchannel0 = mockSubchannels.poll(); + + // child pick_first eagerly connects, so initial state is CONNECTING + inOrder.verify(helper, atLeast(1)).updateBalancingState(eq(CONNECTING), pickerCaptor.capture()); + RoundRobinPicker picker0 = (RoundRobinPicker) pickerCaptor.getValue(); + assertThat(picker0.pickList.get(0)).isInstanceOf(ChildLbPickerEntry.class); + + // Initial eager connection request on 1st subchannel + verify(subchannel0).requestConnection(); + + // Deliver TRANSIENT_FAILURE to subchannel0 to trigger failover to 2nd backend address + deliverSubchannelState( + subchannel0, ConnectivityStateInfo.forTransientFailure(Status.UNAVAILABLE)); + + // Verify 2nd subchannel is created for backendList.get(1) + inOrder.verify(helper).createSubchannel(createSubchannelArgsCaptor.capture()); + CreateSubchannelArgs createSubchannelArgs1 = createSubchannelArgsCaptor.getValue(); + assertThat(createSubchannelArgs1.getAddresses()) + .containsExactly(backendList.get(1)); + + assertThat(mockSubchannels).hasSize(1); + Subchannel subchannel1 = mockSubchannels.poll(); + verify(subchannel1).requestConnection(); + + // 2nd subchannel becomes READY (PickFirstLeafLoadBalancer shuts down subchannel0) + deliverSubchannelState(subchannel1, ConnectivityStateInfo.forNonError(READY)); + verify(subchannel0).shutdown(); + inOrder.verify(helper, atLeast(1)).updateBalancingState(eq(READY), pickerCaptor.capture()); + RoundRobinPicker picker1 = (RoundRobinPicker) pickerCaptor.getValue(); + assertThat(picker1.dropList).containsExactly(null, null); + ChildLbPickerEntry readyEntry = (ChildLbPickerEntry) picker1.pickList.get(0); + assertThat( + readyEntry.getChildPicker().pickSubchannel(mock(PickSubchannelArgs.class)).getSubchannel()) + .isEqualTo(subchannel1); + + // Finally, an LB response arrives, which brings us out of fallback + List backends1 = Arrays.asList( + new ServerEntry("127.0.0.1", 2000, "token0001"), + new ServerEntry("127.0.0.1", 2010, "token0002")); + inOrder.verify(helper, never()) + .updateBalancingState(any(ConnectivityState.class), any(SubchannelPicker.class)); + lbResponseObserver.onNext(buildInitialResponse()); + lbResponseObserver.onNext(buildLbResponse(backends1)); + + // In new PF, active fallback subchannel1 is shut down because its address is not in backends1 + verify(subchannel1).shutdown(); + // And LB state transitions to IDLE with RequestConnectionPicker + inOrder.verify(helper).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 subchannel2 = mockSubchannels.poll(); + verify(subchannel2).requestConnection(); + + inOrder.verify(helper, atLeast(1)) + .updateBalancingState(eq(CONNECTING), any(SubchannelPicker.class)); + + // subchannel2 becomes READY + deliverSubchannelState(subchannel2, ConnectivityStateInfo.forNonError(READY)); + inOrder.verify(helper, atLeast(1)).updateBalancingState(eq(READY), pickerCaptor.capture()); + RoundRobinPicker pickerReady = (RoundRobinPicker) pickerCaptor.getValue(); + ChildLbPickerEntry readyEntry2 = (ChildLbPickerEntry) pickerReady.pickList.get(0); + assertThat( + readyEntry2.getChildPicker().pickSubchannel(mock(PickSubchannelArgs.class)).getSubchannel()) + .isEqualTo(subchannel2); + + // Subchannel goes IDLE, grpclb follows + deliverSubchannelState(subchannel2, ConnectivityStateInfo.forNonError(IDLE)); + inOrder.verify(helper).updateBalancingState(eq(IDLE), pickerCaptor.capture()); + RoundRobinPicker pickerIdle2 = (RoundRobinPicker) pickerCaptor.getValue(); + + // Verify connection is NOT eagerly requested again (still only 1st request from start) + verify(subchannel2, times(1)).requestConnection(); + + // Picking while IDLE triggers a new connection request + PickResult pick2 = pickerIdle2.pickSubchannel(args); + assertThat(pick2.getSubchannel()).isNull(); // BUFFERing while IDLE + verify(subchannel2, times(2)).requestConnection(); + + balancer.requestConnection(); + verify(subchannel2, times(3)).requestConnection(); + + // PICK_FIRST doesn't use subchannelPool + verify(subchannelPool, never()) + .takeOrCreateSubchannel(any(EquivalentAddressGroup.class), any(Attributes.class)); + verify(subchannelPool, never()) + .returnSubchannel(any(Subchannel.class), any(ConnectivityStateInfo.class)); + } + @Test public void switchMode() throws Exception { + PickFirstLoadBalancerProviderAccessor.setEnableNewPickFirst(false); InOrder inOrder = inOrder(helper); List grpclbBalancerList = createResolvedBalancerAddresses(1); @@ -2344,12 +2713,101 @@ public void switchMode() throws Exception { .updateBalancingState(eq(CONNECTING), any(SubchannelPicker.class)); } + @Test + public void switchMode_newPickFirst() throws Exception { + PickFirstLoadBalancerProviderAccessor.setEnableNewPickFirst(true); + InOrder inOrder = inOrder(helper); + + List grpclbBalancerList = createResolvedBalancerAddresses(1); + deliverResolvedAddresses( + Collections.emptyList(), + grpclbBalancerList, + GrpclbConfig.create(Mode.ROUND_ROBIN)); + + assertEquals(1, fakeOobChannels.size()); + ManagedChannel oobChannel = fakeOobChannels.poll(); + verify(mockLbService).balanceLoad(lbResponseObserverCaptor.capture()); + StreamObserver lbResponseObserver = lbResponseObserverCaptor.getValue(); + assertEquals(1, lbRequestObservers.size()); + StreamObserver lbRequestObserver = lbRequestObservers.poll(); + verify(lbRequestObserver).onNext( + eq(LoadBalanceRequest.newBuilder().setInitialRequest( + InitialLoadBalanceRequest.newBuilder().setName(SERVICE_AUTHORITY).build()) + .build())); + + // Simulate receiving LB response + List backends1 = Arrays.asList( + new ServerEntry("127.0.0.1", 2000, "token0001"), + new ServerEntry("127.0.0.1", 2010, "token0002")); + + // RR Mode: Ensure no updates before initial response + inOrder.verify(helper, never()) + .updateBalancingState(any(ConnectivityState.class), any(SubchannelPicker.class)); + lbResponseObserver.onNext(buildInitialResponse()); + lbResponseObserver.onNext(buildLbResponse(backends1)); + + // ROUND_ROBIN: create one subchannel per server + verify(subchannelPool).takeOrCreateSubchannel( + eq(new EquivalentAddressGroup(backends1.get(0).addr, LB_BACKEND_ATTRS)), + any(Attributes.class)); + verify(subchannelPool).takeOrCreateSubchannel( + eq(new EquivalentAddressGroup(backends1.get(1).addr, LB_BACKEND_ATTRS)), + any(Attributes.class)); + inOrder.verify(helper).updateBalancingState(eq(CONNECTING), any(SubchannelPicker.class)); + assertEquals(2, mockSubchannels.size()); + Subchannel subchannel1 = mockSubchannels.poll(); + Subchannel subchannel2 = mockSubchannels.poll(); + verify(subchannelPool, never()) + .returnSubchannel(any(Subchannel.class), any(ConnectivityStateInfo.class)); + + // Switch to PICK_FIRST + deliverResolvedAddresses( + Collections.emptyList(), + grpclbBalancerList, GrpclbConfig.create(Mode.PICK_FIRST)); + + // GrpclbState will be shutdown, and a new one will be created + assertThat(oobChannel.isShutdown()).isTrue(); + verify(subchannelPool) + .returnSubchannel(same(subchannel1), eq(ConnectivityStateInfo.forNonError(IDLE))); + verify(subchannelPool) + .returnSubchannel(same(subchannel2), eq(ConnectivityStateInfo.forNonError(IDLE))); + + // A new LB stream is created + assertEquals(1, fakeOobChannels.size()); + verify(mockLbService, times(2)).balanceLoad(lbResponseObserverCaptor.capture()); + lbResponseObserver = lbResponseObserverCaptor.getValue(); + assertEquals(1, lbRequestObservers.size()); + lbRequestObserver = lbRequestObservers.poll(); + verify(lbRequestObserver).onNext( + eq(LoadBalanceRequest.newBuilder().setInitialRequest( + InitialLoadBalanceRequest.newBuilder().setName(SERVICE_AUTHORITY).build()) + .build())); + + // Simulate receiving LB response for PICK_FIRST + inOrder.verify(helper, never()) + .updateBalancingState(any(ConnectivityState.class), any(SubchannelPicker.class)); + lbResponseObserver.onNext(buildInitialResponse()); + lbResponseObserver.onNext(buildLbResponse(backends1)); + + // PICK_FIRST Subchannel: child LB creates it with only the 1st address + inOrder.verify(helper).createSubchannel(createSubchannelArgsCaptor.capture()); + CreateSubchannelArgs createSubchannelArgs = createSubchannelArgsCaptor.getValue(); + assertThat(createSubchannelArgs.getAddresses()) + .containsExactly( + new EquivalentAddressGroup(backends1.get(0).addr, eagAttrsWithToken("token0001"))); + + // Child pick_first eagerly connects, so initial state is CONNECTING (not IDLE) + inOrder.verify(helper, atLeast(1)) + .updateBalancingState(eq(CONNECTING), any(SubchannelPicker.class)); + } + private static Attributes eagAttrsWithToken(String token) { return LB_BACKEND_ATTRS.toBuilder().set(GrpclbConstants.TOKEN_ATTRIBUTE_KEY, token).build(); } @Test public void switchMode_nullLbPolicy() throws Exception { + PickFirstLoadBalancerProviderAccessor.setEnableNewPickFirst(false); InOrder inOrder = inOrder(helper); final List grpclbBalancerList = createResolvedBalancerAddresses(1); @@ -2434,6 +2892,92 @@ public void switchMode_nullLbPolicy() throws Exception { .updateBalancingState(eq(CONNECTING), any(SubchannelPicker.class)); } + @Test + public void switchMode_nullLbPolicy_newPickFirst() throws Exception { + PickFirstLoadBalancerProviderAccessor.setEnableNewPickFirst(true); + InOrder inOrder = inOrder(helper); + + final List grpclbBalancerList = createResolvedBalancerAddresses(1); + deliverResolvedAddresses( + Collections.emptyList(), + grpclbBalancerList); + + assertEquals(1, fakeOobChannels.size()); + ManagedChannel oobChannel = fakeOobChannels.poll(); + verify(mockLbService).balanceLoad(lbResponseObserverCaptor.capture()); + StreamObserver lbResponseObserver = lbResponseObserverCaptor.getValue(); + assertEquals(1, lbRequestObservers.size()); + StreamObserver lbRequestObserver = lbRequestObservers.poll(); + verify(lbRequestObserver).onNext( + eq(LoadBalanceRequest.newBuilder().setInitialRequest( + InitialLoadBalanceRequest.newBuilder().setName(SERVICE_AUTHORITY).build()) + .build())); + + // Simulate receiving LB response (Initial default mode: ROUND_ROBIN) + List backends1 = Arrays.asList( + new ServerEntry("127.0.0.1", 2000, "token0001"), + new ServerEntry("127.0.0.1", 2010, "token0002")); + inOrder.verify(helper, never()) + .updateBalancingState(any(ConnectivityState.class), any(SubchannelPicker.class)); + lbResponseObserver.onNext(buildInitialResponse()); + lbResponseObserver.onNext(buildLbResponse(backends1)); + + // ROUND_ROBIN: create one subchannel per server + verify(subchannelPool).takeOrCreateSubchannel( + eq(new EquivalentAddressGroup(backends1.get(0).addr, LB_BACKEND_ATTRS)), + any(Attributes.class)); + verify(subchannelPool).takeOrCreateSubchannel( + eq(new EquivalentAddressGroup(backends1.get(1).addr, LB_BACKEND_ATTRS)), + any(Attributes.class)); + inOrder.verify(helper).updateBalancingState(eq(CONNECTING), any(SubchannelPicker.class)); + assertEquals(2, mockSubchannels.size()); + Subchannel subchannel1 = mockSubchannels.poll(); + Subchannel subchannel2 = mockSubchannels.poll(); + verify(subchannelPool, never()) + .returnSubchannel(any(Subchannel.class), any(ConnectivityStateInfo.class)); + + // Switch to PICK_FIRST + deliverResolvedAddresses( + Collections.emptyList(), + grpclbBalancerList, + GrpclbConfig.create(Mode.PICK_FIRST)); + + // GrpclbState will be shutdown, and a new one will be created + assertThat(oobChannel.isShutdown()).isTrue(); + verify(subchannelPool) + .returnSubchannel(same(subchannel1), eq(ConnectivityStateInfo.forNonError(IDLE))); + verify(subchannelPool) + .returnSubchannel(same(subchannel2), eq(ConnectivityStateInfo.forNonError(IDLE))); + + // A new LB stream is created + assertEquals(1, fakeOobChannels.size()); + verify(mockLbService, times(2)).balanceLoad(lbResponseObserverCaptor.capture()); + lbResponseObserver = lbResponseObserverCaptor.getValue(); + assertEquals(1, lbRequestObservers.size()); + lbRequestObserver = lbRequestObservers.poll(); + verify(lbRequestObserver).onNext( + eq(LoadBalanceRequest.newBuilder().setInitialRequest( + InitialLoadBalanceRequest.newBuilder().setName(SERVICE_AUTHORITY).build()) + .build())); + + // Simulate receiving LB response for PICK_FIRST + inOrder.verify(helper, never()) + .updateBalancingState(any(ConnectivityState.class), any(SubchannelPicker.class)); + lbResponseObserver.onNext(buildInitialResponse()); + lbResponseObserver.onNext(buildLbResponse(backends1)); + + // PICK_FIRST Subchannel: child LB creates the subchannel with only the 1st address + inOrder.verify(helper).createSubchannel(createSubchannelArgsCaptor.capture()); + CreateSubchannelArgs createSubchannelArgs = createSubchannelArgsCaptor.getValue(); + assertThat(createSubchannelArgs.getAddresses()) + .containsExactly( + new EquivalentAddressGroup(backends1.get(0).addr, eagAttrsWithToken("token0001"))); + + // Child pick_first eagerly connects, so state is CONNECTING (not IDLE) + inOrder.verify(helper, atLeast(1)) + .updateBalancingState(eq(CONNECTING), any(SubchannelPicker.class)); + } + @Test public void switchServiceName() throws Exception { InOrder inOrder = inOrder(helper); From 37c510cadeeba13754085de508dcbb0e931759a5 Mon Sep 17 00:00:00 2001 From: Kannan J Date: Wed, 7 Oct 2026 06:44:15 +0000 Subject: [PATCH 3/7] grpclb: fix test assertions for new pick-first load balancer Update requestConnection invocation counts and intermediate state verifications in GrpclbLoadBalancerTest to match PickFirstLeafLoadBalancer's state tracking behavior when connecting from IDLE. TAG=agy CONV=2b34d6af-c0df-46b4-8326-2dea3641f8cc --- .../java/io/grpc/grpclb/GrpclbLoadBalancerTest.java | 13 +++++++------ 1 file changed, 7 insertions(+), 6 deletions(-) diff --git a/grpclb/src/test/java/io/grpc/grpclb/GrpclbLoadBalancerTest.java b/grpclb/src/test/java/io/grpc/grpclb/GrpclbLoadBalancerTest.java index 9ba8b4d3bf5..927616ccf51 100644 --- a/grpclb/src/test/java/io/grpc/grpclb/GrpclbLoadBalancerTest.java +++ b/grpclb/src/test/java/io/grpc/grpclb/GrpclbLoadBalancerTest.java @@ -2119,7 +2119,9 @@ public void grpclbWorking_pickFirstMode_newPickFirst() throws Exception { assertThat(pick.getSubchannel()).isNull(); verify(subchannel0, times(2)).requestConnection(); balancer.requestConnection(); - verify(subchannel0, times(3)).requestConnection(); + // PickFirstLeafLoadBalancer is already in CONNECTING state, so no additional + // requestConnection() call is made to subchannel0. + verify(subchannel0, times(2)).requestConnection(); // PICK_FIRST doesn't use subchannelPool verify(subchannelPool, never()) @@ -2571,7 +2573,7 @@ private void pickFirstModeFallback_newPickFirst(long timeout) throws Exception { // In new PF, active fallback subchannel1 is shut down because its address is not in backends1 verify(subchannel1).shutdown(); // And LB state transitions to IDLE with RequestConnectionPicker - inOrder.verify(helper).updateBalancingState(eq(IDLE), pickerCaptor.capture()); + inOrder.verify(helper, atLeast(1)).updateBalancingState(eq(IDLE), pickerCaptor.capture()); RoundRobinPicker pickerIdle = (RoundRobinPicker) pickerCaptor.getValue(); ChildLbPickerEntry idleEntry = (ChildLbPickerEntry) pickerIdle.pickList.get(0); @@ -2589,9 +2591,6 @@ private void pickFirstModeFallback_newPickFirst(long timeout) throws Exception { Subchannel subchannel2 = mockSubchannels.poll(); verify(subchannel2).requestConnection(); - inOrder.verify(helper, atLeast(1)) - .updateBalancingState(eq(CONNECTING), any(SubchannelPicker.class)); - // subchannel2 becomes READY deliverSubchannelState(subchannel2, ConnectivityStateInfo.forNonError(READY)); inOrder.verify(helper, atLeast(1)).updateBalancingState(eq(READY), pickerCaptor.capture()); @@ -2615,7 +2614,9 @@ private void pickFirstModeFallback_newPickFirst(long timeout) throws Exception { verify(subchannel2, times(2)).requestConnection(); balancer.requestConnection(); - verify(subchannel2, times(3)).requestConnection(); + // PickFirstLeafLoadBalancer is already in CONNECTING state, so no additional + // requestConnection() call is made to subchannel2. + verify(subchannel2, times(2)).requestConnection(); // PICK_FIRST doesn't use subchannelPool verify(subchannelPool, never()) From e88415f3c91e108ffa11525c5ae26b5e819779b3 Mon Sep 17 00:00:00 2001 From: Kannan J Date: Wed, 7 Oct 2026 08:34:28 +0000 Subject: [PATCH 4/7] rls: duplicate RlsLoadBalancerTest to test old and new pick-first Separate lb_working_withDefaultTarget_rlsResponding into legacy and new pick-first test variants with the appropriate InOrder createSubchannel verifications. TAG=agy CONV=2b34d6af-c0df-46b4-8326-2dea3641f8cc --- .../java/io/grpc/rls/RlsLoadBalancerTest.java | 78 +++++++++++++++++++ 1 file changed, 78 insertions(+) diff --git a/rls/src/test/java/io/grpc/rls/RlsLoadBalancerTest.java b/rls/src/test/java/io/grpc/rls/RlsLoadBalancerTest.java index 79c9a95b119..3806bd72302 100644 --- a/rls/src/test/java/io/grpc/rls/RlsLoadBalancerTest.java +++ b/rls/src/test/java/io/grpc/rls/RlsLoadBalancerTest.java @@ -73,6 +73,8 @@ import io.grpc.inprocess.InProcessServerBuilder; import io.grpc.internal.FakeClock; import io.grpc.internal.JsonParser; +import io.grpc.internal.PickFirstLoadBalancerProvider; +import io.grpc.internal.PickFirstLoadBalancerProviderAccessor; import io.grpc.internal.PickSubchannelArgsImpl; import io.grpc.internal.testing.StreamRecorder; import io.grpc.lookup.v1.RouteLookupServiceGrpc; @@ -118,6 +120,8 @@ public class RlsLoadBalancerTest { public final GrpcCleanupRule grpcCleanupRule = new GrpcCleanupRule(); @Rule public final MockitoRule mocks = MockitoJUnit.rule(); + private final boolean defaultNewPickFirst = + PickFirstLoadBalancerProvider.isEnabledNewPickFirst(); private final RlsLoadBalancerProvider provider = new RlsLoadBalancerProvider(); private final FakeClock fakeClock = new FakeClock(); private final SynchronizationContext syncContext = @@ -201,6 +205,7 @@ public CachingRlsLbClient.Builder get() { @After public void tearDown() { rlsLb.shutdown(); + PickFirstLoadBalancerProviderAccessor.setEnableNewPickFirst(defaultNewPickFirst); } @Test @@ -247,6 +252,7 @@ public void lb_serverStatusCodeConversion() throws Exception { @Test public void lb_working_withDefaultTarget_rlsResponding() throws Exception { + PickFirstLoadBalancerProviderAccessor.setEnableNewPickFirst(false); helper.getSynchronizationContext().execute(() -> { try { deliverResolvedAddresses(); @@ -315,6 +321,78 @@ public void lb_working_withDefaultTarget_rlsResponding() throws Exception { verifyLongCounterAdd("grpc.lb.rls.target_picks", 1, 1, "wilderness", "fail"); } + @Test + public void lb_working_withDefaultTarget_rlsResponding_newPickFirst() throws Exception { + PickFirstLoadBalancerProviderAccessor.setEnableNewPickFirst(true); + helper.getSynchronizationContext().execute(() -> { + try { + deliverResolvedAddresses(); + } catch (Exception e) { + throw new RuntimeException(e); + } + }); + InOrder inOrder = inOrder(helper); + inOrder.verify(helper) + .updateBalancingState(eq(ConnectivityState.CONNECTING), pickerCaptor.capture()); + inOrder.verify(helper).createSubchannel(any(CreateSubchannelArgs.class)); + SubchannelPicker picker = pickerCaptor.getValue(); + // Warm-up pick; will be queued + PickResult res = picker.pickSubchannel(searchSubchannelArgs); + assertThat(res.getStatus().isOk()).isTrue(); + assertThat(res.getSubchannel()).isNull(); + assertThat(res.getDelayType()).isEqualTo("rls_lookup_pending"); + assertThat(res.getDelayReason()).contains("Route Lookup Service query pending"); + // Cache is warm, but still unconnected + res = picker.pickSubchannel(searchSubchannelArgs); + inOrder.verify(helper).createSubchannel(any(CreateSubchannelArgs.class)); + inOrder.verify(helper, atLeast(0)) + .updateBalancingState(eq(ConnectivityState.CONNECTING), any(SubchannelPicker.class)); + inOrder.verify(helper, atLeast(0)).getSynchronizationContext(); + inOrder.verify(helper, atLeast(0)).getScheduledExecutorService(); + inOrder.verify(helper, atLeast(0)).getMetricRecorder(); + inOrder.verify(helper, atLeast(0)).getChannelTarget(); + inOrder.verifyNoMoreInteractions(); + + assertThat(res.getStatus().isOk()).isTrue(); + assertThat(subchannels).hasSize(2); // includes fallback sub-channel + FakeSubchannel searchSubchannel = subchannels.getLast(); + assertThat(subchannelIsReady(searchSubchannel)).isFalse(); + + searchSubchannel.updateState(ConnectivityStateInfo.forNonError(ConnectivityState.READY)); + inOrder.verify(helper) + .updateBalancingState(eq(ConnectivityState.READY), pickerCaptor.capture()); + inOrder.verifyNoMoreInteractions(); + res = picker.pickSubchannel(searchSubchannelArgs); + assertThat(subchannelIsReady(res.getSubchannel())).isTrue(); + assertThat(res.getSubchannel()).isSameInstanceAs(searchSubchannel); + verifyLongCounterAdd("grpc.lb.rls.target_picks", 1, 1, "wilderness", "complete"); + + // rescue should be pending status although the overall channel state is READY + res = picker.pickSubchannel(rescueSubchannelArgs); + inOrder.verify(helper).createSubchannel(any(CreateSubchannelArgs.class)); + // other rls picker itself is ready due to first channel. + assertThat(res.getStatus().isOk()).isTrue(); + assertThat(subchannelIsReady(res.getSubchannel())).isFalse(); + assertThat(subchannels).hasSize(3); // includes fallback sub-channel + FakeSubchannel rescueSubchannel = subchannels.getLast(); + + // search subchannel is down, rescue subchannel is connecting + searchSubchannel.updateState(ConnectivityStateInfo.forTransientFailure(Status.UNAVAILABLE)); + inOrder.verify(helper) + .updateBalancingState(eq(ConnectivityState.CONNECTING), pickerCaptor.capture()); + + rescueSubchannel.updateState(ConnectivityStateInfo.forNonError(ConnectivityState.READY)); + inOrder.verify(helper) + .updateBalancingState(eq(ConnectivityState.READY), pickerCaptor.capture()); + + // search again, verify that it doesn't use fallback, since RLS server responded, even though + // subchannel is in failure mode + res = picker.pickSubchannel(searchSubchannelArgs); + assertThat(res.getStatus().getCode()).isEqualTo(Status.Code.UNAVAILABLE); + assertThat(subchannelIsReady(res.getSubchannel())).isFalse(); + verifyLongCounterAdd("grpc.lb.rls.target_picks", 1, 1, "wilderness", "fail"); + } + @Test public void fallbackWithDelay_succeeds() throws Exception { fakeRlsServerImpl.setResponseDelay(100, TimeUnit.MILLISECONDS); From 474c6d4e1700c8b0afe8c339d53c618381ee08e4 Mon Sep 17 00:00:00 2001 From: Kannan J Date: Wed, 7 Oct 2026 08:42:20 +0000 Subject: [PATCH 5/7] util: duplicate RoundRobinLoadBalancerTest.pickAfterStateChange for new pick-first PickFirstLeafLoadBalancer does not trigger a second refreshNameResolution when transitioning from TRANSIENT_FAILURE to IDLE during backoff. Duplicate pickAfterStateChange to verify both old and new PF behaviors. TAG=agy CONV=2b34d6af-c0df-46b4-8326-2dea3641f8cc --- .../grpc/util/RoundRobinLoadBalancerTest.java | 37 +++++++++++++++++++ 1 file changed, 37 insertions(+) diff --git a/util/src/test/java/io/grpc/util/RoundRobinLoadBalancerTest.java b/util/src/test/java/io/grpc/util/RoundRobinLoadBalancerTest.java index c0f5bbc526f..549569f04a8 100644 --- a/util/src/test/java/io/grpc/util/RoundRobinLoadBalancerTest.java +++ b/util/src/test/java/io/grpc/util/RoundRobinLoadBalancerTest.java @@ -230,6 +230,7 @@ public void pickAfterResolvedUpdatedHosts() throws Exception { @Test public void pickAfterStateChange() throws Exception { + PickFirstLoadBalancerProviderAccessor.setEnableNewPickFirst(false); InOrder inOrder = inOrder(mockHelper); Status addressesAcceptanceStatus = acceptAddresses(Arrays.asList(servers.get(0)), Attributes.EMPTY); @@ -263,6 +264,42 @@ public void pickAfterStateChange() throws Exception { AbstractTestHelper.verifyNoMoreMeaningfulInteractions(mockHelper); } + @Test + public void pickAfterStateChange_newPickFirst() throws Exception { + PickFirstLoadBalancerProviderAccessor.setEnableNewPickFirst(true); + InOrder inOrder = inOrder(mockHelper); + Status addressesAcceptanceStatus = + acceptAddresses(Arrays.asList(servers.get(0)), Attributes.EMPTY); + assertThat(addressesAcceptanceStatus.isOk()).isTrue(); + inOrder.verify(mockHelper).createSubchannel(any(CreateSubchannelArgs.class)); + + // TODO figure out if this method testing the right things + + assertThat(subchannels).hasSize(1); + Subchannel subchannel = subchannels.values().iterator().next(); + + inOrder.verify(mockHelper).updateBalancingState(eq(CONNECTING), eq(EMPTY_PICKER)); + + deliverSubchannelState(subchannel, ConnectivityStateInfo.forNonError(READY)); + inOrder.verify(mockHelper).updateBalancingState(eq(READY), pickerCaptor.capture()); + assertThat(pickerCaptor.getValue()).isInstanceOf(ReadyPicker.class); + + Status error = Status.UNKNOWN.withDescription("¯\\_(ツ)_//¯"); + deliverSubchannelState(subchannel, + ConnectivityStateInfo.forTransientFailure(error)); + AbstractTestHelper.refreshInvokedAndUpdateBS( + inOrder, TRANSIENT_FAILURE, mockHelper, pickerCaptor); + assertThat(pickerCaptor.getValue().pickSubchannel(mockArgs).getStatus()).isEqualTo(error); + + deliverSubchannelState(subchannel, ConnectivityStateInfo.forNonError(IDLE)); + inOrder.verify(mockHelper, never()).refreshNameResolution(); + inOrder.verify(mockHelper, never()) + .updateBalancingState(eq(TRANSIENT_FAILURE), any(SubchannelPicker.class)); + + verify(subchannel, atLeastOnce()).requestConnection(); + AbstractTestHelper.verifyNoMoreMeaningfulInteractions(mockHelper); + } + @Test public void ignoreShutdownSubchannelStateChange() { InOrder inOrder = inOrder(mockHelper); From 7308226d2521903322ab8252ba87811a53993942 Mon Sep 17 00:00:00 2001 From: Kannan J Date: Wed, 7 Oct 2026 09:29:45 +0000 Subject: [PATCH 6/7] xds: Add separate test variants for new PickFirst LB in RingHash and LeastRequest Duplicated failing tests in RingHashLoadBalancerTest and LeastRequestLoadBalancerTest to test legacy PickFirst and new PickFirst behaviors separately under both flag configurations. TAG=agy CONV=2b34d6af-c0df-46b4-8326-2dea3641f8cc --- .../xds/LeastRequestLoadBalancerTest.java | 45 ++- .../io/grpc/xds/RingHashLoadBalancerTest.java | 284 +++++++++++++++++- 2 files changed, 324 insertions(+), 5 deletions(-) diff --git a/xds/src/test/java/io/grpc/xds/LeastRequestLoadBalancerTest.java b/xds/src/test/java/io/grpc/xds/LeastRequestLoadBalancerTest.java index 302faed95a4..ca6b7082ecd 100644 --- a/xds/src/test/java/io/grpc/xds/LeastRequestLoadBalancerTest.java +++ b/xds/src/test/java/io/grpc/xds/LeastRequestLoadBalancerTest.java @@ -62,6 +62,7 @@ import io.grpc.Metadata; import io.grpc.Status; import io.grpc.internal.PickFirstLoadBalancerProvider; +import io.grpc.internal.PickFirstLoadBalancerProviderAccessor; import io.grpc.util.AbstractTestHelper; import io.grpc.util.MultiChildLoadBalancer.ChildLbState; import io.grpc.xds.LeastRequestLoadBalancer.LeastRequestConfig; @@ -100,6 +101,7 @@ public class LeastRequestLoadBalancerTest { private final Map, Subchannel> subchannels = Maps.newLinkedHashMap(); private final Attributes affinity = Attributes.newBuilder().set(MAJOR_KEY, "I got the keys").build(); + private boolean defaultNewPickFirst = PickFirstLoadBalancerProvider.isEnabledNewPickFirst(); @Captor private ArgumentCaptor pickerCaptor; @@ -129,6 +131,7 @@ public void setUp() { @After public void tearDown() throws Exception { + PickFirstLoadBalancerProviderAccessor.setEnableNewPickFirst(defaultNewPickFirst); verifyNoMoreInteractions(mockRandom); verifyNoMoreInteractions(mockArgs); } @@ -231,6 +234,7 @@ private Subchannel getSubchannel(EquivalentAddressGroup removedEag) { @Test public void pickAfterStateChange() throws Exception { + PickFirstLoadBalancerProviderAccessor.setEnableNewPickFirst(false); InOrder inOrder = inOrder(helper); Status addressesAcceptanceStatus = loadBalancer.acceptResolvedAddresses( ResolvedAddresses.newBuilder().setAddresses(servers).setAttributes(Attributes.EMPTY) @@ -261,8 +265,45 @@ public void pickAfterStateChange() throws Exception { assertThat(childLbState.getCurrentState()).isEqualTo(TRANSIENT_FAILURE); assertThat(childLbState.getCurrentPicker().toString()).contains(error.toString()); - int expectedCount = PickFirstLoadBalancerProvider.isEnabledNewPickFirst() ? 1 : 2; - verify(subchannel, times(expectedCount)).requestConnection(); + verify(subchannel, times(2)).requestConnection(); + verify(helper, times(3)).createSubchannel(any(CreateSubchannelArgs.class)); + AbstractTestHelper.verifyNoMoreMeaningfulInteractions(helper); + } + + @Test + public void pickAfterStateChange_newPickFirst() throws Exception { + PickFirstLoadBalancerProviderAccessor.setEnableNewPickFirst(true); + InOrder inOrder = inOrder(helper); + Status addressesAcceptanceStatus = loadBalancer.acceptResolvedAddresses( + ResolvedAddresses.newBuilder().setAddresses(servers).setAttributes(Attributes.EMPTY) + .build()); + assertThat(addressesAcceptanceStatus.isOk()).isTrue(); + ChildLbState childLbState = loadBalancer.getChildLbStates().iterator().next(); + Subchannel subchannel = getSubchannel(servers.get(0)); + + inOrder.verify(helper) + .updateBalancingState(eq(CONNECTING), pickerReturns(PickResult.withNoResult())); + assertThat(childLbState.getCurrentState()).isEqualTo(CONNECTING); + + deliverSubchannelState(subchannel, ConnectivityStateInfo.forNonError(READY)); + inOrder.verify(helper).updateBalancingState(eq(READY), pickerCaptor.capture()); + assertThat(pickerCaptor.getValue()).isInstanceOf(ReadyPicker.class); + assertThat(childLbState.getCurrentState()).isEqualTo(READY); + + Status error = Status.UNKNOWN.withDescription("¯\\_(ツ)_//¯"); + deliverSubchannelState(subchannel, ConnectivityStateInfo.forTransientFailure(error)); + assertThat(childLbState.getCurrentState()).isEqualTo(TRANSIENT_FAILURE); + assertThat(childLbState.getCurrentPicker().toString()).contains(error.toString()); + refreshInvokedAndUpdateBS(inOrder, CONNECTING); + assertThat(pickerCaptor.getValue().pickSubchannel(mockArgs)) + .isEqualTo(PickResult.withNoResult()); + + deliverSubchannelState(subchannel, ConnectivityStateInfo.forNonError(IDLE)); + inOrder.verify(helper, never()).refreshNameResolution(); + assertThat(childLbState.getCurrentState()).isEqualTo(TRANSIENT_FAILURE); + assertThat(childLbState.getCurrentPicker().toString()).contains(error.toString()); + + verify(subchannel, times(1)).requestConnection(); verify(helper, times(3)).createSubchannel(any(CreateSubchannelArgs.class)); AbstractTestHelper.verifyNoMoreMeaningfulInteractions(helper); } diff --git a/xds/src/test/java/io/grpc/xds/RingHashLoadBalancerTest.java b/xds/src/test/java/io/grpc/xds/RingHashLoadBalancerTest.java index da11df24af4..1ced3f49e71 100644 --- a/xds/src/test/java/io/grpc/xds/RingHashLoadBalancerTest.java +++ b/xds/src/test/java/io/grpc/xds/RingHashLoadBalancerTest.java @@ -533,6 +533,7 @@ public void pickWithRandomHash_atLeastOneSubchannelConnecting() { @Test public void pickWithRandomHash_firstSubchannelInTransientFailure_remainingSubchannelsIdle() { + PickFirstLoadBalancerProviderAccessor.setEnableNewPickFirst(false); // Map each server address to exactly one ring entry. RingHashConfig config = new RingHashConfig(3, 3, "dummy-random-hash"); List servers = createWeightedServerAddrs(1, 1, 1); @@ -555,6 +556,31 @@ public void pickWithRandomHash_firstSubchannelInTransientFailure_remainingSubcha verifyConnection(1); } + @Test + public void + pickWithRandomHash_firstSubchannelInTransientFailure_remainingSubchannelsIdle_newPickFirst() { + PickFirstLoadBalancerProviderAccessor.setEnableNewPickFirst(true); + // Map each server address to exactly one ring entry. + RingHashConfig config = new RingHashConfig(3, 3, "dummy-random-hash"); + List servers = createWeightedServerAddrs(1, 1, 1); + initializeLbSubchannels(config, servers); + + // Bring one subchannel to TRANSIENT_FAILURE. + deliverSubchannelUnreachable(getSubChannel(servers.get(0))); + verify(helper).updateBalancingState(eq(CONNECTING), pickerCaptor.capture()); + verifyConnection(1); + + SubchannelPicker picker = pickerCaptor.getValue(); + PickSubchannelArgs args = getDefaultPickSubchannelArgs(hashFunc.hashVoid()); + PickResult result = picker.pickSubchannel(args); + assertThat(result.getStatus().isOk()).isTrue(); + assertThat(result.getSubchannel()).isNull(); // buffer request + assertThat(result.getDelayType()).isEqualTo("connecting"); + assertThat(result.getDelayReason()).isEqualTo("ring_hash: waiting for connection"); + int expected = connectionRequestedQueue.isEmpty() ? 0 : 1; + verifyConnection(expected); + } + private Subchannel getSubChannel(EquivalentAddressGroup eag) { return subchannels.get(Collections.singletonList(eag)); } @@ -631,6 +657,7 @@ private PickSubchannelArgs getDefaultPickSubchannelArgsForServer(int serverid) { @Test public void skipFailingHosts_firstTwoHostsFailed_pickNextFirstReady() { + PickFirstLoadBalancerProviderAccessor.setEnableNewPickFirst(false); // Map each server address to exactly one ring entry. RingHashConfig config = new RingHashConfig(3, 3, ""); List servers = createWeightedServerAddrs(1, 1, 1); @@ -658,8 +685,73 @@ public void skipFailingHosts_firstTwoHostsFailed_pickNextFirstReady() { verifyConnection(2); PickResult result = pickerCaptor.getValue().pickSubchannel(args); // activate last subchannel assertThat(result.getStatus().isOk()).isTrue(); - int expectedCount = PickFirstLoadBalancerProvider.isEnabledNewPickFirst() ? 0 : 1; - verifyConnection(expectedCount); + verifyConnection(1); + + deliverSubchannelState( + getSubchannel(servers, 0), + ConnectivityStateInfo.forTransientFailure( + Status.PERMISSION_DENIED.withDescription("permission denied again"))); + verify(helper, times(2)).updateBalancingState(eq(TRANSIENT_FAILURE), pickerCaptor.capture()); + result = pickerCaptor.getValue().pickSubchannel(args); + assertThat(result.getStatus().isOk()).isFalse(); // fail the RPC + assertThat(result.getStatus().getCode()) + .isEqualTo(Code.UNAVAILABLE); // with error status for the original server hit by hash + assertThat(result.getStatus().getDescription()).isEqualTo("unreachable"); + + // Now connecting to server1. + deliverSubchannelState(getSubchannel(servers, 1), CSI_CONNECTING); + + reset(helper); + + result = pickerCaptor.getValue().pickSubchannel(args); + assertThat(result.getStatus().isOk()).isFalse(); // fail the RPC + assertThat(result.getStatus().getCode()) + .isEqualTo(Code.UNAVAILABLE); // with error status for the original server hit by hash + assertThat(result.getStatus().getDescription()).isEqualTo("unreachable"); + + // Simulate server1 becomes READY. + deliverSubchannelState(getSubchannel(servers, 1), CSI_READY); + verify(helper).updateBalancingState(eq(READY), pickerCaptor.capture()); + + SubchannelPicker picker = pickerCaptor.getValue(); + result = picker.pickSubchannel(args); + assertThat(result.getStatus().isOk()).isTrue(); // succeed + assertThat(result.getSubchannel().getAddresses()).isEqualTo(servers.get(1)); // with server1 + assertThat(picker.pickSubchannel(getDefaultPickSubchannelArgsForServer(0))).isEqualTo(result); + assertThat(picker.pickSubchannel(getDefaultPickSubchannelArgsForServer(2))).isEqualTo(result); + } + + @Test + public void skipFailingHosts_firstTwoHostsFailed_pickNextFirstReady_newPickFirst() { + PickFirstLoadBalancerProviderAccessor.setEnableNewPickFirst(true); + // Map each server address to exactly one ring entry. + RingHashConfig config = new RingHashConfig(3, 3, ""); + List servers = createWeightedServerAddrs(1, 1, 1); + + initializeLbSubchannels(config, servers); + + // ring: + // "FakeSocketAddress-server0_0" + // "FakeSocketAddress-server1_0" + // "FakeSocketAddress-server2_0" + + long rpcHash = hashFunc.hashAsciiString("FakeSocketAddress-server1_0"); + PickSubchannelArgs args = getDefaultPickSubchannelArgs(rpcHash); + + // Bring down server0 and server2 to force trying server1. + deliverSubchannelState( + getSubchannel(servers, 1), + ConnectivityStateInfo.forTransientFailure( + Status.UNAVAILABLE.withDescription("unreachable"))); + deliverSubchannelState( + getSubchannel(servers, 2), + ConnectivityStateInfo.forTransientFailure( + Status.PERMISSION_DENIED.withDescription("permission denied"))); + verify(helper).updateBalancingState(eq(TRANSIENT_FAILURE), pickerCaptor.capture()); + verifyConnection(1); + PickResult result = pickerCaptor.getValue().pickSubchannel(args); // activate last subchannel + assertThat(result.getStatus().isOk()).isTrue(); + verifyConnection(0); deliverSubchannelState( getSubchannel(servers, 0), @@ -714,6 +806,7 @@ public void removingAddressShutdownSubchannel() { @Test public void allSubchannelsInTransientFailure() { + PickFirstLoadBalancerProviderAccessor.setEnableNewPickFirst(false); // Map each server address to exactly one ring entry. RingHashConfig config = new RingHashConfig(3, 3, ""); List servers = createWeightedServerAddrs(1, 1, 1); @@ -739,6 +832,34 @@ public void allSubchannelsInTransientFailure() { verifyConnection(0); // TF has already started taking care of this, pick doesn't need to } + @Test + public void allSubchannelsInTransientFailure_newPickFirst() { + PickFirstLoadBalancerProviderAccessor.setEnableNewPickFirst(true); + // Map each server address to exactly one ring entry. + RingHashConfig config = new RingHashConfig(3, 3, ""); + List servers = createWeightedServerAddrs(1, 1, 1); + initializeLbSubchannels(config, servers); + + // Bring all subchannels to TRANSIENT_FAILURE. + for (Subchannel subchannel : subchannels.values()) { + deliverSubchannelState(subchannel, ConnectivityStateInfo.forTransientFailure( + Status.UNAVAILABLE.withDescription( + subchannel.getAddresses().getAddresses() + " unreachable"))); + } + verify(helper, atLeastOnce()) + .updateBalancingState(eq(TRANSIENT_FAILURE), pickerCaptor.capture()); + verifyConnection(1); + + // Picking subchannel triggers connection. RPC hash hits server0. + PickSubchannelArgs args = getDefaultPickSubchannelArgsForServer(0); + PickResult result = pickerCaptor.getValue().pickSubchannel(args); + assertThat(result.getStatus().isOk()).isFalse(); + assertThat(result.getStatus().getCode()).isEqualTo(Code.UNAVAILABLE); + assertThat(result.getStatus().getDescription()) + .isEqualTo("[FakeSocketAddress-server0] unreachable"); + verifyConnection(0); // TF has already started taking care of this, pick doesn't need to + } + @Test public void firstSubchannelIdle() { // Map each server address to exactly one ring entry. @@ -787,6 +908,7 @@ private Subchannel getSubchannel(List servers, int serve @Test public void firstSubchannelFailure() { + PickFirstLoadBalancerProviderAccessor.setEnableNewPickFirst(false); // Map each server address to exactly one ring entry. RingHashConfig config = new RingHashConfig(3, 3, ""); List servers = createWeightedServerAddrs(1, 1, 1); @@ -816,8 +938,40 @@ public void firstSubchannelFailure() { verify(subchannelList.get(2), never()).requestConnection(); // Not one of the first 2 } + @Test + public void firstSubchannelFailure_newPickFirst() { + PickFirstLoadBalancerProviderAccessor.setEnableNewPickFirst(true); + // Map each server address to exactly one ring entry. + RingHashConfig config = new RingHashConfig(3, 3, ""); + List servers = createWeightedServerAddrs(1, 1, 1); + + List subchannelList = + initializeLbSubchannels(config, servers, RESET_SUBCHANNEL_MOCKS); + + // ring: + // "FakeSocketAddress-server1_0" + // "FakeSocketAddress-server0_0" + // "FakeSocketAddress-server2_0" + + deliverSubchannelState(subchannelList.get(0), + ConnectivityStateInfo.forTransientFailure( + Status.UNAVAILABLE.withDescription("unreachable"))); + verify(helper).updateBalancingState(eq(CONNECTING), pickerCaptor.capture()); + verifyConnection(1); + + PickSubchannelArgs args = getDefaultPickSubchannelArgs(hashFunc.hashVoid()); + SubchannelPicker picker1 = pickerCaptor.getValue(); + PickResult result = picker1.pickSubchannel(args); + assertThat(result.getStatus().isOk()).isTrue(); + assertThat(result.getSubchannel()).isNull(); + verify(subchannelList.get(0), never()).requestConnection(); + verify(subchannelList.get(1), never()).requestConnection(); + verify(subchannelList.get(2), never()).requestConnection(); + } + @Test public void secondSubchannelConnecting() { + PickFirstLoadBalancerProviderAccessor.setEnableNewPickFirst(false); // Map each server address to exactly one ring entry. RingHashConfig config = new RingHashConfig(3, 3, ""); List servers = createWeightedServerAddrs(1, 1, 1); @@ -847,8 +1001,41 @@ public void secondSubchannelConnecting() { verifyConnection(1); } + @Test + public void secondSubchannelConnecting_newPickFirst() { + PickFirstLoadBalancerProviderAccessor.setEnableNewPickFirst(true); + // Map each server address to exactly one ring entry. + RingHashConfig config = new RingHashConfig(3, 3, ""); + List servers = createWeightedServerAddrs(1, 1, 1); + + initializeLbSubchannels(config, servers); + + // ring: + // "FakeSocketAddress-server1_0" + // "FakeSocketAddress-server0_0" + // "FakeSocketAddress-server2_0" + + Subchannel firstSubchannel = getSubchannel(servers, 0); + deliverSubchannelUnreachable(firstSubchannel); + verifyConnection(1); + + deliverSubchannelState(getSubchannel(servers, 2), CSI_CONNECTING); + verify(helper, times(2)).updateBalancingState(eq(CONNECTING), pickerCaptor.capture()); + verifyConnection(0); + + // Picking subchannel when idle triggers connection. + deliverSubchannelState(getSubchannel(servers, 2), + ConnectivityStateInfo.forNonError(IDLE)); + verifyConnection(0); + PickSubchannelArgs args = getDefaultPickSubchannelArgs(hashFunc.hashVoid()); + PickResult result = pickerCaptor.getValue().pickSubchannel(args); + assertThat(result.getStatus().isOk()).isTrue(); + verifyConnection(0); + } + @Test public void secondSubchannelFailure() { + PickFirstLoadBalancerProviderAccessor.setEnableNewPickFirst(false); // Map each server address to exactly one ring entry. RingHashConfig config = new RingHashConfig(3, 3, ""); List servers = createWeightedServerAddrs(1, 1, 1); @@ -874,8 +1061,37 @@ public void secondSubchannelFailure() { verifyConnection(1); } + @Test + public void secondSubchannelFailure_newPickFirst() { + PickFirstLoadBalancerProviderAccessor.setEnableNewPickFirst(true); + // Map each server address to exactly one ring entry. + RingHashConfig config = new RingHashConfig(3, 3, ""); + List servers = createWeightedServerAddrs(1, 1, 1); + + initializeLbSubchannels(config, servers); + + // ring: + // "FakeSocketAddress-server1_0" + // "FakeSocketAddress-server0_0" + // "FakeSocketAddress-server2_0" + + Subchannel firstSubchannel = getSubchannel(servers, 0); + deliverSubchannelUnreachable(firstSubchannel); + deliverSubchannelUnreachable(getSubchannel(servers, 2)); + verify(helper).updateBalancingState(eq(TRANSIENT_FAILURE), pickerCaptor.capture()); + verifyConnection(1); + + // Picking subchannel triggers connection. + PickSubchannelArgs args = getDefaultPickSubchannelArgs(hashFunc.hashVoid()); + PickResult result = pickerCaptor.getValue().pickSubchannel(args); + assertThat(result.getStatus().isOk()).isTrue(); + verify(getSubchannel(servers, 1), never()).requestConnection(); + verifyConnection(0); + } + @Test public void thirdSubchannelConnecting() { + PickFirstLoadBalancerProviderAccessor.setEnableNewPickFirst(false); // Map each server address to exactly one ring entry. RingHashConfig config = new RingHashConfig(3, 3, ""); List servers = createWeightedServerAddrs(1, 1, 1); @@ -903,8 +1119,39 @@ public void thirdSubchannelConnecting() { verifyConnection(0); } + @Test + public void thirdSubchannelConnecting_newPickFirst() { + PickFirstLoadBalancerProviderAccessor.setEnableNewPickFirst(true); + // Map each server address to exactly one ring entry. + RingHashConfig config = new RingHashConfig(3, 3, ""); + List servers = createWeightedServerAddrs(1, 1, 1); + + initializeLbSubchannels(config, servers); + + // ring: + // "FakeSocketAddress-server1_0" + // "FakeSocketAddress-server0_0" + // "FakeSocketAddress-server2_0" + + Subchannel firstSubchannel = getSubchannel(servers, 0); + + deliverSubchannelUnreachable(firstSubchannel); + deliverSubchannelUnreachable(getSubchannel(servers, 2)); + deliverSubchannelState(getSubchannel(servers, 1), CSI_CONNECTING); + verify(helper, atLeastOnce()) + .updateBalancingState(eq(TRANSIENT_FAILURE), pickerCaptor.capture()); + verifyConnection(1); + + // Picking subchannel should not trigger connection per gRFC A61. + PickSubchannelArgs args = getDefaultPickSubchannelArgs(hashFunc.hashVoid()); + PickResult result = pickerCaptor.getValue().pickSubchannel(args); + assertThat(result.getStatus().isOk()).isTrue(); + verifyConnection(0); + } + @Test public void stickyTransientFailure() { + PickFirstLoadBalancerProviderAccessor.setEnableNewPickFirst(false); // Map each server address to exactly one ring entry. RingHashConfig config = new RingHashConfig(3, 3, ""); List servers = createWeightedServerAddrs(1, 1, 1); @@ -923,7 +1170,7 @@ public void stickyTransientFailure() { // Should not have called updateBalancingState on the helper again because PickFirst is // shielding the higher level from the state change. verify(helper, never()).updateBalancingState(any(), any()); - verifyConnection(PickFirstLoadBalancerProvider.isEnabledNewPickFirst() ? 0 : 1); + verifyConnection(1); // Picking subchannel triggers connection on second address. RPC hash hits server0. PickSubchannelArgs args = getDefaultPickSubchannelArgs(hashFunc.hashVoid()); @@ -933,6 +1180,37 @@ public void stickyTransientFailure() { verify(getSubchannel(servers, 2), never()).requestConnection(); } + @Test + public void stickyTransientFailure_newPickFirst() { + PickFirstLoadBalancerProviderAccessor.setEnableNewPickFirst(true); + // Map each server address to exactly one ring entry. + RingHashConfig config = new RingHashConfig(3, 3, ""); + List servers = createWeightedServerAddrs(1, 1, 1); + + initializeLbSubchannels(config, servers); + + // Bring one subchannel to TRANSIENT_FAILURE. + Subchannel firstSubchannel = getSubchannel(servers, 0); + deliverSubchannelUnreachable(firstSubchannel); + + verify(helper).updateBalancingState(eq(CONNECTING), pickerCaptor.capture()); + verifyConnection(1); + + reset(helper); + deliverSubchannelState(firstSubchannel, ConnectivityStateInfo.forNonError(IDLE)); + // Should not have called updateBalancingState on the helper again because PickFirst is + // shielding the higher level from the state change. + verify(helper, never()).updateBalancingState(any(), any()); + verifyConnection(0); + + // Picking subchannel triggers connection on second address. RPC hash hits server0. + PickSubchannelArgs args = getDefaultPickSubchannelArgs(hashFunc.hashVoid()); + PickResult result = pickerCaptor.getValue().pickSubchannel(args); + assertThat(result.getStatus().isOk()).isTrue(); + verify(getSubchannel(servers, 1), never()).requestConnection(); + verify(getSubchannel(servers, 2), never()).requestConnection(); + } + @Test public void largeWeights() { RingHashConfig config = new RingHashConfig(10000, 100000, ""); // large ring From b6c36a68044e34b7ce3e0cca14c9ce11e75f925b Mon Sep 17 00:00:00 2001 From: Kannan J Date: Fri, 9 Oct 2026 05:49:26 +0000 Subject: [PATCH 7/7] test: Use in-place conditional assertions for new PickFirst LB Refactor test suites across grpclb, rls, util, and xds to remove duplicated _newPickFirst test variants in favor of in-place conditional assertions using PickFirstLoadBalancerProvider.isEnabledNewPickFirst(). TAG=agy CONV=2b34d6af-c0df-46b4-8326-2dea3641f8cc --- .../PickFirstLoadBalancerProvider.java | 2 +- .../grpc/grpclb/GrpclbLoadBalancerTest.java | 712 +++--------------- .../java/io/grpc/rls/RlsLoadBalancerTest.java | 80 +- .../grpc/util/RoundRobinLoadBalancerTest.java | 43 +- .../xds/LeastRequestLoadBalancerTest.java | 49 +- .../io/grpc/xds/RingHashLoadBalancerTest.java | 314 +------- 6 files changed, 158 insertions(+), 1042 deletions(-) diff --git a/core/src/main/java/io/grpc/internal/PickFirstLoadBalancerProvider.java b/core/src/main/java/io/grpc/internal/PickFirstLoadBalancerProvider.java index 9a204a6d740..83b3fb7d8e6 100644 --- a/core/src/main/java/io/grpc/internal/PickFirstLoadBalancerProvider.java +++ b/core/src/main/java/io/grpc/internal/PickFirstLoadBalancerProvider.java @@ -37,7 +37,7 @@ public final class PickFirstLoadBalancerProvider extends LoadBalancerProvider { private static final String SHUFFLE_ADDRESS_LIST_KEY = "shuffleAddressList"; static boolean enableNewPickFirst = - GrpcUtil.getFlag("GRPC_EXPERIMENTAL_ENABLE_NEW_PICK_FIRST", true); + GrpcUtil.getFlag("GRPC_EXPERIMENTAL_ENABLE_NEW_PICK_FIRST", false); public static boolean isEnabledHappyEyeballs() { diff --git a/grpclb/src/test/java/io/grpc/grpclb/GrpclbLoadBalancerTest.java b/grpclb/src/test/java/io/grpc/grpclb/GrpclbLoadBalancerTest.java index 927616ccf51..fab813711c2 100644 --- a/grpclb/src/test/java/io/grpc/grpclb/GrpclbLoadBalancerTest.java +++ b/grpclb/src/test/java/io/grpc/grpclb/GrpclbLoadBalancerTest.java @@ -83,7 +83,6 @@ import io.grpc.internal.BackoffPolicy; import io.grpc.internal.FakeClock; import io.grpc.internal.PickFirstLoadBalancerProvider; -import io.grpc.internal.PickFirstLoadBalancerProviderAccessor; import io.grpc.lb.v1.ClientStats; import io.grpc.lb.v1.ClientStatsPerToken; import io.grpc.lb.v1.FallbackResponse; @@ -207,7 +206,6 @@ public void uncaughtException(Thread t, Throwable e) { private GrpclbLoadBalancer balancer; private final ArgumentCaptor createSubchannelArgsCaptor = ArgumentCaptor.forClass(CreateSubchannelArgs.class); - private final boolean defaultNewPickFirst = PickFirstLoadBalancerProvider.isEnabledNewPickFirst(); @Before public void setUp() throws Exception { @@ -267,7 +265,6 @@ public void run() { // No timer should linger after shutdown assertThat(fakeClock.getPendingTasks()).isEmpty(); } finally { - PickFirstLoadBalancerProviderAccessor.setEnableNewPickFirst(defaultNewPickFirst); if (fakeLbServer != null) { fakeLbServer.shutdownNow(); } @@ -1894,7 +1891,6 @@ public void grpclbBalancerStreamClosedAndRetried() throws Exception { @Test public void grpclbWorking_pickFirstMode() throws Exception { - PickFirstLoadBalancerProviderAccessor.setEnableNewPickFirst(false); InOrder inOrder = inOrder(helper); List grpclbBalancerList = createResolvedBalancerAddresses(1); @@ -1926,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()); @@ -1944,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)); @@ -1982,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 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( @@ -2009,119 +2017,8 @@ public void grpclbWorking_pickFirstMode() throws Exception { assertThat(pick.getSubchannel()).isNull(); verify(subchannel, times(2)).requestConnection(); balancer.requestConnection(); - verify(subchannel, times(3)).requestConnection(); - - // PICK_FIRST doesn't use subchannelPool - verify(subchannelPool, never()) - .takeOrCreateSubchannel(any(EquivalentAddressGroup.class), any(Attributes.class)); - verify(subchannelPool, never()) - .returnSubchannel(any(Subchannel.class), any(ConnectivityStateInfo.class)); - } - - @Test - public void grpclbWorking_pickFirstMode_newPickFirst() throws Exception { - PickFirstLoadBalancerProviderAccessor.setEnableNewPickFirst(true); - InOrder inOrder = inOrder(helper); - - List grpclbBalancerList = createResolvedBalancerAddresses(1); - - deliverResolvedAddresses( - Collections.emptyList(), - grpclbBalancerList, - GrpclbConfig.create(Mode.PICK_FIRST)); - - assertEquals(1, fakeOobChannels.size()); - verify(mockLbService).balanceLoad(lbResponseObserverCaptor.capture()); - StreamObserver lbResponseObserver = lbResponseObserverCaptor.getValue(); - assertEquals(1, lbRequestObservers.size()); - StreamObserver lbRequestObserver = lbRequestObservers.poll(); - verify(lbRequestObserver).onNext( - eq(LoadBalanceRequest.newBuilder().setInitialRequest( - InitialLoadBalanceRequest.newBuilder().setName(SERVICE_AUTHORITY).build()) - .build())); - - // Simulate receiving LB response - List backends1 = Arrays.asList( - new ServerEntry("127.0.0.1", 2000, "token0001"), - new ServerEntry("127.0.0.1", 2010, "token0002")); - inOrder.verify(helper, never()) - .updateBalancingState(any(ConnectivityState.class), any(SubchannelPicker.class)); - lbResponseObserver.onNext(buildInitialResponse()); - lbResponseObserver.onNext(buildLbResponse(backends1)); - - // With delegation, the child pick_first creates the subchannel for the 1st address - inOrder.verify(helper).createSubchannel(createSubchannelArgsCaptor.capture()); - CreateSubchannelArgs createSubchannelArgs0 = createSubchannelArgsCaptor.getValue(); - assertThat(createSubchannelArgs0.getAddresses()) - .containsExactly( - new EquivalentAddressGroup(backends1.get(0).addr, eagAttrsWithToken("token0001"))); - - // Child pick_first eagerly connects, so we start in CONNECTING - inOrder.verify(helper, atLeast(1)).updateBalancingState(eq(CONNECTING), pickerCaptor.capture()); - RoundRobinPicker picker0 = (RoundRobinPicker) pickerCaptor.getValue(); - // Only one subchannel is created by the child LB initially - assertThat(mockSubchannels).hasSize(1); - Subchannel subchannel0 = mockSubchannels.poll(); - assertThat(picker0.dropList).containsExactly(null, null); - assertThat(picker0.pickList).hasSize(1); - assertThat(picker0.pickList.get(0)).isInstanceOf(ChildLbPickerEntry.class); - - // Child pick_first eagerly calls requestConnection() - verify(subchannel0).requestConnection(); - - // READY on subchannel0 - deliverSubchannelState(subchannel0, ConnectivityStateInfo.forNonError(READY)); - inOrder.verify(helper, atLeast(1)).updateBalancingState(eq(READY), pickerCaptor.capture()); - RoundRobinPicker picker1 = (RoundRobinPicker) pickerCaptor.getValue(); - assertThat(picker1.dropList).containsExactly(null, null); - ChildLbPickerEntry readyEntry = (ChildLbPickerEntry) picker1.pickList.get(0); - PickResult readyResult = - readyEntry.getChildPicker().pickSubchannel(mock(PickSubchannelArgs.class)); - assertThat(readyResult.getSubchannel()).isEqualTo(subchannel0); - - // New server list with drops - List backends2 = Arrays.asList( - new ServerEntry("127.0.0.1", 2000, "token0001"), - new ServerEntry("token0003"), // drop - new ServerEntry("127.0.0.1", 2020, "token0004")); - inOrder.verify(helper, never()) - .updateBalancingState(any(ConnectivityState.class), any(SubchannelPicker.class)); - lbResponseObserver.onNext(buildLbResponse(backends2)); - - // Verify child LB updates subchannel0 in-place with backends2.get(0) attributes - inOrder.verify(helper, never()).createSubchannel(any(CreateSubchannelArgs.class)); - verify(helper, times(1)).createSubchannel(any(CreateSubchannelArgs.class)); - assertThat(mockSubchannels).isEmpty(); - - // The child LB policy calls updateAddresses on subchannel0 with ONLY the current address - verify(subchannel0).updateAddresses( - eq(Collections.singletonList( - new EquivalentAddressGroup(backends2.get(0).addr, eagAttrsWithToken("token0001"))))); - inOrder.verify(helper, atLeast(1)).updateBalancingState(eq(READY), pickerCaptor.capture()); - RoundRobinPicker picker2 = (RoundRobinPicker) pickerCaptor.getValue(); - assertThat(picker2.dropList).containsExactly( - null, new DropEntry(getLoadRecorder(), "token0003"), null); - ChildLbPickerEntry updatedEntry = (ChildLbPickerEntry) picker2.pickList.get(0); - PickResult updatedResult = - updatedEntry.getChildPicker().pickSubchannel(mock(PickSubchannelArgs.class)); - assertThat(updatedResult.getSubchannel()).isEqualTo(subchannel0); - - // Subchannel goes IDLE, grpclb state should follow - deliverSubchannelState(subchannel0, ConnectivityStateInfo.forNonError(IDLE)); - inOrder.verify(helper).updateBalancingState(eq(IDLE), pickerCaptor.capture()); - RoundRobinPicker picker3 = (RoundRobinPicker) pickerCaptor.getValue(); - - // No new connection request should have happened yet (beyond the first eager one) - verify(subchannel0, times(1)).requestConnection(); - PickSubchannelArgs args = mock(PickSubchannelArgs.class); - PickResult pick = picker3.pickSubchannel(args); - // Child pick_first picker returns withNoResult() when IDLE and requests connection - assertThat(pick.getSubchannel()).isNull(); - verify(subchannel0, times(2)).requestConnection(); - balancer.requestConnection(); - // PickFirstLeafLoadBalancer is already in CONNECTING state, so no additional - // requestConnection() call is made to subchannel0. - verify(subchannel0, times(2)).requestConnection(); + int expectedRequests = PickFirstLoadBalancerProvider.isEnabledNewPickFirst() ? 2 : 3; + verify(subchannel, times(expectedRequests)).requestConnection(); // PICK_FIRST doesn't use subchannelPool verify(subchannelPool, never()) @@ -2132,7 +2029,6 @@ public void grpclbWorking_pickFirstMode_newPickFirst() throws Exception { @Test public void grpclbWorking_pickFirstMode_lbSendsEmptyAddress() throws Exception { - PickFirstLoadBalancerProviderAccessor.setEnableNewPickFirst(false); InOrder inOrder = inOrder(helper); List grpclbBalancerList = createResolvedBalancerAddresses(1); @@ -2163,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()); @@ -2222,119 +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)); - assertThat(mockSubchannels).hasSize(1); - Subchannel subchannel2 = mockSubchannels.poll(); - inOrder.verify(helper, atLeast(1)).updateBalancingState(eq(CONNECTING), pickerCaptor.capture()); - - // Subchannel became READY - deliverSubchannelState(subchannel2, ConnectivityStateInfo.forNonError(READY)); - inOrder.verify(helper, atLeast(1)).updateBalancingState(eq(READY), pickerCaptor.capture()); - RoundRobinPicker pickerFinal = (RoundRobinPicker) pickerCaptor.getValue(); - assertThat(pickerFinal.dropList).containsExactly( - null, new DropEntry(getLoadRecorder(), "token0003"), null); - ChildLbPickerEntry finalEntry = (ChildLbPickerEntry) pickerFinal.pickList.get(0); - assertThat( - finalEntry.getChildPicker().pickSubchannel(mock(PickSubchannelArgs.class)).getSubchannel()) - .isEqualTo(subchannel2); - } - - @Test - public void grpclbWorking_pickFirstMode_lbSendsEmptyAddress_newPickFirst() throws Exception { - PickFirstLoadBalancerProviderAccessor.setEnableNewPickFirst(true); - InOrder inOrder = inOrder(helper); - - List grpclbBalancerList = createResolvedBalancerAddresses(1); - deliverResolvedAddresses( - Collections.emptyList(), - grpclbBalancerList, - GrpclbConfig.create(Mode.PICK_FIRST)); - - assertEquals(1, fakeOobChannels.size()); - verify(mockLbService).balanceLoad(lbResponseObserverCaptor.capture()); - StreamObserver lbResponseObserver = lbResponseObserverCaptor.getValue(); - assertEquals(1, lbRequestObservers.size()); - StreamObserver lbRequestObserver = lbRequestObservers.poll(); - verify(lbRequestObserver).onNext( - eq(LoadBalanceRequest.newBuilder().setInitialRequest( - InitialLoadBalanceRequest.newBuilder().setName(SERVICE_AUTHORITY).build()) - .build())); - - // Simulate receiving LB response - List backends1 = Arrays.asList( - new ServerEntry("127.0.0.1", 2000, "token0001"), - new ServerEntry("127.0.0.1", 2010, "token0002")); - inOrder.verify(helper, never()) - .updateBalancingState(any(ConnectivityState.class), any(SubchannelPicker.class)); - lbResponseObserver.onNext(buildInitialResponse()); - lbResponseObserver.onNext(buildLbResponse(backends1)); - - // The child pick_first creates the first subchannel with only the first address - inOrder.verify(helper).createSubchannel(createSubchannelArgsCaptor.capture()); - CreateSubchannelArgs createSubchannelArgs = createSubchannelArgsCaptor.getValue(); - assertThat(createSubchannelArgs.getAddresses()) - .containsExactly( - new EquivalentAddressGroup(backends1.get(0).addr, eagAttrsWithToken("token0001"))); - - // Child pick_first eagerly connects, so initial state is CONNECTING - inOrder.verify(helper, atLeast(1)).updateBalancingState(eq(CONNECTING), pickerCaptor.capture()); - RoundRobinPicker picker0 = (RoundRobinPicker) pickerCaptor.getValue(); - // Verify subchannel creation by child LB - assertThat(mockSubchannels).hasSize(1); - Subchannel subchannel = mockSubchannels.poll(); - assertThat(picker0.dropList).containsExactly(null, null); - assertThat(picker0.pickList).hasSize(1); - assertThat(picker0.pickList.get(0)).isInstanceOf(ChildLbPickerEntry.class); - - // Child pick_first eagerly calls requestConnection() - verify(subchannel).requestConnection(); - - // READY - deliverSubchannelState(subchannel, ConnectivityStateInfo.forNonError(READY)); - inOrder.verify(helper, atLeast(1)).updateBalancingState(eq(READY), pickerCaptor.capture()); - RoundRobinPicker pickerReady = (RoundRobinPicker) pickerCaptor.getValue(); - // Verify the subchannel in the delegated picker - ChildLbPickerEntry readyEntry = (ChildLbPickerEntry) pickerReady.pickList.get(0); - assertThat( - readyEntry.getChildPicker().pickSubchannel(mock(PickSubchannelArgs.class)).getSubchannel()) - .isEqualTo(subchannel); - inOrder.verify(helper, never()) - .updateBalancingState(any(ConnectivityState.class), any(SubchannelPicker.class)); - - // Empty addresses from LB - child LB is shutdown - lbResponseObserver.onNext(buildLbResponse(Collections.emptyList())); - - // Child LB is shutdown (which shuts down its subchannel) - // createSubchannel() has ever been called only once - inOrder.verify(helper, never()).createSubchannel(any(CreateSubchannelArgs.class)); - assertThat(mockSubchannels).isEmpty(); - verify(subchannel).shutdown(); - - // RPC error status includes message of no backends provided by balancer - inOrder.verify(helper, atLeast(1)) - .updateBalancingState(eq(TRANSIENT_FAILURE), pickerCaptor.capture()); - RoundRobinPicker errorPicker = (RoundRobinPicker) pickerCaptor.getValue(); - assertThat(errorPicker.pickList) - .containsExactly(new ErrorEntry(GrpclbState.NO_AVAILABLE_BACKENDS_STATUS)); - - lbResponseObserver.onNext(buildLbResponse(Collections.emptyList())); - - // Test recover from new LB response with addresses - // New server list with drops - List backends2 = Arrays.asList( - new ServerEntry("127.0.0.1", 2000, "token0001"), - new ServerEntry("token0003"), // drop - new ServerEntry("127.0.0.1", 2020, "token0004")); - inOrder.verify(helper, never()) - .updateBalancingState(any(ConnectivityState.class), any(SubchannelPicker.class)); - lbResponseObserver.onNext(buildLbResponse(backends2)); - - // A NEW child LB and NEW subchannel are created upon recovery with only backends2.get(0) inOrder.verify(helper).createSubchannel(createSubchannelArgsCaptor.capture()); CreateSubchannelArgs createSubchannelArgs2 = createSubchannelArgsCaptor.getValue(); - assertThat(createSubchannelArgs2.getAddresses()) - .containsExactly( - new EquivalentAddressGroup(backends2.get(0).addr, eagAttrsWithToken("token0001"))); + 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()); @@ -2381,30 +2182,15 @@ private void subtestShutdownWithoutSubchannel(GrpclbConfig grpclbConfig) { @Test public void pickFirstMode_defaultTimeout_fallback() throws Exception { - PickFirstLoadBalancerProviderAccessor.setEnableNewPickFirst(false); pickFirstModeFallback(GrpclbState.FALLBACK_TIMEOUT_MS); } - @Test - public void pickFirstMode_defaultTimeout_fallback_newPickFirst() throws Exception { - PickFirstLoadBalancerProviderAccessor.setEnableNewPickFirst(true); - pickFirstModeFallback_newPickFirst(GrpclbState.FALLBACK_TIMEOUT_MS); - } - @Test public void pickFirstMode_serviceConfigTimeout_fallback() throws Exception { - PickFirstLoadBalancerProviderAccessor.setEnableNewPickFirst(false); pickFirstModeFallback(12345); } - @Test - public void pickFirstMode_serviceConfigTimeout_fallback_newPickFirst() throws Exception { - PickFirstLoadBalancerProviderAccessor.setEnableNewPickFirst(true); - pickFirstModeFallback_newPickFirst(12345); - } - private void pickFirstModeFallback(long timeout) throws Exception { - PickFirstLoadBalancerProviderAccessor.setEnableNewPickFirst(false); InOrder inOrder = inOrder(helper); // Name resolver returns balancer and backend addresses @@ -2425,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(); @@ -2457,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); @@ -2492,131 +2311,8 @@ private void pickFirstModeFallback(long timeout) throws Exception { verify(subchannel, times(2)).requestConnection(); balancer.requestConnection(); - verify(subchannel, times(3)).requestConnection(); - - // PICK_FIRST doesn't use subchannelPool - verify(subchannelPool, never()) - .takeOrCreateSubchannel(any(EquivalentAddressGroup.class), any(Attributes.class)); - verify(subchannelPool, never()) - .returnSubchannel(any(Subchannel.class), any(ConnectivityStateInfo.class)); - } - - private void pickFirstModeFallback_newPickFirst(long timeout) throws Exception { - InOrder inOrder = inOrder(helper); - - // Name resolver returns balancer and backend addresses - List backendList = createResolvedBackendAddresses(2); - List grpclbBalancerList = createResolvedBalancerAddresses(1); - deliverResolvedAddresses( - backendList, grpclbBalancerList, GrpclbConfig.create(Mode.PICK_FIRST, null, timeout)); - - // Attempted to connect to balancer - assertEquals(1, fakeOobChannels.size()); - verify(mockLbService).balanceLoad(lbResponseObserverCaptor.capture()); - StreamObserver lbResponseObserver = lbResponseObserverCaptor.getValue(); - assertEquals(1, lbRequestObservers.size()); - - // Fallback timer expires with no response - fakeClock.forwardTime(timeout, TimeUnit.MILLISECONDS); - - // Entering fallback mode - child LB creates 1st subchannel for 1st backend address - inOrder.verify(helper).createSubchannel(createSubchannelArgsCaptor.capture()); - CreateSubchannelArgs createSubchannelArgs0 = createSubchannelArgsCaptor.getValue(); - assertThat(createSubchannelArgs0.getAddresses()) - .containsExactly(backendList.get(0)); - - assertThat(mockSubchannels).hasSize(1); - Subchannel subchannel0 = mockSubchannels.poll(); - - // child pick_first eagerly connects, so initial state is CONNECTING - inOrder.verify(helper, atLeast(1)).updateBalancingState(eq(CONNECTING), pickerCaptor.capture()); - RoundRobinPicker picker0 = (RoundRobinPicker) pickerCaptor.getValue(); - assertThat(picker0.pickList.get(0)).isInstanceOf(ChildLbPickerEntry.class); - - // Initial eager connection request on 1st subchannel - verify(subchannel0).requestConnection(); - - // Deliver TRANSIENT_FAILURE to subchannel0 to trigger failover to 2nd backend address - deliverSubchannelState( - subchannel0, ConnectivityStateInfo.forTransientFailure(Status.UNAVAILABLE)); - - // Verify 2nd subchannel is created for backendList.get(1) - inOrder.verify(helper).createSubchannel(createSubchannelArgsCaptor.capture()); - CreateSubchannelArgs createSubchannelArgs1 = createSubchannelArgsCaptor.getValue(); - assertThat(createSubchannelArgs1.getAddresses()) - .containsExactly(backendList.get(1)); - - assertThat(mockSubchannels).hasSize(1); - Subchannel subchannel1 = mockSubchannels.poll(); - verify(subchannel1).requestConnection(); - - // 2nd subchannel becomes READY (PickFirstLeafLoadBalancer shuts down subchannel0) - deliverSubchannelState(subchannel1, ConnectivityStateInfo.forNonError(READY)); - verify(subchannel0).shutdown(); - inOrder.verify(helper, atLeast(1)).updateBalancingState(eq(READY), pickerCaptor.capture()); - RoundRobinPicker picker1 = (RoundRobinPicker) pickerCaptor.getValue(); - assertThat(picker1.dropList).containsExactly(null, null); - ChildLbPickerEntry readyEntry = (ChildLbPickerEntry) picker1.pickList.get(0); - assertThat( - readyEntry.getChildPicker().pickSubchannel(mock(PickSubchannelArgs.class)).getSubchannel()) - .isEqualTo(subchannel1); - - // Finally, an LB response arrives, which brings us out of fallback - List backends1 = Arrays.asList( - new ServerEntry("127.0.0.1", 2000, "token0001"), - new ServerEntry("127.0.0.1", 2010, "token0002")); - inOrder.verify(helper, never()) - .updateBalancingState(any(ConnectivityState.class), any(SubchannelPicker.class)); - lbResponseObserver.onNext(buildInitialResponse()); - lbResponseObserver.onNext(buildLbResponse(backends1)); - - // In new PF, active fallback subchannel1 is shut down because its address is not in backends1 - verify(subchannel1).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 subchannel2 = mockSubchannels.poll(); - verify(subchannel2).requestConnection(); - - // subchannel2 becomes READY - deliverSubchannelState(subchannel2, ConnectivityStateInfo.forNonError(READY)); - inOrder.verify(helper, atLeast(1)).updateBalancingState(eq(READY), pickerCaptor.capture()); - RoundRobinPicker pickerReady = (RoundRobinPicker) pickerCaptor.getValue(); - ChildLbPickerEntry readyEntry2 = (ChildLbPickerEntry) pickerReady.pickList.get(0); - assertThat( - readyEntry2.getChildPicker().pickSubchannel(mock(PickSubchannelArgs.class)).getSubchannel()) - .isEqualTo(subchannel2); - - // Subchannel goes IDLE, grpclb follows - deliverSubchannelState(subchannel2, ConnectivityStateInfo.forNonError(IDLE)); - inOrder.verify(helper).updateBalancingState(eq(IDLE), pickerCaptor.capture()); - RoundRobinPicker pickerIdle2 = (RoundRobinPicker) pickerCaptor.getValue(); - - // Verify connection is NOT eagerly requested again (still only 1st request from start) - verify(subchannel2, times(1)).requestConnection(); - - // Picking while IDLE triggers a new connection request - PickResult pick2 = pickerIdle2.pickSubchannel(args); - assertThat(pick2.getSubchannel()).isNull(); // BUFFERing while IDLE - verify(subchannel2, times(2)).requestConnection(); - - balancer.requestConnection(); - // PickFirstLeafLoadBalancer is already in CONNECTING state, so no additional - // requestConnection() call is made to subchannel2. - verify(subchannel2, times(2)).requestConnection(); + int expectedRequests = PickFirstLoadBalancerProvider.isEnabledNewPickFirst() ? 2 : 3; + verify(subchannel, times(expectedRequests)).requestConnection(); // PICK_FIRST doesn't use subchannelPool verify(subchannelPool, never()) @@ -2627,7 +2323,6 @@ private void pickFirstModeFallback_newPickFirst(long timeout) throws Exception { @Test public void switchMode() throws Exception { - PickFirstLoadBalancerProviderAccessor.setEnableNewPickFirst(false); InOrder inOrder = inOrder(helper); List grpclbBalancerList = createResolvedBalancerAddresses(1); @@ -2704,98 +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"))); - - // Child pick_first eagerly connects, so initial state is CONNECTING (not IDLE) - inOrder.verify(helper, atLeast(1)) - .updateBalancingState(eq(CONNECTING), any(SubchannelPicker.class)); - } - - @Test - public void switchMode_newPickFirst() throws Exception { - PickFirstLoadBalancerProviderAccessor.setEnableNewPickFirst(true); - InOrder inOrder = inOrder(helper); - - List grpclbBalancerList = createResolvedBalancerAddresses(1); - deliverResolvedAddresses( - Collections.emptyList(), - grpclbBalancerList, - GrpclbConfig.create(Mode.ROUND_ROBIN)); - - assertEquals(1, fakeOobChannels.size()); - ManagedChannel oobChannel = fakeOobChannels.poll(); - verify(mockLbService).balanceLoad(lbResponseObserverCaptor.capture()); - StreamObserver lbResponseObserver = lbResponseObserverCaptor.getValue(); - assertEquals(1, lbRequestObservers.size()); - StreamObserver lbRequestObserver = lbRequestObservers.poll(); - verify(lbRequestObserver).onNext( - eq(LoadBalanceRequest.newBuilder().setInitialRequest( - InitialLoadBalanceRequest.newBuilder().setName(SERVICE_AUTHORITY).build()) - .build())); - - // Simulate receiving LB response - List backends1 = Arrays.asList( - new ServerEntry("127.0.0.1", 2000, "token0001"), - new ServerEntry("127.0.0.1", 2010, "token0002")); - - // RR Mode: Ensure no updates before initial response - inOrder.verify(helper, never()) - .updateBalancingState(any(ConnectivityState.class), any(SubchannelPicker.class)); - lbResponseObserver.onNext(buildInitialResponse()); - lbResponseObserver.onNext(buildLbResponse(backends1)); - - // ROUND_ROBIN: create one subchannel per server - verify(subchannelPool).takeOrCreateSubchannel( - eq(new EquivalentAddressGroup(backends1.get(0).addr, LB_BACKEND_ATTRS)), - any(Attributes.class)); - verify(subchannelPool).takeOrCreateSubchannel( - eq(new EquivalentAddressGroup(backends1.get(1).addr, LB_BACKEND_ATTRS)), - any(Attributes.class)); - inOrder.verify(helper).updateBalancingState(eq(CONNECTING), any(SubchannelPicker.class)); - assertEquals(2, mockSubchannels.size()); - Subchannel subchannel1 = mockSubchannels.poll(); - Subchannel subchannel2 = mockSubchannels.poll(); - verify(subchannelPool, never()) - .returnSubchannel(any(Subchannel.class), any(ConnectivityStateInfo.class)); - - // Switch to PICK_FIRST - deliverResolvedAddresses( - Collections.emptyList(), - grpclbBalancerList, GrpclbConfig.create(Mode.PICK_FIRST)); - - // GrpclbState will be shutdown, and a new one will be created - assertThat(oobChannel.isShutdown()).isTrue(); - verify(subchannelPool) - .returnSubchannel(same(subchannel1), eq(ConnectivityStateInfo.forNonError(IDLE))); - verify(subchannelPool) - .returnSubchannel(same(subchannel2), eq(ConnectivityStateInfo.forNonError(IDLE))); - - // A new LB stream is created - assertEquals(1, fakeOobChannels.size()); - verify(mockLbService, times(2)).balanceLoad(lbResponseObserverCaptor.capture()); - lbResponseObserver = lbResponseObserverCaptor.getValue(); - assertEquals(1, lbRequestObservers.size()); - lbRequestObserver = lbRequestObservers.poll(); - verify(lbRequestObserver).onNext( - eq(LoadBalanceRequest.newBuilder().setInitialRequest( - InitialLoadBalanceRequest.newBuilder().setName(SERVICE_AUTHORITY).build()) - .build())); - - // Simulate receiving LB response for PICK_FIRST - inOrder.verify(helper, never()) - .updateBalancingState(any(ConnectivityState.class), any(SubchannelPicker.class)); - lbResponseObserver.onNext(buildInitialResponse()); - lbResponseObserver.onNext(buildLbResponse(backends1)); - - // PICK_FIRST Subchannel: child LB creates it with only the 1st address - inOrder.verify(helper).createSubchannel(createSubchannelArgsCaptor.capture()); - CreateSubchannelArgs createSubchannelArgs = createSubchannelArgsCaptor.getValue(); - assertThat(createSubchannelArgs.getAddresses()) - .containsExactly( - new EquivalentAddressGroup(backends1.get(0).addr, eagAttrsWithToken("token0001"))); + 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)) @@ -2808,7 +2421,6 @@ private static Attributes eagAttrsWithToken(String token) { @Test public void switchMode_nullLbPolicy() throws Exception { - PickFirstLoadBalancerProviderAccessor.setEnableNewPickFirst(false); InOrder inOrder = inOrder(helper); final List grpclbBalancerList = createResolvedBalancerAddresses(1); @@ -2883,96 +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"))); - - // Child pick_first eagerly connects, so state is CONNECTING (not IDLE) - inOrder.verify(helper, atLeast(1)) - .updateBalancingState(eq(CONNECTING), any(SubchannelPicker.class)); - } - - @Test - public void switchMode_nullLbPolicy_newPickFirst() throws Exception { - PickFirstLoadBalancerProviderAccessor.setEnableNewPickFirst(true); - InOrder inOrder = inOrder(helper); - - final List grpclbBalancerList = createResolvedBalancerAddresses(1); - deliverResolvedAddresses( - Collections.emptyList(), - grpclbBalancerList); - - assertEquals(1, fakeOobChannels.size()); - ManagedChannel oobChannel = fakeOobChannels.poll(); - verify(mockLbService).balanceLoad(lbResponseObserverCaptor.capture()); - StreamObserver lbResponseObserver = lbResponseObserverCaptor.getValue(); - assertEquals(1, lbRequestObservers.size()); - StreamObserver lbRequestObserver = lbRequestObservers.poll(); - verify(lbRequestObserver).onNext( - eq(LoadBalanceRequest.newBuilder().setInitialRequest( - InitialLoadBalanceRequest.newBuilder().setName(SERVICE_AUTHORITY).build()) - .build())); - - // Simulate receiving LB response (Initial default mode: ROUND_ROBIN) - List backends1 = Arrays.asList( - new ServerEntry("127.0.0.1", 2000, "token0001"), - new ServerEntry("127.0.0.1", 2010, "token0002")); - inOrder.verify(helper, never()) - .updateBalancingState(any(ConnectivityState.class), any(SubchannelPicker.class)); - lbResponseObserver.onNext(buildInitialResponse()); - lbResponseObserver.onNext(buildLbResponse(backends1)); - - // ROUND_ROBIN: create one subchannel per server - verify(subchannelPool).takeOrCreateSubchannel( - eq(new EquivalentAddressGroup(backends1.get(0).addr, LB_BACKEND_ATTRS)), - any(Attributes.class)); - verify(subchannelPool).takeOrCreateSubchannel( - eq(new EquivalentAddressGroup(backends1.get(1).addr, LB_BACKEND_ATTRS)), - any(Attributes.class)); - inOrder.verify(helper).updateBalancingState(eq(CONNECTING), any(SubchannelPicker.class)); - assertEquals(2, mockSubchannels.size()); - Subchannel subchannel1 = mockSubchannels.poll(); - Subchannel subchannel2 = mockSubchannels.poll(); - verify(subchannelPool, never()) - .returnSubchannel(any(Subchannel.class), any(ConnectivityStateInfo.class)); - - // Switch to PICK_FIRST - deliverResolvedAddresses( - Collections.emptyList(), - grpclbBalancerList, - GrpclbConfig.create(Mode.PICK_FIRST)); - - // GrpclbState will be shutdown, and a new one will be created - assertThat(oobChannel.isShutdown()).isTrue(); - verify(subchannelPool) - .returnSubchannel(same(subchannel1), eq(ConnectivityStateInfo.forNonError(IDLE))); - verify(subchannelPool) - .returnSubchannel(same(subchannel2), eq(ConnectivityStateInfo.forNonError(IDLE))); - - // A new LB stream is created - assertEquals(1, fakeOobChannels.size()); - verify(mockLbService, times(2)).balanceLoad(lbResponseObserverCaptor.capture()); - lbResponseObserver = lbResponseObserverCaptor.getValue(); - assertEquals(1, lbRequestObservers.size()); - lbRequestObserver = lbRequestObservers.poll(); - verify(lbRequestObserver).onNext( - eq(LoadBalanceRequest.newBuilder().setInitialRequest( - InitialLoadBalanceRequest.newBuilder().setName(SERVICE_AUTHORITY).build()) - .build())); - - // Simulate receiving LB response for PICK_FIRST - inOrder.verify(helper, never()) - .updateBalancingState(any(ConnectivityState.class), any(SubchannelPicker.class)); - lbResponseObserver.onNext(buildInitialResponse()); - lbResponseObserver.onNext(buildLbResponse(backends1)); - - // PICK_FIRST Subchannel: child LB creates the subchannel with only the 1st address - inOrder.verify(helper).createSubchannel(createSubchannelArgsCaptor.capture()); - CreateSubchannelArgs createSubchannelArgs = createSubchannelArgsCaptor.getValue(); - assertThat(createSubchannelArgs.getAddresses()) - .containsExactly( - new EquivalentAddressGroup(backends1.get(0).addr, eagAttrsWithToken("token0001"))); + 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)) diff --git a/rls/src/test/java/io/grpc/rls/RlsLoadBalancerTest.java b/rls/src/test/java/io/grpc/rls/RlsLoadBalancerTest.java index 3806bd72302..714f23559d8 100644 --- a/rls/src/test/java/io/grpc/rls/RlsLoadBalancerTest.java +++ b/rls/src/test/java/io/grpc/rls/RlsLoadBalancerTest.java @@ -74,7 +74,6 @@ import io.grpc.internal.FakeClock; import io.grpc.internal.JsonParser; import io.grpc.internal.PickFirstLoadBalancerProvider; -import io.grpc.internal.PickFirstLoadBalancerProviderAccessor; import io.grpc.internal.PickSubchannelArgsImpl; import io.grpc.internal.testing.StreamRecorder; import io.grpc.lookup.v1.RouteLookupServiceGrpc; @@ -120,8 +119,6 @@ public class RlsLoadBalancerTest { public final GrpcCleanupRule grpcCleanupRule = new GrpcCleanupRule(); @Rule public final MockitoRule mocks = MockitoJUnit.rule(); - private final boolean defaultNewPickFirst = - PickFirstLoadBalancerProvider.isEnabledNewPickFirst(); private final RlsLoadBalancerProvider provider = new RlsLoadBalancerProvider(); private final FakeClock fakeClock = new FakeClock(); private final SynchronizationContext syncContext = @@ -205,7 +202,6 @@ public CachingRlsLbClient.Builder get() { @After public void tearDown() { rlsLb.shutdown(); - PickFirstLoadBalancerProviderAccessor.setEnableNewPickFirst(defaultNewPickFirst); } @Test @@ -252,7 +248,6 @@ public void lb_serverStatusCodeConversion() throws Exception { @Test public void lb_working_withDefaultTarget_rlsResponding() throws Exception { - PickFirstLoadBalancerProviderAccessor.setEnableNewPickFirst(false); helper.getSynchronizationContext().execute(() -> { try { deliverResolvedAddresses(); @@ -263,78 +258,9 @@ public void lb_working_withDefaultTarget_rlsResponding() throws Exception { InOrder inOrder = inOrder(helper); inOrder.verify(helper) .updateBalancingState(eq(ConnectivityState.CONNECTING), pickerCaptor.capture()); - SubchannelPicker picker = pickerCaptor.getValue(); - // Warm-up pick; will be queued - PickResult res = picker.pickSubchannel(searchSubchannelArgs); - assertThat(res.getStatus().isOk()).isTrue(); - assertThat(res.getSubchannel()).isNull(); - assertThat(res.getDelayType()).isEqualTo("rls_lookup_pending"); - assertThat(res.getDelayReason()).contains("Route Lookup Service query pending"); - // Cache is warm, but still unconnected - res = picker.pickSubchannel(searchSubchannelArgs); - inOrder.verify(helper).createSubchannel(any(CreateSubchannelArgs.class)); - inOrder.verify(helper, atLeast(0)) - .updateBalancingState(eq(ConnectivityState.CONNECTING), any(SubchannelPicker.class)); - inOrder.verify(helper, atLeast(0)).getSynchronizationContext(); - inOrder.verify(helper, atLeast(0)).getScheduledExecutorService(); - inOrder.verify(helper, atLeast(0)).getMetricRecorder(); - inOrder.verify(helper, atLeast(0)).getChannelTarget(); - inOrder.verifyNoMoreInteractions(); - - assertThat(res.getStatus().isOk()).isTrue(); - assertThat(subchannels).hasSize(2); // includes fallback sub-channel - FakeSubchannel searchSubchannel = subchannels.getLast(); - assertThat(subchannelIsReady(searchSubchannel)).isFalse(); - - searchSubchannel.updateState(ConnectivityStateInfo.forNonError(ConnectivityState.READY)); - inOrder.verify(helper) - .updateBalancingState(eq(ConnectivityState.READY), pickerCaptor.capture()); - inOrder.verifyNoMoreInteractions(); - res = picker.pickSubchannel(searchSubchannelArgs); - assertThat(subchannelIsReady(res.getSubchannel())).isTrue(); - assertThat(res.getSubchannel()).isSameInstanceAs(searchSubchannel); - verifyLongCounterAdd("grpc.lb.rls.target_picks", 1, 1, "wilderness", "complete"); - - // rescue should be pending status although the overall channel state is READY - res = picker.pickSubchannel(rescueSubchannelArgs); - inOrder.verify(helper).createSubchannel(any(CreateSubchannelArgs.class)); - // other rls picker itself is ready due to first channel. - assertThat(res.getStatus().isOk()).isTrue(); - assertThat(subchannelIsReady(res.getSubchannel())).isFalse(); - assertThat(subchannels).hasSize(3); // includes fallback sub-channel - FakeSubchannel rescueSubchannel = subchannels.getLast(); - - // search subchannel is down, rescue subchannel is connecting - searchSubchannel.updateState(ConnectivityStateInfo.forTransientFailure(Status.UNAVAILABLE)); - inOrder.verify(helper) - .updateBalancingState(eq(ConnectivityState.CONNECTING), pickerCaptor.capture()); - - rescueSubchannel.updateState(ConnectivityStateInfo.forNonError(ConnectivityState.READY)); - inOrder.verify(helper) - .updateBalancingState(eq(ConnectivityState.READY), pickerCaptor.capture()); - - // search again, verify that it doesn't use fallback, since RLS server responded, even though - // subchannel is in failure mode - res = picker.pickSubchannel(searchSubchannelArgs); - assertThat(res.getStatus().getCode()).isEqualTo(Status.Code.UNAVAILABLE); - assertThat(subchannelIsReady(res.getSubchannel())).isFalse(); - verifyLongCounterAdd("grpc.lb.rls.target_picks", 1, 1, "wilderness", "fail"); - } - - @Test - public void lb_working_withDefaultTarget_rlsResponding_newPickFirst() throws Exception { - PickFirstLoadBalancerProviderAccessor.setEnableNewPickFirst(true); - helper.getSynchronizationContext().execute(() -> { - try { - deliverResolvedAddresses(); - } catch (Exception e) { - throw new RuntimeException(e); - } - }); - InOrder inOrder = inOrder(helper); - inOrder.verify(helper) - .updateBalancingState(eq(ConnectivityState.CONNECTING), pickerCaptor.capture()); - inOrder.verify(helper).createSubchannel(any(CreateSubchannelArgs.class)); + 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); diff --git a/util/src/test/java/io/grpc/util/RoundRobinLoadBalancerTest.java b/util/src/test/java/io/grpc/util/RoundRobinLoadBalancerTest.java index 549569f04a8..06473117a26 100644 --- a/util/src/test/java/io/grpc/util/RoundRobinLoadBalancerTest.java +++ b/util/src/test/java/io/grpc/util/RoundRobinLoadBalancerTest.java @@ -230,7 +230,6 @@ public void pickAfterResolvedUpdatedHosts() throws Exception { @Test public void pickAfterStateChange() throws Exception { - PickFirstLoadBalancerProviderAccessor.setEnableNewPickFirst(false); InOrder inOrder = inOrder(mockHelper); Status addressesAcceptanceStatus = acceptAddresses(Arrays.asList(servers.get(0)), Attributes.EMPTY); @@ -256,43 +255,11 @@ public void pickAfterStateChange() throws Exception { assertThat(pickerCaptor.getValue().pickSubchannel(mockArgs).getStatus()).isEqualTo(error); deliverSubchannelState(subchannel, ConnectivityStateInfo.forNonError(IDLE)); - inOrder.verify(mockHelper).refreshNameResolution(); - inOrder.verify(mockHelper, never()) - .updateBalancingState(eq(TRANSIENT_FAILURE), any(SubchannelPicker.class)); - - verify(subchannel, atLeastOnce()).requestConnection(); - AbstractTestHelper.verifyNoMoreMeaningfulInteractions(mockHelper); - } - - @Test - public void pickAfterStateChange_newPickFirst() throws Exception { - PickFirstLoadBalancerProviderAccessor.setEnableNewPickFirst(true); - InOrder inOrder = inOrder(mockHelper); - Status addressesAcceptanceStatus = - acceptAddresses(Arrays.asList(servers.get(0)), Attributes.EMPTY); - assertThat(addressesAcceptanceStatus.isOk()).isTrue(); - inOrder.verify(mockHelper).createSubchannel(any(CreateSubchannelArgs.class)); - - // TODO figure out if this method testing the right things - - assertThat(subchannels).hasSize(1); - Subchannel subchannel = subchannels.values().iterator().next(); - - inOrder.verify(mockHelper).updateBalancingState(eq(CONNECTING), eq(EMPTY_PICKER)); - - deliverSubchannelState(subchannel, ConnectivityStateInfo.forNonError(READY)); - inOrder.verify(mockHelper).updateBalancingState(eq(READY), pickerCaptor.capture()); - assertThat(pickerCaptor.getValue()).isInstanceOf(ReadyPicker.class); - - Status error = Status.UNKNOWN.withDescription("¯\\_(ツ)_//¯"); - deliverSubchannelState(subchannel, - ConnectivityStateInfo.forTransientFailure(error)); - AbstractTestHelper.refreshInvokedAndUpdateBS( - inOrder, TRANSIENT_FAILURE, mockHelper, pickerCaptor); - assertThat(pickerCaptor.getValue().pickSubchannel(mockArgs).getStatus()).isEqualTo(error); - - deliverSubchannelState(subchannel, ConnectivityStateInfo.forNonError(IDLE)); - inOrder.verify(mockHelper, never()).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)); diff --git a/xds/src/test/java/io/grpc/xds/LeastRequestLoadBalancerTest.java b/xds/src/test/java/io/grpc/xds/LeastRequestLoadBalancerTest.java index ca6b7082ecd..71d4bdb224b 100644 --- a/xds/src/test/java/io/grpc/xds/LeastRequestLoadBalancerTest.java +++ b/xds/src/test/java/io/grpc/xds/LeastRequestLoadBalancerTest.java @@ -62,7 +62,6 @@ import io.grpc.Metadata; import io.grpc.Status; import io.grpc.internal.PickFirstLoadBalancerProvider; -import io.grpc.internal.PickFirstLoadBalancerProviderAccessor; import io.grpc.util.AbstractTestHelper; import io.grpc.util.MultiChildLoadBalancer.ChildLbState; import io.grpc.xds.LeastRequestLoadBalancer.LeastRequestConfig; @@ -101,7 +100,6 @@ public class LeastRequestLoadBalancerTest { private final Map, Subchannel> subchannels = Maps.newLinkedHashMap(); private final Attributes affinity = Attributes.newBuilder().set(MAJOR_KEY, "I got the keys").build(); - private boolean defaultNewPickFirst = PickFirstLoadBalancerProvider.isEnabledNewPickFirst(); @Captor private ArgumentCaptor pickerCaptor; @@ -131,7 +129,6 @@ public void setUp() { @After public void tearDown() throws Exception { - PickFirstLoadBalancerProviderAccessor.setEnableNewPickFirst(defaultNewPickFirst); verifyNoMoreInteractions(mockRandom); verifyNoMoreInteractions(mockArgs); } @@ -234,7 +231,6 @@ private Subchannel getSubchannel(EquivalentAddressGroup removedEag) { @Test public void pickAfterStateChange() throws Exception { - PickFirstLoadBalancerProviderAccessor.setEnableNewPickFirst(false); InOrder inOrder = inOrder(helper); Status addressesAcceptanceStatus = loadBalancer.acceptResolvedAddresses( ResolvedAddresses.newBuilder().setAddresses(servers).setAttributes(Attributes.EMPTY) @@ -261,49 +257,14 @@ public void pickAfterStateChange() throws Exception { .isEqualTo(PickResult.withNoResult()); deliverSubchannelState(subchannel, ConnectivityStateInfo.forNonError(IDLE)); - inOrder.verify(helper).refreshNameResolution(); - assertThat(childLbState.getCurrentState()).isEqualTo(TRANSIENT_FAILURE); - assertThat(childLbState.getCurrentPicker().toString()).contains(error.toString()); - - verify(subchannel, times(2)).requestConnection(); - verify(helper, times(3)).createSubchannel(any(CreateSubchannelArgs.class)); - AbstractTestHelper.verifyNoMoreMeaningfulInteractions(helper); - } - - @Test - public void pickAfterStateChange_newPickFirst() throws Exception { - PickFirstLoadBalancerProviderAccessor.setEnableNewPickFirst(true); - InOrder inOrder = inOrder(helper); - Status addressesAcceptanceStatus = loadBalancer.acceptResolvedAddresses( - ResolvedAddresses.newBuilder().setAddresses(servers).setAttributes(Attributes.EMPTY) - .build()); - assertThat(addressesAcceptanceStatus.isOk()).isTrue(); - ChildLbState childLbState = loadBalancer.getChildLbStates().iterator().next(); - Subchannel subchannel = getSubchannel(servers.get(0)); - - inOrder.verify(helper) - .updateBalancingState(eq(CONNECTING), pickerReturns(PickResult.withNoResult())); - assertThat(childLbState.getCurrentState()).isEqualTo(CONNECTING); - - deliverSubchannelState(subchannel, ConnectivityStateInfo.forNonError(READY)); - inOrder.verify(helper).updateBalancingState(eq(READY), pickerCaptor.capture()); - assertThat(pickerCaptor.getValue()).isInstanceOf(ReadyPicker.class); - assertThat(childLbState.getCurrentState()).isEqualTo(READY); - - Status error = Status.UNKNOWN.withDescription("¯\\_(ツ)_//¯"); - deliverSubchannelState(subchannel, ConnectivityStateInfo.forTransientFailure(error)); - assertThat(childLbState.getCurrentState()).isEqualTo(TRANSIENT_FAILURE); - assertThat(childLbState.getCurrentPicker().toString()).contains(error.toString()); - refreshInvokedAndUpdateBS(inOrder, CONNECTING); - assertThat(pickerCaptor.getValue().pickSubchannel(mockArgs)) - .isEqualTo(PickResult.withNoResult()); - - deliverSubchannelState(subchannel, ConnectivityStateInfo.forNonError(IDLE)); - inOrder.verify(helper, never()).refreshNameResolution(); + if (!PickFirstLoadBalancerProvider.isEnabledNewPickFirst()) { + inOrder.verify(helper).refreshNameResolution(); + } assertThat(childLbState.getCurrentState()).isEqualTo(TRANSIENT_FAILURE); assertThat(childLbState.getCurrentPicker().toString()).contains(error.toString()); - verify(subchannel, times(1)).requestConnection(); + int expectedCount = PickFirstLoadBalancerProvider.isEnabledNewPickFirst() ? 1 : 2; + verify(subchannel, times(expectedCount)).requestConnection(); verify(helper, times(3)).createSubchannel(any(CreateSubchannelArgs.class)); AbstractTestHelper.verifyNoMoreMeaningfulInteractions(helper); } diff --git a/xds/src/test/java/io/grpc/xds/RingHashLoadBalancerTest.java b/xds/src/test/java/io/grpc/xds/RingHashLoadBalancerTest.java index 1ced3f49e71..9c7f8dc62e0 100644 --- a/xds/src/test/java/io/grpc/xds/RingHashLoadBalancerTest.java +++ b/xds/src/test/java/io/grpc/xds/RingHashLoadBalancerTest.java @@ -533,7 +533,6 @@ public void pickWithRandomHash_atLeastOneSubchannelConnecting() { @Test public void pickWithRandomHash_firstSubchannelInTransientFailure_remainingSubchannelsIdle() { - PickFirstLoadBalancerProviderAccessor.setEnableNewPickFirst(false); // Map each server address to exactly one ring entry. RingHashConfig config = new RingHashConfig(3, 3, "dummy-random-hash"); List servers = createWeightedServerAddrs(1, 1, 1); @@ -553,32 +552,10 @@ public void pickWithRandomHash_firstSubchannelInTransientFailure_remainingSubcha assertThat(result.getSubchannel()).isNull(); // buffer request assertThat(result.getDelayType()).isEqualTo("connecting"); assertThat(result.getDelayReason()).isEqualTo("ring_hash: waiting for connection"); - verifyConnection(1); - } - - @Test - public void - pickWithRandomHash_firstSubchannelInTransientFailure_remainingSubchannelsIdle_newPickFirst() { - PickFirstLoadBalancerProviderAccessor.setEnableNewPickFirst(true); - // Map each server address to exactly one ring entry. - RingHashConfig config = new RingHashConfig(3, 3, "dummy-random-hash"); - List servers = createWeightedServerAddrs(1, 1, 1); - initializeLbSubchannels(config, servers); - - // Bring one subchannel to TRANSIENT_FAILURE. - deliverSubchannelUnreachable(getSubChannel(servers.get(0))); - verify(helper).updateBalancingState(eq(CONNECTING), pickerCaptor.capture()); - verifyConnection(1); - - SubchannelPicker picker = pickerCaptor.getValue(); - PickSubchannelArgs args = getDefaultPickSubchannelArgs(hashFunc.hashVoid()); - PickResult result = picker.pickSubchannel(args); - assertThat(result.getStatus().isOk()).isTrue(); - assertThat(result.getSubchannel()).isNull(); // buffer request - assertThat(result.getDelayType()).isEqualTo("connecting"); - assertThat(result.getDelayReason()).isEqualTo("ring_hash: waiting for connection"); - int expected = connectionRequestedQueue.isEmpty() ? 0 : 1; - verifyConnection(expected); + int expectedCount = + PickFirstLoadBalancerProvider.isEnabledNewPickFirst() + ? (connectionRequestedQueue.isEmpty() ? 0 : 1) : 1; + verifyConnection(expectedCount); } private Subchannel getSubChannel(EquivalentAddressGroup eag) { @@ -657,73 +634,6 @@ private PickSubchannelArgs getDefaultPickSubchannelArgsForServer(int serverid) { @Test public void skipFailingHosts_firstTwoHostsFailed_pickNextFirstReady() { - PickFirstLoadBalancerProviderAccessor.setEnableNewPickFirst(false); - // Map each server address to exactly one ring entry. - RingHashConfig config = new RingHashConfig(3, 3, ""); - List servers = createWeightedServerAddrs(1, 1, 1); - - initializeLbSubchannels(config, servers); - - // ring: - // "FakeSocketAddress-server0_0" - // "FakeSocketAddress-server1_0" - // "FakeSocketAddress-server2_0" - - long rpcHash = hashFunc.hashAsciiString("FakeSocketAddress-server1_0"); - PickSubchannelArgs args = getDefaultPickSubchannelArgs(rpcHash); - - // Bring down server0 and server2 to force trying server1. - deliverSubchannelState( - getSubchannel(servers, 1), - ConnectivityStateInfo.forTransientFailure( - Status.UNAVAILABLE.withDescription("unreachable"))); - deliverSubchannelState( - getSubchannel(servers, 2), - ConnectivityStateInfo.forTransientFailure( - Status.PERMISSION_DENIED.withDescription("permission denied"))); - verify(helper).updateBalancingState(eq(TRANSIENT_FAILURE), pickerCaptor.capture()); - verifyConnection(2); - PickResult result = pickerCaptor.getValue().pickSubchannel(args); // activate last subchannel - assertThat(result.getStatus().isOk()).isTrue(); - verifyConnection(1); - - deliverSubchannelState( - getSubchannel(servers, 0), - ConnectivityStateInfo.forTransientFailure( - Status.PERMISSION_DENIED.withDescription("permission denied again"))); - verify(helper, times(2)).updateBalancingState(eq(TRANSIENT_FAILURE), pickerCaptor.capture()); - result = pickerCaptor.getValue().pickSubchannel(args); - assertThat(result.getStatus().isOk()).isFalse(); // fail the RPC - assertThat(result.getStatus().getCode()) - .isEqualTo(Code.UNAVAILABLE); // with error status for the original server hit by hash - assertThat(result.getStatus().getDescription()).isEqualTo("unreachable"); - - // Now connecting to server1. - deliverSubchannelState(getSubchannel(servers, 1), CSI_CONNECTING); - - reset(helper); - - result = pickerCaptor.getValue().pickSubchannel(args); - assertThat(result.getStatus().isOk()).isFalse(); // fail the RPC - assertThat(result.getStatus().getCode()) - .isEqualTo(Code.UNAVAILABLE); // with error status for the original server hit by hash - assertThat(result.getStatus().getDescription()).isEqualTo("unreachable"); - - // Simulate server1 becomes READY. - deliverSubchannelState(getSubchannel(servers, 1), CSI_READY); - verify(helper).updateBalancingState(eq(READY), pickerCaptor.capture()); - - SubchannelPicker picker = pickerCaptor.getValue(); - result = picker.pickSubchannel(args); - assertThat(result.getStatus().isOk()).isTrue(); // succeed - assertThat(result.getSubchannel().getAddresses()).isEqualTo(servers.get(1)); // with server1 - assertThat(picker.pickSubchannel(getDefaultPickSubchannelArgsForServer(0))).isEqualTo(result); - assertThat(picker.pickSubchannel(getDefaultPickSubchannelArgsForServer(2))).isEqualTo(result); - } - - @Test - public void skipFailingHosts_firstTwoHostsFailed_pickNextFirstReady_newPickFirst() { - PickFirstLoadBalancerProviderAccessor.setEnableNewPickFirst(true); // Map each server address to exactly one ring entry. RingHashConfig config = new RingHashConfig(3, 3, ""); List servers = createWeightedServerAddrs(1, 1, 1); @@ -748,10 +658,10 @@ public void skipFailingHosts_firstTwoHostsFailed_pickNextFirstReady_newPickFirst ConnectivityStateInfo.forTransientFailure( Status.PERMISSION_DENIED.withDescription("permission denied"))); verify(helper).updateBalancingState(eq(TRANSIENT_FAILURE), pickerCaptor.capture()); - verifyConnection(1); + verifyConnection(PickFirstLoadBalancerProvider.isEnabledNewPickFirst() ? 1 : 2); PickResult result = pickerCaptor.getValue().pickSubchannel(args); // activate last subchannel assertThat(result.getStatus().isOk()).isTrue(); - verifyConnection(0); + verifyConnection(PickFirstLoadBalancerProvider.isEnabledNewPickFirst() ? 0 : 1); deliverSubchannelState( getSubchannel(servers, 0), @@ -806,35 +716,6 @@ public void removingAddressShutdownSubchannel() { @Test public void allSubchannelsInTransientFailure() { - PickFirstLoadBalancerProviderAccessor.setEnableNewPickFirst(false); - // Map each server address to exactly one ring entry. - RingHashConfig config = new RingHashConfig(3, 3, ""); - List servers = createWeightedServerAddrs(1, 1, 1); - initializeLbSubchannels(config, servers); - - // Bring all subchannels to TRANSIENT_FAILURE. - for (Subchannel subchannel : subchannels.values()) { - deliverSubchannelState(subchannel, ConnectivityStateInfo.forTransientFailure( - Status.UNAVAILABLE.withDescription( - subchannel.getAddresses().getAddresses() + " unreachable"))); - } - verify(helper, atLeastOnce()) - .updateBalancingState(eq(TRANSIENT_FAILURE), pickerCaptor.capture()); - verifyConnection(2); - - // Picking subchannel triggers connection. RPC hash hits server0. - PickSubchannelArgs args = getDefaultPickSubchannelArgsForServer(0); - PickResult result = pickerCaptor.getValue().pickSubchannel(args); - assertThat(result.getStatus().isOk()).isFalse(); - assertThat(result.getStatus().getCode()).isEqualTo(Code.UNAVAILABLE); - assertThat(result.getStatus().getDescription()) - .isEqualTo("[FakeSocketAddress-server0] unreachable"); - verifyConnection(0); // TF has already started taking care of this, pick doesn't need to - } - - @Test - public void allSubchannelsInTransientFailure_newPickFirst() { - PickFirstLoadBalancerProviderAccessor.setEnableNewPickFirst(true); // Map each server address to exactly one ring entry. RingHashConfig config = new RingHashConfig(3, 3, ""); List servers = createWeightedServerAddrs(1, 1, 1); @@ -848,7 +729,7 @@ public void allSubchannelsInTransientFailure_newPickFirst() { } verify(helper, atLeastOnce()) .updateBalancingState(eq(TRANSIENT_FAILURE), pickerCaptor.capture()); - verifyConnection(1); + verifyConnection(PickFirstLoadBalancerProvider.isEnabledNewPickFirst() ? 1 : 2); // Picking subchannel triggers connection. RPC hash hits server0. PickSubchannelArgs args = getDefaultPickSubchannelArgsForServer(0); @@ -908,7 +789,6 @@ private Subchannel getSubchannel(List servers, int serve @Test public void firstSubchannelFailure() { - PickFirstLoadBalancerProviderAccessor.setEnableNewPickFirst(false); // Map each server address to exactly one ring entry. RingHashConfig config = new RingHashConfig(3, 3, ""); List servers = createWeightedServerAddrs(1, 1, 1); @@ -934,76 +814,14 @@ public void firstSubchannelFailure() { assertThat(result.getStatus().isOk()).isTrue(); assertThat(result.getSubchannel()).isNull(); verify(subchannelList.get(0), never()).requestConnection(); // In TF - verify(subchannelList.get(1)).requestConnection(); + verify(subchannelList.get(1), + PickFirstLoadBalancerProvider.isEnabledNewPickFirst() ? never() : times(1)) + .requestConnection(); verify(subchannelList.get(2), never()).requestConnection(); // Not one of the first 2 } - @Test - public void firstSubchannelFailure_newPickFirst() { - PickFirstLoadBalancerProviderAccessor.setEnableNewPickFirst(true); - // Map each server address to exactly one ring entry. - RingHashConfig config = new RingHashConfig(3, 3, ""); - List servers = createWeightedServerAddrs(1, 1, 1); - - List subchannelList = - initializeLbSubchannels(config, servers, RESET_SUBCHANNEL_MOCKS); - - // ring: - // "FakeSocketAddress-server1_0" - // "FakeSocketAddress-server0_0" - // "FakeSocketAddress-server2_0" - - deliverSubchannelState(subchannelList.get(0), - ConnectivityStateInfo.forTransientFailure( - Status.UNAVAILABLE.withDescription("unreachable"))); - verify(helper).updateBalancingState(eq(CONNECTING), pickerCaptor.capture()); - verifyConnection(1); - - PickSubchannelArgs args = getDefaultPickSubchannelArgs(hashFunc.hashVoid()); - SubchannelPicker picker1 = pickerCaptor.getValue(); - PickResult result = picker1.pickSubchannel(args); - assertThat(result.getStatus().isOk()).isTrue(); - assertThat(result.getSubchannel()).isNull(); - verify(subchannelList.get(0), never()).requestConnection(); - verify(subchannelList.get(1), never()).requestConnection(); - verify(subchannelList.get(2), never()).requestConnection(); - } - @Test public void secondSubchannelConnecting() { - PickFirstLoadBalancerProviderAccessor.setEnableNewPickFirst(false); - // Map each server address to exactly one ring entry. - RingHashConfig config = new RingHashConfig(3, 3, ""); - List servers = createWeightedServerAddrs(1, 1, 1); - - initializeLbSubchannels(config, servers); - - // ring: - // "FakeSocketAddress-server1_0" - // "FakeSocketAddress-server0_0" - // "FakeSocketAddress-server2_0" - - Subchannel firstSubchannel = getSubchannel(servers, 0); - deliverSubchannelUnreachable(firstSubchannel); - verifyConnection(1); - - deliverSubchannelState(getSubchannel(servers, 2), CSI_CONNECTING); - verify(helper, times(2)).updateBalancingState(eq(CONNECTING), pickerCaptor.capture()); - verifyConnection(0); - - // Picking subchannel when idle triggers connection. - deliverSubchannelState(getSubchannel(servers, 2), - ConnectivityStateInfo.forNonError(IDLE)); - verifyConnection(1); - PickSubchannelArgs args = getDefaultPickSubchannelArgs(hashFunc.hashVoid()); - PickResult result = pickerCaptor.getValue().pickSubchannel(args); - assertThat(result.getStatus().isOk()).isTrue(); - verifyConnection(1); - } - - @Test - public void secondSubchannelConnecting_newPickFirst() { - PickFirstLoadBalancerProviderAccessor.setEnableNewPickFirst(true); // Map each server address to exactly one ring entry. RingHashConfig config = new RingHashConfig(3, 3, ""); List servers = createWeightedServerAddrs(1, 1, 1); @@ -1026,16 +844,15 @@ public void secondSubchannelConnecting_newPickFirst() { // Picking subchannel when idle triggers connection. deliverSubchannelState(getSubchannel(servers, 2), ConnectivityStateInfo.forNonError(IDLE)); - verifyConnection(0); + verifyConnection(PickFirstLoadBalancerProvider.isEnabledNewPickFirst() ? 0 : 1); PickSubchannelArgs args = getDefaultPickSubchannelArgs(hashFunc.hashVoid()); PickResult result = pickerCaptor.getValue().pickSubchannel(args); assertThat(result.getStatus().isOk()).isTrue(); - verifyConnection(0); + verifyConnection(PickFirstLoadBalancerProvider.isEnabledNewPickFirst() ? 0 : 1); } @Test public void secondSubchannelFailure() { - PickFirstLoadBalancerProviderAccessor.setEnableNewPickFirst(false); // Map each server address to exactly one ring entry. RingHashConfig config = new RingHashConfig(3, 3, ""); List servers = createWeightedServerAddrs(1, 1, 1); @@ -1051,47 +868,20 @@ public void secondSubchannelFailure() { deliverSubchannelUnreachable(firstSubchannel); deliverSubchannelUnreachable(getSubchannel(servers, 2)); verify(helper).updateBalancingState(eq(TRANSIENT_FAILURE), pickerCaptor.capture()); - verifyConnection(2); + verifyConnection(PickFirstLoadBalancerProvider.isEnabledNewPickFirst() ? 1 : 2); // Picking subchannel triggers connection. PickSubchannelArgs args = getDefaultPickSubchannelArgs(hashFunc.hashVoid()); PickResult result = pickerCaptor.getValue().pickSubchannel(args); assertThat(result.getStatus().isOk()).isTrue(); - verify(getSubchannel(servers, 1)).requestConnection(); - verifyConnection(1); - } - - @Test - public void secondSubchannelFailure_newPickFirst() { - PickFirstLoadBalancerProviderAccessor.setEnableNewPickFirst(true); - // Map each server address to exactly one ring entry. - RingHashConfig config = new RingHashConfig(3, 3, ""); - List servers = createWeightedServerAddrs(1, 1, 1); - - initializeLbSubchannels(config, servers); - - // ring: - // "FakeSocketAddress-server1_0" - // "FakeSocketAddress-server0_0" - // "FakeSocketAddress-server2_0" - - Subchannel firstSubchannel = getSubchannel(servers, 0); - deliverSubchannelUnreachable(firstSubchannel); - deliverSubchannelUnreachable(getSubchannel(servers, 2)); - verify(helper).updateBalancingState(eq(TRANSIENT_FAILURE), pickerCaptor.capture()); - verifyConnection(1); - - // Picking subchannel triggers connection. - PickSubchannelArgs args = getDefaultPickSubchannelArgs(hashFunc.hashVoid()); - PickResult result = pickerCaptor.getValue().pickSubchannel(args); - assertThat(result.getStatus().isOk()).isTrue(); - verify(getSubchannel(servers, 1), never()).requestConnection(); - verifyConnection(0); + verify(getSubchannel(servers, 1), + PickFirstLoadBalancerProvider.isEnabledNewPickFirst() ? never() : times(1)) + .requestConnection(); + verifyConnection(PickFirstLoadBalancerProvider.isEnabledNewPickFirst() ? 0 : 1); } @Test public void thirdSubchannelConnecting() { - PickFirstLoadBalancerProviderAccessor.setEnableNewPickFirst(false); // Map each server address to exactly one ring entry. RingHashConfig config = new RingHashConfig(3, 3, ""); List servers = createWeightedServerAddrs(1, 1, 1); @@ -1110,37 +900,7 @@ public void thirdSubchannelConnecting() { deliverSubchannelState(getSubchannel(servers, 1), CSI_CONNECTING); verify(helper, atLeastOnce()) .updateBalancingState(eq(TRANSIENT_FAILURE), pickerCaptor.capture()); - verifyConnection(2); - - // Picking subchannel should not trigger connection per gRFC A61. - PickSubchannelArgs args = getDefaultPickSubchannelArgs(hashFunc.hashVoid()); - PickResult result = pickerCaptor.getValue().pickSubchannel(args); - assertThat(result.getStatus().isOk()).isTrue(); - verifyConnection(0); - } - - @Test - public void thirdSubchannelConnecting_newPickFirst() { - PickFirstLoadBalancerProviderAccessor.setEnableNewPickFirst(true); - // Map each server address to exactly one ring entry. - RingHashConfig config = new RingHashConfig(3, 3, ""); - List servers = createWeightedServerAddrs(1, 1, 1); - - initializeLbSubchannels(config, servers); - - // ring: - // "FakeSocketAddress-server1_0" - // "FakeSocketAddress-server0_0" - // "FakeSocketAddress-server2_0" - - Subchannel firstSubchannel = getSubchannel(servers, 0); - - deliverSubchannelUnreachable(firstSubchannel); - deliverSubchannelUnreachable(getSubchannel(servers, 2)); - deliverSubchannelState(getSubchannel(servers, 1), CSI_CONNECTING); - verify(helper, atLeastOnce()) - .updateBalancingState(eq(TRANSIENT_FAILURE), pickerCaptor.capture()); - verifyConnection(1); + verifyConnection(PickFirstLoadBalancerProvider.isEnabledNewPickFirst() ? 1 : 2); // Picking subchannel should not trigger connection per gRFC A61. PickSubchannelArgs args = getDefaultPickSubchannelArgs(hashFunc.hashVoid()); @@ -1151,38 +911,6 @@ public void thirdSubchannelConnecting_newPickFirst() { @Test public void stickyTransientFailure() { - PickFirstLoadBalancerProviderAccessor.setEnableNewPickFirst(false); - // Map each server address to exactly one ring entry. - RingHashConfig config = new RingHashConfig(3, 3, ""); - List servers = createWeightedServerAddrs(1, 1, 1); - - initializeLbSubchannels(config, servers); - - // Bring one subchannel to TRANSIENT_FAILURE. - Subchannel firstSubchannel = getSubchannel(servers, 0); - deliverSubchannelUnreachable(firstSubchannel); - - verify(helper).updateBalancingState(eq(CONNECTING), pickerCaptor.capture()); - verifyConnection(1); - - reset(helper); - deliverSubchannelState(firstSubchannel, ConnectivityStateInfo.forNonError(IDLE)); - // Should not have called updateBalancingState on the helper again because PickFirst is - // shielding the higher level from the state change. - verify(helper, never()).updateBalancingState(any(), any()); - verifyConnection(1); - - // Picking subchannel triggers connection on second address. RPC hash hits server0. - PickSubchannelArgs args = getDefaultPickSubchannelArgs(hashFunc.hashVoid()); - PickResult result = pickerCaptor.getValue().pickSubchannel(args); - assertThat(result.getStatus().isOk()).isTrue(); - verify(getSubchannel(servers, 1)).requestConnection(); - verify(getSubchannel(servers, 2), never()).requestConnection(); - } - - @Test - public void stickyTransientFailure_newPickFirst() { - PickFirstLoadBalancerProviderAccessor.setEnableNewPickFirst(true); // Map each server address to exactly one ring entry. RingHashConfig config = new RingHashConfig(3, 3, ""); List servers = createWeightedServerAddrs(1, 1, 1); @@ -1201,13 +929,15 @@ public void stickyTransientFailure_newPickFirst() { // Should not have called updateBalancingState on the helper again because PickFirst is // shielding the higher level from the state change. verify(helper, never()).updateBalancingState(any(), any()); - verifyConnection(0); + verifyConnection(PickFirstLoadBalancerProvider.isEnabledNewPickFirst() ? 0 : 1); // Picking subchannel triggers connection on second address. RPC hash hits server0. PickSubchannelArgs args = getDefaultPickSubchannelArgs(hashFunc.hashVoid()); PickResult result = pickerCaptor.getValue().pickSubchannel(args); assertThat(result.getStatus().isOk()).isTrue(); - verify(getSubchannel(servers, 1), never()).requestConnection(); + verify(getSubchannel(servers, 1), + PickFirstLoadBalancerProvider.isEnabledNewPickFirst() ? never() : times(1)) + .requestConnection(); verify(getSubchannel(servers, 2), never()).requestConnection(); }