chore: Move channel priming out of InstantiatingGrpcTransportProvider (#2644)

* chore: Move channel priming out of InstantiatingGrpcTransportProvider

Change-Id: I7214aa3016bd7e7f7f167c64cbaa04134b54a352

* delete comment

Change-Id: I26d9b4929be6ba941eb88da05244fc031a747b68

* fixes

Change-Id: Ib414d250f9f1ce0555450e81a706f6c0619d80f2

* Update channel priming annotation

Co-authored-by: Mattie Fu <[email protected]>

* Update google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/stub/BigtableClientContext.java

Co-authored-by: Mattie Fu <[email protected]>

* fixes

Change-Id: I2f27c788652c9e5e12ea470355098c73b29ab40a

* add no op channel primer

Change-Id: I343b3dc44d12638be2af32398b6aaa8529c7ebfb

* add noop channel primer test

Change-Id: I2a11cd44ddd9caaea4db918f66aa7313a31a0057

* lint

Change-Id: Ifc564bcd182d4ea18e8e63ce78787281d589d8fd

* feat/add-async-channel-ping

Change-Id: I3e390b6f7a6c9beaee52d74f37ef557629af7759

* lint

Change-Id: I5e2b13674c7fc945d5567b7f86478c847017aa41

* fixes

Change-Id: Ibaa32ee3425b23596a3141bb4e29cfb068caf0ee

* fixes

Change-Id: I3d43864fee0b05cc64f27aabcb2ffaf84c9c0ab8

* minor change to test

Change-Id: I63f9f0e930f95d48ce56962e49d6bdca0fc5918d

* clirr-ignored-differences.xml

Change-Id: Ia4e25a7c80256762e241f30a585056aa78fb2878

* increase timeout for testing purposes

Change-Id: I96fb32e098862546aae0de498637459c5a9e197d

* increase offset jitter to make more resilient to flakes

Change-Id: I8754683cae05ed9fd0267a1fd16746756450c15e

* more fault tolerance

Change-Id: I997c3752d8b5ae13084762a7eb3fc6b042ead206

* revert flaky test changes

Change-Id: I8e0ad310444c754df71dc61782b23d007bbfdaba

---------

Co-authored-by: Mattie Fu <[email protected]>

Author: nicholsl

google-cloud-bigtable/clirr-ignored-differences.xml

@@ -382,4 +382,18 @@
         <method>*</method>
         <to>*</to>
     </difference>
+    <difference>
+        <!-- InternalApi was updated -->
+        <differenceType>7004</differenceType>
+        <className>com/google/cloud/bigtable/gaxx/grpc/BigtableChannelPool</className>
+        <method>*create*</method>
+        <to>*</to>
+    </difference>
+    <difference>
+        <!-- InternalApi was updated -->
+        <differenceType>7004</differenceType>
+        <className>com/google/cloud/bigtable/gaxx/grpc/BigtableTransportChannelProvider</className>
+        <method>*create*</method>
+        <to>*</to>
+    </difference>
 </differences>

google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/stub/BigtableChannelPrimer.java

@@ -15,7 +15,7 @@
  */
 package com.google.cloud.bigtable.data.v2.stub;
 
-import com.google.api.core.BetaApi;
+import com.google.api.core.InternalApi;
 import com.google.api.core.SettableApiFuture;
 import com.google.api.gax.grpc.ChannelPrimer;
 import com.google.auth.Credentials;
@@ -41,13 +41,13 @@
 
 /**
  * A channel warmer that ensures that a Bigtable channel is ready to be used before being added to
- * the active {@link com.google.api.gax.grpc.ChannelPool}.
+ * the active {@link com.google.cloud.bigtable.gaxx.grpc.BigtableChannelPool}.
  *
  * <p>This implementation is subject to change in the future, but currently it will prime the
  * channel by sending a ReadRow request for a hardcoded, non-existent row key.
  */
-@BetaApi("Channel priming is not currently stable and might change in the future")
-class BigtableChannelPrimer implements ChannelPrimer {
+@InternalApi
+public class BigtableChannelPrimer implements ChannelPrimer {
   private static Logger LOG = Logger.getLogger(BigtableChannelPrimer.class.toString());
 
   static final Metadata.Key<String> REQUEST_PARAMS =
@@ -96,43 +96,12 @@ public void primeChannel(ManagedChannel managedChannel) {
   }
 
   private void primeChannelUnsafe(ManagedChannel managedChannel) throws IOException {
-    sendPrimeRequests(managedChannel);
+    sendPrimeRequestsBlocking(managedChannel);
   }
 
-  private void sendPrimeRequests(ManagedChannel managedChannel) {
+  private void sendPrimeRequestsBlocking(ManagedChannel managedChannel) {
     try {
-      ClientCall<PingAndWarmRequest, PingAndWarmResponse> clientCall =
-          managedChannel.newCall(
-              BigtableGrpc.getPingAndWarmMethod(),
-              CallOptions.DEFAULT
-                  .withCallCredentials(callCredentials)
-                  .withDeadline(Deadline.after(1, TimeUnit.MINUTES)));
-
-      SettableApiFuture<PingAndWarmResponse> future = SettableApiFuture.create();
-      clientCall.start(
-          new ClientCall.Listener<PingAndWarmResponse>() {
-            PingAndWarmResponse response;
-
-            @Override
-            public void onMessage(PingAndWarmResponse message) {
-              response = message;
-            }
-
-            @Override
-            public void onClose(Status status, Metadata trailers) {
-              if (status.isOk()) {
-                future.set(response);
-              } else {
-                future.setException(status.asException());
-              }
-            }
-          },
-          createMetadata(headers, request));
-      clientCall.sendMessage(request);
-      clientCall.halfClose();
-      clientCall.request(Integer.MAX_VALUE);
-
-      future.get(1, TimeUnit.MINUTES);
+      sendPrimeRequestsAsync(managedChannel).get(1, TimeUnit.MINUTES);
     } catch (Throwable e) {
       // TODO: Not sure if we should swallow the error here. We are pre-emptively swapping
       // channels if the new
@@ -141,6 +110,53 @@ public void onClose(Status status, Metadata trailers) {
     }
   }
 
+  public SettableApiFuture<PingAndWarmResponse> sendPrimeRequestsAsync(
+      ManagedChannel managedChannel) {
+    ClientCall<PingAndWarmRequest, PingAndWarmResponse> clientCall =
+        managedChannel.newCall(
+            BigtableGrpc.getPingAndWarmMethod(),
+            CallOptions.DEFAULT
+                .withCallCredentials(callCredentials)
+                .withDeadline(Deadline.after(1, TimeUnit.MINUTES)));
+
+    SettableApiFuture<PingAndWarmResponse> future = SettableApiFuture.create();
+    clientCall.start(
+        new ClientCall.Listener<PingAndWarmResponse>() {
+          private PingAndWarmResponse response;
+
+          @Override
+          public void onMessage(PingAndWarmResponse message) {
+            response = message;
+          }
+
+          @Override
+          public void onClose(Status status, Metadata trailers) {
+            if (status.isOk()) {
+              future.set(response);
+            } else {
+              // Propagate the gRPC error to the future.
+              future.setException(status.asException(trailers));
+            }
+          }
+        },
+        createMetadata(headers, request));
+
+    try {
+      // Send the request message.
+      clientCall.sendMessage(request);
+      // Signal that no more messages will be sent.
+      clientCall.halfClose();
+      // Request the response from the server.
+      clientCall.request(Integer.MAX_VALUE);
+    } catch (Throwable t) {
+      // If sending fails, cancel the call and notify the future.
+      clientCall.cancel("Failed to send priming request", t);
+      future.setException(t);
+    }
+
+    return future;
+  }
+
   private static Metadata createMetadata(Map<String, String> headers, PingAndWarmRequest request) {
     Metadata metadata = new Metadata();
 

google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/stub/BigtableClientContext.java

@@ -20,6 +20,7 @@
 import com.google.api.gax.core.BackgroundResource;
 import com.google.api.gax.core.CredentialsProvider;
 import com.google.api.gax.core.FixedCredentialsProvider;
+import com.google.api.gax.grpc.ChannelPrimer;
 import com.google.api.gax.grpc.InstantiatingGrpcChannelProvider;
 import com.google.api.gax.rpc.ClientContext;
 import com.google.auth.Credentials;
@@ -121,20 +122,22 @@ public static BigtableClientContext create(EnhancedBigtableStubSettings settings
         setupCookieHolder(transportProvider);
       }
 
+      ChannelPrimer channelPrimer = NoOpChannelPrimer.create();
+
       // Inject channel priming if enabled
       if (builder.isRefreshingChannel()) {
-        transportProvider.setChannelPrimer(
+        channelPrimer =
             BigtableChannelPrimer.create(
                 builder.getProjectId(),
                 builder.getInstanceId(),
                 builder.getAppProfileId(),
                 credentials,
-                builder.getHeaderProvider().getHeaders()));
+                builder.getHeaderProvider().getHeaders());
       }
 
       BigtableTransportChannelProvider btTransportProvider =
           BigtableTransportChannelProvider.create(
-              (InstantiatingGrpcChannelProvider) transportProvider.build());
+              (InstantiatingGrpcChannelProvider) transportProvider.build(), channelPrimer);
 
       builder.setTransportChannelProvider(btTransportProvider);
     }

google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/stub/NoOpChannelPrimer.java

@@ -0,0 +1,34 @@
+/*
+ * Copyright 2025 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.bigtable.data.v2.stub;
+
+import com.google.api.core.InternalApi;
+import com.google.api.gax.grpc.ChannelPrimer;
+import io.grpc.ManagedChannel;
+
+@InternalApi
+public class NoOpChannelPrimer implements ChannelPrimer {
+  static NoOpChannelPrimer create() {
+    return new NoOpChannelPrimer();
+  }
+
+  private NoOpChannelPrimer() {}
+
+  @Override
+  public void primeChannel(ManagedChannel managedChannel) {
+    // No op
+  }
+}

google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/gaxx/grpc/BigtableChannelPool.java

@@ -17,6 +17,7 @@
 
 import com.google.api.core.InternalApi;
 import com.google.api.gax.grpc.ChannelFactory;
+import com.google.api.gax.grpc.ChannelPrimer;
 import com.google.common.annotations.VisibleForTesting;
 import com.google.common.base.Preconditions;
 import com.google.common.collect.ImmutableList;
@@ -60,6 +61,8 @@ public class BigtableChannelPool extends ManagedChannel {
 
   private final BigtableChannelPoolSettings settings;
   private final ChannelFactory channelFactory;
+
+  private final ChannelPrimer channelPrimer;
   private final ScheduledExecutorService executor;
 
   private final Object entryWriteLock = new Object();
@@ -68,9 +71,12 @@ public class BigtableChannelPool extends ManagedChannel {
   private final String authority;
 
   public static BigtableChannelPool create(
-      BigtableChannelPoolSettings settings, ChannelFactory channelFactory) throws IOException {
+      BigtableChannelPoolSettings settings,
+      ChannelFactory channelFactory,
+      ChannelPrimer channelPrimer)
+      throws IOException {
     return new BigtableChannelPool(
-        settings, channelFactory, Executors.newSingleThreadScheduledExecutor());
+        settings, channelFactory, channelPrimer, Executors.newSingleThreadScheduledExecutor());
   }
 
   /**
@@ -84,15 +90,19 @@ public static BigtableChannelPool create(
   BigtableChannelPool(
       BigtableChannelPoolSettings settings,
       ChannelFactory channelFactory,
+      ChannelPrimer channelPrimer,
       ScheduledExecutorService executor)
       throws IOException {
     this.settings = settings;
     this.channelFactory = channelFactory;
+    this.channelPrimer = channelPrimer;
 
     ImmutableList.Builder<Entry> initialListBuilder = ImmutableList.builder();
 
     for (int i = 0; i < settings.getInitialChannelCount(); i++) {
-      initialListBuilder.add(new Entry(channelFactory.createSingleChannel()));
+      ManagedChannel newChannel = channelFactory.createSingleChannel();
+      channelPrimer.primeChannel(newChannel);
+      initialListBuilder.add(new Entry(newChannel));
     }
 
     entries.set(initialListBuilder.build());
@@ -316,7 +326,9 @@ private void expand(int desiredSize) {
 
     for (int i = 0; i < desiredSize - localEntries.size(); i++) {
       try {
-        newEntries.add(new Entry(channelFactory.createSingleChannel()));
+        ManagedChannel newChannel = channelFactory.createSingleChannel();
+        this.channelPrimer.primeChannel(newChannel);
+        newEntries.add(new Entry(newChannel));
       } catch (IOException e) {
         LOG.log(Level.WARNING, "Failed to add channel", e);
       }
@@ -354,7 +366,9 @@ void refresh() {
 
       for (int i = 0; i < newEntries.size(); i++) {
         try {
-          newEntries.set(i, new Entry(channelFactory.createSingleChannel()));
+          ManagedChannel newChannel = channelFactory.createSingleChannel();
+          this.channelPrimer.primeChannel(newChannel);
+          newEntries.set(i, new Entry(newChannel));
         } catch (IOException e) {
           LOG.log(Level.WARNING, "Failed to refresh channel, leaving old channel", e);
         }

google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/gaxx/grpc/BigtableTransportChannelProvider.java

@@ -18,6 +18,7 @@
 import com.google.api.core.InternalApi;
 import com.google.api.gax.grpc.ChannelFactory;
 import com.google.api.gax.grpc.ChannelPoolSettings;
+import com.google.api.gax.grpc.ChannelPrimer;
 import com.google.api.gax.grpc.GrpcTransportChannel;
 import com.google.api.gax.grpc.InstantiatingGrpcChannelProvider;
 import com.google.api.gax.rpc.TransportChannel;
@@ -38,10 +39,13 @@
 public final class BigtableTransportChannelProvider implements TransportChannelProvider {
 
   private final InstantiatingGrpcChannelProvider delegate;
+  private final ChannelPrimer channelPrimer;
 
   private BigtableTransportChannelProvider(
-      InstantiatingGrpcChannelProvider instantiatingGrpcChannelProvider) {
+      InstantiatingGrpcChannelProvider instantiatingGrpcChannelProvider,
+      ChannelPrimer channelPrimer) {
     delegate = Preconditions.checkNotNull(instantiatingGrpcChannelProvider);
+    this.channelPrimer = channelPrimer;
   }
 
   @Override
@@ -63,7 +67,7 @@ public BigtableTransportChannelProvider withExecutor(ScheduledExecutorService ex
   public BigtableTransportChannelProvider withExecutor(Executor executor) {
     InstantiatingGrpcChannelProvider newChannelProvider =
         (InstantiatingGrpcChannelProvider) delegate.withExecutor(executor);
-    return new BigtableTransportChannelProvider(newChannelProvider);
+    return new BigtableTransportChannelProvider(newChannelProvider, channelPrimer);
   }
 
   @Override
@@ -75,7 +79,7 @@ public boolean needsHeaders() {
   public BigtableTransportChannelProvider withHeaders(Map<String, String> headers) {
     InstantiatingGrpcChannelProvider newChannelProvider =
         (InstantiatingGrpcChannelProvider) delegate.withHeaders(headers);
-    return new BigtableTransportChannelProvider(newChannelProvider);
+    return new BigtableTransportChannelProvider(newChannelProvider, channelPrimer);
   }
 
   @Override
@@ -87,7 +91,7 @@ public boolean needsEndpoint() {
   public TransportChannelProvider withEndpoint(String endpoint) {
     InstantiatingGrpcChannelProvider newChannelProvider =
         (InstantiatingGrpcChannelProvider) delegate.withEndpoint(endpoint);
-    return new BigtableTransportChannelProvider(newChannelProvider);
+    return new BigtableTransportChannelProvider(newChannelProvider, channelPrimer);
   }
 
   @Deprecated
@@ -101,7 +105,7 @@ public boolean acceptsPoolSize() {
   public TransportChannelProvider withPoolSize(int size) {
     InstantiatingGrpcChannelProvider newChannelProvider =
         (InstantiatingGrpcChannelProvider) delegate.withPoolSize(size);
-    return new BigtableTransportChannelProvider(newChannelProvider);
+    return new BigtableTransportChannelProvider(newChannelProvider, channelPrimer);
   }
 
   /** Expected to only be called once when BigtableClientContext is created */
@@ -130,7 +134,8 @@ public TransportChannel getTransportChannel() throws IOException {
     BigtableChannelPoolSettings btPoolSettings =
         BigtableChannelPoolSettings.copyFrom(delegate.getChannelPoolSettings());
 
-    BigtableChannelPool btChannelPool = BigtableChannelPool.create(btPoolSettings, channelFactory);
+    BigtableChannelPool btChannelPool =
+        BigtableChannelPool.create(btPoolSettings, channelFactory, channelPrimer);
 
     return GrpcTransportChannel.create(btChannelPool);
   }
@@ -149,12 +154,13 @@ public boolean needsCredentials() {
   public TransportChannelProvider withCredentials(Credentials credentials) {
     InstantiatingGrpcChannelProvider newChannelProvider =
         (InstantiatingGrpcChannelProvider) delegate.withCredentials(credentials);
-    return new BigtableTransportChannelProvider(newChannelProvider);
+    return new BigtableTransportChannelProvider(newChannelProvider, channelPrimer);
   }
 
   /** Creates a BigtableTransportChannelProvider. */
   public static BigtableTransportChannelProvider create(
-      InstantiatingGrpcChannelProvider instantiatingGrpcChannelProvider) {
-    return new BigtableTransportChannelProvider(instantiatingGrpcChannelProvider);
+      InstantiatingGrpcChannelProvider instantiatingGrpcChannelProvider,
+      ChannelPrimer channelPrimer) {
+    return new BigtableTransportChannelProvider(instantiatingGrpcChannelProvider, channelPrimer);
   }
 }