Threading refactor (8/8): reuse user-configured executor for user-callback dispatch

java-bigtable/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/internal/api/Client.java

@@ -49,6 +49,7 @@
 import java.util.List;
 import java.util.Set;
 import java.util.WeakHashMap;
+import java.util.concurrent.Executor;
 import java.util.concurrent.ExecutorService;
 import java.util.concurrent.Executors;
 import java.util.concurrent.ScheduledExecutorService;
@@ -87,11 +88,12 @@ public class Client implements AutoCloseable {
   private final FeatureFlags featureFlags;
   private final ClientInfo clientInfo;
   private final Resource<ScheduledExecutorService> backgroundExecutor;
-  // Drains the per-op SerializingExecutor. Cached pool so a blocked user callback does not starve
-  // heartbeats, retry delays, or other vRPCs (which all run on backgroundExecutor).
-  // TODO: source from the gax TransportChannelProvider so transport and user-callback dispatch
-  // share the same pool. Blocked on missing APIs to extract the configured executor from gax.
-  private final Resource<ExecutorService> userCallbackExecutor;
+  // Drains the per-op SerializingExecutor. Sourced from the gax TransportChannelProvider /
+  // StubSettings on the compat path (see ShimImpl.selectUserCallbackExecutor) so a user-configured
+  // transport executor is reused for callback dispatch; otherwise a dedicated cached pool keeps a
+  // blocked user callback from starving heartbeats, retry delays, or other vRPCs (which all run on
+  // backgroundExecutor).
+  private final Resource<Executor> userCallbackExecutor;
   // Hashed-wheel timer for heartbeat / deadline / watchdog / retry scheduling. Built over
   // backgroundExecutor (the timer's tick thread dispatches bodies onto it). Single tick thread per
   // Client, shared across every SessionPoolImpl.
@@ -187,7 +189,8 @@ public static Client create(ClientSettings settings) throws IOException {
         Resource.createOwned(metrics, metrics::close),
         Resource.createOwned(configManager, configManager::close),
         Resource.createOwned(backgroundExecutor, backgroundExecutor::shutdown),
-        Resource.createOwned(userCallbackExecutor, () -> shutdownAndAwait(userCallbackExecutor)));
+        Resource.<Executor>createOwned(
+            userCallbackExecutor, () -> shutdownAndAwait(userCallbackExecutor)));
   }
 
   public Client(
@@ -197,7 +200,7 @@ public Client(
       Resource<Metrics> metrics,
       Resource<ClientConfigurationManager> configManager,
       Resource<ScheduledExecutorService> bgExecutor,
-      Resource<ExecutorService> userCallbackExecutor)
+      Resource<Executor> userCallbackExecutor)
       throws IOException {
     this.featureFlags = featureFlags;
     this.clientInfo = clientInfo;

java-bigtable/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/internal/compat/ShimImpl.java

@@ -15,6 +15,7 @@
  */
 package com.google.cloud.bigtable.data.v2.internal.compat;
 
+import com.google.api.gax.core.ExecutorProvider;
 import com.google.api.gax.grpc.ChannelPoolSettings;
 import com.google.api.gax.grpc.InstantiatingGrpcChannelProvider;
 import com.google.api.gax.rpc.StubSettings;
@@ -45,6 +46,7 @@
 import com.google.cloud.bigtable.data.v2.models.RowAdapter;
 import com.google.cloud.bigtable.data.v2.models.RowMutation;
 import com.google.cloud.bigtable.data.v2.stub.MetadataExtractorInterceptor;
+import com.google.common.annotations.VisibleForTesting;
 import com.google.common.base.Strings;
 import com.google.common.collect.ImmutableMap;
 import com.google.common.util.concurrent.ThreadFactoryBuilder;
@@ -55,6 +57,7 @@
 import java.io.IOException;
 import java.time.Duration;
 import java.util.Optional;
+import java.util.concurrent.Executor;
 import java.util.concurrent.ExecutorService;
 import java.util.concurrent.Executors;
 import java.util.concurrent.ScheduledExecutorService;
@@ -163,12 +166,9 @@ public static Shim create(
       featureFlags = featureFlags.toBuilder().setSessionsRequired(true).build();
     }
 
-    ExecutorService userCallbackExecutor =
-        Executors.newCachedThreadPool(
-            new ThreadFactoryBuilder()
-                .setNameFormat("bigtable-callback-shim-%d")
-                .setDaemon(true)
-                .build());
+    Resource<Executor> userCallbackExecutor =
+        selectUserCallbackExecutor(
+            stubSettings.getTransportChannelProvider(), stubSettings.getExecutorProvider());
 
     Client client =
         new Client(
@@ -178,8 +178,7 @@ public static Shim create(
             Resource.createShared(metrics),
             Resource.createShared(configManager),
             Resource.createShared(bgExecutor),
-            Resource.createOwned(
-                userCallbackExecutor, () -> Client.shutdownAndAwait(userCallbackExecutor)));
+            userCallbackExecutor);
 
     return new ShimImpl(configManager, client);
   }
@@ -192,6 +191,34 @@ public ShimImpl(ClientConfigurationManager configManager, Client client) {
     this.mutateRowShim = new MutateRowShim(client);
   }
 
