grpclb: pick_first delegation #12568

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Jump to bottom

Draft

shivaspeaks wants to merge 2 commits into grpc:master from shivaspeaks:port-grpclb

+235 −230

Member

shivaspeaks commented Dec 15, 2025

Summary of Changes
This pull request refactors the grpclb load balancer's PICK_FIRST mode to delegate its logic to a standard pick_first load balancing policy.

The key changes are as follows:

grpclb/build.gradle

Added dependency on grpc-util module to access ForwardingLoadBalancerHelper

grpclb/src/main/java/io/grpc/grpclb/GrpclbState.java

New imports:
LoadBalancer, LoadBalancerProvider, LoadBalancerRegistry, ResolvedAddresses, FixedResultPicker, ForwardingLoadBalancerHelper
New fields for PICK_FIRST delegation:
- pickFirstLbProvider - Provider for creating child pick_first LB
- pickFirstLb - The child LoadBalancer instance
- pickFirstLbState / pickFirstLbPicker - Track child LB's state and picker
- currentPickFirstLoadRecorder - Load recorder for token attachment
Key behavioral changes:
- updateServerList() PICK_FIRST case: Instead of creating a single subchannel, it now:
  - Creates the child pick_first LB once and then updates it with new addresses on subsequent updates.
  - Passes addresses to child LB via acceptResolvedAddresses()
- maybeUpdatePicker() PICK_FIRST case: Uses child LB's state and picker wrapped with ChildLbPickerEntry
- RoundRobinEntry.picked() signature change: Changed from picked(Metadata) to picked(PickSubchannelArgs) to allow child picker delegation
- New ChildLbPickerEntry class: Wraps child LB's picker and attaches TokenAttachingTracerFactory for token propagation
- New PickFirstLbHelper class: Forwarding helper that intercepts updateBalancingState() to store child state and trigger grpclb picker updates
- Updated shutdown(), requestConnection(), maybeUseFallbackBackends(): Handle the new child LB delegation model

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

Updated tests to reflect the new delegation behavior:
- Initial state is now CONNECTING (not IDLE) since standard pick_first eagerly connects
- Tests now verify the child LB is created only once and then handles address updates internally.
- Adjusted verification expectations for the new flow
Key Behavioral Changes:
- Delegation to child LB: The grpclb state no longer directly manages subchannels for PICK_FIRST mode; the child pick_first LB handles this internally.


          grpclb: pick_first delegation

bd7aa4c

shivaspeaks requested review from ejona86 and kannanjgithub

December 15, 2025 07:18

kannanjgithub requested changes

View reviewed changes

Contributor

kannanjgithub left a comment

Review comments.

grpclb/src/test/java/io/grpc/grpclb/GrpclbLoadBalancerTest.java Show resolved Hide resolved

grpclb/src/test/java/io/grpc/grpclb/GrpclbLoadBalancerTest.java Show resolved Hide resolved

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

    
                  inOrder.verify(helper).updateBalancingState(eq(READY), pickerCaptor.capture());

                  RoundRobinPicker picker3 = (RoundRobinPicker) pickerCaptor.getValue();

                  assertThat(picker3.dropList).containsExactly(null, null);

                  assertThat(picker3.pickList).containsExactly(

Contributor

kannanjgithub Dec 16, 2025

Similarly here, don't remove the assertion for the pickList with the fake subchannel, assert using accessor in ChildLbPickerEntry for use by the test.

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

Comment on lines 1945 to 1952

    
                  // TRANSIENT_FAILURE

                  Status error = Status.UNAVAILABLE.withDescription("Simulated connection error");

                  deliverSubchannelState(subchannel, ConnectivityStateInfo.forTransientFailure(error));

                  inOrder.verify(helper).updateBalancingState(eq(TRANSIENT_FAILURE), pickerCaptor.capture());

                  RoundRobinPicker picker2 = (RoundRobinPicker) pickerCaptor.getValue();

                  assertThat(picker2.dropList).containsExactly(null, null);

                  assertThat(picker2.pickList).containsExactly(new ErrorEntry(error));

Contributor

kannanjgithub Dec 16, 2025

Do not remove the testing for simulated transient failure.
Now we cannot directly do

assertThat(picker2.pickList).containsExactly(new ErrorEntry(error));

as the ErorrEntry is wrapped in a ChildLbPickerEntry. Create accessor for the wrapped entry for use in the unit test so we can assert using that.

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

Comment on lines -2219 to 2230

    
                  assertThat(mockSubchannels).isEmpty();

                  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).updateBalancingState(eq(READY), pickerCaptor.capture());

                  RoundRobinPicker picker2 = (RoundRobinPicker) pickerCaptor.getValue();

                  assertThat(picker2.dropList).containsExactly(null, null);

                  assertThat(picker2.pickList).containsExactly(

                      new BackendEntry(subchannel, new TokenAttachingTracerFactory(getLoadRecorder())));

                  // the same child LB is updated with new addresses. We expect

                  // only one subchannel to be created for the whole test. The child LB will call

                  // updateAddresses() internally, so we remove the verifications for recreation.

                  verify(helper, times(1)).createSubchannel(any(CreateSubchannelArgs.class));

