From bf0003b1996b10fd66388d36a28342bed9884b3c Mon Sep 17 00:00:00 2001 From: Sushan Bhattarai Date: Tue, 10 Mar 2026 02:45:21 -0400 Subject: [PATCH 1/8] int commit --- .../data/v2/stub/BigtableChannelPrimer.java | 15 +- .../v2/stub/EnhancedBigtableStubSettings.java | 52 ++++-- .../data/v2/stub/NoOpChannelPrimer.java | 12 +- .../BigtableTransportChannelProvider.java | 156 +++++++++++++----- .../bigtable/gaxx/grpc/ChannelPrimer.java | 14 +- .../gaxx/grpc/DirectAccessChecker.java | 26 +++ .../gaxx/grpc/UnaryDirectAccessChecker.java | 58 +++++++ 7 files changed, 260 insertions(+), 73 deletions(-) create mode 100644 google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/gaxx/grpc/DirectAccessChecker.java create mode 100644 google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/gaxx/grpc/UnaryDirectAccessChecker.java diff --git a/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/stub/BigtableChannelPrimer.java b/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/stub/BigtableChannelPrimer.java index e770e04ccb..85adcf5772 100644 --- a/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/stub/BigtableChannelPrimer.java +++ b/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/stub/BigtableChannelPrimer.java @@ -26,6 +26,7 @@ import com.google.cloud.bigtable.gaxx.grpc.ChannelPrimer; import io.grpc.CallCredentials; import io.grpc.CallOptions; +import io.grpc.Channel; import io.grpc.ClientCall; import io.grpc.Deadline; import io.grpc.ManagedChannel; @@ -88,21 +89,21 @@ static BigtableChannelPrimer create( } @Override - public void primeChannel(ManagedChannel managedChannel) { + public void primeChannel(Channel channel) { try { - primeChannelUnsafe(managedChannel); + primeChannelUnsafe(channel); } catch (IOException | RuntimeException e) { LOG.log(Level.WARNING, "Unexpected error while trying to prime a channel", e); } } - private void primeChannelUnsafe(ManagedChannel managedChannel) throws IOException { - sendPrimeRequestsBlocking(managedChannel); + private void primeChannelUnsafe(Channel channel) throws IOException { + sendPrimeRequestsBlocking(channel); } - private void sendPrimeRequestsBlocking(ManagedChannel managedChannel) { + private void sendPrimeRequestsBlocking(Channel channel) { try { - sendPrimeRequestsAsync(managedChannel).get(1, TimeUnit.MINUTES); + sendPrimeRequestsAsync(channel).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 @@ -112,7 +113,7 @@ private void sendPrimeRequestsBlocking(ManagedChannel managedChannel) { } @Override - public ApiFuture sendPrimeRequestsAsync(ManagedChannel managedChannel) { + public ApiFuture sendPrimeRequestsAsync(Channel managedChannel) { ClientCall clientCall = managedChannel.newCall( BigtableGrpc.getPingAndWarmMethod(), diff --git a/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/stub/EnhancedBigtableStubSettings.java b/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/stub/EnhancedBigtableStubSettings.java index db33e93fec..7ced90f8ef 100644 --- a/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/stub/EnhancedBigtableStubSettings.java +++ b/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/stub/EnhancedBigtableStubSettings.java @@ -242,6 +242,25 @@ public ClientOperationSettings getPerOpSettings() { return perOpSettings; } + /** Applies common pool, message size, and keep-alive settings to the provided builder. */ + private static InstantiatingGrpcChannelProvider.Builder commonTraits( + InstantiatingGrpcChannelProvider.Builder builder) { + return builder + .setChannelPoolSettings( + ChannelPoolSettings.builder() + .setInitialChannelCount(10) + .setMinRpcsPerChannel(1) + // Keep it conservative as we scale the channel size every 1min + // and delta is 2 channels. + .setMaxRpcsPerChannel(25) + .setPreemptiveRefreshEnabled(true) + .build()) + .setMaxInboundMessageSize(MAX_MESSAGE_SIZE) + .setKeepAliveTime(Duration.ofSeconds(30)) // sends ping in this interval + .setKeepAliveTimeout( + Duration.ofSeconds(10)); // wait this long before considering the connection dead + } + /** Returns a builder for the default ChannelProvider for this service. */ public static InstantiatingGrpcChannelProvider.Builder defaultGrpcTransportProviderBuilder() { InstantiatingGrpcChannelProvider.Builder grpcTransportProviderBuilder = @@ -261,22 +280,27 @@ public static InstantiatingGrpcChannelProvider.Builder defaultGrpcTransportProvi Collections.singletonList(InstantiatingGrpcChannelProvider.HardBoundTokenTypes.ALTS)); } } - return grpcTransportProviderBuilder - .setChannelPoolSettings( - ChannelPoolSettings.builder() - .setInitialChannelCount(10) - .setMinRpcsPerChannel(1) - // Keep it conservative as we scale the channel size every 1min - // and delta is 2 channels. - .setMaxRpcsPerChannel(25) - .setPreemptiveRefreshEnabled(true) - .build()) - .setMaxInboundMessageSize(MAX_MESSAGE_SIZE) - .setKeepAliveTime(Duration.ofSeconds(30)) // sends ping in this interval - .setKeepAliveTimeout( - Duration.ofSeconds(10)); // wait this long before considering the connection dead + return commonTraits(grpcTransportProviderBuilder); } + /** Applies Direct Access traits (DirectPath & ALTS) to an existing builder. */ + public static InstantiatingGrpcChannelProvider.Builder applyDirectAccessTraits( + InstantiatingGrpcChannelProvider.Builder builder) { + + builder + .setAttemptDirectPathXds() + .setAttemptDirectPath(true) + .setAllowNonDefaultServiceAccount(true); + + if (!DIRECT_PATH_BOUND_TOKEN_DISABLED) { + builder.setAllowHardBoundTokenTypes( + Collections.singletonList(InstantiatingGrpcChannelProvider.HardBoundTokenTypes.ALTS)); + } + + return builder; + } + + @SuppressWarnings("WeakerAccess") public static TransportChannelProvider defaultTransportChannelProvider() { return defaultGrpcTransportProviderBuilder().build(); diff --git a/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/stub/NoOpChannelPrimer.java b/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/stub/NoOpChannelPrimer.java index 3cb98d9dee..86c564bcc0 100644 --- a/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/stub/NoOpChannelPrimer.java +++ b/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/stub/NoOpChannelPrimer.java @@ -16,11 +16,11 @@ package com.google.cloud.bigtable.data.v2.stub; import com.google.api.core.ApiFuture; +import com.google.api.core.ApiFutures; import com.google.api.core.InternalApi; -import com.google.api.core.SettableApiFuture; import com.google.bigtable.v2.PingAndWarmResponse; import com.google.cloud.bigtable.gaxx.grpc.ChannelPrimer; -import io.grpc.ManagedChannel; +import io.grpc.Channel; @InternalApi public class NoOpChannelPrimer implements ChannelPrimer { @@ -31,14 +31,12 @@ static NoOpChannelPrimer create() { private NoOpChannelPrimer() {} @Override - public void primeChannel(ManagedChannel channel) { + public void primeChannel(Channel channel) { // No op } @Override - public ApiFuture sendPrimeRequestsAsync(ManagedChannel channel) { - SettableApiFuture future = SettableApiFuture.create(); - future.set(PingAndWarmResponse.getDefaultInstance()); - return future; + public ApiFuture sendPrimeRequestsAsync(Channel channel) { + return ApiFutures.immediateFuture(PingAndWarmResponse.getDefaultInstance()); } } diff --git a/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/gaxx/grpc/BigtableTransportChannelProvider.java b/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/gaxx/grpc/BigtableTransportChannelProvider.java index b75f145108..facb28dcf8 100644 --- a/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/gaxx/grpc/BigtableTransportChannelProvider.java +++ b/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/gaxx/grpc/BigtableTransportChannelProvider.java @@ -16,7 +16,6 @@ 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; @@ -24,12 +23,19 @@ import com.google.api.gax.rpc.TransportChannelProvider; import com.google.auth.Credentials; import com.google.cloud.bigtable.data.v2.internal.csm.tracers.ChannelPoolMetricsTracer; +import com.google.cloud.bigtable.data.v2.internal.csm.tracers.DirectPathCompatibleTracer; +import com.google.cloud.bigtable.data.v2.stub.BigtableChannelFactory; +import com.google.cloud.bigtable.data.v2.stub.DirectAccessChecker; +import com.google.cloud.bigtable.data.v2.stub.EnhancedBigtableStubSettings; +import com.google.cloud.bigtable.data.v2.stub.UnaryDirectAccessChecker; 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; +import java.util.logging.Level; +import java.util.logging.Logger; import javax.annotation.Nullable; /** @@ -38,20 +44,25 @@ */ @InternalApi public final class BigtableTransportChannelProvider implements TransportChannelProvider { + private static final Logger LOG = + Logger.getLogger(BigtableTransportChannelProvider.class.getName()); private final InstantiatingGrpcChannelProvider delegate; private final ChannelPrimer channelPrimer; @Nullable private final ChannelPoolMetricsTracer channelPoolMetricsTracer; @Nullable private final ScheduledExecutorService backgroundExecutor; + DirectPathCompatibleTracer directPathCompatibleTracer; private BigtableTransportChannelProvider( - InstantiatingGrpcChannelProvider instantiatingGrpcChannelProvider, - ChannelPrimer channelPrimer, - @Nullable ChannelPoolMetricsTracer channelPoolMetricsTracer, - @Nullable ScheduledExecutorService backgroundExecutor) { + InstantiatingGrpcChannelProvider instantiatingGrpcChannelProvider, + ChannelPrimer channelPrimer, + @Nullable ChannelPoolMetricsTracer channelPoolMetricsTracer, + @Nullable ScheduledExecutorService backgroundExecutor, + DirectPathCompatibleTracer directPathCompatibleTracer) { delegate = Preconditions.checkNotNull(instantiatingGrpcChannelProvider); this.channelPrimer = channelPrimer; this.channelPoolMetricsTracer = channelPoolMetricsTracer; this.backgroundExecutor = backgroundExecutor; + this.directPathCompatibleTracer = directPathCompatibleTracer; } @Override @@ -66,7 +77,6 @@ public boolean needsExecutor() { } @Override - @Deprecated public BigtableTransportChannelProvider withExecutor(ScheduledExecutorService executor) { return withExecutor((Executor) executor); } @@ -76,9 +86,13 @@ public BigtableTransportChannelProvider withExecutor(ScheduledExecutorService ex @Override public BigtableTransportChannelProvider withExecutor(Executor executor) { InstantiatingGrpcChannelProvider newChannelProvider = - (InstantiatingGrpcChannelProvider) delegate.withExecutor(executor); + (InstantiatingGrpcChannelProvider) delegate.withExecutor(executor); return new BigtableTransportChannelProvider( - newChannelProvider, channelPrimer, channelPoolMetricsTracer, backgroundExecutor); + newChannelProvider, + channelPrimer, + channelPoolMetricsTracer, + backgroundExecutor, + directPathCompatibleTracer); } @Override @@ -89,9 +103,13 @@ public boolean needsBackgroundExecutor() { @Override public TransportChannelProvider withBackgroundExecutor(ScheduledExecutorService executor) { InstantiatingGrpcChannelProvider newChannelProvider = - (InstantiatingGrpcChannelProvider) delegate.withBackgroundExecutor(executor); + (InstantiatingGrpcChannelProvider) delegate.withBackgroundExecutor(executor); return new BigtableTransportChannelProvider( - newChannelProvider, channelPrimer, channelPoolMetricsTracer, executor); + newChannelProvider, + channelPrimer, + channelPoolMetricsTracer, + executor, + directPathCompatibleTracer); } @Override @@ -102,9 +120,13 @@ public boolean needsHeaders() { @Override public BigtableTransportChannelProvider withHeaders(Map headers) { InstantiatingGrpcChannelProvider newChannelProvider = - (InstantiatingGrpcChannelProvider) delegate.withHeaders(headers); + (InstantiatingGrpcChannelProvider) delegate.withHeaders(headers); return new BigtableTransportChannelProvider( - newChannelProvider, channelPrimer, channelPoolMetricsTracer, backgroundExecutor); + newChannelProvider, + channelPrimer, + channelPoolMetricsTracer, + backgroundExecutor, + directPathCompatibleTracer); } @Override @@ -115,9 +137,13 @@ public boolean needsEndpoint() { @Override public TransportChannelProvider withEndpoint(String endpoint) { InstantiatingGrpcChannelProvider newChannelProvider = - (InstantiatingGrpcChannelProvider) delegate.withEndpoint(endpoint); + (InstantiatingGrpcChannelProvider) delegate.withEndpoint(endpoint); return new BigtableTransportChannelProvider( - newChannelProvider, channelPrimer, channelPoolMetricsTracer, backgroundExecutor); + newChannelProvider, + channelPrimer, + channelPoolMetricsTracer, + backgroundExecutor, + directPathCompatibleTracer); } @Deprecated @@ -130,14 +156,50 @@ public boolean acceptsPoolSize() { @Override public TransportChannelProvider withPoolSize(int size) { InstantiatingGrpcChannelProvider newChannelProvider = - (InstantiatingGrpcChannelProvider) delegate.withPoolSize(size); + (InstantiatingGrpcChannelProvider) delegate.withPoolSize(size); return new BigtableTransportChannelProvider( - newChannelProvider, channelPrimer, channelPoolMetricsTracer, backgroundExecutor); + newChannelProvider, + channelPrimer, + channelPoolMetricsTracer, + backgroundExecutor, + directPathCompatibleTracer); } + // We need this for direct access checker. + /** Expected to only be called once when BigtableClientContext is created */ @Override public TransportChannel getTransportChannel() throws IOException { + InstantiatingGrpcChannelProvider directAccessProvider = + EnhancedBigtableStubSettings.applyDirectAccessTraits(delegate.toBuilder()) + .setChannelPoolSettings(ChannelPoolSettings.staticallySized(1)) + .build(); + + BigtableChannelFactory maybeDirectAccessChannelFactory = + () -> { + GrpcTransportChannel channel = + (GrpcTransportChannel) directAccessProvider.getTransportChannel(); + return (ManagedChannel) channel.getChannel(); + }; + + DirectAccessChecker directAccessChecker = UnaryDirectAccessChecker.create(channelPrimer); + boolean isDirectAccessEligible = false; + + try { + isDirectAccessEligible = + directAccessChecker.check(maybeDirectAccessChannelFactory, directPathCompatibleTracer); + } catch (Exception e) { + LOG.log(Level.FINE, "Client is not direct access eligible, using standard transport.", e); + } + + InstantiatingGrpcChannelProvider selectedProvider; + + if (isDirectAccessEligible) { + selectedProvider = directAccessProvider; + } else { + selectedProvider = delegate; + } + // This provider's main purpose is to replace the default GAX ChannelPool // with a custom BigtableChannelPool, reusing the delegate's configuration. @@ -145,30 +207,32 @@ public TransportChannel getTransportChannel() throws IOException { // 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); - } - }; + selectedProvider.toBuilder() + .setChannelPoolSettings(ChannelPoolSettings.staticallySized(1)) + .build(); + + BigtableChannelFactory 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()); + BigtableChannelPoolSettings.copyFrom(delegate.getChannelPoolSettings()); BigtableChannelPool btChannelPool = - BigtableChannelPool.create( - btPoolSettings, channelFactory, channelPrimer, backgroundExecutor); + BigtableChannelPool.create( + btPoolSettings, channelFactory, channelPrimer, backgroundExecutor); if (channelPoolMetricsTracer != null) { channelPoolMetricsTracer.registerChannelInsightsProvider(btChannelPool); channelPoolMetricsTracer.registerLoadBalancingStrategy( - btPoolSettings.getLoadBalancingStrategy()); + btPoolSettings.getLoadBalancingStrategy()); } return GrpcTransportChannel.create(btChannelPool); @@ -187,21 +251,27 @@ public boolean needsCredentials() { @Override public TransportChannelProvider withCredentials(Credentials credentials) { InstantiatingGrpcChannelProvider newChannelProvider = - (InstantiatingGrpcChannelProvider) delegate.withCredentials(credentials); + (InstantiatingGrpcChannelProvider) delegate.withCredentials(credentials); return new BigtableTransportChannelProvider( - newChannelProvider, channelPrimer, channelPoolMetricsTracer, backgroundExecutor); + newChannelProvider, + channelPrimer, + channelPoolMetricsTracer, + backgroundExecutor, + directPathCompatibleTracer); } /** Creates a BigtableTransportChannelProvider. */ public static BigtableTransportChannelProvider create( - InstantiatingGrpcChannelProvider instantiatingGrpcChannelProvider, - ChannelPrimer channelPrimer, - ChannelPoolMetricsTracer outstandingRpcsMetricTracker, - ScheduledExecutorService backgroundExecutor) { + InstantiatingGrpcChannelProvider instantiatingGrpcChannelProvider, + ChannelPrimer channelPrimer, + ChannelPoolMetricsTracer outstandingRpcsMetricTracker, + ScheduledExecutorService backgroundExecutor, + DirectPathCompatibleTracer directPathCompatibleTracer) { return new BigtableTransportChannelProvider( - instantiatingGrpcChannelProvider, - channelPrimer, - outstandingRpcsMetricTracker, - backgroundExecutor); + instantiatingGrpcChannelProvider, + channelPrimer, + outstandingRpcsMetricTracker, + backgroundExecutor, + directPathCompatibleTracer); } -} +} \ No newline at end of file diff --git a/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/gaxx/grpc/ChannelPrimer.java b/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/gaxx/grpc/ChannelPrimer.java index ea7cc70175..6234b4410c 100644 --- a/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/gaxx/grpc/ChannelPrimer.java +++ b/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/gaxx/grpc/ChannelPrimer.java @@ -18,11 +18,21 @@ import com.google.api.core.ApiFuture; import com.google.api.core.InternalApi; import com.google.bigtable.v2.PingAndWarmResponse; +import io.grpc.Channel; import io.grpc.ManagedChannel; @InternalApi("For internal use by google-cloud-java clients only") public interface ChannelPrimer { - void primeChannel(ManagedChannel channel); + @Deprecated + default void primeChannel(ManagedChannel channel) { + primeChannel((Channel) channel); + } - ApiFuture sendPrimeRequestsAsync(ManagedChannel channel); + void primeChannel(Channel channel); + + @Deprecated + default ApiFuture sendPrimeRequestsAsync(ManagedChannel channel) { + return sendPrimeRequestsAsync((Channel) channel); + } + ApiFuture sendPrimeRequestsAsync(Channel channel); } diff --git a/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/gaxx/grpc/DirectAccessChecker.java b/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/gaxx/grpc/DirectAccessChecker.java new file mode 100644 index 0000000000..5ab0109b25 --- /dev/null +++ b/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/gaxx/grpc/DirectAccessChecker.java @@ -0,0 +1,26 @@ +/* + * 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.gaxx.grpc; + +import com.google.api.core.InternalApi; +import io.grpc.Channel; + +@InternalApi +/* Evaluates whether a given channel supports Direct Access. */ +public interface DirectAccessChecker { + /// Performs a request on the provided channel to check for Direct Access eligibility. + boolean check(Channel channel ); +} \ No newline at end of file diff --git a/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/gaxx/grpc/UnaryDirectAccessChecker.java b/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/gaxx/grpc/UnaryDirectAccessChecker.java new file mode 100644 index 0000000000..138dbd245c --- /dev/null +++ b/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/gaxx/grpc/UnaryDirectAccessChecker.java @@ -0,0 +1,58 @@ +/* + * 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.gaxx.grpc; + +import com.google.api.core.InternalApi; +import com.google.bigtable.v2.PeerInfo; +import com.google.cloud.bigtable.data.v2.stub.MetadataExtractorInterceptor; +import io.grpc.Channel; +import io.grpc.ClientInterceptors; + +import java.util.Optional; + +/** + * Evaluates whether a given channel has Direct Access (DirectPath) routing + * by executing a RPC and inspecting the response headers. + */ +@InternalApi +public class UnaryDirectAccessChecker implements DirectAccessChecker { + + private final ChannelPrimer channelPrimer; + + private UnaryDirectAccessChecker(ChannelPrimer channelPrimer) { + this.channelPrimer = channelPrimer; + } + + public static UnaryDirectAccessChecker create(ChannelPrimer channelPrimer) { + return new UnaryDirectAccessChecker(channelPrimer); + } + + @Override + public boolean check(Channel channel) { + MetadataExtractorInterceptor interceptor = new MetadataExtractorInterceptor(); + Channel interceptedChannel = ClientInterceptors.intercept(channel, interceptor); + channelPrimer.primeChannel(interceptedChannel); + + // Extract the sideband data populated by the interceptor + MetadataExtractorInterceptor.SidebandData sidebandData = interceptor.getSidebandData(); + + return Optional.of(sidebandData) + .map(MetadataExtractorInterceptor.SidebandData::getPeerInfo) + .map(PeerInfo::getTransportType) + .map(type -> type == PeerInfo.TransportType.TRANSPORT_TYPE_DIRECT_ACCESS) + .orElse(false); + } +} \ No newline at end of file From a6910b290aa2fed094c63c4cc31cb48b47832e00 Mon Sep 17 00:00:00 2001 From: Sushan Bhattarai Date: Tue, 10 Mar 2026 16:41:46 -0400 Subject: [PATCH 2/8] fix --- .../cloud/bigtable/data/v2/stub/BigtableChannelPrimer.java | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/stub/BigtableChannelPrimer.java b/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/stub/BigtableChannelPrimer.java index 85adcf5772..73c89e3778 100644 --- a/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/stub/BigtableChannelPrimer.java +++ b/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/stub/BigtableChannelPrimer.java @@ -29,10 +29,10 @@ import io.grpc.Channel; import io.grpc.ClientCall; import io.grpc.Deadline; -import io.grpc.ManagedChannel; import io.grpc.Metadata; import io.grpc.Status; import io.grpc.auth.MoreCallCredentials; + import java.io.IOException; import java.io.UnsupportedEncodingException; import java.net.URLEncoder; From 60bc8f4c9c3d0a4e7ff733866b42e0a477d51fb1 Mon Sep 17 00:00:00 2001 From: Sushan Bhattarai Date: Thu, 12 Mar 2026 03:15:32 -0400 Subject: [PATCH 3/8] fix --- .../data/v2/internal/csm/Metrics.java | 4 + .../data/v2/internal/csm/MetricsImpl.java | 11 + .../DefaultDirectPathCompatibleTracer.java | 42 ++ .../tracers/DirectPathCompatibleTracer.java | 39 ++ .../v2/stub/BigtableChannelFactory.java} | 12 +- .../data/v2/stub/BigtableClientContext.java | 3 +- .../data/v2/stub/DirectAccessChecker.java | 34 ++ .../v2/stub/EnhancedBigtableStubSettings.java | 8 +- .../v2/stub/MetadataExtractorInterceptor.java | 34 ++ .../v2/stub/UnaryDirectAccessChecker.java | 84 ++++ .../gaxx/grpc/BigtableChannelPool.java | 8 +- .../BigtableTransportChannelProvider.java | 446 +++++++++--------- .../gaxx/grpc/UnaryDirectAccessChecker.java | 58 --- .../v2/stub/UnaryDirectAccessCheckerTest.java | 134 ++++++ .../gaxx/grpc/BigtableChannelPoolTest.java | 4 +- 15 files changed, 625 insertions(+), 296 deletions(-) create mode 100644 google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/internal/csm/tracers/DefaultDirectPathCompatibleTracer.java create mode 100644 google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/internal/csm/tracers/DirectPathCompatibleTracer.java rename google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/{gaxx/grpc/DirectAccessChecker.java => data/v2/stub/BigtableChannelFactory.java} (68%) create mode 100644 google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/stub/DirectAccessChecker.java create mode 100644 google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/stub/UnaryDirectAccessChecker.java delete mode 100644 google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/gaxx/grpc/UnaryDirectAccessChecker.java create mode 100644 google-cloud-bigtable/src/test/java/com/google/cloud/bigtable/data/v2/stub/UnaryDirectAccessCheckerTest.java diff --git a/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/internal/csm/Metrics.java b/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/internal/csm/Metrics.java index 7df665c673..823458e0a4 100644 --- a/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/internal/csm/Metrics.java +++ b/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/internal/csm/Metrics.java @@ -18,6 +18,7 @@ import com.google.api.gax.tracing.ApiTracerFactory; import com.google.cloud.bigtable.data.v2.internal.csm.attributes.ClientInfo; import com.google.cloud.bigtable.data.v2.internal.csm.tracers.ChannelPoolMetricsTracer; +import com.google.cloud.bigtable.data.v2.internal.csm.tracers.DirectPathCompatibleTracer; import io.grpc.ManagedChannelBuilder; import java.io.Closeable; import java.io.IOException; @@ -31,6 +32,9 @@ public interface Metrics extends Closeable { @Nullable ChannelPoolMetricsTracer getChannelPoolMetricsTracer(); + @Nullable + DirectPathCompatibleTracer getDirectPathCompatibleTracer(); + void start(); @Override diff --git a/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/internal/csm/MetricsImpl.java b/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/internal/csm/MetricsImpl.java index f0efac7e96..5d4c8d0e58 100644 --- a/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/internal/csm/MetricsImpl.java +++ b/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/internal/csm/MetricsImpl.java @@ -31,6 +31,8 @@ import com.google.cloud.bigtable.data.v2.internal.csm.tracers.BuiltinMetricsTracerFactory; import com.google.cloud.bigtable.data.v2.internal.csm.tracers.ChannelPoolMetricsTracer; import com.google.cloud.bigtable.data.v2.internal.csm.tracers.CompositeTracerFactory; +import com.google.cloud.bigtable.data.v2.internal.csm.tracers.DefaultDirectPathCompatibleTracer; +import com.google.cloud.bigtable.data.v2.internal.csm.tracers.DirectPathCompatibleTracer; import com.google.cloud.bigtable.data.v2.internal.csm.tracers.Pacemaker; import com.google.common.base.Preconditions; import com.google.common.collect.ImmutableList; @@ -67,6 +69,7 @@ public class MetricsImpl implements Metrics, Closeable { @Nullable private final GrpcOpenTelemetry grpcOtel; @Nullable private final ChannelPoolMetricsTracer channelPoolMetricsTracer; + @Nullable private final DirectPathCompatibleTracer directPathCompatibleTracer; @Nullable private final Pacemaker pacemaker; private final List> tasks = new ArrayList<>(); @@ -94,6 +97,7 @@ public MetricsImpl( this.internalRecorder = metricRegistry.newRecorderRegistry(internalOtel.getMeterProvider()); this.pacemaker = new Pacemaker(internalRecorder, clientInfo, "background"); this.channelPoolMetricsTracer = new ChannelPoolMetricsTracer(internalRecorder, clientInfo); + this.directPathCompatibleTracer = new DefaultDirectPathCompatibleTracer(clientInfo, internalRecorder); this.grpcOtel = GrpcOpenTelemetry.newBuilder() .sdk(internalOtel) @@ -109,6 +113,7 @@ public MetricsImpl( this.grpcOtel = null; this.pacemaker = null; this.channelPoolMetricsTracer = null; + this.directPathCompatibleTracer = null; } if (userOtel != null) { @@ -171,6 +176,12 @@ public ChannelPoolMetricsTracer getChannelPoolMetricsTracer() { return channelPoolMetricsTracer; } + @Nullable + @Override + public DirectPathCompatibleTracer getDirectPathCompatibleTracer() { + return directPathCompatibleTracer; + } + public static OpenTelemetrySdk createBuiltinOtel( MetricRegistry metricRegistry, ClientInfo clientInfo, diff --git a/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/internal/csm/tracers/DefaultDirectPathCompatibleTracer.java b/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/internal/csm/tracers/DefaultDirectPathCompatibleTracer.java new file mode 100644 index 0000000000..041b226ab2 --- /dev/null +++ b/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/internal/csm/tracers/DefaultDirectPathCompatibleTracer.java @@ -0,0 +1,42 @@ +/* + * 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.csm.tracers; + +import com.google.api.core.InternalApi; +import com.google.cloud.bigtable.data.v2.internal.csm.MetricRegistry; +import com.google.cloud.bigtable.data.v2.internal.csm.attributes.ClientInfo; + +@InternalApi +public class DefaultDirectPathCompatibleTracer implements DirectPathCompatibleTracer { + private final ClientInfo clientInfo; + private final MetricRegistry.RecorderRegistry recorder; + + public DefaultDirectPathCompatibleTracer( + ClientInfo clientInfo, MetricRegistry.RecorderRegistry recorder) { + this.clientInfo = clientInfo; + this.recorder = recorder; + } + + @Override + public void recordSuccess(String ipPreference) { + recorder.dpCompatGuage.recordSuccess(clientInfo, ipPreference); + } + + @Override + public void recordFailure(String reason) { + recorder.dpCompatGuage.recordFailure(clientInfo, reason); + } +} diff --git a/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/internal/csm/tracers/DirectPathCompatibleTracer.java b/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/internal/csm/tracers/DirectPathCompatibleTracer.java new file mode 100644 index 0000000000..0a385494cf --- /dev/null +++ b/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/internal/csm/tracers/DirectPathCompatibleTracer.java @@ -0,0 +1,39 @@ +/* + * 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.csm.tracers; + +import com.google.api.core.InternalApi; + +/** + * Interface for recording DirectPath/DirectAccess eligibility metrics. + */ +@InternalApi +public interface DirectPathCompatibleTracer { + + /** + * Records that the environment is eligible and successfully connected via DirectPath. + * + * @param ipPreference The IP preference used (e.g., "ipv6"). + */ + void recordSuccess(String ipPreference); + + /** + * Records that the environment is not eligible or failed to connect via DirectPath. + * + * @param reason The reason for the failure (e.g., "routing_check_failed"). + */ + void recordFailure(String reason); +} \ No newline at end of file diff --git a/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/gaxx/grpc/DirectAccessChecker.java b/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/stub/BigtableChannelFactory.java similarity index 68% rename from google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/gaxx/grpc/DirectAccessChecker.java rename to google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/stub/BigtableChannelFactory.java index 5ab0109b25..a4ce11b58a 100644 --- a/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/gaxx/grpc/DirectAccessChecker.java +++ b/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/stub/BigtableChannelFactory.java @@ -13,14 +13,14 @@ * See the License for the specific language governing permissions and * limitations under the License. */ -package com.google.cloud.bigtable.gaxx.grpc; +package com.google.cloud.bigtable.data.v2.stub; import com.google.api.core.InternalApi; -import io.grpc.Channel; +import io.grpc.ManagedChannel; + +import java.io.IOException; @InternalApi -/* Evaluates whether a given channel supports Direct Access. */ -public interface DirectAccessChecker { - /// Performs a request on the provided channel to check for Direct Access eligibility. - boolean check(Channel channel ); +public interface BigtableChannelFactory { + ManagedChannel createSingleChannel() throws IOException; } \ No newline at end of file diff --git a/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/stub/BigtableClientContext.java b/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/stub/BigtableClientContext.java index 47bb9ccff9..3bebe1844d 100644 --- a/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/stub/BigtableClientContext.java +++ b/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/stub/BigtableClientContext.java @@ -166,7 +166,8 @@ public static BigtableClientContext create( transportProvider.build(), channelPrimer, metrics.getChannelPoolMetricsTracer(), - backgroundExecutor); + backgroundExecutor, metrics.getDirectPathCompatibleTracer() + ); builder.setTransportChannelProvider(btTransportProvider); } diff --git a/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/stub/DirectAccessChecker.java b/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/stub/DirectAccessChecker.java new file mode 100644 index 0000000000..281ab7b48f --- /dev/null +++ b/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/stub/DirectAccessChecker.java @@ -0,0 +1,34 @@ +/* + * 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.stub; + +import com.google.api.core.InternalApi; +import com.google.cloud.bigtable.data.v2.internal.csm.tracers.DirectPathCompatibleTracer; +import io.grpc.Channel; + +import javax.annotation.Nullable; + +@InternalApi +/* Evaluates whether a given channel supports Direct Access. */ +public interface DirectAccessChecker { + /** + * Evaluates if Direct Access is available by creating a test channel. + * + * @param channelFactory A factory to create the test channel + * @return true if the channel is eligible for Direct Access + */ + boolean check(BigtableChannelFactory channelFactory, @Nullable DirectPathCompatibleTracer tracer); +} \ No newline at end of file diff --git a/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/stub/EnhancedBigtableStubSettings.java b/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/stub/EnhancedBigtableStubSettings.java index 7ced90f8ef..47937216da 100644 --- a/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/stub/EnhancedBigtableStubSettings.java +++ b/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/stub/EnhancedBigtableStubSettings.java @@ -634,12 +634,16 @@ private Builder() { perOpSettings = new ClientOperationSettings.Builder(); + // Note: RouteLookup evaluates and returns directpath targets + // only if Traffic Director sends the request (with grpc as target type) + // For GFE/CFE, sending setDirectAccessRequested + // is fine as GFE/CFE sends with gslb target type featureFlags = FeatureFlags.newBuilder() .setReverseScans(true) .setLastScannedRowResponses(true) - .setDirectAccessRequested(DIRECT_PATH_ENABLED) - .setTrafficDirectorEnabled(DIRECT_PATH_ENABLED) + .setDirectAccessRequested(true) + .setTrafficDirectorEnabled(true) .setPeerInfo(true); } diff --git a/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/stub/MetadataExtractorInterceptor.java b/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/stub/MetadataExtractorInterceptor.java index d38f164fc6..145e4064ea 100644 --- a/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/stub/MetadataExtractorInterceptor.java +++ b/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/stub/MetadataExtractorInterceptor.java @@ -29,10 +29,16 @@ import io.grpc.ClientInterceptors; import io.grpc.ForwardingClientCall; import io.grpc.ForwardingClientCallListener; +import io.grpc.Grpc; import io.grpc.Metadata; import io.grpc.MethodDescriptor; import io.grpc.Status; import io.grpc.alts.AltsContextUtil; + +import java.net.Inet4Address; +import java.net.Inet6Address; +import java.net.InetSocketAddress; +import java.net.SocketAddress; import java.time.Duration; import java.util.Base64; import java.util.regex.Matcher; @@ -85,6 +91,12 @@ public SidebandData getSidebandData() { } public static class SidebandData { + public enum IpProtocol { + IPV4, + IPV6, + UNKNOWN + } + private static final CallOptions.Key KEY = CallOptions.Key.create("bigtable-sideband"); @@ -105,6 +117,7 @@ public static SidebandData from(CallOptions callOptions) { @Nullable private volatile ResponseParams responseParams; @Nullable private volatile PeerInfo peerInfo; @Nullable private volatile Duration gfeTiming; + @Nullable private volatile IpProtocol ipProtocol; @Nullable public ResponseParams getResponseParams() { @@ -121,16 +134,23 @@ public Duration getGfeTiming() { return gfeTiming; } + @Nullable + public IpProtocol getIpProtocol() { + return ipProtocol; + } + private void reset() { responseParams = null; peerInfo = null; gfeTiming = null; + ipProtocol = IpProtocol.UNKNOWN; } void onResponseHeaders(Metadata md, Attributes attributes) { responseParams = extractResponseParams(md); gfeTiming = extractGfeLatency(md); peerInfo = extractPeerInfo(md, gfeTiming, attributes); + ipProtocol = extractIpProtocol(attributes); } void onClose(Status status, Metadata trailers) { @@ -139,6 +159,20 @@ void onClose(Status status, Metadata trailers) { } } + @Nullable + private static IpProtocol extractIpProtocol(Attributes attributes) { + SocketAddress remoteAddr = attributes.get(Grpc.TRANSPORT_ATTR_REMOTE_ADDR); + if (remoteAddr instanceof InetSocketAddress) { + InetSocketAddress inetAddr = (InetSocketAddress) remoteAddr; + if (inetAddr.getAddress() instanceof Inet4Address) { + return IpProtocol.IPV4; + } else if (inetAddr.getAddress() instanceof Inet6Address) { + return IpProtocol.IPV6; + } + } + return IpProtocol.UNKNOWN; + } + @Nullable private static Duration extractGfeLatency(Metadata metadata) { String serverTiming = metadata.get(SERVER_TIMING_HEADER_KEY); diff --git a/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/stub/UnaryDirectAccessChecker.java b/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/stub/UnaryDirectAccessChecker.java new file mode 100644 index 0000000000..edb842235f --- /dev/null +++ b/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/stub/UnaryDirectAccessChecker.java @@ -0,0 +1,84 @@ +/* + * 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.stub; + +import com.google.api.core.InternalApi; +import com.google.bigtable.v2.PeerInfo; +import com.google.cloud.bigtable.data.v2.internal.csm.tracers.DirectPathCompatibleTracer; +import com.google.cloud.bigtable.gaxx.grpc.ChannelPrimer; +import io.grpc.Channel; +import io.grpc.ClientInterceptors; +import io.grpc.ManagedChannel; + +import javax.annotation.Nullable; +import java.util.Optional; +import java.util.logging.Level; +import java.util.logging.Logger; + +/** + * Evaluates whether a given channel has Direct Access (DirectPath) routing + * by executing a RPC and inspecting the response headers. + */ +@InternalApi +public class UnaryDirectAccessChecker implements DirectAccessChecker { + private static final Logger LOG = Logger.getLogger(UnaryDirectAccessChecker.class.getName()); + private final ChannelPrimer channelPrimer; + + private UnaryDirectAccessChecker(ChannelPrimer channelPrimer) { + this.channelPrimer = channelPrimer; + } + + public static UnaryDirectAccessChecker create(ChannelPrimer channelPrimer) { + return new UnaryDirectAccessChecker(channelPrimer); + } + + @Override + public boolean check(BigtableChannelFactory channelFactory, @Nullable DirectPathCompatibleTracer tracer) { + ManagedChannel channel = null; + try { + channel = channelFactory.createSingleChannel(); + MetadataExtractorInterceptor interceptor = new MetadataExtractorInterceptor(); + Channel interceptedChannel = ClientInterceptors.intercept(channel, interceptor); + channelPrimer.primeChannel(interceptedChannel); + + // Extract the sideband data populated by the interceptor + MetadataExtractorInterceptor.SidebandData sidebandData = interceptor.getSidebandData(); + + boolean isEligible = Optional.ofNullable(sidebandData) + .map(MetadataExtractorInterceptor.SidebandData::getPeerInfo) + .map(PeerInfo::getTransportType) + .map(type -> type == PeerInfo.TransportType.TRANSPORT_TYPE_DIRECT_ACCESS) + .orElse(false); + + if (isEligible && tracer != null) { + String ipProtocolStr = Optional.ofNullable(sidebandData) + .map(MetadataExtractorInterceptor.SidebandData::getIpProtocol) + .map(String::valueOf) + .map(String::toLowerCase) + .orElse("unknown"); + tracer.recordSuccess(ipProtocolStr); + } + return isEligible; + } catch (Exception e) { + LOG.log(Level.FINE, "Failed to evaluate direct access eligibility.", e); + return false; + } finally { + if (channel != null) { + channel.shutdownNow(); + } + } + } +} \ No newline at end of file diff --git a/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/gaxx/grpc/BigtableChannelPool.java b/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/gaxx/grpc/BigtableChannelPool.java index 188dc83ac0..0298b730e7 100644 --- a/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/gaxx/grpc/BigtableChannelPool.java +++ b/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/gaxx/grpc/BigtableChannelPool.java @@ -16,8 +16,8 @@ package com.google.cloud.bigtable.gaxx.grpc; import com.google.api.core.InternalApi; -import com.google.api.gax.grpc.ChannelFactory; import com.google.bigtable.v2.PeerInfo; +import com.google.cloud.bigtable.data.v2.stub.BigtableChannelFactory; import com.google.cloud.bigtable.data.v2.stub.MetadataExtractorInterceptor; import com.google.cloud.bigtable.gaxx.grpc.ChannelPoolHealthChecker.ProbeResult; import com.google.common.annotations.VisibleForTesting; @@ -67,7 +67,7 @@ public class BigtableChannelPool extends ManagedChannel implements BigtableChann private static final java.time.Duration REFRESH_PERIOD = java.time.Duration.ofMinutes(50); private final BigtableChannelPoolSettings settings; - private final ChannelFactory channelFactory; + private final BigtableChannelFactory channelFactory; private final ChannelPrimer channelPrimer; private final Object entryWriteLock = new Object(); @@ -81,7 +81,7 @@ public class BigtableChannelPool extends ManagedChannel implements BigtableChann public static BigtableChannelPool create( BigtableChannelPoolSettings settings, - ChannelFactory channelFactory, + BigtableChannelFactory channelFactory, ChannelPrimer channelPrimer, ScheduledExecutorService backgroundExecutor) throws IOException { @@ -98,7 +98,7 @@ public static BigtableChannelPool create( @VisibleForTesting BigtableChannelPool( BigtableChannelPoolSettings settings, - ChannelFactory channelFactory, + BigtableChannelFactory channelFactory, ChannelPrimer channelPrimer, ScheduledExecutorService executor) throws IOException { diff --git a/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/gaxx/grpc/BigtableTransportChannelProvider.java b/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/gaxx/grpc/BigtableTransportChannelProvider.java index facb28dcf8..2b3cb4022e 100644 --- a/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/gaxx/grpc/BigtableTransportChannelProvider.java +++ b/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/gaxx/grpc/BigtableTransportChannelProvider.java @@ -44,234 +44,234 @@ */ @InternalApi public final class BigtableTransportChannelProvider implements TransportChannelProvider { - private static final Logger LOG = - Logger.getLogger(BigtableTransportChannelProvider.class.getName()); - private final InstantiatingGrpcChannelProvider delegate; - private final ChannelPrimer channelPrimer; - @Nullable private final ChannelPoolMetricsTracer channelPoolMetricsTracer; - @Nullable private final ScheduledExecutorService backgroundExecutor; - DirectPathCompatibleTracer directPathCompatibleTracer; - - private BigtableTransportChannelProvider( - InstantiatingGrpcChannelProvider instantiatingGrpcChannelProvider, - ChannelPrimer channelPrimer, - @Nullable ChannelPoolMetricsTracer channelPoolMetricsTracer, - @Nullable ScheduledExecutorService backgroundExecutor, - DirectPathCompatibleTracer directPathCompatibleTracer) { - delegate = Preconditions.checkNotNull(instantiatingGrpcChannelProvider); - this.channelPrimer = channelPrimer; - this.channelPoolMetricsTracer = channelPoolMetricsTracer; - this.backgroundExecutor = backgroundExecutor; - this.directPathCompatibleTracer = directPathCompatibleTracer; - } - - @Override - public boolean shouldAutoClose() { - return delegate.shouldAutoClose(); - } - - @SuppressWarnings("deprecation") - @Override - public boolean needsExecutor() { - return delegate.needsExecutor(); - } - - @Override - public BigtableTransportChannelProvider withExecutor(ScheduledExecutorService executor) { - return withExecutor((Executor) executor); - } - - // This executor if set is for handling rpc callbacks so we can't use it as the background - // executor - @Override - public BigtableTransportChannelProvider withExecutor(Executor executor) { - InstantiatingGrpcChannelProvider newChannelProvider = - (InstantiatingGrpcChannelProvider) delegate.withExecutor(executor); - return new BigtableTransportChannelProvider( - newChannelProvider, - channelPrimer, - channelPoolMetricsTracer, - backgroundExecutor, - directPathCompatibleTracer); - } - - @Override - public boolean needsBackgroundExecutor() { - return delegate.needsBackgroundExecutor(); - } - - @Override - public TransportChannelProvider withBackgroundExecutor(ScheduledExecutorService executor) { - InstantiatingGrpcChannelProvider newChannelProvider = - (InstantiatingGrpcChannelProvider) delegate.withBackgroundExecutor(executor); - return new BigtableTransportChannelProvider( - newChannelProvider, - channelPrimer, - channelPoolMetricsTracer, - executor, - directPathCompatibleTracer); - } - - @Override - public boolean needsHeaders() { - return delegate.needsHeaders(); - } - - @Override - public BigtableTransportChannelProvider withHeaders(Map headers) { - InstantiatingGrpcChannelProvider newChannelProvider = - (InstantiatingGrpcChannelProvider) delegate.withHeaders(headers); - return new BigtableTransportChannelProvider( - newChannelProvider, - channelPrimer, - channelPoolMetricsTracer, - backgroundExecutor, - directPathCompatibleTracer); - } - - @Override - public boolean needsEndpoint() { - return delegate.needsEndpoint(); - } - - @Override - public TransportChannelProvider withEndpoint(String endpoint) { - InstantiatingGrpcChannelProvider newChannelProvider = - (InstantiatingGrpcChannelProvider) delegate.withEndpoint(endpoint); - return new BigtableTransportChannelProvider( - newChannelProvider, - channelPrimer, - channelPoolMetricsTracer, - backgroundExecutor, - directPathCompatibleTracer); - } - - @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, - channelPrimer, - channelPoolMetricsTracer, - backgroundExecutor, - directPathCompatibleTracer); - } - - // We need this for direct access checker. - - /** Expected to only be called once when BigtableClientContext is created */ - @Override - public TransportChannel getTransportChannel() throws IOException { - InstantiatingGrpcChannelProvider directAccessProvider = - EnhancedBigtableStubSettings.applyDirectAccessTraits(delegate.toBuilder()) - .setChannelPoolSettings(ChannelPoolSettings.staticallySized(1)) - .build(); - - BigtableChannelFactory maybeDirectAccessChannelFactory = - () -> { - GrpcTransportChannel channel = - (GrpcTransportChannel) directAccessProvider.getTransportChannel(); - return (ManagedChannel) channel.getChannel(); - }; - - DirectAccessChecker directAccessChecker = UnaryDirectAccessChecker.create(channelPrimer); - boolean isDirectAccessEligible = false; - - try { - isDirectAccessEligible = - directAccessChecker.check(maybeDirectAccessChannelFactory, directPathCompatibleTracer); - } catch (Exception e) { - LOG.log(Level.FINE, "Client is not direct access eligible, using standard transport.", e); + private static final Logger LOG = + Logger.getLogger(BigtableTransportChannelProvider.class.getName()); + private final InstantiatingGrpcChannelProvider delegate; + private final ChannelPrimer channelPrimer; + @Nullable private final ChannelPoolMetricsTracer channelPoolMetricsTracer; + @Nullable private final ScheduledExecutorService backgroundExecutor; + DirectPathCompatibleTracer directPathCompatibleTracer; + + private BigtableTransportChannelProvider( + InstantiatingGrpcChannelProvider instantiatingGrpcChannelProvider, + ChannelPrimer channelPrimer, + @Nullable ChannelPoolMetricsTracer channelPoolMetricsTracer, + @Nullable ScheduledExecutorService backgroundExecutor, + DirectPathCompatibleTracer directPathCompatibleTracer) { + delegate = Preconditions.checkNotNull(instantiatingGrpcChannelProvider); + this.channelPrimer = channelPrimer; + this.channelPoolMetricsTracer = channelPoolMetricsTracer; + this.backgroundExecutor = backgroundExecutor; + this.directPathCompatibleTracer = directPathCompatibleTracer; } - InstantiatingGrpcChannelProvider selectedProvider; + @Override + public boolean shouldAutoClose() { + return delegate.shouldAutoClose(); + } + + @SuppressWarnings("deprecation") + @Override + public boolean needsExecutor() { + return delegate.needsExecutor(); + } + + @Override + public BigtableTransportChannelProvider withExecutor(ScheduledExecutorService executor) { + return withExecutor((Executor) executor); + } + + // This executor if set is for handling rpc callbacks so we can't use it as the background + // executor + @Override + public BigtableTransportChannelProvider withExecutor(Executor executor) { + InstantiatingGrpcChannelProvider newChannelProvider = + (InstantiatingGrpcChannelProvider) delegate.withExecutor(executor); + return new BigtableTransportChannelProvider( + newChannelProvider, + channelPrimer, + channelPoolMetricsTracer, + backgroundExecutor, + directPathCompatibleTracer); + } + + @Override + public boolean needsBackgroundExecutor() { + return delegate.needsBackgroundExecutor(); + } + + @Override + public TransportChannelProvider withBackgroundExecutor(ScheduledExecutorService executor) { + InstantiatingGrpcChannelProvider newChannelProvider = + (InstantiatingGrpcChannelProvider) delegate.withBackgroundExecutor(executor); + return new BigtableTransportChannelProvider( + newChannelProvider, + channelPrimer, + channelPoolMetricsTracer, + executor, + directPathCompatibleTracer); + } + + @Override + public boolean needsHeaders() { + return delegate.needsHeaders(); + } + + @Override + public BigtableTransportChannelProvider withHeaders(Map headers) { + InstantiatingGrpcChannelProvider newChannelProvider = + (InstantiatingGrpcChannelProvider) delegate.withHeaders(headers); + return new BigtableTransportChannelProvider( + newChannelProvider, + channelPrimer, + channelPoolMetricsTracer, + backgroundExecutor, + directPathCompatibleTracer); + } + + @Override + public boolean needsEndpoint() { + return delegate.needsEndpoint(); + } + + @Override + public TransportChannelProvider withEndpoint(String endpoint) { + InstantiatingGrpcChannelProvider newChannelProvider = + (InstantiatingGrpcChannelProvider) delegate.withEndpoint(endpoint); + return new BigtableTransportChannelProvider( + newChannelProvider, + channelPrimer, + channelPoolMetricsTracer, + backgroundExecutor, + directPathCompatibleTracer); + } + + @Deprecated + @Override + public boolean acceptsPoolSize() { + return delegate.acceptsPoolSize(); + } - if (isDirectAccessEligible) { - selectedProvider = directAccessProvider; - } else { - selectedProvider = delegate; + @Deprecated + @Override + public TransportChannelProvider withPoolSize(int size) { + InstantiatingGrpcChannelProvider newChannelProvider = + (InstantiatingGrpcChannelProvider) delegate.withPoolSize(size); + return new BigtableTransportChannelProvider( + newChannelProvider, + channelPrimer, + channelPoolMetricsTracer, + backgroundExecutor, + directPathCompatibleTracer); } - // 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 = - selectedProvider.toBuilder() - .setChannelPoolSettings(ChannelPoolSettings.staticallySized(1)) - .build(); - - BigtableChannelFactory 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, channelPrimer, backgroundExecutor); - - if (channelPoolMetricsTracer != null) { - channelPoolMetricsTracer.registerChannelInsightsProvider(btChannelPool); - channelPoolMetricsTracer.registerLoadBalancingStrategy( - btPoolSettings.getLoadBalancingStrategy()); + // We need this for direct access checker. + + /** Expected to only be called once when BigtableClientContext is created */ + @Override + public TransportChannel getTransportChannel() throws IOException { + InstantiatingGrpcChannelProvider directAccessProvider = + EnhancedBigtableStubSettings.applyDirectAccessTraits(delegate.toBuilder()) + .setChannelPoolSettings(ChannelPoolSettings.staticallySized(1)) + .build(); + + BigtableChannelFactory maybeDirectAccessChannelFactory = + () -> { + GrpcTransportChannel channel = + (GrpcTransportChannel) directAccessProvider.getTransportChannel(); + return (ManagedChannel) channel.getChannel(); + }; + + DirectAccessChecker directAccessChecker = UnaryDirectAccessChecker.create(channelPrimer); + boolean isDirectAccessEligible = false; + + try { + isDirectAccessEligible = + directAccessChecker.check(maybeDirectAccessChannelFactory, directPathCompatibleTracer); + } catch (Exception e) { + LOG.log(Level.FINE, "Client is not direct access eligible, using standard transport.", e); + } + + InstantiatingGrpcChannelProvider selectedProvider; + + if (isDirectAccessEligible) { + selectedProvider = directAccessProvider; + } else { + selectedProvider = delegate; + } + + // 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 = + selectedProvider.toBuilder() + .setChannelPoolSettings(ChannelPoolSettings.staticallySized(1)) + .build(); + + BigtableChannelFactory 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, channelPrimer, backgroundExecutor); + + if (channelPoolMetricsTracer != null) { + channelPoolMetricsTracer.registerChannelInsightsProvider(btChannelPool); + channelPoolMetricsTracer.registerLoadBalancingStrategy( + btPoolSettings.getLoadBalancingStrategy()); + } + + return GrpcTransportChannel.create(btChannelPool); } - 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, - channelPrimer, - channelPoolMetricsTracer, - backgroundExecutor, - directPathCompatibleTracer); - } - - /** Creates a BigtableTransportChannelProvider. */ - public static BigtableTransportChannelProvider create( - InstantiatingGrpcChannelProvider instantiatingGrpcChannelProvider, - ChannelPrimer channelPrimer, - ChannelPoolMetricsTracer outstandingRpcsMetricTracker, - ScheduledExecutorService backgroundExecutor, - DirectPathCompatibleTracer directPathCompatibleTracer) { - return new BigtableTransportChannelProvider( - instantiatingGrpcChannelProvider, - channelPrimer, - outstandingRpcsMetricTracker, - backgroundExecutor, - directPathCompatibleTracer); - } + @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, + channelPrimer, + channelPoolMetricsTracer, + backgroundExecutor, + directPathCompatibleTracer); + } + + /** Creates a BigtableTransportChannelProvider. */ + public static BigtableTransportChannelProvider create( + InstantiatingGrpcChannelProvider instantiatingGrpcChannelProvider, + ChannelPrimer channelPrimer, + ChannelPoolMetricsTracer outstandingRpcsMetricTracker, + ScheduledExecutorService backgroundExecutor, + DirectPathCompatibleTracer directPathCompatibleTracer) { + return new BigtableTransportChannelProvider( + instantiatingGrpcChannelProvider, + channelPrimer, + outstandingRpcsMetricTracker, + backgroundExecutor, + directPathCompatibleTracer); + } } \ No newline at end of file diff --git a/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/gaxx/grpc/UnaryDirectAccessChecker.java b/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/gaxx/grpc/UnaryDirectAccessChecker.java deleted file mode 100644 index 138dbd245c..0000000000 --- a/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/gaxx/grpc/UnaryDirectAccessChecker.java +++ /dev/null @@ -1,58 +0,0 @@ -/* - * 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.gaxx.grpc; - -import com.google.api.core.InternalApi; -import com.google.bigtable.v2.PeerInfo; -import com.google.cloud.bigtable.data.v2.stub.MetadataExtractorInterceptor; -import io.grpc.Channel; -import io.grpc.ClientInterceptors; - -import java.util.Optional; - -/** - * Evaluates whether a given channel has Direct Access (DirectPath) routing - * by executing a RPC and inspecting the response headers. - */ -@InternalApi -public class UnaryDirectAccessChecker implements DirectAccessChecker { - - private final ChannelPrimer channelPrimer; - - private UnaryDirectAccessChecker(ChannelPrimer channelPrimer) { - this.channelPrimer = channelPrimer; - } - - public static UnaryDirectAccessChecker create(ChannelPrimer channelPrimer) { - return new UnaryDirectAccessChecker(channelPrimer); - } - - @Override - public boolean check(Channel channel) { - MetadataExtractorInterceptor interceptor = new MetadataExtractorInterceptor(); - Channel interceptedChannel = ClientInterceptors.intercept(channel, interceptor); - channelPrimer.primeChannel(interceptedChannel); - - // Extract the sideband data populated by the interceptor - MetadataExtractorInterceptor.SidebandData sidebandData = interceptor.getSidebandData(); - - return Optional.of(sidebandData) - .map(MetadataExtractorInterceptor.SidebandData::getPeerInfo) - .map(PeerInfo::getTransportType) - .map(type -> type == PeerInfo.TransportType.TRANSPORT_TYPE_DIRECT_ACCESS) - .orElse(false); - } -} \ No newline at end of file diff --git a/google-cloud-bigtable/src/test/java/com/google/cloud/bigtable/data/v2/stub/UnaryDirectAccessCheckerTest.java b/google-cloud-bigtable/src/test/java/com/google/cloud/bigtable/data/v2/stub/UnaryDirectAccessCheckerTest.java new file mode 100644 index 0000000000..5894565898 --- /dev/null +++ b/google-cloud-bigtable/src/test/java/com/google/cloud/bigtable/data/v2/stub/UnaryDirectAccessCheckerTest.java @@ -0,0 +1,134 @@ +/* + * 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.stub; + +import static com.google.common.truth.Truth.assertThat; +import static org.mockito.ArgumentMatchers.any; +import static org.mockito.Mockito.*; + +import com.google.bigtable.v2.PeerInfo; +import com.google.cloud.bigtable.data.v2.internal.csm.tracers.DirectPathCompatibleTracer; +import com.google.cloud.bigtable.gaxx.grpc.ChannelPrimer; +import io.grpc.Channel; +import io.grpc.ManagedChannel; +import org.junit.Before; +import org.junit.Rule; +import org.junit.Test; +import org.junit.runner.RunWith; +import org.junit.runners.JUnit4; +import org.mockito.Mock; +import org.mockito.MockedConstruction; +import org.mockito.junit.MockitoJUnit; +import org.mockito.junit.MockitoRule; + +@RunWith(JUnit4.class) +public class UnaryDirectAccessCheckerTest { + + @Rule public final MockitoRule mockito = MockitoJUnit.rule(); + + @Mock private ChannelPrimer mockChannelPrimer; + @Mock private BigtableChannelFactory mockChannelFactory; + @Mock private DirectPathCompatibleTracer mockTracer; + @Mock private ManagedChannel mockChannel; + @Mock private MetadataExtractorInterceptor.SidebandData mockSidebandData; + + private UnaryDirectAccessChecker checker; + + @Before + public void setUp() throws Exception { + checker = UnaryDirectAccessChecker.create(mockChannelPrimer); + when(mockChannelFactory.createSingleChannel()).thenReturn(mockChannel); + } + + @Test + public void testEligibleForDirectAccess() { + PeerInfo peerInfo = PeerInfo.newBuilder() + .setTransportType(PeerInfo.TransportType.TRANSPORT_TYPE_DIRECT_ACCESS) + .build(); + when(mockSidebandData.getPeerInfo()).thenReturn(peerInfo); + + when(mockSidebandData.getIpProtocol()).thenReturn(MetadataExtractorInterceptor.SidebandData.IpProtocol.IPV6); + + try (MockedConstruction ignored = + mockConstruction( + MetadataExtractorInterceptor.class, + (mock, context) -> { + when(mock.getSidebandData()).thenReturn(mockSidebandData); + })) { + + boolean isEligible = checker.check(mockChannelFactory, mockTracer); + + assertThat(isEligible).isFalse(); + verify(mockChannelPrimer).primeChannel(any(Channel.class)); + verify(mockTracer).recordSuccess("ipv6"); + verify(mockChannel).shutdownNow(); + } + } + + @Test + public void testNotEligibleProxiedRouting() { + // 1. Setup sideband data to simulate standard CloudPath routing + PeerInfo peerInfo = PeerInfo.newBuilder() + .setTransportType(PeerInfo.TransportType.TRANSPORT_TYPE_CLOUD_PATH) + .build(); + when(mockSidebandData.getPeerInfo()).thenReturn(peerInfo); + + try (MockedConstruction mocked = + mockConstruction( + MetadataExtractorInterceptor.class, + (mock, context) -> { + when(mock.getSidebandData()).thenReturn(mockSidebandData); + })) { + + boolean isEligible = checker.check(mockChannelFactory, mockTracer); + + assertThat(isEligible).isFalse(); + verifyNoInteractions(mockTracer); + verify(mockChannel).shutdownNow(); + } + } + + @Test + public void testMissingSidebandData() { + // Interceptor failed to capture anything (returns null) + try (MockedConstruction mocked = + mockConstruction( + MetadataExtractorInterceptor.class, + (mock, context) -> { + when(mock.getSidebandData()).thenReturn(null); + })) { + + boolean isEligible = checker.check(mockChannelFactory, mockTracer); + + assertThat(isEligible).isFalse(); + verifyNoInteractions(mockTracer); + verify(mockChannel).shutdownNow(); + } + } + + @Test + public void testExceptionSafetyAndCleanup() { + doThrow(new RuntimeException("Simulated primer failure")) + .when(mockChannelPrimer) + .primeChannel(any(Channel.class)); + + boolean isEligible = checker.check(mockChannelFactory, mockTracer); + + assertThat(isEligible).isFalse(); + verifyNoInteractions(mockTracer); + verify(mockChannel).shutdownNow(); + } +} \ No newline at end of file diff --git a/google-cloud-bigtable/src/test/java/com/google/cloud/bigtable/gaxx/grpc/BigtableChannelPoolTest.java b/google-cloud-bigtable/src/test/java/com/google/cloud/bigtable/gaxx/grpc/BigtableChannelPoolTest.java index d1059c0362..1cd98a92c4 100644 --- a/google-cloud-bigtable/src/test/java/com/google/cloud/bigtable/gaxx/grpc/BigtableChannelPoolTest.java +++ b/google-cloud-bigtable/src/test/java/com/google/cloud/bigtable/gaxx/grpc/BigtableChannelPoolTest.java @@ -19,7 +19,7 @@ import static org.mockito.ArgumentMatchers.any; import static org.mockito.Mockito.*; -import com.google.api.gax.grpc.ChannelFactory; +import com.google.cloud.bigtable.data.v2.stub.BigtableChannelFactory; import com.google.common.collect.Iterables; import io.grpc.CallOptions; import io.grpc.ClientCall; @@ -48,7 +48,7 @@ public class BigtableChannelPoolTest { @Rule public final MockitoRule mockito = MockitoJUnit.rule(); - @Mock private ChannelFactory mockChannelFactory; + @Mock private BigtableChannelFactory mockChannelFactory; @Mock private ChannelPrimer mockChannelPrimer; @Mock private ManagedChannel mockChannel; @Mock private ClientCall mockClientCall; From 547f461ecafddec4ea6b18b37aafdfcfb9042105 Mon Sep 17 00:00:00 2001 From: Sushan Bhattarai Date: Thu, 12 Mar 2026 12:02:27 -0400 Subject: [PATCH 4/8] fix --- .../cloud/bigtable/data/v2/stub/BigtableChannelFactory.java | 2 +- .../google/cloud/bigtable/data/v2/stub/DirectAccessChecker.java | 2 +- .../cloud/bigtable/data/v2/stub/UnaryDirectAccessChecker.java | 2 +- .../bigtable/data/v2/stub/UnaryDirectAccessCheckerTest.java | 2 +- 4 files changed, 4 insertions(+), 4 deletions(-) diff --git a/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/stub/BigtableChannelFactory.java b/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/stub/BigtableChannelFactory.java index a4ce11b58a..80244b05ea 100644 --- a/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/stub/BigtableChannelFactory.java +++ b/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/stub/BigtableChannelFactory.java @@ -23,4 +23,4 @@ @InternalApi public interface BigtableChannelFactory { ManagedChannel createSingleChannel() throws IOException; -} \ No newline at end of file +} diff --git a/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/stub/DirectAccessChecker.java b/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/stub/DirectAccessChecker.java index 281ab7b48f..b919d380bd 100644 --- a/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/stub/DirectAccessChecker.java +++ b/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/stub/DirectAccessChecker.java @@ -31,4 +31,4 @@ public interface DirectAccessChecker { * @return true if the channel is eligible for Direct Access */ boolean check(BigtableChannelFactory channelFactory, @Nullable DirectPathCompatibleTracer tracer); -} \ No newline at end of file +} diff --git a/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/stub/UnaryDirectAccessChecker.java b/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/stub/UnaryDirectAccessChecker.java index edb842235f..0c8469ee58 100644 --- a/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/stub/UnaryDirectAccessChecker.java +++ b/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/stub/UnaryDirectAccessChecker.java @@ -81,4 +81,4 @@ public boolean check(BigtableChannelFactory channelFactory, @Nullable DirectPath } } } -} \ No newline at end of file +} diff --git a/google-cloud-bigtable/src/test/java/com/google/cloud/bigtable/data/v2/stub/UnaryDirectAccessCheckerTest.java b/google-cloud-bigtable/src/test/java/com/google/cloud/bigtable/data/v2/stub/UnaryDirectAccessCheckerTest.java index 5894565898..4684c2a72f 100644 --- a/google-cloud-bigtable/src/test/java/com/google/cloud/bigtable/data/v2/stub/UnaryDirectAccessCheckerTest.java +++ b/google-cloud-bigtable/src/test/java/com/google/cloud/bigtable/data/v2/stub/UnaryDirectAccessCheckerTest.java @@ -131,4 +131,4 @@ public void testExceptionSafetyAndCleanup() { verifyNoInteractions(mockTracer); verify(mockChannel).shutdownNow(); } -} \ No newline at end of file +} From 30c17d562ceaf8a6b2a5aedff2dc99b0939f42eb Mon Sep 17 00:00:00 2001 From: Sushan Bhattarai Date: Thu, 12 Mar 2026 12:13:39 -0400 Subject: [PATCH 5/8] fix --- .../data/v2/internal/csm/MetricsImpl.java | 3 +- .../DefaultDirectPathCompatibleTracer.java | 30 +-- .../tracers/DirectPathCompatibleTracer.java | 30 ++- .../data/v2/stub/BigtableChannelFactory.java | 3 +- .../data/v2/stub/BigtableChannelPrimer.java | 1 - .../data/v2/stub/BigtableClientContext.java | 4 +- .../data/v2/stub/DirectAccessChecker.java | 16 +- .../v2/stub/EnhancedBigtableStubSettings.java | 39 ++-- .../v2/stub/MetadataExtractorInterceptor.java | 1 - .../v2/stub/UnaryDirectAccessChecker.java | 90 ++++----- .../bigtable/gaxx/grpc/ChannelPrimer.java | 1 + .../v2/stub/UnaryDirectAccessCheckerTest.java | 179 +++++++++--------- 12 files changed, 198 insertions(+), 199 deletions(-) diff --git a/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/internal/csm/MetricsImpl.java b/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/internal/csm/MetricsImpl.java index 5d4c8d0e58..5c384a55fb 100644 --- a/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/internal/csm/MetricsImpl.java +++ b/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/internal/csm/MetricsImpl.java @@ -97,7 +97,8 @@ public MetricsImpl( this.internalRecorder = metricRegistry.newRecorderRegistry(internalOtel.getMeterProvider()); this.pacemaker = new Pacemaker(internalRecorder, clientInfo, "background"); this.channelPoolMetricsTracer = new ChannelPoolMetricsTracer(internalRecorder, clientInfo); - this.directPathCompatibleTracer = new DefaultDirectPathCompatibleTracer(clientInfo, internalRecorder); + this.directPathCompatibleTracer = + new DefaultDirectPathCompatibleTracer(clientInfo, internalRecorder); this.grpcOtel = GrpcOpenTelemetry.newBuilder() .sdk(internalOtel) diff --git a/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/internal/csm/tracers/DefaultDirectPathCompatibleTracer.java b/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/internal/csm/tracers/DefaultDirectPathCompatibleTracer.java index 041b226ab2..465e147868 100644 --- a/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/internal/csm/tracers/DefaultDirectPathCompatibleTracer.java +++ b/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/internal/csm/tracers/DefaultDirectPathCompatibleTracer.java @@ -21,22 +21,22 @@ @InternalApi public class DefaultDirectPathCompatibleTracer implements DirectPathCompatibleTracer { - private final ClientInfo clientInfo; - private final MetricRegistry.RecorderRegistry recorder; + private final ClientInfo clientInfo; + private final MetricRegistry.RecorderRegistry recorder; - public DefaultDirectPathCompatibleTracer( - ClientInfo clientInfo, MetricRegistry.RecorderRegistry recorder) { - this.clientInfo = clientInfo; - this.recorder = recorder; - } + public DefaultDirectPathCompatibleTracer( + ClientInfo clientInfo, MetricRegistry.RecorderRegistry recorder) { + this.clientInfo = clientInfo; + this.recorder = recorder; + } - @Override - public void recordSuccess(String ipPreference) { - recorder.dpCompatGuage.recordSuccess(clientInfo, ipPreference); - } + @Override + public void recordSuccess(String ipPreference) { + recorder.dpCompatGuage.recordSuccess(clientInfo, ipPreference); + } - @Override - public void recordFailure(String reason) { - recorder.dpCompatGuage.recordFailure(clientInfo, reason); - } + @Override + public void recordFailure(String reason) { + recorder.dpCompatGuage.recordFailure(clientInfo, reason); + } } diff --git a/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/internal/csm/tracers/DirectPathCompatibleTracer.java b/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/internal/csm/tracers/DirectPathCompatibleTracer.java index 0a385494cf..4776785dc2 100644 --- a/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/internal/csm/tracers/DirectPathCompatibleTracer.java +++ b/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/internal/csm/tracers/DirectPathCompatibleTracer.java @@ -17,23 +17,21 @@ import com.google.api.core.InternalApi; -/** - * Interface for recording DirectPath/DirectAccess eligibility metrics. - */ +/** Interface for recording DirectPath/DirectAccess eligibility metrics. */ @InternalApi public interface DirectPathCompatibleTracer { - /** - * Records that the environment is eligible and successfully connected via DirectPath. - * - * @param ipPreference The IP preference used (e.g., "ipv6"). - */ - void recordSuccess(String ipPreference); + /** + * Records that the environment is eligible and successfully connected via DirectPath. + * + * @param ipPreference The IP preference used (e.g., "ipv6"). + */ + void recordSuccess(String ipPreference); - /** - * Records that the environment is not eligible or failed to connect via DirectPath. - * - * @param reason The reason for the failure (e.g., "routing_check_failed"). - */ - void recordFailure(String reason); -} \ No newline at end of file + /** + * Records that the environment is not eligible or failed to connect via DirectPath. + * + * @param reason The reason for the failure (e.g., "routing_check_failed"). + */ + void recordFailure(String reason); +} diff --git a/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/stub/BigtableChannelFactory.java b/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/stub/BigtableChannelFactory.java index 80244b05ea..ce814e3073 100644 --- a/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/stub/BigtableChannelFactory.java +++ b/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/stub/BigtableChannelFactory.java @@ -17,10 +17,9 @@ import com.google.api.core.InternalApi; import io.grpc.ManagedChannel; - import java.io.IOException; @InternalApi public interface BigtableChannelFactory { - ManagedChannel createSingleChannel() throws IOException; + ManagedChannel createSingleChannel() throws IOException; } diff --git a/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/stub/BigtableChannelPrimer.java b/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/stub/BigtableChannelPrimer.java index 73c89e3778..bcec4c5403 100644 --- a/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/stub/BigtableChannelPrimer.java +++ b/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/stub/BigtableChannelPrimer.java @@ -32,7 +32,6 @@ import io.grpc.Metadata; import io.grpc.Status; import io.grpc.auth.MoreCallCredentials; - import java.io.IOException; import java.io.UnsupportedEncodingException; import java.net.URLEncoder; diff --git a/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/stub/BigtableClientContext.java b/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/stub/BigtableClientContext.java index 3bebe1844d..2680558995 100644 --- a/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/stub/BigtableClientContext.java +++ b/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/stub/BigtableClientContext.java @@ -166,8 +166,8 @@ public static BigtableClientContext create( transportProvider.build(), channelPrimer, metrics.getChannelPoolMetricsTracer(), - backgroundExecutor, metrics.getDirectPathCompatibleTracer() - ); + backgroundExecutor, + metrics.getDirectPathCompatibleTracer()); builder.setTransportChannelProvider(btTransportProvider); } diff --git a/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/stub/DirectAccessChecker.java b/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/stub/DirectAccessChecker.java index b919d380bd..9c5d66c559 100644 --- a/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/stub/DirectAccessChecker.java +++ b/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/stub/DirectAccessChecker.java @@ -17,18 +17,16 @@ import com.google.api.core.InternalApi; import com.google.cloud.bigtable.data.v2.internal.csm.tracers.DirectPathCompatibleTracer; -import io.grpc.Channel; - import javax.annotation.Nullable; @InternalApi /* Evaluates whether a given channel supports Direct Access. */ public interface DirectAccessChecker { - /** - * Evaluates if Direct Access is available by creating a test channel. - * - * @param channelFactory A factory to create the test channel - * @return true if the channel is eligible for Direct Access - */ - boolean check(BigtableChannelFactory channelFactory, @Nullable DirectPathCompatibleTracer tracer); + /** + * Evaluates if Direct Access is available by creating a test channel. + * + * @param channelFactory A factory to create the test channel + * @return true if the channel is eligible for Direct Access + */ + boolean check(BigtableChannelFactory channelFactory, @Nullable DirectPathCompatibleTracer tracer); } diff --git a/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/stub/EnhancedBigtableStubSettings.java b/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/stub/EnhancedBigtableStubSettings.java index 47937216da..88904a59c0 100644 --- a/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/stub/EnhancedBigtableStubSettings.java +++ b/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/stub/EnhancedBigtableStubSettings.java @@ -244,21 +244,21 @@ public ClientOperationSettings getPerOpSettings() { /** Applies common pool, message size, and keep-alive settings to the provided builder. */ private static InstantiatingGrpcChannelProvider.Builder commonTraits( - InstantiatingGrpcChannelProvider.Builder builder) { + InstantiatingGrpcChannelProvider.Builder builder) { return builder - .setChannelPoolSettings( - ChannelPoolSettings.builder() - .setInitialChannelCount(10) - .setMinRpcsPerChannel(1) - // Keep it conservative as we scale the channel size every 1min - // and delta is 2 channels. - .setMaxRpcsPerChannel(25) - .setPreemptiveRefreshEnabled(true) - .build()) - .setMaxInboundMessageSize(MAX_MESSAGE_SIZE) - .setKeepAliveTime(Duration.ofSeconds(30)) // sends ping in this interval - .setKeepAliveTimeout( - Duration.ofSeconds(10)); // wait this long before considering the connection dead + .setChannelPoolSettings( + ChannelPoolSettings.builder() + .setInitialChannelCount(10) + .setMinRpcsPerChannel(1) + // Keep it conservative as we scale the channel size every 1min + // and delta is 2 channels. + .setMaxRpcsPerChannel(25) + .setPreemptiveRefreshEnabled(true) + .build()) + .setMaxInboundMessageSize(MAX_MESSAGE_SIZE) + .setKeepAliveTime(Duration.ofSeconds(30)) // sends ping in this interval + .setKeepAliveTimeout( + Duration.ofSeconds(10)); // wait this long before considering the connection dead } /** Returns a builder for the default ChannelProvider for this service. */ @@ -285,22 +285,21 @@ public static InstantiatingGrpcChannelProvider.Builder defaultGrpcTransportProvi /** Applies Direct Access traits (DirectPath & ALTS) to an existing builder. */ public static InstantiatingGrpcChannelProvider.Builder applyDirectAccessTraits( - InstantiatingGrpcChannelProvider.Builder builder) { + InstantiatingGrpcChannelProvider.Builder builder) { builder - .setAttemptDirectPathXds() - .setAttemptDirectPath(true) - .setAllowNonDefaultServiceAccount(true); + .setAttemptDirectPathXds() + .setAttemptDirectPath(true) + .setAllowNonDefaultServiceAccount(true); if (!DIRECT_PATH_BOUND_TOKEN_DISABLED) { builder.setAllowHardBoundTokenTypes( - Collections.singletonList(InstantiatingGrpcChannelProvider.HardBoundTokenTypes.ALTS)); + Collections.singletonList(InstantiatingGrpcChannelProvider.HardBoundTokenTypes.ALTS)); } return builder; } - @SuppressWarnings("WeakerAccess") public static TransportChannelProvider defaultTransportChannelProvider() { return defaultGrpcTransportProviderBuilder().build(); diff --git a/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/stub/MetadataExtractorInterceptor.java b/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/stub/MetadataExtractorInterceptor.java index 145e4064ea..f2d3a1502f 100644 --- a/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/stub/MetadataExtractorInterceptor.java +++ b/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/stub/MetadataExtractorInterceptor.java @@ -34,7 +34,6 @@ import io.grpc.MethodDescriptor; import io.grpc.Status; import io.grpc.alts.AltsContextUtil; - import java.net.Inet4Address; import java.net.Inet6Address; import java.net.InetSocketAddress; diff --git a/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/stub/UnaryDirectAccessChecker.java b/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/stub/UnaryDirectAccessChecker.java index 0c8469ee58..9549fbe403 100644 --- a/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/stub/UnaryDirectAccessChecker.java +++ b/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/stub/UnaryDirectAccessChecker.java @@ -22,63 +22,65 @@ import io.grpc.Channel; import io.grpc.ClientInterceptors; import io.grpc.ManagedChannel; - -import javax.annotation.Nullable; import java.util.Optional; import java.util.logging.Level; import java.util.logging.Logger; +import javax.annotation.Nullable; /** - * Evaluates whether a given channel has Direct Access (DirectPath) routing - * by executing a RPC and inspecting the response headers. + * Evaluates whether a given channel has Direct Access (DirectPath) routing by executing a RPC and + * inspecting the response headers. */ @InternalApi public class UnaryDirectAccessChecker implements DirectAccessChecker { - private static final Logger LOG = Logger.getLogger(UnaryDirectAccessChecker.class.getName()); - private final ChannelPrimer channelPrimer; + private static final Logger LOG = Logger.getLogger(UnaryDirectAccessChecker.class.getName()); + private final ChannelPrimer channelPrimer; - private UnaryDirectAccessChecker(ChannelPrimer channelPrimer) { - this.channelPrimer = channelPrimer; - } + private UnaryDirectAccessChecker(ChannelPrimer channelPrimer) { + this.channelPrimer = channelPrimer; + } - public static UnaryDirectAccessChecker create(ChannelPrimer channelPrimer) { - return new UnaryDirectAccessChecker(channelPrimer); - } + public static UnaryDirectAccessChecker create(ChannelPrimer channelPrimer) { + return new UnaryDirectAccessChecker(channelPrimer); + } - @Override - public boolean check(BigtableChannelFactory channelFactory, @Nullable DirectPathCompatibleTracer tracer) { - ManagedChannel channel = null; - try { - channel = channelFactory.createSingleChannel(); - MetadataExtractorInterceptor interceptor = new MetadataExtractorInterceptor(); - Channel interceptedChannel = ClientInterceptors.intercept(channel, interceptor); - channelPrimer.primeChannel(interceptedChannel); + @Override + public boolean check( + BigtableChannelFactory channelFactory, @Nullable DirectPathCompatibleTracer tracer) { + ManagedChannel channel = null; + try { + channel = channelFactory.createSingleChannel(); + MetadataExtractorInterceptor interceptor = new MetadataExtractorInterceptor(); + Channel interceptedChannel = ClientInterceptors.intercept(channel, interceptor); + channelPrimer.primeChannel(interceptedChannel); - // Extract the sideband data populated by the interceptor - MetadataExtractorInterceptor.SidebandData sidebandData = interceptor.getSidebandData(); + // Extract the sideband data populated by the interceptor + MetadataExtractorInterceptor.SidebandData sidebandData = interceptor.getSidebandData(); - boolean isEligible = Optional.ofNullable(sidebandData) - .map(MetadataExtractorInterceptor.SidebandData::getPeerInfo) - .map(PeerInfo::getTransportType) - .map(type -> type == PeerInfo.TransportType.TRANSPORT_TYPE_DIRECT_ACCESS) - .orElse(false); + boolean isEligible = + Optional.ofNullable(sidebandData) + .map(MetadataExtractorInterceptor.SidebandData::getPeerInfo) + .map(PeerInfo::getTransportType) + .map(type -> type == PeerInfo.TransportType.TRANSPORT_TYPE_DIRECT_ACCESS) + .orElse(false); - if (isEligible && tracer != null) { - String ipProtocolStr = Optional.ofNullable(sidebandData) - .map(MetadataExtractorInterceptor.SidebandData::getIpProtocol) - .map(String::valueOf) - .map(String::toLowerCase) - .orElse("unknown"); - tracer.recordSuccess(ipProtocolStr); - } - return isEligible; - } catch (Exception e) { - LOG.log(Level.FINE, "Failed to evaluate direct access eligibility.", e); - return false; - } finally { - if (channel != null) { - channel.shutdownNow(); - } - } + if (isEligible && tracer != null) { + String ipProtocolStr = + Optional.ofNullable(sidebandData) + .map(MetadataExtractorInterceptor.SidebandData::getIpProtocol) + .map(String::valueOf) + .map(String::toLowerCase) + .orElse("unknown"); + tracer.recordSuccess(ipProtocolStr); + } + return isEligible; + } catch (Exception e) { + LOG.log(Level.FINE, "Failed to evaluate direct access eligibility.", e); + return false; + } finally { + if (channel != null) { + channel.shutdownNow(); + } } + } } diff --git a/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/gaxx/grpc/ChannelPrimer.java b/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/gaxx/grpc/ChannelPrimer.java index 6234b4410c..b9fb28bcb6 100644 --- a/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/gaxx/grpc/ChannelPrimer.java +++ b/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/gaxx/grpc/ChannelPrimer.java @@ -34,5 +34,6 @@ default void primeChannel(ManagedChannel channel) { default ApiFuture sendPrimeRequestsAsync(ManagedChannel channel) { return sendPrimeRequestsAsync((Channel) channel); } + ApiFuture sendPrimeRequestsAsync(Channel channel); } diff --git a/google-cloud-bigtable/src/test/java/com/google/cloud/bigtable/data/v2/stub/UnaryDirectAccessCheckerTest.java b/google-cloud-bigtable/src/test/java/com/google/cloud/bigtable/data/v2/stub/UnaryDirectAccessCheckerTest.java index 4684c2a72f..275c6313ce 100644 --- a/google-cloud-bigtable/src/test/java/com/google/cloud/bigtable/data/v2/stub/UnaryDirectAccessCheckerTest.java +++ b/google-cloud-bigtable/src/test/java/com/google/cloud/bigtable/data/v2/stub/UnaryDirectAccessCheckerTest.java @@ -37,98 +37,101 @@ @RunWith(JUnit4.class) public class UnaryDirectAccessCheckerTest { - @Rule public final MockitoRule mockito = MockitoJUnit.rule(); - - @Mock private ChannelPrimer mockChannelPrimer; - @Mock private BigtableChannelFactory mockChannelFactory; - @Mock private DirectPathCompatibleTracer mockTracer; - @Mock private ManagedChannel mockChannel; - @Mock private MetadataExtractorInterceptor.SidebandData mockSidebandData; - - private UnaryDirectAccessChecker checker; - - @Before - public void setUp() throws Exception { - checker = UnaryDirectAccessChecker.create(mockChannelPrimer); - when(mockChannelFactory.createSingleChannel()).thenReturn(mockChannel); - } - - @Test - public void testEligibleForDirectAccess() { - PeerInfo peerInfo = PeerInfo.newBuilder() - .setTransportType(PeerInfo.TransportType.TRANSPORT_TYPE_DIRECT_ACCESS) - .build(); - when(mockSidebandData.getPeerInfo()).thenReturn(peerInfo); - - when(mockSidebandData.getIpProtocol()).thenReturn(MetadataExtractorInterceptor.SidebandData.IpProtocol.IPV6); - - try (MockedConstruction ignored = - mockConstruction( - MetadataExtractorInterceptor.class, - (mock, context) -> { - when(mock.getSidebandData()).thenReturn(mockSidebandData); - })) { - - boolean isEligible = checker.check(mockChannelFactory, mockTracer); - - assertThat(isEligible).isFalse(); - verify(mockChannelPrimer).primeChannel(any(Channel.class)); - verify(mockTracer).recordSuccess("ipv6"); - verify(mockChannel).shutdownNow(); - } + @Rule public final MockitoRule mockito = MockitoJUnit.rule(); + + @Mock private ChannelPrimer mockChannelPrimer; + @Mock private BigtableChannelFactory mockChannelFactory; + @Mock private DirectPathCompatibleTracer mockTracer; + @Mock private ManagedChannel mockChannel; + @Mock private MetadataExtractorInterceptor.SidebandData mockSidebandData; + + private UnaryDirectAccessChecker checker; + + @Before + public void setUp() throws Exception { + checker = UnaryDirectAccessChecker.create(mockChannelPrimer); + when(mockChannelFactory.createSingleChannel()).thenReturn(mockChannel); + } + + @Test + public void testEligibleForDirectAccess() { + PeerInfo peerInfo = + PeerInfo.newBuilder() + .setTransportType(PeerInfo.TransportType.TRANSPORT_TYPE_DIRECT_ACCESS) + .build(); + when(mockSidebandData.getPeerInfo()).thenReturn(peerInfo); + + when(mockSidebandData.getIpProtocol()) + .thenReturn(MetadataExtractorInterceptor.SidebandData.IpProtocol.IPV6); + + try (MockedConstruction ignored = + mockConstruction( + MetadataExtractorInterceptor.class, + (mock, context) -> { + when(mock.getSidebandData()).thenReturn(mockSidebandData); + })) { + + boolean isEligible = checker.check(mockChannelFactory, mockTracer); + + assertThat(isEligible).isFalse(); + verify(mockChannelPrimer).primeChannel(any(Channel.class)); + verify(mockTracer).recordSuccess("ipv6"); + verify(mockChannel).shutdownNow(); } - - @Test - public void testNotEligibleProxiedRouting() { - // 1. Setup sideband data to simulate standard CloudPath routing - PeerInfo peerInfo = PeerInfo.newBuilder() - .setTransportType(PeerInfo.TransportType.TRANSPORT_TYPE_CLOUD_PATH) - .build(); - when(mockSidebandData.getPeerInfo()).thenReturn(peerInfo); - - try (MockedConstruction mocked = - mockConstruction( - MetadataExtractorInterceptor.class, - (mock, context) -> { - when(mock.getSidebandData()).thenReturn(mockSidebandData); - })) { - - boolean isEligible = checker.check(mockChannelFactory, mockTracer); - - assertThat(isEligible).isFalse(); - verifyNoInteractions(mockTracer); - verify(mockChannel).shutdownNow(); - } + } + + @Test + public void testNotEligibleProxiedRouting() { + // 1. Setup sideband data to simulate standard CloudPath routing + PeerInfo peerInfo = + PeerInfo.newBuilder() + .setTransportType(PeerInfo.TransportType.TRANSPORT_TYPE_CLOUD_PATH) + .build(); + when(mockSidebandData.getPeerInfo()).thenReturn(peerInfo); + + try (MockedConstruction mocked = + mockConstruction( + MetadataExtractorInterceptor.class, + (mock, context) -> { + when(mock.getSidebandData()).thenReturn(mockSidebandData); + })) { + + boolean isEligible = checker.check(mockChannelFactory, mockTracer); + + assertThat(isEligible).isFalse(); + verifyNoInteractions(mockTracer); + verify(mockChannel).shutdownNow(); } - - @Test - public void testMissingSidebandData() { - // Interceptor failed to capture anything (returns null) - try (MockedConstruction mocked = - mockConstruction( - MetadataExtractorInterceptor.class, - (mock, context) -> { - when(mock.getSidebandData()).thenReturn(null); - })) { - - boolean isEligible = checker.check(mockChannelFactory, mockTracer); - - assertThat(isEligible).isFalse(); - verifyNoInteractions(mockTracer); - verify(mockChannel).shutdownNow(); - } + } + + @Test + public void testMissingSidebandData() { + // Interceptor failed to capture anything (returns null) + try (MockedConstruction mocked = + mockConstruction( + MetadataExtractorInterceptor.class, + (mock, context) -> { + when(mock.getSidebandData()).thenReturn(null); + })) { + + boolean isEligible = checker.check(mockChannelFactory, mockTracer); + + assertThat(isEligible).isFalse(); + verifyNoInteractions(mockTracer); + verify(mockChannel).shutdownNow(); } + } - @Test - public void testExceptionSafetyAndCleanup() { - doThrow(new RuntimeException("Simulated primer failure")) - .when(mockChannelPrimer) - .primeChannel(any(Channel.class)); + @Test + public void testExceptionSafetyAndCleanup() { + doThrow(new RuntimeException("Simulated primer failure")) + .when(mockChannelPrimer) + .primeChannel(any(Channel.class)); - boolean isEligible = checker.check(mockChannelFactory, mockTracer); + boolean isEligible = checker.check(mockChannelFactory, mockTracer); - assertThat(isEligible).isFalse(); - verifyNoInteractions(mockTracer); - verify(mockChannel).shutdownNow(); - } + assertThat(isEligible).isFalse(); + verifyNoInteractions(mockTracer); + verify(mockChannel).shutdownNow(); + } } From f69837e5407ff318e8a41b9d45e9cfd8b0490181 Mon Sep 17 00:00:00 2001 From: Sushan Bhattarai Date: Sat, 21 Mar 2026 00:25:43 -0400 Subject: [PATCH 6/8] address review comments --- .../data/v2/internal/csm/Metrics.java | 1 - .../data/v2/internal/csm/MetricsImpl.java | 3 +- .../NoopDirectPathCompatibleTracer.java | 37 ++ .../dp/ClassicDirectAccessChecker.java} | 38 +- .../dp}/DirectAccessChecker.java | 8 +- ...tory.java => BigtableChannelSupplier.java} | 6 +- .../v2/stub/EnhancedBigtableStubSettings.java | 4 +- .../gaxx/grpc/BigtableChannelPool.java | 21 +- .../BigtableTransportChannelProvider.java | 459 +++++++++--------- .../v2/BigtableDataClientFactoryTest.java | 7 +- .../dp/ClassicDirectAccessCheckerTest.java} | 81 ++-- .../v2/stub/EnhancedBigtableStubTest.java | 3 +- .../gaxx/grpc/BigtableChannelPoolTest.java | 6 +- 13 files changed, 355 insertions(+), 319 deletions(-) create mode 100644 google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/internal/csm/tracers/NoopDirectPathCompatibleTracer.java rename google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/{stub/UnaryDirectAccessChecker.java => internal/dp/ClassicDirectAccessChecker.java} (67%) rename google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/{stub => internal/dp}/DirectAccessChecker.java (78%) rename google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/stub/{BigtableChannelFactory.java => BigtableChannelSupplier.java} (84%) rename google-cloud-bigtable/src/test/java/com/google/cloud/bigtable/data/v2/{stub/UnaryDirectAccessCheckerTest.java => internal/dp/ClassicDirectAccessCheckerTest.java} (58%) diff --git a/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/internal/csm/Metrics.java b/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/internal/csm/Metrics.java index 823458e0a4..6e30d3dd2b 100644 --- a/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/internal/csm/Metrics.java +++ b/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/internal/csm/Metrics.java @@ -32,7 +32,6 @@ public interface Metrics extends Closeable { @Nullable ChannelPoolMetricsTracer getChannelPoolMetricsTracer(); - @Nullable DirectPathCompatibleTracer getDirectPathCompatibleTracer(); void start(); diff --git a/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/internal/csm/MetricsImpl.java b/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/internal/csm/MetricsImpl.java index 5c384a55fb..69391b9796 100644 --- a/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/internal/csm/MetricsImpl.java +++ b/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/internal/csm/MetricsImpl.java @@ -33,6 +33,7 @@ import com.google.cloud.bigtable.data.v2.internal.csm.tracers.CompositeTracerFactory; import com.google.cloud.bigtable.data.v2.internal.csm.tracers.DefaultDirectPathCompatibleTracer; import com.google.cloud.bigtable.data.v2.internal.csm.tracers.DirectPathCompatibleTracer; +import com.google.cloud.bigtable.data.v2.internal.csm.tracers.NoopDirectPathCompatibleTracer; import com.google.cloud.bigtable.data.v2.internal.csm.tracers.Pacemaker; import com.google.common.base.Preconditions; import com.google.common.collect.ImmutableList; @@ -114,7 +115,7 @@ public MetricsImpl( this.grpcOtel = null; this.pacemaker = null; this.channelPoolMetricsTracer = null; - this.directPathCompatibleTracer = null; + this.directPathCompatibleTracer = NoopDirectPathCompatibleTracer.INSTANCE; } if (userOtel != null) { diff --git a/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/internal/csm/tracers/NoopDirectPathCompatibleTracer.java b/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/internal/csm/tracers/NoopDirectPathCompatibleTracer.java new file mode 100644 index 0000000000..bcfb8a48f7 --- /dev/null +++ b/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/internal/csm/tracers/NoopDirectPathCompatibleTracer.java @@ -0,0 +1,37 @@ +/* + * 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.csm.tracers; + +import com.google.api.core.InternalApi; + +@InternalApi +public class NoopDirectPathCompatibleTracer implements DirectPathCompatibleTracer { + + public static final NoopDirectPathCompatibleTracer INSTANCE = + new NoopDirectPathCompatibleTracer(); + + private NoopDirectPathCompatibleTracer() {} + + @Override + public void recordSuccess(String ipPreference) { + // No-op + } + + @Override + public void recordFailure(String reason) { + // No-op + } +} diff --git a/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/stub/UnaryDirectAccessChecker.java b/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/internal/dp/ClassicDirectAccessChecker.java similarity index 67% rename from google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/stub/UnaryDirectAccessChecker.java rename to google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/internal/dp/ClassicDirectAccessChecker.java index 9549fbe403..cfd14741cc 100644 --- a/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/stub/UnaryDirectAccessChecker.java +++ b/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/internal/dp/ClassicDirectAccessChecker.java @@ -13,16 +13,19 @@ * See the License for the specific language governing permissions and * limitations under the License. */ -package com.google.cloud.bigtable.data.v2.stub; +package com.google.cloud.bigtable.data.v2.internal.dp; import com.google.api.core.InternalApi; import com.google.bigtable.v2.PeerInfo; import com.google.cloud.bigtable.data.v2.internal.csm.tracers.DirectPathCompatibleTracer; +import com.google.cloud.bigtable.data.v2.stub.MetadataExtractorInterceptor; import com.google.cloud.bigtable.gaxx.grpc.ChannelPrimer; +import com.google.common.annotations.VisibleForTesting; import io.grpc.Channel; import io.grpc.ClientInterceptors; import io.grpc.ManagedChannel; import java.util.Optional; +import java.util.function.Supplier; import java.util.logging.Level; import java.util.logging.Logger; import javax.annotation.Nullable; @@ -32,25 +35,30 @@ * inspecting the response headers. */ @InternalApi -public class UnaryDirectAccessChecker implements DirectAccessChecker { - private static final Logger LOG = Logger.getLogger(UnaryDirectAccessChecker.class.getName()); +public class ClassicDirectAccessChecker implements DirectAccessChecker { + private static final Logger LOG = Logger.getLogger(ClassicDirectAccessChecker.class.getName()); private final ChannelPrimer channelPrimer; - private UnaryDirectAccessChecker(ChannelPrimer channelPrimer) { + private ClassicDirectAccessChecker(ChannelPrimer channelPrimer) { this.channelPrimer = channelPrimer; } - public static UnaryDirectAccessChecker create(ChannelPrimer channelPrimer) { - return new UnaryDirectAccessChecker(channelPrimer); + public static ClassicDirectAccessChecker create(ChannelPrimer channelPrimer) { + return new ClassicDirectAccessChecker(channelPrimer); + } + + @VisibleForTesting + MetadataExtractorInterceptor createInterceptor() { + return new MetadataExtractorInterceptor(); } @Override public boolean check( - BigtableChannelFactory channelFactory, @Nullable DirectPathCompatibleTracer tracer) { + Supplier channelSupplier, @Nullable DirectPathCompatibleTracer tracer) { ManagedChannel channel = null; try { - channel = channelFactory.createSingleChannel(); - MetadataExtractorInterceptor interceptor = new MetadataExtractorInterceptor(); + channel = channelSupplier.get(); + MetadataExtractorInterceptor interceptor = createInterceptor(); Channel interceptedChannel = ClientInterceptors.intercept(channel, interceptor); channelPrimer.primeChannel(interceptedChannel); @@ -64,16 +72,16 @@ public boolean check( .map(type -> type == PeerInfo.TransportType.TRANSPORT_TYPE_DIRECT_ACCESS) .orElse(false); - if (isEligible && tracer != null) { + if (isEligible) { String ipProtocolStr = - Optional.ofNullable(sidebandData) - .map(MetadataExtractorInterceptor.SidebandData::getIpProtocol) - .map(String::valueOf) - .map(String::toLowerCase) - .orElse("unknown"); + sidebandData.getIpProtocol() != null + ? sidebandData.getIpProtocol().toString().toLowerCase() + : "unknown"; tracer.recordSuccess(ipProtocolStr); } + return isEligible; + } catch (Exception e) { LOG.log(Level.FINE, "Failed to evaluate direct access eligibility.", e); return false; diff --git a/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/stub/DirectAccessChecker.java b/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/internal/dp/DirectAccessChecker.java similarity index 78% rename from google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/stub/DirectAccessChecker.java rename to google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/internal/dp/DirectAccessChecker.java index 9c5d66c559..beab21e8d4 100644 --- a/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/stub/DirectAccessChecker.java +++ b/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/internal/dp/DirectAccessChecker.java @@ -13,10 +13,12 @@ * See the License for the specific language governing permissions and * limitations under the License. */ -package com.google.cloud.bigtable.data.v2.stub; +package com.google.cloud.bigtable.data.v2.internal.dp; import com.google.api.core.InternalApi; import com.google.cloud.bigtable.data.v2.internal.csm.tracers.DirectPathCompatibleTracer; +import io.grpc.ManagedChannel; +import java.util.function.Supplier; import javax.annotation.Nullable; @InternalApi @@ -25,8 +27,8 @@ public interface DirectAccessChecker { /** * Evaluates if Direct Access is available by creating a test channel. * - * @param channelFactory A factory to create the test channel + * @param supplier A supplier to create maybe direct access channel * @return true if the channel is eligible for Direct Access */ - boolean check(BigtableChannelFactory channelFactory, @Nullable DirectPathCompatibleTracer tracer); + boolean check(Supplier supplier, @Nullable DirectPathCompatibleTracer tracer); } diff --git a/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/stub/BigtableChannelFactory.java b/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/stub/BigtableChannelSupplier.java similarity index 84% rename from google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/stub/BigtableChannelFactory.java rename to google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/stub/BigtableChannelSupplier.java index ce814e3073..24aa062e14 100644 --- a/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/stub/BigtableChannelFactory.java +++ b/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/stub/BigtableChannelSupplier.java @@ -17,9 +17,7 @@ import com.google.api.core.InternalApi; import io.grpc.ManagedChannel; -import java.io.IOException; +import java.util.function.Supplier; @InternalApi -public interface BigtableChannelFactory { - ManagedChannel createSingleChannel() throws IOException; -} +public interface BigtableChannelSupplier extends Supplier {} diff --git a/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/stub/EnhancedBigtableStubSettings.java b/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/stub/EnhancedBigtableStubSettings.java index 88904a59c0..76c3e49b51 100644 --- a/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/stub/EnhancedBigtableStubSettings.java +++ b/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/stub/EnhancedBigtableStubSettings.java @@ -641,8 +641,8 @@ private Builder() { FeatureFlags.newBuilder() .setReverseScans(true) .setLastScannedRowResponses(true) - .setDirectAccessRequested(true) - .setTrafficDirectorEnabled(true) + .setDirectAccessRequested(DIRECT_PATH_ENABLED) + .setTrafficDirectorEnabled(DIRECT_PATH_ENABLED) .setPeerInfo(true); } diff --git a/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/gaxx/grpc/BigtableChannelPool.java b/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/gaxx/grpc/BigtableChannelPool.java index 0298b730e7..4a127a5502 100644 --- a/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/gaxx/grpc/BigtableChannelPool.java +++ b/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/gaxx/grpc/BigtableChannelPool.java @@ -17,7 +17,6 @@ import com.google.api.core.InternalApi; import com.google.bigtable.v2.PeerInfo; -import com.google.cloud.bigtable.data.v2.stub.BigtableChannelFactory; import com.google.cloud.bigtable.data.v2.stub.MetadataExtractorInterceptor; import com.google.cloud.bigtable.gaxx.grpc.ChannelPoolHealthChecker.ProbeResult; import com.google.common.annotations.VisibleForTesting; @@ -67,7 +66,7 @@ public class BigtableChannelPool extends ManagedChannel implements BigtableChann private static final java.time.Duration REFRESH_PERIOD = java.time.Duration.ofMinutes(50); private final BigtableChannelPoolSettings settings; - private final BigtableChannelFactory channelFactory; + private final Supplier channelSupplier; private final ChannelPrimer channelPrimer; private final Object entryWriteLock = new Object(); @@ -81,11 +80,11 @@ public class BigtableChannelPool extends ManagedChannel implements BigtableChann public static BigtableChannelPool create( BigtableChannelPoolSettings settings, - BigtableChannelFactory channelFactory, + Supplier channelSupplier, ChannelPrimer channelPrimer, ScheduledExecutorService backgroundExecutor) throws IOException { - return new BigtableChannelPool(settings, channelFactory, channelPrimer, backgroundExecutor); + return new BigtableChannelPool(settings, channelSupplier, channelPrimer, backgroundExecutor); } /** @@ -98,12 +97,12 @@ public static BigtableChannelPool create( @VisibleForTesting BigtableChannelPool( BigtableChannelPoolSettings settings, - BigtableChannelFactory channelFactory, + Supplier channelSupplier, ChannelPrimer channelPrimer, ScheduledExecutorService executor) throws IOException { this.settings = settings; - this.channelFactory = channelFactory; + this.channelSupplier = channelSupplier; this.channelPrimer = channelPrimer; Clock systemClock = Clock.systemUTC(); ChannelPoolHealthChecker channelPoolHealthChecker = @@ -113,7 +112,7 @@ public static BigtableChannelPool create( ImmutableList.Builder initialListBuilder = ImmutableList.builder(); for (int i = 0; i < settings.getInitialChannelCount(); i++) { - ManagedChannel newChannel = channelFactory.createSingleChannel(); + ManagedChannel newChannel = channelSupplier.get(); channelPrimer.primeChannel(newChannel); initialListBuilder.add(new Entry(newChannel)); } @@ -419,10 +418,10 @@ private void expand(int desiredSize) { for (int i = 0; i < desiredSize - localEntries.size(); i++) { try { - ManagedChannel newChannel = channelFactory.createSingleChannel(); + ManagedChannel newChannel = channelSupplier.get(); this.channelPrimer.primeChannel(newChannel); newEntries.add(new Entry(newChannel)); - } catch (IOException e) { + } catch (Exception e) { LOG.log(Level.WARNING, "Failed to add channel", e); } } @@ -459,10 +458,10 @@ void refresh() { for (int i = 0; i < newEntries.size(); i++) { try { - ManagedChannel newChannel = channelFactory.createSingleChannel(); + ManagedChannel newChannel = channelSupplier.get(); this.channelPrimer.primeChannel(newChannel); newEntries.set(i, new Entry(newChannel)); - } catch (IOException e) { + } catch (Exception e) { LOG.log(Level.WARNING, "Failed to refresh channel, leaving old channel", e); } } diff --git a/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/gaxx/grpc/BigtableTransportChannelProvider.java b/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/gaxx/grpc/BigtableTransportChannelProvider.java index 2b3cb4022e..8f3342bba4 100644 --- a/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/gaxx/grpc/BigtableTransportChannelProvider.java +++ b/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/gaxx/grpc/BigtableTransportChannelProvider.java @@ -24,16 +24,17 @@ import com.google.auth.Credentials; import com.google.cloud.bigtable.data.v2.internal.csm.tracers.ChannelPoolMetricsTracer; import com.google.cloud.bigtable.data.v2.internal.csm.tracers.DirectPathCompatibleTracer; -import com.google.cloud.bigtable.data.v2.stub.BigtableChannelFactory; -import com.google.cloud.bigtable.data.v2.stub.DirectAccessChecker; +import com.google.cloud.bigtable.data.v2.internal.dp.ClassicDirectAccessChecker; +import com.google.cloud.bigtable.data.v2.internal.dp.DirectAccessChecker; +import com.google.cloud.bigtable.data.v2.stub.BigtableChannelSupplier; import com.google.cloud.bigtable.data.v2.stub.EnhancedBigtableStubSettings; -import com.google.cloud.bigtable.data.v2.stub.UnaryDirectAccessChecker; 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; +import java.util.function.Supplier; import java.util.logging.Level; import java.util.logging.Logger; import javax.annotation.Nullable; @@ -44,234 +45,238 @@ */ @InternalApi public final class BigtableTransportChannelProvider implements TransportChannelProvider { - private static final Logger LOG = - Logger.getLogger(BigtableTransportChannelProvider.class.getName()); - private final InstantiatingGrpcChannelProvider delegate; - private final ChannelPrimer channelPrimer; - @Nullable private final ChannelPoolMetricsTracer channelPoolMetricsTracer; - @Nullable private final ScheduledExecutorService backgroundExecutor; - DirectPathCompatibleTracer directPathCompatibleTracer; - - private BigtableTransportChannelProvider( - InstantiatingGrpcChannelProvider instantiatingGrpcChannelProvider, - ChannelPrimer channelPrimer, - @Nullable ChannelPoolMetricsTracer channelPoolMetricsTracer, - @Nullable ScheduledExecutorService backgroundExecutor, - DirectPathCompatibleTracer directPathCompatibleTracer) { - delegate = Preconditions.checkNotNull(instantiatingGrpcChannelProvider); - this.channelPrimer = channelPrimer; - this.channelPoolMetricsTracer = channelPoolMetricsTracer; - this.backgroundExecutor = backgroundExecutor; - this.directPathCompatibleTracer = directPathCompatibleTracer; + private static final Logger LOG = + Logger.getLogger(BigtableTransportChannelProvider.class.getName()); + private final InstantiatingGrpcChannelProvider delegate; + private final ChannelPrimer channelPrimer; + @Nullable private final ChannelPoolMetricsTracer channelPoolMetricsTracer; + @Nullable private final ScheduledExecutorService backgroundExecutor; + DirectPathCompatibleTracer directPathCompatibleTracer; + + private BigtableTransportChannelProvider( + InstantiatingGrpcChannelProvider instantiatingGrpcChannelProvider, + ChannelPrimer channelPrimer, + @Nullable ChannelPoolMetricsTracer channelPoolMetricsTracer, + @Nullable ScheduledExecutorService backgroundExecutor, + DirectPathCompatibleTracer directPathCompatibleTracer) { + delegate = Preconditions.checkNotNull(instantiatingGrpcChannelProvider); + this.channelPrimer = channelPrimer; + this.channelPoolMetricsTracer = channelPoolMetricsTracer; + this.backgroundExecutor = backgroundExecutor; + this.directPathCompatibleTracer = directPathCompatibleTracer; + } + + @Override + public boolean shouldAutoClose() { + return delegate.shouldAutoClose(); + } + + @SuppressWarnings("deprecation") + @Override + public boolean needsExecutor() { + return delegate.needsExecutor(); + } + + @Override + public BigtableTransportChannelProvider withExecutor(ScheduledExecutorService executor) { + return withExecutor((Executor) executor); + } + + // This executor if set is for handling rpc callbacks so we can't use it as the background + // executor + @Override + public BigtableTransportChannelProvider withExecutor(Executor executor) { + InstantiatingGrpcChannelProvider newChannelProvider = + (InstantiatingGrpcChannelProvider) delegate.withExecutor(executor); + return new BigtableTransportChannelProvider( + newChannelProvider, + channelPrimer, + channelPoolMetricsTracer, + backgroundExecutor, + directPathCompatibleTracer); + } + + @Override + public boolean needsBackgroundExecutor() { + return delegate.needsBackgroundExecutor(); + } + + @Override + public TransportChannelProvider withBackgroundExecutor(ScheduledExecutorService executor) { + InstantiatingGrpcChannelProvider newChannelProvider = + (InstantiatingGrpcChannelProvider) delegate.withBackgroundExecutor(executor); + return new BigtableTransportChannelProvider( + newChannelProvider, + channelPrimer, + channelPoolMetricsTracer, + executor, + directPathCompatibleTracer); + } + + @Override + public boolean needsHeaders() { + return delegate.needsHeaders(); + } + + @Override + public BigtableTransportChannelProvider withHeaders(Map headers) { + InstantiatingGrpcChannelProvider newChannelProvider = + (InstantiatingGrpcChannelProvider) delegate.withHeaders(headers); + return new BigtableTransportChannelProvider( + newChannelProvider, + channelPrimer, + channelPoolMetricsTracer, + backgroundExecutor, + directPathCompatibleTracer); + } + + @Override + public boolean needsEndpoint() { + return delegate.needsEndpoint(); + } + + @Override + public TransportChannelProvider withEndpoint(String endpoint) { + InstantiatingGrpcChannelProvider newChannelProvider = + (InstantiatingGrpcChannelProvider) delegate.withEndpoint(endpoint); + return new BigtableTransportChannelProvider( + newChannelProvider, + channelPrimer, + channelPoolMetricsTracer, + backgroundExecutor, + directPathCompatibleTracer); + } + + @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, + channelPrimer, + channelPoolMetricsTracer, + backgroundExecutor, + directPathCompatibleTracer); + } + + // We need this for direct access checker. + + /** Expected to only be called once when BigtableClientContext is created */ + @Override + public TransportChannel getTransportChannel() throws IOException { + InstantiatingGrpcChannelProvider directAccessProvider = + EnhancedBigtableStubSettings.applyDirectAccessTraits(delegate.toBuilder()) + .setChannelPoolSettings(ChannelPoolSettings.staticallySized(1)) + .build(); + + Supplier maybeDirectAccessChannelFactory = + () -> { + try { + GrpcTransportChannel channel = + (GrpcTransportChannel) directAccessProvider.getTransportChannel(); + return (ManagedChannel) channel.getChannel(); + } catch (IOException e) { + throw new java.io.UncheckedIOException(e); + } + }; + + DirectAccessChecker directAccessChecker = ClassicDirectAccessChecker.create(channelPrimer); + boolean isDirectAccessEligible = false; + + try { + isDirectAccessEligible = + directAccessChecker.check(maybeDirectAccessChannelFactory, directPathCompatibleTracer); + } catch (Exception e) { + LOG.log(Level.FINE, "Client is not direct access eligible, using standard transport.", e); } - @Override - public boolean shouldAutoClose() { - return delegate.shouldAutoClose(); - } - - @SuppressWarnings("deprecation") - @Override - public boolean needsExecutor() { - return delegate.needsExecutor(); - } - - @Override - public BigtableTransportChannelProvider withExecutor(ScheduledExecutorService executor) { - return withExecutor((Executor) executor); - } - - // This executor if set is for handling rpc callbacks so we can't use it as the background - // executor - @Override - public BigtableTransportChannelProvider withExecutor(Executor executor) { - InstantiatingGrpcChannelProvider newChannelProvider = - (InstantiatingGrpcChannelProvider) delegate.withExecutor(executor); - return new BigtableTransportChannelProvider( - newChannelProvider, - channelPrimer, - channelPoolMetricsTracer, - backgroundExecutor, - directPathCompatibleTracer); - } - - @Override - public boolean needsBackgroundExecutor() { - return delegate.needsBackgroundExecutor(); - } - - @Override - public TransportChannelProvider withBackgroundExecutor(ScheduledExecutorService executor) { - InstantiatingGrpcChannelProvider newChannelProvider = - (InstantiatingGrpcChannelProvider) delegate.withBackgroundExecutor(executor); - return new BigtableTransportChannelProvider( - newChannelProvider, - channelPrimer, - channelPoolMetricsTracer, - executor, - directPathCompatibleTracer); - } - - @Override - public boolean needsHeaders() { - return delegate.needsHeaders(); - } - - @Override - public BigtableTransportChannelProvider withHeaders(Map headers) { - InstantiatingGrpcChannelProvider newChannelProvider = - (InstantiatingGrpcChannelProvider) delegate.withHeaders(headers); - return new BigtableTransportChannelProvider( - newChannelProvider, - channelPrimer, - channelPoolMetricsTracer, - backgroundExecutor, - directPathCompatibleTracer); - } - - @Override - public boolean needsEndpoint() { - return delegate.needsEndpoint(); - } - - @Override - public TransportChannelProvider withEndpoint(String endpoint) { - InstantiatingGrpcChannelProvider newChannelProvider = - (InstantiatingGrpcChannelProvider) delegate.withEndpoint(endpoint); - return new BigtableTransportChannelProvider( - newChannelProvider, - channelPrimer, - channelPoolMetricsTracer, - backgroundExecutor, - directPathCompatibleTracer); - } - - @Deprecated - @Override - public boolean acceptsPoolSize() { - return delegate.acceptsPoolSize(); - } + InstantiatingGrpcChannelProvider selectedProvider; - @Deprecated - @Override - public TransportChannelProvider withPoolSize(int size) { - InstantiatingGrpcChannelProvider newChannelProvider = - (InstantiatingGrpcChannelProvider) delegate.withPoolSize(size); - return new BigtableTransportChannelProvider( - newChannelProvider, - channelPrimer, - channelPoolMetricsTracer, - backgroundExecutor, - directPathCompatibleTracer); + if (isDirectAccessEligible) { + selectedProvider = directAccessProvider; + } else { + selectedProvider = delegate; } - // We need this for direct access checker. - - /** Expected to only be called once when BigtableClientContext is created */ - @Override - public TransportChannel getTransportChannel() throws IOException { - InstantiatingGrpcChannelProvider directAccessProvider = - EnhancedBigtableStubSettings.applyDirectAccessTraits(delegate.toBuilder()) - .setChannelPoolSettings(ChannelPoolSettings.staticallySized(1)) - .build(); - - BigtableChannelFactory maybeDirectAccessChannelFactory = - () -> { - GrpcTransportChannel channel = - (GrpcTransportChannel) directAccessProvider.getTransportChannel(); - return (ManagedChannel) channel.getChannel(); - }; - - DirectAccessChecker directAccessChecker = UnaryDirectAccessChecker.create(channelPrimer); - boolean isDirectAccessEligible = false; - - try { - isDirectAccessEligible = - directAccessChecker.check(maybeDirectAccessChannelFactory, directPathCompatibleTracer); - } catch (Exception e) { - LOG.log(Level.FINE, "Client is not direct access eligible, using standard transport.", e); - } - - InstantiatingGrpcChannelProvider selectedProvider; - - if (isDirectAccessEligible) { - selectedProvider = directAccessProvider; - } else { - selectedProvider = delegate; - } - - // 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 = - selectedProvider.toBuilder() - .setChannelPoolSettings(ChannelPoolSettings.staticallySized(1)) - .build(); - - BigtableChannelFactory 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, channelPrimer, backgroundExecutor); - - if (channelPoolMetricsTracer != null) { - channelPoolMetricsTracer.registerChannelInsightsProvider(btChannelPool); - channelPoolMetricsTracer.registerLoadBalancingStrategy( - btPoolSettings.getLoadBalancingStrategy()); - } - - return GrpcTransportChannel.create(btChannelPool); + // 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 = + selectedProvider.toBuilder() + .setChannelPoolSettings(ChannelPoolSettings.staticallySized(1)) + .build(); + + BigtableChannelSupplier 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, channelPrimer, backgroundExecutor); + + if (channelPoolMetricsTracer != null) { + channelPoolMetricsTracer.registerChannelInsightsProvider(btChannelPool); + channelPoolMetricsTracer.registerLoadBalancingStrategy( + btPoolSettings.getLoadBalancingStrategy()); } - @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, - channelPrimer, - channelPoolMetricsTracer, - backgroundExecutor, - directPathCompatibleTracer); - } - - /** Creates a BigtableTransportChannelProvider. */ - public static BigtableTransportChannelProvider create( - InstantiatingGrpcChannelProvider instantiatingGrpcChannelProvider, - ChannelPrimer channelPrimer, - ChannelPoolMetricsTracer outstandingRpcsMetricTracker, - ScheduledExecutorService backgroundExecutor, - DirectPathCompatibleTracer directPathCompatibleTracer) { - return new BigtableTransportChannelProvider( - instantiatingGrpcChannelProvider, - channelPrimer, - outstandingRpcsMetricTracker, - backgroundExecutor, - directPathCompatibleTracer); - } -} \ No newline at end of file + 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, + channelPrimer, + channelPoolMetricsTracer, + backgroundExecutor, + directPathCompatibleTracer); + } + + /** Creates a BigtableTransportChannelProvider. */ + public static BigtableTransportChannelProvider create( + InstantiatingGrpcChannelProvider instantiatingGrpcChannelProvider, + ChannelPrimer channelPrimer, + ChannelPoolMetricsTracer outstandingRpcsMetricTracker, + ScheduledExecutorService backgroundExecutor, + DirectPathCompatibleTracer directPathCompatibleTracer) { + return new BigtableTransportChannelProvider( + instantiatingGrpcChannelProvider, + channelPrimer, + outstandingRpcsMetricTracker, + backgroundExecutor, + directPathCompatibleTracer); + } +} diff --git a/google-cloud-bigtable/src/test/java/com/google/cloud/bigtable/data/v2/BigtableDataClientFactoryTest.java b/google-cloud-bigtable/src/test/java/com/google/cloud/bigtable/data/v2/BigtableDataClientFactoryTest.java index f5c9472b47..cbe6d45457 100644 --- a/google-cloud-bigtable/src/test/java/com/google/cloud/bigtable/data/v2/BigtableDataClientFactoryTest.java +++ b/google-cloud-bigtable/src/test/java/com/google/cloud/bigtable/data/v2/BigtableDataClientFactoryTest.java @@ -285,8 +285,9 @@ public void testCreateWithRefreshingChannel() throws Exception { Mockito.verify(credentialsProvider, Mockito.times(2)).getCredentials(); Mockito.verify(executorProvider, Mockito.times(1)).getExecutor(); Mockito.verify(watchdogProvider, Mockito.times(1)).getWatchdog(); - - assertThat(warmedChannels).hasSize(poolSize); + // +1 accounts for the temporary channel created by UnaryDirectAccessChecker + int expectedChannelCount = poolSize + 1; + assertThat(warmedChannels).hasSize(expectedChannelCount); assertThat(warmedChannels.values()).doesNotContain(false); // Wait for all the connections to close asynchronously @@ -294,7 +295,7 @@ public void testCreateWithRefreshingChannel() throws Exception { long sleepTimeMs = 1000; Thread.sleep(sleepTimeMs); // Verify that all the channels are closed - assertThat(terminateAttributes).hasSize(poolSize); + assertThat(terminateAttributes).hasSize(expectedChannelCount); } @Test diff --git a/google-cloud-bigtable/src/test/java/com/google/cloud/bigtable/data/v2/stub/UnaryDirectAccessCheckerTest.java b/google-cloud-bigtable/src/test/java/com/google/cloud/bigtable/data/v2/internal/dp/ClassicDirectAccessCheckerTest.java similarity index 58% rename from google-cloud-bigtable/src/test/java/com/google/cloud/bigtable/data/v2/stub/UnaryDirectAccessCheckerTest.java rename to google-cloud-bigtable/src/test/java/com/google/cloud/bigtable/data/v2/internal/dp/ClassicDirectAccessCheckerTest.java index 275c6313ce..39fcfb6a00 100644 --- a/google-cloud-bigtable/src/test/java/com/google/cloud/bigtable/data/v2/stub/UnaryDirectAccessCheckerTest.java +++ b/google-cloud-bigtable/src/test/java/com/google/cloud/bigtable/data/v2/internal/dp/ClassicDirectAccessCheckerTest.java @@ -13,7 +13,7 @@ * See the License for the specific language governing permissions and * limitations under the License. */ -package com.google.cloud.bigtable.data.v2.stub; +package com.google.cloud.bigtable.data.v2.internal.dp; import static com.google.common.truth.Truth.assertThat; import static org.mockito.ArgumentMatchers.any; @@ -21,36 +21,45 @@ import com.google.bigtable.v2.PeerInfo; import com.google.cloud.bigtable.data.v2.internal.csm.tracers.DirectPathCompatibleTracer; +import com.google.cloud.bigtable.data.v2.stub.MetadataExtractorInterceptor; import com.google.cloud.bigtable.gaxx.grpc.ChannelPrimer; import io.grpc.Channel; import io.grpc.ManagedChannel; +import java.util.function.Supplier; import org.junit.Before; import org.junit.Rule; import org.junit.Test; import org.junit.runner.RunWith; import org.junit.runners.JUnit4; import org.mockito.Mock; -import org.mockito.MockedConstruction; import org.mockito.junit.MockitoJUnit; import org.mockito.junit.MockitoRule; @RunWith(JUnit4.class) -public class UnaryDirectAccessCheckerTest { +public class ClassicDirectAccessCheckerTest { @Rule public final MockitoRule mockito = MockitoJUnit.rule(); @Mock private ChannelPrimer mockChannelPrimer; - @Mock private BigtableChannelFactory mockChannelFactory; + @Mock private Supplier mockChannelFactory; @Mock private DirectPathCompatibleTracer mockTracer; @Mock private ManagedChannel mockChannel; + @Mock private MetadataExtractorInterceptor mockInterceptor; @Mock private MetadataExtractorInterceptor.SidebandData mockSidebandData; - private UnaryDirectAccessChecker checker; + private ClassicDirectAccessChecker checker; @Before public void setUp() throws Exception { - checker = UnaryDirectAccessChecker.create(mockChannelPrimer); - when(mockChannelFactory.createSingleChannel()).thenReturn(mockChannel); + // 1. Create a SPY of the checker so we can override just one method + checker = spy(ClassicDirectAccessChecker.create(mockChannelPrimer)); + + // 2. Tell the spy to return our mock interceptor instead of calling new() + doReturn(mockInterceptor).when(checker).createInterceptor(); + + // 3. Setup standard mocks + when(mockChannelFactory.get()).thenReturn(mockChannel); + when(mockInterceptor.getSidebandData()).thenReturn(mockSidebandData); } @Test @@ -60,66 +69,42 @@ public void testEligibleForDirectAccess() { .setTransportType(PeerInfo.TransportType.TRANSPORT_TYPE_DIRECT_ACCESS) .build(); when(mockSidebandData.getPeerInfo()).thenReturn(peerInfo); - when(mockSidebandData.getIpProtocol()) .thenReturn(MetadataExtractorInterceptor.SidebandData.IpProtocol.IPV6); - try (MockedConstruction ignored = - mockConstruction( - MetadataExtractorInterceptor.class, - (mock, context) -> { - when(mock.getSidebandData()).thenReturn(mockSidebandData); - })) { - - boolean isEligible = checker.check(mockChannelFactory, mockTracer); + boolean isEligible = checker.check(mockChannelFactory, mockTracer); - assertThat(isEligible).isFalse(); - verify(mockChannelPrimer).primeChannel(any(Channel.class)); - verify(mockTracer).recordSuccess("ipv6"); - verify(mockChannel).shutdownNow(); - } + assertThat(isEligible).isTrue(); + verify(mockChannelPrimer).primeChannel(any(Channel.class)); + verify(mockTracer).recordSuccess("ipv6"); + verify(mockChannel).shutdownNow(); } @Test public void testNotEligibleProxiedRouting() { - // 1. Setup sideband data to simulate standard CloudPath routing PeerInfo peerInfo = PeerInfo.newBuilder() .setTransportType(PeerInfo.TransportType.TRANSPORT_TYPE_CLOUD_PATH) .build(); when(mockSidebandData.getPeerInfo()).thenReturn(peerInfo); - try (MockedConstruction mocked = - mockConstruction( - MetadataExtractorInterceptor.class, - (mock, context) -> { - when(mock.getSidebandData()).thenReturn(mockSidebandData); - })) { - - boolean isEligible = checker.check(mockChannelFactory, mockTracer); + boolean isEligible = checker.check(mockChannelFactory, mockTracer); - assertThat(isEligible).isFalse(); - verifyNoInteractions(mockTracer); - verify(mockChannel).shutdownNow(); - } + assertThat(isEligible).isFalse(); + verifyNoInteractions(mockTracer); + verify(mockChannel).shutdownNow(); } @Test public void testMissingSidebandData() { - // Interceptor failed to capture anything (returns null) - try (MockedConstruction mocked = - mockConstruction( - MetadataExtractorInterceptor.class, - (mock, context) -> { - when(mock.getSidebandData()).thenReturn(null); - })) { - - boolean isEligible = checker.check(mockChannelFactory, mockTracer); - - assertThat(isEligible).isFalse(); - verifyNoInteractions(mockTracer); - verify(mockChannel).shutdownNow(); - } + // Override the Before setup to return null for this specific test + when(mockInterceptor.getSidebandData()).thenReturn(null); + + boolean isEligible = checker.check(mockChannelFactory, mockTracer); + + assertThat(isEligible).isFalse(); + verifyNoInteractions(mockTracer); + verify(mockChannel).shutdownNow(); } @Test diff --git a/google-cloud-bigtable/src/test/java/com/google/cloud/bigtable/data/v2/stub/EnhancedBigtableStubTest.java b/google-cloud-bigtable/src/test/java/com/google/cloud/bigtable/data/v2/stub/EnhancedBigtableStubTest.java index a9eda04356..1037aa1c15 100644 --- a/google-cloud-bigtable/src/test/java/com/google/cloud/bigtable/data/v2/stub/EnhancedBigtableStubTest.java +++ b/google-cloud-bigtable/src/test/java/com/google/cloud/bigtable/data/v2/stub/EnhancedBigtableStubTest.java @@ -552,7 +552,8 @@ public void testChannelPrimerConfigured() throws IOException { defaultSettings.toBuilder().setRefreshingChannel(true).build(); try (EnhancedBigtableStub ignored = EnhancedBigtableStub.create(settings)) { - assertThat(fakeDataService.pingRequests).hasSize(1); + // direct access checker ping + assertThat(fakeDataService.pingRequests).hasSize(2); } } diff --git a/google-cloud-bigtable/src/test/java/com/google/cloud/bigtable/gaxx/grpc/BigtableChannelPoolTest.java b/google-cloud-bigtable/src/test/java/com/google/cloud/bigtable/gaxx/grpc/BigtableChannelPoolTest.java index 1cd98a92c4..6719a70d89 100644 --- a/google-cloud-bigtable/src/test/java/com/google/cloud/bigtable/gaxx/grpc/BigtableChannelPoolTest.java +++ b/google-cloud-bigtable/src/test/java/com/google/cloud/bigtable/gaxx/grpc/BigtableChannelPoolTest.java @@ -19,7 +19,6 @@ import static org.mockito.ArgumentMatchers.any; import static org.mockito.Mockito.*; -import com.google.cloud.bigtable.data.v2.stub.BigtableChannelFactory; import com.google.common.collect.Iterables; import io.grpc.CallOptions; import io.grpc.ClientCall; @@ -32,6 +31,7 @@ import java.util.List; import java.util.concurrent.Executors; import java.util.concurrent.ScheduledExecutorService; +import java.util.function.Supplier; import org.junit.Before; import org.junit.Rule; import org.junit.Test; @@ -48,7 +48,7 @@ public class BigtableChannelPoolTest { @Rule public final MockitoRule mockito = MockitoJUnit.rule(); - @Mock private BigtableChannelFactory mockChannelFactory; + @Mock private Supplier mockChannelFactory; @Mock private ChannelPrimer mockChannelPrimer; @Mock private ManagedChannel mockChannel; @Mock private ClientCall mockClientCall; @@ -75,7 +75,7 @@ public String parse(InputStream stream) { @Before public void setUp() throws IOException { - when(mockChannelFactory.createSingleChannel()).thenReturn(mockChannel); + when(mockChannelFactory.get()).thenReturn(mockChannel); when(mockChannel.newCall( ArgumentMatchers.>any(), any(CallOptions.class))) .thenReturn(mockClientCall); From 1ca692cfc49525b261d396c79b8c4e0c7da762ab Mon Sep 17 00:00:00 2001 From: Sushan Bhattarai Date: Sat, 21 Mar 2026 00:27:24 -0400 Subject: [PATCH 7/8] address review comments --- .../bigtable/gaxx/grpc/BigtableTransportChannelProvider.java | 1 + 1 file changed, 1 insertion(+) diff --git a/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/gaxx/grpc/BigtableTransportChannelProvider.java b/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/gaxx/grpc/BigtableTransportChannelProvider.java index 8f3342bba4..0b7c44cff4 100644 --- a/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/gaxx/grpc/BigtableTransportChannelProvider.java +++ b/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/gaxx/grpc/BigtableTransportChannelProvider.java @@ -78,6 +78,7 @@ public boolean needsExecutor() { } @Override + @Deprecated public BigtableTransportChannelProvider withExecutor(ScheduledExecutorService executor) { return withExecutor((Executor) executor); } From cc306bc27fba5dfa40a986d97a311bb5bc02d8dc Mon Sep 17 00:00:00 2001 From: Sushan Bhattarai Date: Sat, 21 Mar 2026 00:48:47 -0400 Subject: [PATCH 8/8] address review comments --- .../bigtable/data/v2/stub/EnhancedBigtableStubSettings.java | 2 ++ 1 file changed, 2 insertions(+) diff --git a/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/stub/EnhancedBigtableStubSettings.java b/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/stub/EnhancedBigtableStubSettings.java index 76c3e49b51..b7189ecc04 100644 --- a/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/stub/EnhancedBigtableStubSettings.java +++ b/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/stub/EnhancedBigtableStubSettings.java @@ -637,6 +637,8 @@ private Builder() { // only if Traffic Director sends the request (with grpc as target type) // For GFE/CFE, sending setDirectAccessRequested // is fine as GFE/CFE sends with gslb target type + // TODO: flip the bit setDirectAccessRequested and setTrafficDirectorEnabled once we make + // client compatible by default. featureFlags = FeatureFlags.newBuilder() .setReverseScans(true)