+  // If the user configured an executor — either via InstantiatingGrpcChannelProvider#setExecutor
+  // or the legacy StubSettings#setExecutorProvider — reuse it for user-callback dispatch so the
+  // transport pool and callback pool are the same. Ownership: transport-set is borrowed by
+  // convention; legacy provider honors shouldAutoClose(). Otherwise allocate a dedicated cached
+  // pool that this Client owns.
+  @VisibleForTesting
+  static Resource<Executor> selectUserCallbackExecutor(
+      TransportChannelProvider transportProvider, @Nullable ExecutorProvider executorProvider) {
+    Executor transportExecutor = transportProvider.getExecutor();
+    if (transportExecutor != null) {
+      return Resource.createShared(transportExecutor);
+    }
+    if (executorProvider != null) {
+      ScheduledExecutorService executor = executorProvider.getExecutor();
+      if (executorProvider.shouldAutoClose()) {
+        return Resource.createOwned(executor, () -> Client.shutdownAndAwait(executor));
+      }
+      return Resource.createShared(executor);
+    }
+    ExecutorService owned =
+        Executors.newCachedThreadPool(
+            new ThreadFactoryBuilder()
+                .setNameFormat("bigtable-callback-shim-%d")
+                .setDaemon(true)
+                .build());
+    return Resource.createOwned(owned, () -> Client.shutdownAndAwait(owned));
+  }
+
   private static ChannelProvider configureChannelProvider(
       ChannelProvider channelProvider, Metrics metrics) {
     return new ConfiguredChannelProvider(

java-bigtable/google-cloud-bigtable/src/test/java/com/google/cloud/bigtable/data/v2/internal/compat/ShimImplTest.java

@@ -0,0 +1,97 @@
+/*
+ * 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.bigtable.data.v2.internal.compat;
+
+import static com.google.common.truth.Truth.assertThat;
+import static org.mockito.Mockito.mock;
+import static org.mockito.Mockito.never;
+import static org.mockito.Mockito.verify;
+
+import com.google.api.gax.core.FixedExecutorProvider;
+import com.google.api.gax.grpc.InstantiatingGrpcChannelProvider;
+import com.google.cloud.bigtable.data.v2.internal.api.Client.Resource;
+import java.util.concurrent.CountDownLatch;
+import java.util.concurrent.Executor;
+import java.util.concurrent.ExecutorService;
+import java.util.concurrent.ScheduledExecutorService;
+import java.util.concurrent.TimeUnit;
+import org.junit.Test;
+
+public class ShimImplTest {
+
+  private static final InstantiatingGrpcChannelProvider EMPTY_TRANSPORT =
+      InstantiatingGrpcChannelProvider.newBuilder().build();
+
+  @Test
+  public void selectUserCallbackExecutor_transportSet_returnsSharedTransportExecutor() {
+    Executor userExec = mock(Executor.class);
+    InstantiatingGrpcChannelProvider transport =
+        InstantiatingGrpcChannelProvider.newBuilder().setExecutor(userExec).build();
+
+    Resource<Executor> resource = ShimImpl.selectUserCallbackExecutor(transport, null);
+
+    assertThat(resource.get()).isSameInstanceAs(userExec);
+    // Shared resources have a no-op closer; close() must not throw.
+    resource.close();
+  }
+
+  @Test
+  public void selectUserCallbackExecutor_legacyNoAutoClose_returnsSharedExecutor() {
+    ScheduledExecutorService userExec = mock(ScheduledExecutorService.class);
+
+    Resource<Executor> resource =
+        ShimImpl.selectUserCallbackExecutor(
+            EMPTY_TRANSPORT, FixedExecutorProvider.create(userExec, /* shouldAutoClose= */ false));
+
+    assertThat(resource.get()).isSameInstanceAs(userExec);
+    resource.close();
+    verify(userExec, never()).shutdown();
+  }
+
+  @Test
+  public void selectUserCallbackExecutor_legacyAutoClose_returnsOwnedExecutor() {
+    ScheduledExecutorService userExec = mock(ScheduledExecutorService.class);
+
+    Resource<Executor> resource =
+        ShimImpl.selectUserCallbackExecutor(
+            EMPTY_TRANSPORT, FixedExecutorProvider.create(userExec, /* shouldAutoClose= */ true));
+
+    assertThat(resource.get()).isSameInstanceAs(userExec);
+    resource.close();
+    verify(userExec).shutdown();
+  }
+
+  @Test
+  public void selectUserCallbackExecutor_unset_allocatesOwnedCachedPool() throws Exception {
+    Resource<Executor> resource = ShimImpl.selectUserCallbackExecutor(EMPTY_TRANSPORT, null);
+
+    Executor exec = resource.get();
+    assertThat(exec).isInstanceOf(ExecutorService.class);
+
+    String[] threadName = new String[1];
+    CountDownLatch done = new CountDownLatch(1);
+    exec.execute(
+        () -> {
+          threadName[0] = Thread.currentThread().getName();
+          done.countDown();
+        });
+    assertThat(done.await(5, TimeUnit.SECONDS)).isTrue();
+    assertThat(threadName[0]).startsWith("bigtable-callback-shim-");
+
+    resource.close();
+    assertThat(((ExecutorService) exec).isShutdown()).isTrue();
+  }
+}