chore: integrate channel pool fixes

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

@@ -1,5 +1,5 @@
 /*
- * Copyright 2024 Google LLC
+ * 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.
@@ -32,6 +32,7 @@
 import com.google.cloud.bigtable.data.v2.stub.metrics.ErrorCountPerConnectionMetricTracker;
 import com.google.cloud.bigtable.data.v2.stub.metrics.MetricsProvider;
 import com.google.cloud.bigtable.data.v2.stub.metrics.NoopMetricsProvider;
+import com.google.cloud.bigtable.gaxx.grpc.BigtableTransportChannelProvider;
 import io.grpc.ManagedChannelBuilder;
 import io.grpc.opentelemetry.GrpcOpenTelemetry;
 import io.opentelemetry.api.OpenTelemetry;
@@ -131,7 +132,11 @@ public static BigtableClientContext create(EnhancedBigtableStubSettings settings
                 builder.getHeaderProvider().getHeaders()));
       }
 
-      builder.setTransportChannelProvider(transportProvider.build());
+      BigtableTransportChannelProvider btTransportProvider =
+          BigtableTransportChannelProvider.create(
+              (InstantiatingGrpcChannelProvider) transportProvider.build());
+
+      builder.setTransportChannelProvider(btTransportProvider);
     }
 
     ClientContext clientContext = ClientContext.create(builder.build());

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

@@ -46,14 +46,13 @@
 /**
  * A {@link ManagedChannel} that will send requests round-robin via a set of channels.
  *
- * <p>In addition to spreading requests over a set of child connections, the pool will also actively
- * manage the lifecycle of the channels. Currently lifecycle management is limited to pre-emptively
- * replacing channels every hour. In the future it will dynamically size the pool based on number of
- * outstanding requests.
+ * <p>Spreads over a set of child connections, and actively manages lifecycle of connections.
+ * Dynamically resizes pool based on number of outstanding connections.
  *
- * <p>Package-private for internal use.
+ * <p>Internal API
  */
-class BigtableChannelPool extends ManagedChannel {
+@InternalApi
+public class BigtableChannelPool extends ManagedChannel {
   @VisibleForTesting
   static final Logger LOG = Logger.getLogger(BigtableChannelPool.class.getName());
 
@@ -68,7 +67,7 @@ class BigtableChannelPool extends ManagedChannel {
   private final AtomicInteger indexTicker = new AtomicInteger();
   private final String authority;
 
-  static BigtableChannelPool create(
+  public static BigtableChannelPool create(
       BigtableChannelPoolSettings settings, ChannelFactory channelFactory) throws IOException {
     return new BigtableChannelPool(
         settings, channelFactory, Executors.newSingleThreadScheduledExecutor());

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

@@ -16,6 +16,7 @@
 package com.google.cloud.bigtable.gaxx.grpc;
 
 import com.google.api.core.BetaApi;
+import com.google.api.gax.grpc.ChannelPoolSettings;
 import com.google.auto.value.AutoValue;
 import com.google.common.base.Preconditions;
 import java.time.Duration;
@@ -112,6 +113,17 @@ boolean isStaticSize() {
 
   public abstract Builder toBuilder();
 
+  public static BigtableChannelPoolSettings copyFrom(ChannelPoolSettings externalSettings) {
+    return BigtableChannelPoolSettings.builder()
+        .setMinRpcsPerChannel(externalSettings.getMinRpcsPerChannel())
+        .setMaxRpcsPerChannel(externalSettings.getMaxRpcsPerChannel())
+        .setMinChannelCount(externalSettings.getMinChannelCount())
+        .setMaxChannelCount(externalSettings.getMaxChannelCount())
+        .setInitialChannelCount(externalSettings.getInitialChannelCount())
+        .setPreemptiveRefreshEnabled(externalSettings.isPreemptiveRefreshEnabled())
+        .build();
+  }
+
   public static BigtableChannelPoolSettings staticallySized(int size) {
     return builder()
         .setInitialChannelCount(size)

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

@@ -0,0 +1,157 @@
+/*
+ * 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.gaxx.grpc;
+
+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.GrpcTransportChannel;
+import com.google.api.gax.grpc.InstantiatingGrpcChannelProvider;
+import com.google.api.gax.rpc.TransportChannel;
+import com.google.api.gax.rpc.TransportChannelProvider;
+import com.google.auth.Credentials;
+import com.google.common.base.Preconditions;
+import io.grpc.ManagedChannel;
+import java.io.IOException;
+import java.util.Map;
+import java.util.concurrent.Executor;
+import java.util.concurrent.ScheduledExecutorService;
+
+/** An instance of TransportChannelProvider that provides a TransportChannel through a supplied InstantiatingGrpcChannelProvider. */
+@InternalApi
+public final class BigtableTransportChannelProvider implements TransportChannelProvider {
+
+  private final InstantiatingGrpcChannelProvider delegate;
+
+  private BigtableTransportChannelProvider(
+      InstantiatingGrpcChannelProvider instantiatingGrpcChannelProvider) {
+    delegate = Preconditions.checkNotNull(instantiatingGrpcChannelProvider);
+  }
+
+  @Override
+  public boolean shouldAutoClose() {
+    return delegate.shouldAutoClose();
+  }
+
+  @Override
+  public boolean needsExecutor() {
+    return delegate.needsExecutor();
+  }
+
+  @Override
+  public BigtableTransportChannelProvider withExecutor(ScheduledExecutorService executor) {
+    return withExecutor((Executor) executor);
+  }
+
+  @Override
+  public BigtableTransportChannelProvider withExecutor(Executor executor) {
+    InstantiatingGrpcChannelProvider newChannelProvider =
+        (InstantiatingGrpcChannelProvider) delegate.withExecutor(executor);
+    return new BigtableTransportChannelProvider(newChannelProvider);
+  }
+
+  @Override
+  public boolean needsHeaders() {
+    return delegate.needsHeaders();
+  }
+
+  @Override
+  public BigtableTransportChannelProvider withHeaders(Map<String, String> headers) {
+    InstantiatingGrpcChannelProvider newChannelProvider =
+        (InstantiatingGrpcChannelProvider) delegate.withHeaders(headers);
+    return new BigtableTransportChannelProvider(newChannelProvider);
+  }
+
+  @Override
+  public boolean needsEndpoint() {
+    return delegate.needsEndpoint();
+  }
+
+  @Override
+  public TransportChannelProvider withEndpoint(String endpoint) {
+    InstantiatingGrpcChannelProvider newChannelProvider =
+        (InstantiatingGrpcChannelProvider) delegate.withEndpoint(endpoint);
+    return new BigtableTransportChannelProvider(newChannelProvider);
+  }
+
+  @Deprecated
+  @Override
+  public boolean acceptsPoolSize() {
+    return delegate.acceptsPoolSize();
+  }
+
+  @Deprecated
+  @Override
+  public TransportChannelProvider withPoolSize(int size) {
+    InstantiatingGrpcChannelProvider newChannelProvider =
+        (InstantiatingGrpcChannelProvider) delegate.withPoolSize(size);
+    return new BigtableTransportChannelProvider(newChannelProvider);
+  }
+
+  /**Expected to only be called once when BigtableClientContext is created*/
+  @Override
+  public TransportChannel getTransportChannel() throws IOException {
+    // This provider's main purpose is to replace the default GAX ChannelPool
+    // with a custom BigtableChannelPool, reusing the delegate's configuration.
+
+    // To create our pool, we need a factory for raw gRPC channels.
+    // We achieve this by configuring our delegate to not use its own pooling
+    // (by setting pool size to 1) and then calling getTransportChannel() on it.
+    InstantiatingGrpcChannelProvider singleChannelProvider =
+        delegate.toBuilder().setChannelPoolSettings(ChannelPoolSettings.staticallySized(1)).build();
+
+    ChannelFactory channelFactory =
+        () -> {
+          try {
+            GrpcTransportChannel channel =
+                (GrpcTransportChannel) singleChannelProvider.getTransportChannel();
+            return (ManagedChannel) channel.getChannel();
+          } catch (IOException e) {
+            throw new java.io.UncheckedIOException(e);
+          }
+        };
+
+    BigtableChannelPoolSettings btPoolSettings =
+        BigtableChannelPoolSettings.copyFrom(delegate.getChannelPoolSettings());
+
+    BigtableChannelPool btChannelPool = BigtableChannelPool.create(btPoolSettings, channelFactory);
+
+    return GrpcTransportChannel.create(btChannelPool);
+  }
+
+  @Override
+  public String getTransportName() {
+    return "bigtable";
+  }
+
+  @Override
+  public boolean needsCredentials() {
+    return delegate.needsCredentials();
+  }
+
+  @Override
+  public TransportChannelProvider withCredentials(Credentials credentials) {
+    InstantiatingGrpcChannelProvider newChannelProvider =
+        (InstantiatingGrpcChannelProvider) delegate.withCredentials(credentials);
+    return new BigtableTransportChannelProvider(newChannelProvider);
+  }
+
+  /** Creates a BigtableTransportChannelProvider. */
+  public static BigtableTransportChannelProvider create(
+      InstantiatingGrpcChannelProvider instantiatingGrpcChannelProvider) {
+    return new BigtableTransportChannelProvider(instantiatingGrpcChannelProvider);
+  }
+}

google-cloud-bigtable/src/test/java/com/google/cloud/bigtable/gaxx/grpc/BigtableChannelPoolSettingsTest.java

@@ -0,0 +1,88 @@
+package com.google.cloud.bigtable.gaxx.grpc;
+
+import static com.google.common.truth.Truth.assertThat;
+
+import com.google.api.gax.grpc.ChannelPoolSettings;
+import com.google.common.collect.ImmutableSet;
+import java.lang.reflect.Method;
+import java.lang.reflect.Modifier;
+import java.util.Arrays;
+import java.util.Set;
+import java.util.stream.Collectors;
+import org.junit.Test;
+import org.junit.runner.RunWith;
+import org.junit.runners.JUnit4;
+
+@RunWith(JUnit4.class)
+public class BigtableChannelPoolSettingsTest {
+
+  @Test
+  public void testToBigtableChannelPoolSettingsAllFieldsSetCopiesCorrectly() throws Exception {
+    ChannelPoolSettings originalSettings =
+        ChannelPoolSettings.builder()
+            .setMinRpcsPerChannel(10)
+            .setMaxRpcsPerChannel(50)
+            .setMinChannelCount(5)
+            .setMaxChannelCount(100)
+            .setInitialChannelCount(20)
+            .setPreemptiveRefreshEnabled(true)
+            .build();
+
+    BigtableChannelPoolSettings copiedSettings =
+        BigtableChannelPoolSettings.copyFrom(originalSettings);
+    assertSettingsCopiedCorrectly(originalSettings, copiedSettings);
+  }
+
+  @Test
+  public void testToBigtableChannelPoolSettingsDefaultValuesCopiesCorrectly() throws Exception {
+    ChannelPoolSettings originalSettings = ChannelPoolSettings.builder().build();
+    BigtableChannelPoolSettings copiedSettings =
+        BigtableChannelPoolSettings.copyFrom(originalSettings);
+    assertSettingsCopiedCorrectly(originalSettings, copiedSettings);
+  }
+
+  private void assertSettingsCopiedCorrectly(
+      ChannelPoolSettings originalSettings, BigtableChannelPoolSettings copiedSettings)
+      throws Exception {
+
+    Set<String> supportedGetters =
+        ImmutableSet.of(
+            "getMinRpcsPerChannel",
+            "getMaxRpcsPerChannel",
+            "getMinChannelCount",
+            "getMaxChannelCount",
+            "getInitialChannelCount",
+            "isPreemptiveRefreshEnabled",
+            "isStaticSize");
+
+    Set<String> actualGetters =
+        Arrays.stream(ChannelPoolSettings.class.getDeclaredMethods())
+            .filter(
+                method ->
+                    Modifier.isPublic(method.getModifiers())
+                        && Modifier.isAbstract(method.getModifiers())
+                        && (method.getName().startsWith("get")
+                            || method.getName().startsWith("is")))
+            .map(Method::getName)
+            .collect(Collectors.toSet());
+
+    // If this fails then we need to add support for the additional attributes on the gax
+    // ChannelPool by updating the BigtableChannelPoolSettings.copyFrom method
+    assertThat(supportedGetters).containsAtLeastElementsIn(actualGetters);
+
+    assertThat(originalSettings.getInitialChannelCount())
+        .isEqualTo(copiedSettings.getInitialChannelCount());
+    assertThat(originalSettings.getMaxChannelCount())
+        .isEqualTo(copiedSettings.getMaxChannelCount());
+    assertThat(originalSettings.getMinChannelCount())
+        .isEqualTo(copiedSettings.getMinChannelCount());
+    assertThat(originalSettings.getMaxRpcsPerChannel())
+        .isEqualTo(copiedSettings.getMaxRpcsPerChannel());
+    assertThat(originalSettings.getMinRpcsPerChannel())
+        .isEqualTo(copiedSettings.getMinRpcsPerChannel());
+    assertThat(originalSettings.getInitialChannelCount())
+        .isEqualTo(copiedSettings.getInitialChannelCount());
+    assertThat(originalSettings.isPreemptiveRefreshEnabled())
+        .isEqualTo(copiedSettings.isPreemptiveRefreshEnabled());
+  }
+}