core,api,xds: Implement load balancing policy delay plumbing

AgraVator · AgraVator · commit b50ca84e30c1 · 2026-05-13T22:15:34.000+05:30
This commit implements the plumbing required to propagate delay reason tokens from load balancing policies up to the transport layer and tracers, as specified in the LB policy delay design.
diff --git a/api/src/main/java/io/grpc/ClientStreamTracer.java b/api/src/main/java/io/grpc/ClientStreamTracer.java
@@ -57,6 +57,23 @@ public void streamCreated(@Grpc.TransportAttr Attributes transportAttrs, Metadat
   public void createPendingStream() {
   }
 
+  /**
+   * A delay segment started with a specific reason during load balancing.
+   *
+   * @param reasonToken the reason for the delay, e.g., "pick_first:connecting"
+   * @since 1.82.0
+   */
+  public void delayStarted(String reasonToken) {
+  }
+
+  /**
+   * The current delay segment ended.
+   *
+   * @since 1.82.0
+   */
+  public void delayEnded() {
+  }
+
   /**
    * Headers has been sent to the socket.
    */
diff --git a/api/src/main/java/io/grpc/LoadBalancer.java b/api/src/main/java/io/grpc/LoadBalancer.java
@@ -549,25 +549,30 @@ public static final class PickResult {
     // True if the result is created by withDrop()
     private final boolean drop;
     @Nullable private final String authorityOverride;
+    @Nullable private final String delayReasonToken;
 
     private PickResult(
         @Nullable Subchannel subchannel, @Nullable ClientStreamTracer.Factory streamTracerFactory,
         Status status, boolean drop) {
-      this.subchannel = subchannel;
-      this.streamTracerFactory = streamTracerFactory;
-      this.status = checkNotNull(status, "status");
-      this.drop = drop;
-      this.authorityOverride = null;
+      this(subchannel, streamTracerFactory, status, drop, null, null);
     }
 
     private PickResult(
         @Nullable Subchannel subchannel, @Nullable ClientStreamTracer.Factory streamTracerFactory,
         Status status, boolean drop, @Nullable String authorityOverride) {
+      this(subchannel, streamTracerFactory, status, drop, authorityOverride, null);
+    }
+
+    private PickResult(
+        @Nullable Subchannel subchannel, @Nullable ClientStreamTracer.Factory streamTracerFactory,
+        Status status, boolean drop, @Nullable String authorityOverride,
+        @Nullable String delayReasonToken) {
       this.subchannel = subchannel;
       this.streamTracerFactory = streamTracerFactory;
       this.status = checkNotNull(status, "status");
       this.drop = drop;
       this.authorityOverride = authorityOverride;
+      this.delayReasonToken = delayReasonToken;
     }
 
     /**
@@ -727,6 +732,22 @@ public static PickResult withNoResult() {
       return NO_RESULT;
     }
 
+    /**
+     * No decision could be made.  The RPC will stay buffered with a specific reason.
+     *
+     * @since 1.82.0
+     */
+    public static PickResult withNoResult(String delayReasonToken) {
+      Preconditions.checkNotNull(delayReasonToken, "delayReasonToken");
+      return new PickResult(null, null, Status.OK, false, null, delayReasonToken);
+    }
+
+    /** Returns the delay reason token if any. */
+    @Nullable
+    public String getDelayReasonToken() {
+      return delayReasonToken;
+    }
+
     /** Returns the authority override if any. */
     @ExperimentalApi("https://github.com/grpc/grpc-java/issues/11656")
     @Nullable
diff --git a/core/src/main/java/io/grpc/internal/DelayedClientTransport.java b/core/src/main/java/io/grpc/internal/DelayedClientTransport.java
@@ -157,7 +157,8 @@ public final ClientStream newStream(
         synchronized (lock) {
           PickerState newerState = pickerState;
           if (state == newerState) {
-            return createPendingStream(args, tracers, pickResult);
+            String token = pickResult != null ? pickResult.getDelayReasonToken() : null;
+            return createPendingStream(args, tracers, pickResult, token);
           }
           state = newerState;
         }
@@ -173,8 +174,8 @@ public final ClientStream newStream(
    */
   @GuardedBy("lock")
   private PendingStream createPendingStream(PickSubchannelArgs args, ClientStreamTracer[] tracers,
-      PickResult pickResult) {
-    PendingStream pendingStream = new PendingStream(args, tracers);
+      PickResult pickResult, @Nullable String delayReasonToken) {
+    PendingStream pendingStream = new PendingStream(args, tracers, delayReasonToken);
     if (args.getCallOptions().isWaitForReady() && pickResult != null && pickResult.hasResult()) {
       pendingStream.lastPickStatus = pickResult.getStatus();
     }
@@ -303,6 +304,7 @@ final void reprocess(@Nullable SubchannelPicker picker) {
       final ClientTransport transport = GrpcUtil.getTransportFromPickResult(pickResult,
           callOptions.isWaitForReady());
       if (transport != null) {
+        stream.endDelay();
         Executor executor = defaultAppExecutor;
         // createRealStream may be expensive. It will start real streams on the transport. If
         // there are pending requests, they will be serialized too, which may be expensive. Since
@@ -315,7 +317,9 @@ final void reprocess(@Nullable SubchannelPicker picker) {
           executor.execute(runnable);
         }
         toRemove.add(stream);
-      }  // else: stay pending
+      } else { // stay pending
+        stream.updateDelayReason(pickResult.getDelayReasonToken());
+      }
     }
 
     synchronized (lock) {
@@ -361,11 +365,43 @@ private class PendingStream extends DelayedStream {
     private final Context context = Context.current();
     private final ClientStreamTracer[] tracers;
     private volatile Status lastPickStatus;
+    @Nullable private String delayReasonToken;
 
-    private PendingStream(PickSubchannelArgs args, ClientStreamTracer[] tracers) {
+    private PendingStream(PickSubchannelArgs args, ClientStreamTracer[] tracers, @Nullable String initialToken) {
       super("connecting_and_lb");
       this.args = args;
       this.tracers = tracers;
+      this.delayReasonToken = initialToken;
+      if (initialToken != null) {
+        for (ClientStreamTracer tracer : tracers) {
+          tracer.delayStarted(initialToken);
+        }
+      }
+    }
+
+    void updateDelayReason(String newToken) {
+      if (!java.util.Objects.equals(delayReasonToken, newToken)) {
+        if (delayReasonToken != null) {
+          for (ClientStreamTracer tracer : tracers) {
+            tracer.delayEnded();
+          }
+        }
+        delayReasonToken = newToken;
+        if (newToken != null) {
+          for (ClientStreamTracer tracer : tracers) {
+            tracer.delayStarted(newToken);
+          }
+        }
+      }
+    }
+
+    void endDelay() {
+      if (delayReasonToken != null) {
+        for (ClientStreamTracer tracer : tracers) {
+          tracer.delayEnded();
+        }
+        delayReasonToken = null;
+      }
     }
 
     /** Runnable may be null. */
@@ -391,6 +427,7 @@ private Runnable createRealStream(ClientTransport transport, String authorityOve
 
     @Override
     public void cancel(Status reason) {
+      endDelay();
       super.cancel(reason);
       synchronized (lock) {
         if (reportTransportTerminated != null) {
diff --git a/core/src/main/java/io/grpc/internal/PickFirstLoadBalancer.java b/core/src/main/java/io/grpc/internal/PickFirstLoadBalancer.java
@@ -38,6 +38,7 @@
  * list and sticking to the first that works.
  */
 final class PickFirstLoadBalancer extends LoadBalancer {
+  private static final PickResult CONNECTING_RESULT = PickResult.withNoResult("pick_first:connecting");
   private final Helper helper;
   private Subchannel subchannel;
   private ConnectivityState currentState = IDLE;
@@ -83,7 +84,7 @@ public void onSubchannelState(ConnectivityStateInfo stateInfo) {
 
       // The channel state does not get updated when doing name resolving today, so for the moment
       // let LB report CONNECTION and call subchannel.requestConnection() immediately.
-      updateBalancingState(CONNECTING, new FixedResultPicker(PickResult.withNoResult()));
+      updateBalancingState(CONNECTING, new FixedResultPicker(CONNECTING_RESULT));
       subchannel.requestConnection();
     } else {
       subchannel.updateAddresses(servers);
@@ -135,7 +136,7 @@ private void processSubchannelState(Subchannel subchannel, ConnectivityStateInfo
       case CONNECTING:
         // It's safe to use RequestConnectionPicker here, so when coming from IDLE we could leave
         // the current picker in-place. But ignoring the potential optimization is simpler.
-        picker = new FixedResultPicker(PickResult.withNoResult());
+        picker = new FixedResultPicker(CONNECTING_RESULT);
         break;
       case READY:
         picker = new FixedResultPicker(PickResult.withSubchannel(subchannel));
diff --git a/core/src/test/java/io/grpc/internal/DelayedClientTransportTest.java b/core/src/test/java/io/grpc/internal/DelayedClientTransportTest.java
@@ -772,6 +772,54 @@ public void pendingStream_appendTimeoutInsight_waitForReady_withLastPickFailure(
             + " connecting_and_lb_delay=[0-9]+ns, was_still_waiting]");
   }
 
+  @Test
+  public void streamDelayMetrics() {
+    ClientStreamTracer mockTracer = mock(ClientStreamTracer.class);
+    ClientStreamTracer[] customTracers = new ClientStreamTracer[] { mockTracer };
+    
+    SubchannelPicker connectingPicker = mock(SubchannelPicker.class);
+    when(connectingPicker.pickSubchannel(any(PickSubchannelArgs.class)))
+        .thenReturn(PickResult.withNoResult("pick_first:connecting"));
+        
+    delayedTransport.reprocess(connectingPicker);
+    ClientStream stream = delayedTransport.newStream(method, headers, callOptions, customTracers);
+    
+    InOrder inOrder = inOrder(mockTracer);
+    inOrder.verify(mockTracer).delayStarted("pick_first:connecting");
+    
+    SubchannelPicker customDelayPicker = mock(SubchannelPicker.class);
+    when(customDelayPicker.pickSubchannel(any(PickSubchannelArgs.class)))
+        .thenReturn(PickResult.withNoResult("rls:lookup_pending"));
+        
+    delayedTransport.reprocess(customDelayPicker);
+    
+    inOrder.verify(mockTracer).delayEnded();
+    inOrder.verify(mockTracer).delayStarted("rls:lookup_pending");
+    
+    delayedTransport.reprocess(mockPicker);
+    
+    inOrder.verify(mockTracer).delayEnded();
+  }
+
+  @Test
+  public void streamDelayMetrics_cancelled() {
+    ClientStreamTracer mockTracer = mock(ClientStreamTracer.class);
+    ClientStreamTracer[] customTracers = new ClientStreamTracer[] { mockTracer };
+    
+    SubchannelPicker connectingPicker = mock(SubchannelPicker.class);
+    when(connectingPicker.pickSubchannel(any(PickSubchannelArgs.class)))
+        .thenReturn(PickResult.withNoResult("pick_first:connecting"));
+        
+    delayedTransport.reprocess(connectingPicker);
+    ClientStream stream = delayedTransport.newStream(method, headers, callOptions, customTracers);
+    
+    verify(mockTracer).delayStarted("pick_first:connecting");
+    
+    stream.cancel(Status.CANCELLED);
+    
+    verify(mockTracer).delayEnded();
+  }
+
   private static TransportProvider newTransportProvider(final ClientTransport transport) {
     return new TransportProvider() {
       @Override
diff --git a/core/src/test/java/io/grpc/internal/PickFirstLoadBalancerTest.java b/core/src/test/java/io/grpc/internal/PickFirstLoadBalancerTest.java
@@ -147,8 +147,9 @@ public void pickAfterResolved() throws Exception {
     verify(mockSubchannel).requestConnection();
 
     // Calling pickSubchannel() twice gave the same result
-    assertEquals(pickerCaptor.getValue().pickSubchannel(mockArgs),
-        pickerCaptor.getValue().pickSubchannel(mockArgs));
+    PickResult result = pickerCaptor.getValue().pickSubchannel(mockArgs);
+    assertThat(result.getDelayReasonToken()).isEqualTo("pick_first:connecting");
+    assertEquals(result, pickerCaptor.getValue().pickSubchannel(mockArgs));
 
     verifyNoMoreInteractions(mockHelper);
   }
diff --git a/rls/src/main/java/io/grpc/rls/CachingRlsLbClient.java b/rls/src/main/java/io/grpc/rls/CachingRlsLbClient.java
@@ -1050,7 +1050,7 @@ public PickResult pickSubchannel(PickSubchannelArgs args) {
             convertRlsServerStatus(response.getStatus(),
                 lbPolicyConfig.getRouteLookupConfig().lookupService()));
       } else {
-        return PickResult.withNoResult();
+        return PickResult.withNoResult("rls:lookup_pending");
       }
     }
 
diff --git a/rls/src/test/java/io/grpc/rls/RlsLoadBalancerTest.java b/rls/src/test/java/io/grpc/rls/RlsLoadBalancerTest.java
@@ -262,6 +262,7 @@ public void lb_working_withDefaultTarget_rlsResponding() throws Exception {
     PickResult res = picker.pickSubchannel(searchSubchannelArgs);
     assertThat(res.getStatus().isOk()).isTrue();
     assertThat(res.getSubchannel()).isNull();
+    assertThat(res.getDelayReasonToken()).isEqualTo("rls:lookup_pending");
     // Cache is warm, but still unconnected
     res = picker.pickSubchannel(searchSubchannelArgs);
     inOrder.verify(helper).createSubchannel(any(CreateSubchannelArgs.class));
@@ -493,6 +494,7 @@ public void lb_working_withoutDefaultTarget() throws Exception {
     PickResult res = picker.pickSubchannel(searchSubchannelArgs);
     assertThat(res.getStatus().isOk()).isTrue();
     assertThat(res.getSubchannel()).isNull();
+    assertThat(res.getDelayReasonToken()).isEqualTo("rls:lookup_pending");
     // Cache is warm, but still unconnected
     res = picker.pickSubchannel(searchSubchannelArgs);
     inOrder.verify(helper).createSubchannel(any(CreateSubchannelArgs.class));
diff --git a/util/src/main/java/io/grpc/util/RoundRobinLoadBalancer.java b/util/src/main/java/io/grpc/util/RoundRobinLoadBalancer.java
@@ -41,8 +41,9 @@
  * EquivalentAddressGroup}s from the {@link NameResolver}.
  */
 final class RoundRobinLoadBalancer extends MultiChildLoadBalancer {
+  private static final PickResult CONNECTING_RESULT = PickResult.withNoResult("round_robin:connecting");
   private final AtomicInteger sequence = new AtomicInteger(new Random().nextInt());
-  private SubchannelPicker currentPicker = new FixedResultPicker(PickResult.withNoResult());
+  private SubchannelPicker currentPicker = new FixedResultPicker(CONNECTING_RESULT);
 
   public RoundRobinLoadBalancer(Helper helper) {
     super(helper);
@@ -68,7 +69,7 @@ protected void updateOverallBalancingState() {
       }
 
       if (isConnecting) {
-        updateBalancingState(CONNECTING, new FixedResultPicker(PickResult.withNoResult()));
+        updateBalancingState(CONNECTING, new FixedResultPicker(CONNECTING_RESULT));
       } else {
         updateBalancingState(TRANSIENT_FAILURE, createReadyPicker(getChildLbStates()));
       }
diff --git a/util/src/test/java/io/grpc/util/RoundRobinLoadBalancerTest.java b/util/src/test/java/io/grpc/util/RoundRobinLoadBalancerTest.java
@@ -86,7 +86,7 @@
 public class RoundRobinLoadBalancerTest {
   private static final Attributes.Key<String> MAJOR_KEY = Attributes.Key.create("major-key");
   private static final SubchannelPicker EMPTY_PICKER =
-      new FixedResultPicker(PickResult.withNoResult());
+      new FixedResultPicker(PickResult.withNoResult("round_robin:connecting"));
 
   @Rule public final MockitoRule mocks = MockitoJUnit.rule();
 
diff --git a/xds/src/main/java/io/grpc/xds/CdsLoadBalancer2.java b/xds/src/main/java/io/grpc/xds/CdsLoadBalancer2.java
@@ -119,6 +119,7 @@ public Status acceptResolvedAddresses(ResolvedAddresses resolvedAddresses) {
             errorPrefix() + "Unable to find non-dynamic cluster"));
       }
       // The dynamic cluster must not have loaded yet
+      helper.updateBalancingState(CONNECTING, new FixedResultPicker(PickResult.withNoResult("cds:discovery_pending")));
       return Status.OK;
     }
     if (!clusterConfigOr.hasValue()) {
diff --git a/xds/src/main/java/io/grpc/xds/PriorityLoadBalancer.java b/xds/src/main/java/io/grpc/xds/PriorityLoadBalancer.java
@@ -322,7 +322,17 @@ public void updateBalancingState(final ConnectivityState newState,
         }
         ConnectivityState oldState = connectivityState;
         connectivityState = newState;
-        picker = newPicker;
+        picker = new SubchannelPicker() {
+          @Override
+          public PickResult pickSubchannel(PickSubchannelArgs args) {
+            PickResult childResult = newPicker.pickSubchannel(args);
+            if (!childResult.hasResult() && childResult.getDelayReasonToken() != null) {
+              return PickResult.withNoResult(
+                  "priority_" + priority + ":" + childResult.getDelayReasonToken());
+            }
+            return childResult;
+          }
+        };
 
         if (deletionTimer != null && deletionTimer.isPending()) {
           return;
diff --git a/xds/src/main/java/io/grpc/xds/RingHashLoadBalancer.java b/xds/src/main/java/io/grpc/xds/RingHashLoadBalancer.java
@@ -356,6 +356,8 @@ public static EquivalentAddressGroup stripAttrs(EquivalentAddressGroup eag) {
   }
 
   private static final class RingHashPicker extends SubchannelPicker {
+    private static final PickResult RING_HASH_CONNECTING_RESULT = 
+        PickResult.withNoResult("ring_hash:connecting");
     private final SynchronizationContext syncContext;
     private final List<RingEntry> ring;
     // Avoid synchronization between pickSubchannel and subchannel's connectivity state change,
@@ -453,7 +455,7 @@ public PickResult pickSubchannel(PickSubchannelArgs args) {
           // RPCs can be buffered if the next subchannel is pending (per A62). Otherwise, RPCs
           // are failed unless there is a READY connection.
           if (subchannelView.connectivityState == CONNECTING) {
-            return PickResult.withNoResult();
+            return RING_HASH_CONNECTING_RESULT;
           }
 
           if (subchannelView.connectivityState == IDLE) {
@@ -463,7 +465,7 @@ public PickResult pickSubchannel(PickSubchannelArgs args) {
               }
             });
 
-            return PickResult.withNoResult(); // Indicates that this should be retried after backoff
+            return RING_HASH_CONNECTING_RESULT; // Indicates that this should be retried after backoff
           }
         }
       } else {
@@ -487,7 +489,7 @@ public PickResult pickSubchannel(PickSubchannelArgs args) {
           }
         }
         if (requestedConnection) {
-          return PickResult.withNoResult();
+          return RING_HASH_CONNECTING_RESULT;
         }
       }
 
diff --git a/xds/src/test/java/io/grpc/xds/CdsLoadBalancer2Test.java b/xds/src/test/java/io/grpc/xds/CdsLoadBalancer2Test.java
diff --git a/xds/src/test/java/io/grpc/xds/PriorityLoadBalancerTest.java b/xds/src/test/java/io/grpc/xds/PriorityLoadBalancerTest.java
diff --git a/xds/src/test/java/io/grpc/xds/RingHashLoadBalancerTest.java b/xds/src/test/java/io/grpc/xds/RingHashLoadBalancerTest.java

Original file line number	Diff line number	Diff line change
`@@ -1050,7 +1050,7 @@ public PickResult pickSubchannel(PickSubchannelArgs args) {`
`1050`	`1050`	`convertRlsServerStatus(response.getStatus(),`
`1051`	`1051`	`lbPolicyConfig.getRouteLookupConfig().lookupService()));`
`1052`	`1052`	`} else {`
`1053`		`- return PickResult.withNoResult();`
	`1053`	`+ return PickResult.withNoResult("rls:lookup_pending");`
`1054`	`1054`	`}`
`1055`	`1055`	`}`
`1056`	`1056`