Contributor

kannanjgithub Dec 16, 2025

This PF behavior may be done by a different code now but the behavior should still be the same. Why remove the test for it?


          address comments

434c07b

kannanjgithub reviewed

View reviewed changes

Contributor

kannanjgithub left a comment

Some deleted code still needs to be restored.

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

    
                  verify(subchannel).requestConnection();

                  assertThat(picker0.dropList).containsExactly(null, null);

                  assertThat(picker0.pickList).hasSize(1);

                  assertThat(picker0.pickList.get(0)).isInstanceOf(ChildLbPickerEntry.class);

Contributor

kannanjgithub Dec 17, 2025

Now that you have added an accessor in ChildLbPickerEntry we can retrieve that subchannel and do the assertion for it (like you already do below).

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

    
                      readyEntry.getChildPicker().pickSubchannel(mock(PickSubchannelArgs.class));

                  assertThat(readyResult.getSubchannel()).isEqualTo(subchannel);

                  // TRANSIENT_FAILURE

Contributor

kannanjgithub Dec 17, 2025

nit: Add back the comment.

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

    
                }

                @Test

                public void grpclbWorking_pickFirstMode_lbSendsEmptyAddress() throws Exception {

Contributor

kannanjgithub Dec 17, 2025

The removed test code from this test needs to be restored as well.

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

    
                  // Only one subchannel is created

                  // Child pick_first eagerly connects

                  verify(helper, atLeast(1)).updateBalancingState(eq(CONNECTING), any());

Contributor

kannanjgithub Dec 17, 2025

capture the picker and do assertions.

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

Comment on lines -2053 to -2071

    
                  assertThat(picker0.dropList).containsExactly(null, null);

                  assertThat(picker0.pickList).containsExactly(new IdleSubchannelEntry(subchannel, syncContext));

                  // PICK_FIRST doesn't eagerly connect

                  verify(subchannel, never()).requestConnection();

                  // CONNECTING

                  deliverSubchannelState(subchannel, ConnectivityStateInfo.forNonError(CONNECTING));

                  inOrder.verify(helper).updateBalancingState(eq(CONNECTING), pickerCaptor.capture());

                  RoundRobinPicker picker1 = (RoundRobinPicker) pickerCaptor.getValue();

                  assertThat(picker1.dropList).containsExactly(null, null);

                  assertThat(picker1.pickList).containsExactly(BUFFER_ENTRY);

                  // TRANSIENT_FAILURE

                  Status error = Status.UNAVAILABLE.withDescription("Simulated connection error");

                  deliverSubchannelState(subchannel, ConnectivityStateInfo.forTransientFailure(error));

                  inOrder.verify(helper).updateBalancingState(eq(TRANSIENT_FAILURE), pickerCaptor.capture());

                  RoundRobinPicker picker2 = (RoundRobinPicker) pickerCaptor.getValue();

                  assertThat(picker2.dropList).containsExactly(null, null);

Contributor

kannanjgithub Dec 17, 2025

These deleted lines need to be restored as well similar to the previous test.

Contributor

kannanjgithub Dec 17, 2025

Same for other deleted code below.

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

    
                  // new addresses will be updated to the existing subchannel

                  inOrder.verify(helper, times(1)).createSubchannel(any(CreateSubchannelArgs.class));

                  inOrder.verify(helper).updateBalancingState(eq(IDLE), pickerCaptor.capture());

Contributor

kannanjgithub Dec 17, 2025

Restore this statement with the CONNECTING state.

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet