diff --git a/java-bigtable/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/internal/api/Client.java b/java-bigtable/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/internal/api/Client.java index 8bd9998e86f5..bec88617c992 100644 --- a/java-bigtable/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/internal/api/Client.java +++ b/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 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 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 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.createOwned( + userCallbackExecutor, () -> shutdownAndAwait(userCallbackExecutor))); } public Client( @@ -197,7 +200,7 @@ public Client( Resource metrics, Resource configManager, Resource bgExecutor, - Resource userCallbackExecutor) + Resource userCallbackExecutor) throws IOException { this.featureFlags = featureFlags; this.clientInfo = clientInfo; diff --git a/java-bigtable/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/internal/compat/ShimImpl.java b/java-bigtable/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/internal/compat/ShimImpl.java index 24ae2e546939..e51f65b828a9 100644 --- a/java-bigtable/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/internal/compat/ShimImpl.java +++ b/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 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 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( diff --git a/java-bigtable/google-cloud-bigtable/src/test/java/com/google/cloud/bigtable/data/v2/internal/compat/ShimImplTest.java b/java-bigtable/google-cloud-bigtable/src/test/java/com/google/cloud/bigtable/data/v2/internal/compat/ShimImplTest.java new file mode 100644 index 000000000000..1dc064855d16 --- /dev/null +++ b/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 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 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 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 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(); + } +}