fix(grpc-gcp): align the dynamic channel pool with Go DCP to remove one-hot-channel skew

Author: rahul2393

grpc-gcp-java/pom.xml

@@ -63,6 +63,7 @@
     <project.reporting.outputEncoding>UTF-8</project.reporting.outputEncoding>
     <site.installationModule>grpc-gcp</site.installationModule>
     <api-common.version>2.68.0-SNAPSHOT</api-common.version><!-- {x-version-update:api-common:current} -->
+    <awaitility.version>4.3.0</awaitility.version>
     <auto-value.version>1.11.0</auto-value.version>
     <error-prone-annotations.version>2.48.0</error-prone-annotations.version>
     <google-http-client.version>2.2.0</google-http-client.version>
@@ -172,6 +173,12 @@
       <version>${junit.version}</version>
       <scope>test</scope>
     </dependency>
+    <dependency>
+      <groupId>org.awaitility</groupId>
+      <artifactId>awaitility</artifactId>
+      <version>${awaitility.version}</version>
+      <scope>test</scope>
+    </dependency>
     <dependency>
       <groupId>com.google.truth</groupId>
       <artifactId>truth</artifactId>

grpc-gcp-java/src/main/java/com/google/cloud/grpc/GcpChannelPrimer.java

@@ -0,0 +1,32 @@
+/*
+ * Copyright 2026 Google LLC
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *     https://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package com.google.cloud.grpc;
+
+import com.google.common.util.concurrent.ListenableFuture;
+import io.grpc.ManagedChannel;
+
+/** Primes a newly built delegate channel before a dynamic pool publishes it for request picking. */
+@FunctionalInterface
+public interface GcpChannelPrimer {
+
+  /**
+   * Issues a cheap end-to-end RPC on {@code channel} so its connection is warm before real traffic.
+   * For example, a Cloud Spanner implementation can execute {@code SELECT 1}. Return a failed
+   * future to reject and close the channel.
+   */
+  ListenableFuture<Void> prime(ManagedChannel channel);
+}

grpc-gcp-java/src/main/java/com/google/cloud/grpc/GcpClientCall.java

@@ -30,7 +30,7 @@
 import java.util.Collections;
 import java.util.List;
 import java.util.Queue;
-import java.util.concurrent.atomic.AtomicBoolean;
+import java.util.concurrent.atomic.AtomicInteger;
 import javax.annotation.Nullable;
 import javax.annotation.concurrent.GuardedBy;
 
