From bd07668b3f7c2b627504d6abaed39dfab7dc01a4 Mon Sep 17 00:00:00 2001 From: Arushi Garg Date: Wed, 10 Dec 2025 21:24:50 +0530 Subject: [PATCH 1/6] testing grpclb changes --- grpclb/build.gradle | 1 + .../main/java/io/grpc/grpclb/GrpclbState.java | 224 ++++++++++++------ .../grpc/grpclb/GrpclbLoadBalancerTest.java | 218 +++++------------ 3 files changed, 209 insertions(+), 234 deletions(-) diff --git a/grpclb/build.gradle b/grpclb/build.gradle index f543e0d71fc..e8896604f03 100644 --- a/grpclb/build.gradle +++ b/grpclb/build.gradle @@ -19,6 +19,7 @@ dependencies { implementation project(':grpc-core'), project(':grpc-protobuf'), project(':grpc-stub'), + project(':grpc-util'), libraries.guava, libraries.protobuf.java, libraries.protobuf.java.util diff --git a/grpclb/src/main/java/io/grpc/grpclb/GrpclbState.java b/grpclb/src/main/java/io/grpc/grpclb/GrpclbState.java index 49b74645ec8..bf5208140cb 100644 --- a/grpclb/src/main/java/io/grpc/grpclb/GrpclbState.java +++ b/grpclb/src/main/java/io/grpc/grpclb/GrpclbState.java @@ -37,13 +37,18 @@ import io.grpc.ConnectivityStateInfo; import io.grpc.Context; import io.grpc.EquivalentAddressGroup; +import io.grpc.LoadBalancer; import io.grpc.LoadBalancer.CreateSubchannelArgs; +import io.grpc.LoadBalancer.FixedResultPicker; import io.grpc.LoadBalancer.Helper; import io.grpc.LoadBalancer.PickResult; import io.grpc.LoadBalancer.PickSubchannelArgs; +import io.grpc.LoadBalancer.ResolvedAddresses; import io.grpc.LoadBalancer.Subchannel; import io.grpc.LoadBalancer.SubchannelPicker; import io.grpc.LoadBalancer.SubchannelStateListener; +import io.grpc.LoadBalancerProvider; +import io.grpc.LoadBalancerRegistry; import io.grpc.ManagedChannel; import io.grpc.Metadata; import io.grpc.Status; @@ -52,6 +57,7 @@ import io.grpc.grpclb.SubchannelPool.PooledSubchannelStateListener; import io.grpc.internal.BackoffPolicy; import io.grpc.internal.TimeProvider; +import io.grpc.util.ForwardingLoadBalancerHelper; import io.grpc.lb.v1.ClientStats; import io.grpc.lb.v1.InitialLoadBalanceRequest; import io.grpc.lb.v1.InitialLoadBalanceResponse; @@ -119,7 +125,7 @@ final class GrpclbState { @VisibleForTesting static final RoundRobinEntry BUFFER_ENTRY = new RoundRobinEntry() { @Override - public PickResult picked(Metadata headers) { + public PickResult picked(PickSubchannelArgs args) { return PickResult.withNoResult(); } @@ -187,6 +193,15 @@ enum Mode { new RoundRobinPicker(Collections.emptyList(), Arrays.asList(BUFFER_ENTRY)); private boolean requestConnectionPending; + // Child LoadBalancer and state for PICK_FIRST mode delegation. + private final LoadBalancerProvider pickFirstLbProvider; + @Nullable + private LoadBalancer pickFirstLb; + private ConnectivityState pickFirstLbState = CONNECTING; + private SubchannelPicker pickFirstLbPicker = new FixedResultPicker(PickResult.withNoResult()); + @Nullable + private GrpclbClientLoadRecorder currentPickFirstLoadRecorder; + GrpclbState( GrpclbConfig config, Helper helper, @@ -212,6 +227,8 @@ public void onSubchannelState( } else { this.subchannelPool = null; } + this.pickFirstLbProvider = + LoadBalancerRegistry.getDefaultRegistry().getProvider("pick_first"); this.time = checkNotNull(time, "time provider"); this.stopwatch = checkNotNull(stopwatch, "stopwatch"); this.timerService = checkNotNull(helper.getScheduledExecutorService(), "timerService"); @@ -309,6 +326,12 @@ void handleAddresses( void requestConnection() { requestConnectionPending = true; + // For PICK_FIRST mode with delegation, forward to the child LB. + if (config.getMode() == Mode.PICK_FIRST && pickFirstLb != null) { + pickFirstLb.requestConnection(); + requestConnectionPending = false; + return; + } for (RoundRobinEntry entry : currentPicker.pickList) { if (entry instanceof IdleSubchannelEntry) { ((IdleSubchannelEntry) entry).subchannel.requestConnection(); @@ -323,15 +346,23 @@ private void maybeUseFallbackBackends() { } // Balancer RPC should have either been broken or timed out. checkState(fallbackReason != null, "no reason to fallback"); - for (Subchannel subchannel : subchannels.values()) { - ConnectivityStateInfo stateInfo = subchannel.getAttributes().get(STATE_INFO).get(); - if (stateInfo.getState() == READY) { + // For PICK_FIRST mode with delegation, check the child LB's state. + if (config.getMode() == Mode.PICK_FIRST) { + if (pickFirstLb != null && pickFirstLbState == READY) { return; } - // If we do have balancer-provided backends, use one of its error in the error message if - // fail to fallback. - if (stateInfo.getState() == TRANSIENT_FAILURE) { - fallbackReason = stateInfo.getStatus(); + // For PICK_FIRST, we don't have individual subchannel states to use as fallback reason. + } else { + for (Subchannel subchannel : subchannels.values()) { + ConnectivityStateInfo stateInfo = subchannel.getAttributes().get(STATE_INFO).get(); + if (stateInfo.getState() == READY) { + return; + } + // If we do have balancer-provided backends, use one of its error in the error message if + // fail to fallback. + if (stateInfo.getState() == TRANSIENT_FAILURE) { + fallbackReason = stateInfo.getStatus(); + } } } // Fallback conditions met @@ -438,9 +469,10 @@ void shutdown() { subchannelPool.clear(); break; case PICK_FIRST: - if (!subchannels.isEmpty()) { - checkState(subchannels.size() == 1, "Excessive Subchannels: %s", subchannels); - subchannels.values().iterator().next().shutdown(); + // Shutdown the child pick_first LB which manages its own subchannels. + if (pickFirstLb != null) { + pickFirstLb.shutdown(); + pickFirstLb = null; } break; default: @@ -517,22 +549,17 @@ private void updateServerList( subchannels = Collections.unmodifiableMap(newSubchannelMap); break; case PICK_FIRST: - checkState(subchannels.size() <= 1, "Unexpected Subchannel count: %s", subchannels); - final Subchannel subchannel; + // Delegate to child pick_first LB for address management. + // Shutdown existing child LB if addresses become empty. if (newBackendAddrList.isEmpty()) { - if (subchannels.size() == 1) { - subchannel = subchannels.values().iterator().next(); - subchannel.shutdown(); - subchannels = Collections.emptyMap(); + if (pickFirstLb != null) { + pickFirstLb.shutdown(); + pickFirstLb = null; } break; } List eagList = new ArrayList<>(); - // Because for PICK_FIRST, we create a single Subchannel for all addresses, we have to - // attach the tokens to the EAG attributes and use TokenAttachingLoadRecorder to put them on - // headers. - // - // The PICK_FIRST code path doesn't cache Subchannels. + // Attach tokens to EAG attributes for TokenAttachingTracerFactory to retrieve. for (BackendAddressGroup bag : newBackendAddrList) { EquivalentAddressGroup origEag = bag.getAddresses(); Attributes eagAttrs = origEag.getAttributes(); @@ -542,30 +569,24 @@ private void updateServerList( } eagList.add(new EquivalentAddressGroup(origEag.getAddresses(), eagAttrs)); } - if (subchannels.isEmpty()) { - subchannel = - helper.createSubchannel( - CreateSubchannelArgs.newBuilder() - .setAddresses(eagList) - .setAttributes(createSubchannelAttrs()) - .build()); - subchannel.start(new SubchannelStateListener() { - @Override - public void onSubchannelState(ConnectivityStateInfo newState) { - handleSubchannelState(subchannel, newState); - } - }); - if (requestConnectionPending) { - subchannel.requestConnection(); - requestConnectionPending = false; - } - } else { - subchannel = subchannels.values().iterator().next(); - subchannel.updateAddresses(eagList); + // Always shutdown and recreate the child LB when addresses change to avoid + // calling Subchannel.updateAddresses(). This ensures we use the new dualstack- + // compatible path where the child LB creates fresh subchannels. + if (pickFirstLb != null) { + pickFirstLb.shutdown(); + } + pickFirstLb = pickFirstLbProvider.newLoadBalancer(new PickFirstLbHelper()); + // Pass addresses to child LB. + pickFirstLb.acceptResolvedAddresses( + ResolvedAddresses.newBuilder() + .setAddresses(eagList) + .build()); + if (requestConnectionPending) { + pickFirstLb.requestConnection(); + requestConnectionPending = false; } - subchannels = Collections.singletonMap(eagList, subchannel); - newBackendList.add( - new BackendEntry(subchannel, new TokenAttachingTracerFactory(loadRecorder))); + // Store the load recorder for token attachment. + currentPickFirstLoadRecorder = loadRecorder; break; default: throw new AssertionError("Missing case for " + config.getMode()); @@ -842,7 +863,11 @@ private void cleanUp() { private void maybeUpdatePicker() { List pickList; ConnectivityState state; - if (backendList.isEmpty()) { + // For PICK_FIRST mode with delegation, check if child LB exists instead of backendList. + boolean hasBackends = config.getMode() == Mode.PICK_FIRST + ? pickFirstLb != null + : !backendList.isEmpty(); + if (!hasBackends) { // Note balancer (is working) may enforce using fallback backends, and that fallback may // fail. So we should check if currently in fallback first. if (usingFallbackBackends) { @@ -894,26 +919,12 @@ private void maybeUpdatePicker() { } break; case PICK_FIRST: { - checkState(backendList.size() == 1, "Excessive backend entries: %s", backendList); - BackendEntry onlyEntry = backendList.get(0); - ConnectivityStateInfo stateInfo = - onlyEntry.subchannel.getAttributes().get(STATE_INFO).get(); - state = stateInfo.getState(); - switch (state) { - case READY: - pickList = Collections.singletonList(onlyEntry); - break; - case TRANSIENT_FAILURE: - pickList = - Collections.singletonList(new ErrorEntry(stateInfo.getStatus())); - break; - case CONNECTING: - pickList = Collections.singletonList(BUFFER_ENTRY); - break; - default: - pickList = Collections.singletonList( - new IdleSubchannelEntry(onlyEntry.subchannel, syncContext)); - } + // Use child LB's state and picker. Wrap the picker for token attachment. + state = pickFirstLbState; + TokenAttachingTracerFactory tracerFactory = + new TokenAttachingTracerFactory(currentPickFirstLoadRecorder); + pickList = Collections.singletonList( + new ChildLbPickerEntry(pickFirstLbPicker, tracerFactory)); break; } default: @@ -983,7 +994,7 @@ public boolean equals(Object other) { @VisibleForTesting interface RoundRobinEntry { - PickResult picked(Metadata headers); + PickResult picked(PickSubchannelArgs args); } @VisibleForTesting @@ -1024,7 +1035,8 @@ static final class BackendEntry implements RoundRobinEntry { } @Override - public PickResult picked(Metadata headers) { + public PickResult picked(PickSubchannelArgs args) { + Metadata headers = args.getHeaders(); headers.discardAll(GrpclbConstants.TOKEN_METADATA_KEY); if (token != null) { headers.put(GrpclbConstants.TOKEN_METADATA_KEY, token); @@ -1065,7 +1077,7 @@ static final class IdleSubchannelEntry implements RoundRobinEntry { } @Override - public PickResult picked(Metadata headers) { + public PickResult picked(PickSubchannelArgs args) { if (connectionRequested.compareAndSet(false, true)) { syncContext.execute(new Runnable() { @Override @@ -1108,7 +1120,7 @@ static final class ErrorEntry implements RoundRobinEntry { } @Override - public PickResult picked(Metadata headers) { + public PickResult picked(PickSubchannelArgs args) { return result; } @@ -1132,6 +1144,52 @@ public String toString() { } } + /** + * Entry that wraps a child LB's picker for PICK_FIRST mode delegation. + * Attaches TokenAttachingTracerFactory to the pick result for token propagation. + */ + @VisibleForTesting + static final class ChildLbPickerEntry implements RoundRobinEntry { + private final SubchannelPicker childPicker; + private final TokenAttachingTracerFactory tracerFactory; + + ChildLbPickerEntry(SubchannelPicker childPicker, TokenAttachingTracerFactory tracerFactory) { + this.childPicker = checkNotNull(childPicker, "childPicker"); + this.tracerFactory = checkNotNull(tracerFactory, "tracerFactory"); + } + + @Override + public PickResult picked(PickSubchannelArgs args) { + PickResult childResult = childPicker.pickSubchannel(args); + if (childResult.getSubchannel() == null) { + // No subchannel (e.g., buffer, error), return as-is. + return childResult; + } + // Wrap the pick result to attach tokens via the tracer factory. + return PickResult.withSubchannel(childResult.getSubchannel(), tracerFactory); + } + + @Override + public int hashCode() { + return Objects.hashCode(childPicker, tracerFactory); + } + + @Override + public boolean equals(Object other) { + if (!(other instanceof ChildLbPickerEntry)) { + return false; + } + ChildLbPickerEntry that = (ChildLbPickerEntry) other; + return Objects.equal(childPicker, that.childPicker) + && Objects.equal(tracerFactory, that.tracerFactory); + } + + @Override + public String toString() { + return "ChildLbPickerEntry(" + childPicker + ")"; + } + } + @VisibleForTesting static final class RoundRobinPicker extends SubchannelPicker { @VisibleForTesting @@ -1174,7 +1232,7 @@ public PickResult pickSubchannel(PickSubchannelArgs args) { if (pickIndex == pickList.size()) { pickIndex = 0; } - return pick.picked(args.getHeaders()); + return pick.picked(args); } } @@ -1189,4 +1247,28 @@ public String toString() { return MoreObjects.toStringHelper(RoundRobinPicker.class).toString(); } } + + /** + * Helper for the child pick_first LB in PICK_FIRST mode. Intercepts updateBalancingState() + * to store state and trigger the grpclb picker update with drops and token attachment. + */ + private final class PickFirstLbHelper extends ForwardingLoadBalancerHelper { + + @Override + protected Helper delegate() { + return helper; + } + + @Override + public void updateBalancingState(ConnectivityState newState, SubchannelPicker newPicker) { + pickFirstLbState = newState; + pickFirstLbPicker = newPicker; + // Trigger name resolution refresh on TRANSIENT_FAILURE or IDLE, similar to ROUND_ROBIN. + if (newState == TRANSIENT_FAILURE || newState == IDLE) { + helper.refreshNameResolution(); + } + maybeUseFallbackBackends(); + maybeUpdatePicker(); + } + } } diff --git a/grpclb/src/test/java/io/grpc/grpclb/GrpclbLoadBalancerTest.java b/grpclb/src/test/java/io/grpc/grpclb/GrpclbLoadBalancerTest.java index e489129676a..62eb151651e 100644 --- a/grpclb/src/test/java/io/grpc/grpclb/GrpclbLoadBalancerTest.java +++ b/grpclb/src/test/java/io/grpc/grpclb/GrpclbLoadBalancerTest.java @@ -779,7 +779,9 @@ public void receiveNoBackendAndBalancerAddress() { verify(helper).updateBalancingState(eq(TRANSIENT_FAILURE), pickerCaptor.capture()); RoundRobinPicker picker = (RoundRobinPicker) pickerCaptor.getValue(); assertThat(picker.dropList).isEmpty(); - Status error = Iterables.getOnlyElement(picker.pickList).picked(new Metadata()).getStatus(); + PickSubchannelArgs args = mock(PickSubchannelArgs.class); + when(args.getHeaders()).thenReturn(new Metadata()); + Status error = Iterables.getOnlyElement(picker.pickList).picked(args).getStatus(); assertThat(error.getCode()).isEqualTo(Code.UNAVAILABLE); assertThat(error.getDescription()).isEqualTo("No backend or balancer addresses found"); } @@ -1915,90 +1917,47 @@ public void grpclbWorking_pickFirstMode() throws Exception { lbResponseObserver.onNext(buildInitialResponse()); lbResponseObserver.onNext(buildLbResponse(backends1)); - inOrder.verify(helper).createSubchannel(createSubchannelArgsCaptor.capture()); + // With delegation to child pick_first LB, subchannel is created by the child. + // The child pick_first eagerly connects, so initial state is CONNECTING (not IDLE). + verify(helper, atLeast(1)).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"))); - // Initially IDLE - inOrder.verify(helper).updateBalancingState(eq(IDLE), pickerCaptor.capture()); + // Child pick_first eagerly connects, so we start in CONNECTING + verify(helper, atLeast(1)).updateBalancingState(eq(CONNECTING), pickerCaptor.capture()); RoundRobinPicker picker0 = (RoundRobinPicker) pickerCaptor.getValue(); - // Only one subchannel is created + // One subchannel is created by the child LB assertThat(mockSubchannels).hasSize(1); Subchannel subchannel = mockSubchannels.poll(); 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); - assertThat(picker2.pickList).containsExactly(new ErrorEntry(error)); // READY deliverSubchannelState(subchannel, ConnectivityStateInfo.forNonError(READY)); - inOrder.verify(helper).updateBalancingState(eq(READY), pickerCaptor.capture()); - RoundRobinPicker picker3 = (RoundRobinPicker) pickerCaptor.getValue(); - assertThat(picker3.dropList).containsExactly(null, null); - assertThat(picker3.pickList).containsExactly( - new BackendEntry(subchannel, new TokenAttachingTracerFactory(getLoadRecorder()))); - + verify(helper, atLeast(1)).updateBalancingState(eq(READY), pickerCaptor.capture()); + RoundRobinPicker picker1 = (RoundRobinPicker) pickerCaptor.getValue(); + assertThat(picker1.dropList).containsExactly(null, null); - // New server list with drops + // New server list with drops - child LB is recreated (no updateAddresses) 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)); - // new addresses will be updated to the existing subchannel - // createSubchannel() has ever been called only once - verify(helper, times(1)).createSubchannel(any(CreateSubchannelArgs.class)); - assertThat(mockSubchannels).isEmpty(); - verify(subchannel).updateAddresses( - eq(Arrays.asList( - new EquivalentAddressGroup(backends2.get(0).addr, eagAttrsWithToken("token0001")), - new EquivalentAddressGroup(backends2.get(2).addr, - eagAttrsWithToken("token0004"))))); - inOrder.verify(helper).updateBalancingState(eq(READY), pickerCaptor.capture()); - RoundRobinPicker picker4 = (RoundRobinPicker) pickerCaptor.getValue(); - assertThat(picker4.dropList).containsExactly( - null, new DropEntry(getLoadRecorder(), "token0003"), null); - assertThat(picker4.pickList).containsExactly( - new BackendEntry(subchannel, new TokenAttachingTracerFactory(getLoadRecorder()))); - - // Subchannel goes IDLE, but PICK_FIRST will not try to reconnect - deliverSubchannelState(subchannel, ConnectivityStateInfo.forNonError(IDLE)); - inOrder.verify(helper).updateBalancingState(eq(IDLE), pickerCaptor.capture()); - RoundRobinPicker picker5 = (RoundRobinPicker) pickerCaptor.getValue(); - verify(subchannel, never()).requestConnection(); + // With delegation, child LB is recreated so a new subchannel is created. + // No updateAddresses() call - this is the key behavioral change. + verify(subchannel, never()).updateAddresses(any()); + verify(helper, atLeast(2)).createSubchannel(any(CreateSubchannelArgs.class)); - // ... until it's selected - PickSubchannelArgs args = mock(PickSubchannelArgs.class); - PickResult pick = picker5.pickSubchannel(args); - assertThat(pick).isSameInstanceAs(PickResult.withNoResult()); - verify(subchannel).requestConnection(); - - // ... or requested by application - balancer.requestConnection(); - verify(subchannel, times(2)).requestConnection(); + // Verify drop list is updated after new child LB updates state + verify(helper, atLeast(1)).updateBalancingState(any(), pickerCaptor.capture()); + RoundRobinPicker picker2 = (RoundRobinPicker) pickerCaptor.getValue(); + assertThat(picker2.dropList).containsExactly( + null, new DropEntry(getLoadRecorder(), "token0003"), null); // PICK_FIRST doesn't use subchannelPool verify(subchannelPool, never()) @@ -2009,8 +1968,6 @@ public void grpclbWorking_pickFirstMode() throws Exception { @Test public void grpclbWorking_pickFirstMode_lbSendsEmptyAddress() throws Exception { - InOrder inOrder = inOrder(helper); - List grpclbBalancerList = createResolvedBalancerAddresses(1); deliverResolvedAddresses( Collections.emptyList(), @@ -2031,96 +1988,56 @@ public void grpclbWorking_pickFirstMode_lbSendsEmptyAddress() throws Exception { 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)); - inOrder.verify(helper).createSubchannel(createSubchannelArgsCaptor.capture()); + // With delegation to child pick_first LB + verify(helper, atLeast(1)).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"))); - // Initially IDLE - inOrder.verify(helper).updateBalancingState(eq(IDLE), pickerCaptor.capture()); - RoundRobinPicker picker0 = (RoundRobinPicker) pickerCaptor.getValue(); - - // Only one subchannel is created + // Child pick_first eagerly connects + verify(helper, atLeast(1)).updateBalancingState(eq(CONNECTING), any()); assertThat(mockSubchannels).hasSize(1); Subchannel subchannel = mockSubchannels.poll(); - 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); - assertThat(picker2.pickList).containsExactly(new ErrorEntry(error)); // READY deliverSubchannelState(subchannel, ConnectivityStateInfo.forNonError(READY)); - inOrder.verify(helper).updateBalancingState(eq(READY), pickerCaptor.capture()); - RoundRobinPicker picker3 = (RoundRobinPicker) pickerCaptor.getValue(); - assertThat(picker3.dropList).containsExactly(null, null); - assertThat(picker3.pickList).containsExactly( - new BackendEntry(subchannel, new TokenAttachingTracerFactory(getLoadRecorder()))); - - inOrder.verify(helper, never()) - .updateBalancingState(any(ConnectivityState.class), any(SubchannelPicker.class)); + verify(helper, atLeast(1)).updateBalancingState(eq(READY), pickerCaptor.capture()); - // Empty addresses from LB + // Empty addresses from LB - child LB is shutdown lbResponseObserver.onNext(buildLbResponse(Collections.emptyList())); - // new addresses will be updated to the existing subchannel - // createSubchannel() has ever been called only once - inOrder.verify(helper, never()).createSubchannel(any(CreateSubchannelArgs.class)); - assertThat(mockSubchannels).isEmpty(); + // Child LB is shutdown (which shuts down its subchannel) verify(subchannel).shutdown(); // RPC error status includes message of no backends provided by balancer - inOrder.verify(helper).updateBalancingState(eq(TRANSIENT_FAILURE), pickerCaptor.capture()); + 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)); - // 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()); - subchannel = mockSubchannels.poll(); + // A new child LB is created with new subchannel + verify(helper, atLeast(2)).createSubchannel(any(CreateSubchannelArgs.class)); + assertThat(mockSubchannels).hasSize(1); + Subchannel subchannel2 = mockSubchannels.poll(); // Subchannel became READY - deliverSubchannelState(subchannel, ConnectivityStateInfo.forNonError(CONNECTING)); - deliverSubchannelState(subchannel, ConnectivityStateInfo.forNonError(READY)); - inOrder.verify(helper).updateBalancingState(eq(READY), pickerCaptor.capture()); + deliverSubchannelState(subchannel2, ConnectivityStateInfo.forNonError(READY)); + verify(helper, atLeast(2)).updateBalancingState(eq(READY), pickerCaptor.capture()); RoundRobinPicker picker4 = (RoundRobinPicker) pickerCaptor.getValue(); - assertThat(picker4.pickList).containsExactly( - new BackendEntry(subchannel, new TokenAttachingTracerFactory(getLoadRecorder()))); + assertThat(picker4.dropList).containsExactly( + null, new DropEntry(getLoadRecorder(), "token0003"), null); } @Test @@ -2162,8 +2079,6 @@ public void pickFirstMode_serviceConfigTimeout_fallback() throws Exception { } private void pickFirstModeFallback(long timeout) throws Exception { - InOrder inOrder = inOrder(helper); - // Name resolver returns balancer and backend addresses List backendList = createResolvedBackendAddresses(2); List grpclbBalancerList = createResolvedBalancerAddresses(1); @@ -2179,8 +2094,8 @@ private void pickFirstModeFallback(long timeout) throws Exception { // Fallback timer expires with no response fakeClock.forwardTime(timeout, TimeUnit.MILLISECONDS); - // Entering fallback mode - inOrder.verify(helper).createSubchannel(createSubchannelArgsCaptor.capture()); + // Entering fallback mode - child LB is created for fallback backends + verify(helper, atLeast(1)).createSubchannel(createSubchannelArgsCaptor.capture()); CreateSubchannelArgs createSubchannelArgs = createSubchannelArgsCaptor.getValue(); assertThat(createSubchannelArgs.getAddresses()) .containsExactly(backendList.get(0), backendList.get(1)); @@ -2188,45 +2103,24 @@ private void pickFirstModeFallback(long timeout) throws Exception { assertThat(mockSubchannels).hasSize(1); Subchannel subchannel = mockSubchannels.poll(); - // Initially IDLE - inOrder.verify(helper).updateBalancingState(eq(IDLE), pickerCaptor.capture()); - RoundRobinPicker picker0 = (RoundRobinPicker) pickerCaptor.getValue(); + // Child pick_first eagerly connects + verify(helper, atLeast(1)).updateBalancingState(eq(CONNECTING), any()); // READY deliverSubchannelState(subchannel, ConnectivityStateInfo.forNonError(READY)); - inOrder.verify(helper).updateBalancingState(eq(READY), pickerCaptor.capture()); - RoundRobinPicker picker1 = (RoundRobinPicker) pickerCaptor.getValue(); - assertThat(picker1.dropList).containsExactly(null, null); - assertThat(picker1.pickList).containsExactly( - new BackendEntry(subchannel, new TokenAttachingTracerFactory(null))); - - assertThat(picker0.dropList).containsExactly(null, null); - assertThat(picker0.pickList).containsExactly(new IdleSubchannelEntry(subchannel, syncContext)); - + verify(helper, atLeast(1)).updateBalancingState(eq(READY), pickerCaptor.capture()); // Finally, an LB response, 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)); - // new addresses will be updated to the existing subchannel - // createSubchannel() has ever been called only once - inOrder.verify(helper, never()).createSubchannel(any(CreateSubchannelArgs.class)); - 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()))); + // With delegation, child LB is recreated (no updateAddresses) + verify(subchannel, never()).updateAddresses(any()); + // New subchannel is created for LB-provided backends + verify(helper, atLeast(2)).createSubchannel(any(CreateSubchannelArgs.class)); // PICK_FIRST doesn't use subchannelPool verify(subchannelPool, never()) @@ -2304,20 +2198,19 @@ public void switchMode() throws Exception { .build())); // Simulate receiving LB response - inOrder.verify(helper, never()) - .updateBalancingState(any(ConnectivityState.class), any(SubchannelPicker.class)); lbResponseObserver.onNext(buildInitialResponse()); lbResponseObserver.onNext(buildLbResponse(backends1)); - // PICK_FIRST Subchannel - inOrder.verify(helper).createSubchannel(createSubchannelArgsCaptor.capture()); + // PICK_FIRST - with delegation, child LB creates subchannel + verify(helper, atLeast(1)).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"))); - inOrder.verify(helper).updateBalancingState(eq(IDLE), any(SubchannelPicker.class)); + // Child pick_first eagerly connects, so initial state is CONNECTING + verify(helper, atLeast(1)).updateBalancingState(eq(CONNECTING), any(SubchannelPicker.class)); } private static Attributes eagAttrsWithToken(String token) { @@ -2392,20 +2285,19 @@ public void switchMode_nullLbPolicy() throws Exception { .build())); // Simulate receiving LB response - inOrder.verify(helper, never()) - .updateBalancingState(any(ConnectivityState.class), any(SubchannelPicker.class)); lbResponseObserver.onNext(buildInitialResponse()); lbResponseObserver.onNext(buildLbResponse(backends1)); - // PICK_FIRST Subchannel - inOrder.verify(helper).createSubchannel(createSubchannelArgsCaptor.capture()); + // PICK_FIRST - with delegation, child LB creates subchannel + verify(helper, atLeast(1)).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"))); - inOrder.verify(helper).updateBalancingState(eq(IDLE), any(SubchannelPicker.class)); + // Child pick_first eagerly connects, so initial state is CONNECTING + verify(helper, atLeast(1)).updateBalancingState(eq(CONNECTING), any(SubchannelPicker.class)); } @Test From 5dab3a7dab50c57e81a282576e959bb4e46f4253 Mon Sep 17 00:00:00 2001 From: MV Shiva Prasad Date: Thu, 11 Dec 2025 10:59:52 +0530 Subject: [PATCH 2/6] checkstyle and comment --- grpclb/src/main/java/io/grpc/grpclb/GrpclbState.java | 7 +++---- 1 file changed, 3 insertions(+), 4 deletions(-) diff --git a/grpclb/src/main/java/io/grpc/grpclb/GrpclbState.java b/grpclb/src/main/java/io/grpc/grpclb/GrpclbState.java index bf5208140cb..ab954ccb190 100644 --- a/grpclb/src/main/java/io/grpc/grpclb/GrpclbState.java +++ b/grpclb/src/main/java/io/grpc/grpclb/GrpclbState.java @@ -38,7 +38,6 @@ import io.grpc.Context; import io.grpc.EquivalentAddressGroup; import io.grpc.LoadBalancer; -import io.grpc.LoadBalancer.CreateSubchannelArgs; import io.grpc.LoadBalancer.FixedResultPicker; import io.grpc.LoadBalancer.Helper; import io.grpc.LoadBalancer.PickResult; @@ -46,7 +45,6 @@ import io.grpc.LoadBalancer.ResolvedAddresses; import io.grpc.LoadBalancer.Subchannel; import io.grpc.LoadBalancer.SubchannelPicker; -import io.grpc.LoadBalancer.SubchannelStateListener; import io.grpc.LoadBalancerProvider; import io.grpc.LoadBalancerRegistry; import io.grpc.ManagedChannel; @@ -57,7 +55,6 @@ import io.grpc.grpclb.SubchannelPool.PooledSubchannelStateListener; import io.grpc.internal.BackoffPolicy; import io.grpc.internal.TimeProvider; -import io.grpc.util.ForwardingLoadBalancerHelper; import io.grpc.lb.v1.ClientStats; import io.grpc.lb.v1.InitialLoadBalanceRequest; import io.grpc.lb.v1.InitialLoadBalanceResponse; @@ -68,6 +65,7 @@ import io.grpc.lb.v1.Server; import io.grpc.lb.v1.ServerList; import io.grpc.stub.StreamObserver; +import io.grpc.util.ForwardingLoadBalancerHelper; import java.net.InetAddress; import java.net.InetSocketAddress; import java.net.UnknownHostException; @@ -1166,7 +1164,8 @@ public PickResult picked(PickSubchannelArgs args) { return childResult; } // Wrap the pick result to attach tokens via the tracer factory. - return PickResult.withSubchannel(childResult.getSubchannel(), tracerFactory); + return PickResult.withSubchannel( + childResult.getSubchannel(), tracerFactory, childResult.getAuthorityOverride()); } @Override From 8f0a01167e1cdaeab49d63cddc60df14c3bcd994 Mon Sep 17 00:00:00 2001 From: MV Shiva Prasad Date: Thu, 11 Dec 2025 11:33:54 +0530 Subject: [PATCH 3/6] add util in bazel --- grpclb/BUILD.bazel | 1 + 1 file changed, 1 insertion(+) diff --git a/grpclb/BUILD.bazel b/grpclb/BUILD.bazel index 4612968ebcd..ca9975b7ce6 100644 --- a/grpclb/BUILD.bazel +++ b/grpclb/BUILD.bazel @@ -17,6 +17,7 @@ java_library( "//context", "//core:internal", "//stub", + "//util", "@com_google_protobuf//:protobuf_java_util", "@io_grpc_grpc_proto//:grpclb_load_balancer_java_proto", artifact("com.google.code.findbugs:jsr305"), From edd38c5f3d03245137244f2fee277caa6a2e3517 Mon Sep 17 00:00:00 2001 From: MV Shiva Prasad Date: Thu, 11 Dec 2025 11:55:50 +0530 Subject: [PATCH 4/6] non null check --- grpclb/src/main/java/io/grpc/grpclb/GrpclbState.java | 8 ++++++-- 1 file changed, 6 insertions(+), 2 deletions(-) diff --git a/grpclb/src/main/java/io/grpc/grpclb/GrpclbState.java b/grpclb/src/main/java/io/grpc/grpclb/GrpclbState.java index ab954ccb190..fde52ebdeb0 100644 --- a/grpclb/src/main/java/io/grpc/grpclb/GrpclbState.java +++ b/grpclb/src/main/java/io/grpc/grpclb/GrpclbState.java @@ -225,8 +225,9 @@ public void onSubchannelState( } else { this.subchannelPool = null; } - this.pickFirstLbProvider = - LoadBalancerRegistry.getDefaultRegistry().getProvider("pick_first"); + this.pickFirstLbProvider = checkNotNull( + LoadBalancerRegistry.getDefaultRegistry().getProvider("pick_first"), + "pick_first balancer not available"); this.time = checkNotNull(time, "time provider"); this.stopwatch = checkNotNull(stopwatch, "stopwatch"); this.timerService = checkNotNull(helper.getScheduledExecutorService(), "timerService"); @@ -574,6 +575,9 @@ private void updateServerList( pickFirstLb.shutdown(); } pickFirstLb = pickFirstLbProvider.newLoadBalancer(new PickFirstLbHelper()); + // Reset the child LB state, since we created a new one. + pickFirstLbState = CONNECTING; + pickFirstLbPicker = new FixedResultPicker(PickResult.withNoResult()); // Pass addresses to child LB. pickFirstLb.acceptResolvedAddresses( ResolvedAddresses.newBuilder() From b4eb6ae73867f4f2942966c3cfdfa666381e281f Mon Sep 17 00:00:00 2001 From: MV Shiva Prasad Date: Thu, 11 Dec 2025 12:31:42 +0530 Subject: [PATCH 5/6] simple delegate, not trying to avoid updateAddresses() within child lb --- cronet/build.gradle | 2 +- .../src/main/java/io/grpc/grpclb/GrpclbState.java | 13 ++++--------- .../io/grpc/grpclb/GrpclbLoadBalancerTest.java | 15 +++++++-------- 3 files changed, 12 insertions(+), 18 deletions(-) diff --git a/cronet/build.gradle b/cronet/build.gradle index d6d773a97e4..3d4a68bad81 100644 --- a/cronet/build.gradle +++ b/cronet/build.gradle @@ -58,7 +58,7 @@ dependencies { task javadocs(type: Javadoc) { source = android.sourceSets.main.java.srcDirs - classpath += files(android.getBootClasspath()) + // classpath += files(android.getBootClasspath()) classpath += files({ android.libraryVariants.collect { variant -> variant.javaCompileProvider.get().classpath diff --git a/grpclb/src/main/java/io/grpc/grpclb/GrpclbState.java b/grpclb/src/main/java/io/grpc/grpclb/GrpclbState.java index fde52ebdeb0..4ae21651491 100644 --- a/grpclb/src/main/java/io/grpc/grpclb/GrpclbState.java +++ b/grpclb/src/main/java/io/grpc/grpclb/GrpclbState.java @@ -568,16 +568,11 @@ private void updateServerList( } eagList.add(new EquivalentAddressGroup(origEag.getAddresses(), eagAttrs)); } - // Always shutdown and recreate the child LB when addresses change to avoid - // calling Subchannel.updateAddresses(). This ensures we use the new dualstack- - // compatible path where the child LB creates fresh subchannels. - if (pickFirstLb != null) { - pickFirstLb.shutdown(); + + if (pickFirstLb == null) { + pickFirstLb = pickFirstLbProvider.newLoadBalancer(new PickFirstLbHelper()); } - pickFirstLb = pickFirstLbProvider.newLoadBalancer(new PickFirstLbHelper()); - // Reset the child LB state, since we created a new one. - pickFirstLbState = CONNECTING; - pickFirstLbPicker = new FixedResultPicker(PickResult.withNoResult()); + // Pass addresses to child LB. pickFirstLb.acceptResolvedAddresses( ResolvedAddresses.newBuilder() diff --git a/grpclb/src/test/java/io/grpc/grpclb/GrpclbLoadBalancerTest.java b/grpclb/src/test/java/io/grpc/grpclb/GrpclbLoadBalancerTest.java index 62eb151651e..2f0ca4228e2 100644 --- a/grpclb/src/test/java/io/grpc/grpclb/GrpclbLoadBalancerTest.java +++ b/grpclb/src/test/java/io/grpc/grpclb/GrpclbLoadBalancerTest.java @@ -1948,10 +1948,9 @@ public void grpclbWorking_pickFirstMode() throws Exception { new ServerEntry("127.0.0.1", 2020, "token0004")); lbResponseObserver.onNext(buildLbResponse(backends2)); - // With delegation, child LB is recreated so a new subchannel is created. - // No updateAddresses() call - this is the key behavioral change. - verify(subchannel, never()).updateAddresses(any()); - verify(helper, atLeast(2)).createSubchannel(any(CreateSubchannelArgs.class)); + // the child LB is not recreated. A single subchannel is created + // for the lifetime of the test. The child LB will internally call updateAddresses. + verify(helper, times(1)).createSubchannel(any(CreateSubchannelArgs.class)); // Verify drop list is updated after new child LB updates state verify(helper, atLeast(1)).updateBalancingState(any(), pickerCaptor.capture()); @@ -2117,10 +2116,10 @@ private void pickFirstModeFallback(long timeout) throws Exception { lbResponseObserver.onNext(buildInitialResponse()); lbResponseObserver.onNext(buildLbResponse(backends1)); - // With delegation, child LB is recreated (no updateAddresses) - verify(subchannel, never()).updateAddresses(any()); - // New subchannel is created for LB-provided backends - verify(helper, atLeast(2)).createSubchannel(any(CreateSubchannelArgs.class)); + // 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)); // PICK_FIRST doesn't use subchannelPool verify(subchannelPool, never()) From 6e7ebd37552ef739491ddd85ea42bbafce530ce8 Mon Sep 17 00:00:00 2001 From: MV Shiva Prasad Date: Thu, 11 Dec 2025 12:37:13 +0530 Subject: [PATCH 6/6] comments --- cronet/build.gradle | 2 +- grpclb/src/test/java/io/grpc/grpclb/GrpclbLoadBalancerTest.java | 2 +- 2 files changed, 2 insertions(+), 2 deletions(-) diff --git a/cronet/build.gradle b/cronet/build.gradle index 3d4a68bad81..d6d773a97e4 100644 --- a/cronet/build.gradle +++ b/cronet/build.gradle @@ -58,7 +58,7 @@ dependencies { task javadocs(type: Javadoc) { source = android.sourceSets.main.java.srcDirs - // classpath += files(android.getBootClasspath()) + classpath += files(android.getBootClasspath()) classpath += files({ android.libraryVariants.collect { variant -> variant.javaCompileProvider.get().classpath diff --git a/grpclb/src/test/java/io/grpc/grpclb/GrpclbLoadBalancerTest.java b/grpclb/src/test/java/io/grpc/grpclb/GrpclbLoadBalancerTest.java index 2f0ca4228e2..0c994934d1d 100644 --- a/grpclb/src/test/java/io/grpc/grpclb/GrpclbLoadBalancerTest.java +++ b/grpclb/src/test/java/io/grpc/grpclb/GrpclbLoadBalancerTest.java @@ -1941,7 +1941,7 @@ public void grpclbWorking_pickFirstMode() throws Exception { RoundRobinPicker picker1 = (RoundRobinPicker) pickerCaptor.getValue(); assertThat(picker1.dropList).containsExactly(null, null); - // New server list with drops - child LB is recreated (no updateAddresses) + // New server list with drops List backends2 = Arrays.asList( new ServerEntry("127.0.0.1", 2000, "token0001"), new ServerEntry("token0003"), // drop