@@ -52,14 +52,21 @@ public class GcpClientCall<ReqT, RespT> extends ClientCall<ReqT, RespT> {
   private ClientCall<ReqT, RespT> delegateCall = null;
   private List<String> keys = null;
   private boolean received = false;
-  private final AtomicBoolean decremented = new AtomicBoolean(false);
+  // 0 = not counted, 1 = counted, 2 = finished.
+  private final AtomicInteger countState = new AtomicInteger();
 
   @GuardedBy("this")
   private final Queue<Runnable> calls = new ArrayDeque<>();
 
   @GuardedBy("this")
   private boolean started;
 
+  @GuardedBy("this")
+  private boolean cancelQueued;
+
+  @GuardedBy("this")
+  private boolean cancelled;
+
   private long startNanos = 0;
 
   protected GcpClientCall(
@@ -90,7 +97,22 @@ public void setMessageCompression(boolean enabled) {
 
   @Override
   public void cancel(@Nullable String message, @Nullable Throwable cause) {
-    checkSendMessage(() -> checkedCancel(message, cause));
+    synchronized (this) {
+      if (cancelQueued || cancelled) {
+        return;
+      }
+      cancelQueued = true;
+      Runnable cancelCall =
+          () -> {
+            cancelled = true;
+            checkedCancel(message, cause);
+          };
+      if (started) {
+        cancelCall.run();
+      } else {
+        calls.add(cancelCall);
+      }
+    }
   }
 
   @Override
@@ -104,6 +126,7 @@ public void halfClose() {
    */
   @Override
   public void sendMessage(ReqT message) {
+    boolean send;
     synchronized (this) {
       if (!started) {
         startNanos = System.nanoTime();
@@ -123,17 +146,29 @@ public void sendMessage(ReqT message) {
           delegateChannelRef = delegateChannel.getChannelRef(key);
         }
         delegateChannelRef.activeStreamsCountIncr();
-
-        // Create the client call and do the previous operations.
-        delegateCall = delegateChannelRef.getChannel().newCall(methodDescriptor, callOptions);
-        for (Runnable call : calls) {
-          call.run();
+        countState.set(1);
+
+        try {
+          // Create the client call and do the previous operations.
+          CallOptions callOptionsWithChannelId =
+              callOptions.withOption(GcpManagedChannel.CHANNEL_ID_KEY, delegateChannelRef.getId());
+          delegateCall =
+              delegateChannelRef.getChannel().newCall(methodDescriptor, callOptionsWithChannelId);
+          for (Runnable call : calls) {
+            call.run();
+          }
+        } catch (RuntimeException | Error failure) {
+          finishCount(Status.fromThrowable(failure), true);
+          throw failure;
         }
         calls.clear();
         started = true;
       }
+      send = !cancelled;
+    }
+    if (send) {
+      delegateCall.sendMessage(message);
     }
-    delegateCall.sendMessage(message);
   }
 
   /** Calls that send exactly one message should not check this method. */
@@ -162,14 +197,21 @@ public String toString() {
   }
 
   private void checkedCancel(@Nullable String message, @Nullable Throwable cause) {
-    if (!decremented.getAndSet(true)) {
-      delegateChannelRef.activeStreamsCountDecr(startNanos, Status.CANCELLED, true);
-    }
+    finishCount(Status.CANCELLED, true);
     delegateCall.cancel(message, cause);
   }
 
+  private void finishCount(Status status, boolean fromClientSide) {
+    if (countState.compareAndSet(1, 2)) {
+      delegateChannelRef.activeStreamsCountDecr(startNanos, status, fromClientSide);
+    }
+  }
+
   private void checkSendMessage(Runnable call) {
     synchronized (this) {
+      if (cancelQueued || cancelled) {
+        return;
+      }
       if (started) {
         call.run();
       } else {
@@ -185,9 +227,7 @@ private Listener<RespT> getListener(final Listener<RespT> responseListener) {
       // Decrement the stream number by one when the call is closed.
       @Override
       public void onClose(Status status, Metadata trailers) {
-        if (!decremented.getAndSet(true)) {
-          delegateChannelRef.activeStreamsCountDecr(startNanos, status, false);
-        }
+        finishCount(status, false);
         // If the operation completed successfully, bind/unbind the affinity key.
         if (keys != null && status.getCode() == Status.Code.OK) {
           if (affinity.getCommand() == AffinityConfig.Command.UNBIND) {
@@ -219,7 +259,8 @@ public void onMessage(RespT message) {
    * A simple wrapper of ClientCall.
    *
    * <p>It defines the callback function to manage the number of active streams of a ChannelRef
-   * everytime a call is started/closed.
+   * every time a call is created/closed. Stream capacity is reserved in the constructor, before
+   * {@link #start(Listener, Metadata)}, and remains reserved until close or cancel.
    */
   public static class SimpleGcpClientCall<ReqT, RespT> extends ForwardingClientCall<ReqT, RespT> {
 
@@ -230,7 +271,13 @@ public static class SimpleGcpClientCall<ReqT, RespT> extends ForwardingClientCal
     private final boolean unbindOnComplete;
     private long startNanos = 0;
 
-    private final AtomicBoolean decremented = new AtomicBoolean(false);
+    private final Object countLock = new Object();
+
+    @GuardedBy("countLock")
+    private boolean counted;
+
+    @GuardedBy("countLock")
+    private boolean finished;
 
     protected SimpleGcpClientCall(
         GcpManagedChannel delegateChannel,
@@ -244,8 +291,18 @@ protected SimpleGcpClientCall(
       // Set the actual channel ID in callOptions so downstream interceptors can access it.
       CallOptions callOptionsWithChannelId =
           callOptions.withOption(GcpManagedChannel.CHANNEL_ID_KEY, channelRef.getId());
-      this.delegateCall =
-          channelRef.getChannel().newCall(methodDescriptor, callOptionsWithChannelId);
+      startNanos = System.nanoTime();
+      synchronized (countLock) {
+        channelRef.activeStreamsCountIncr();
+        counted = true;
+      }
+      try {
+        this.delegateCall =
+            channelRef.getChannel().newCall(methodDescriptor, callOptionsWithChannelId);
+      } catch (RuntimeException | Error failure) {
+        finishCount(Status.fromThrowable(failure), true);
+        throw failure;
+      }
     }
 
     @Override
@@ -255,16 +312,12 @@ protected ClientCall<ReqT, RespT> delegate() {
 
     @Override
     public void start(Listener<RespT> responseListener, Metadata headers) {
-      startNanos = System.nanoTime();
-
       Listener<RespT> listener =
           new ForwardingClientCallListener.SimpleForwardingClientCallListener<RespT>(
               responseListener) {
             @Override
             public void onClose(Status status, Metadata trailers) {
-              if (!decremented.getAndSet(true)) {
-                channelRef.activeStreamsCountDecr(startNanos, status, false);
-              }
+              finishCount(status, false);
               // Unbind the affinity key when the caller explicitly requests it
               // (e.g., on terminal RPCs like Commit or Rollback) to prevent
               // unbounded growth of the affinity map.
@@ -281,20 +334,34 @@ public void onMessage(RespT message) {
             }
           };
 
-      channelRef.activeStreamsCountIncr();
-      delegateCall.start(listener, headers);
+      try {
+        delegateCall.start(listener, headers);
+      } catch (RuntimeException | Error failure) {
+        finishCount(Status.fromThrowable(failure), true);
+        throw failure;
+      }
     }
 
     @Override
     public void cancel(String message, Throwable cause) {
-      if (!decremented.getAndSet(true)) {
-        channelRef.activeStreamsCountDecr(startNanos, Status.CANCELLED, true);
-      }
+      finishCount(Status.CANCELLED, true);
       // Always unbind on cancel — the transaction is being abandoned.
       if (affinityKey != null) {
         delegateChannel.unbind(Collections.singletonList(affinityKey));
       }
       delegateCall.cancel(message, cause);
     }
+
+    private void finishCount(Status status, boolean fromClientSide) {
+      synchronized (countLock) {
+        if (finished) {
+          return;
+        }
+        finished = true;
+        if (counted) {
+          channelRef.activeStreamsCountDecr(startNanos, status, fromClientSide);
+        }
+      }
+    }
   }
 }