Skip to content
This repository was archived by the owner on May 8, 2026. It is now read-only.

Commit cce0540

Browse files
sushanbsushanb
committed
init commit
1 parent 2933ef6 commit cce0540

7 files changed

Lines changed: 228 additions & 39 deletions

File tree

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

Lines changed: 8 additions & 7 deletions
Original file line numberDiff line numberDiff line change
@@ -26,6 +26,7 @@
2626
import com.google.cloud.bigtable.gaxx.grpc.ChannelPrimer;
2727
import io.grpc.CallCredentials;
2828
import io.grpc.CallOptions;
29+
import io.grpc.Channel;
2930
import io.grpc.ClientCall;
3031
import io.grpc.Deadline;
3132
import io.grpc.ManagedChannel;
@@ -88,21 +89,21 @@ static BigtableChannelPrimer create(
8889
}
8990

9091
@Override
91-
public void primeChannel(ManagedChannel managedChannel) {
92+
public void primeChannel(Channel channel) {
9293
try {
93-
primeChannelUnsafe(managedChannel);
94+
primeChannelUnsafe(channel);
9495
} catch (IOException | RuntimeException e) {
9596
LOG.log(Level.WARNING, "Unexpected error while trying to prime a channel", e);
9697
}
9798
}
9899

99-
private void primeChannelUnsafe(ManagedChannel managedChannel) throws IOException {
100-
sendPrimeRequestsBlocking(managedChannel);
100+
private void primeChannelUnsafe(Channel channel) throws IOException {
101+
sendPrimeRequestsBlocking(channel);
101102
}
102103

103-
private void sendPrimeRequestsBlocking(ManagedChannel managedChannel) {
104+
private void sendPrimeRequestsBlocking(Channel channel) {
104105
try {
105-
sendPrimeRequestsAsync(managedChannel).get(1, TimeUnit.MINUTES);
106+
sendPrimeRequestsAsync(channel).get(1, TimeUnit.MINUTES);
106107
} catch (Throwable e) {
107108
// TODO: Not sure if we should swallow the error here. We are pre-emptively swapping
108109
// channels if the new
@@ -111,7 +112,7 @@ private void sendPrimeRequestsBlocking(ManagedChannel managedChannel) {
111112
}
112113
}
113114

114-
public ApiFuture<PingAndWarmResponse> sendPrimeRequestsAsync(ManagedChannel managedChannel) {
115+
public ApiFuture<PingAndWarmResponse> sendPrimeRequestsAsync(Channel managedChannel) {
115116
ClientCall<PingAndWarmRequest, PingAndWarmResponse> clientCall =
116117
managedChannel.newCall(
117118
BigtableGrpc.getPingAndWarmMethod(),

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

Lines changed: 38 additions & 14 deletions
Original file line numberDiff line numberDiff line change
@@ -242,6 +242,25 @@ public ClientOperationSettings getPerOpSettings() {
242242
return perOpSettings;
243243
}
244244

245+
/** Applies common pool, message size, and keep-alive settings to the provided builder. */
246+
private static InstantiatingGrpcChannelProvider.Builder commonTraits(
247+
InstantiatingGrpcChannelProvider.Builder builder) {
248+
return builder
249+
.setChannelPoolSettings(
250+
ChannelPoolSettings.builder()
251+
.setInitialChannelCount(10)
252+
.setMinRpcsPerChannel(1)
253+
// Keep it conservative as we scale the channel size every 1min
254+
// and delta is 2 channels.
255+
.setMaxRpcsPerChannel(25)
256+
.setPreemptiveRefreshEnabled(true)
257+
.build())
258+
.setMaxInboundMessageSize(MAX_MESSAGE_SIZE)
259+
.setKeepAliveTime(Duration.ofSeconds(30)) // sends ping in this interval
260+
.setKeepAliveTimeout(
261+
Duration.ofSeconds(10)); // wait this long before considering the connection dead
262+
}
263+
245264
/** Returns a builder for the default ChannelProvider for this service. */
246265
public static InstantiatingGrpcChannelProvider.Builder defaultGrpcTransportProviderBuilder() {
247266
InstantiatingGrpcChannelProvider.Builder grpcTransportProviderBuilder =
@@ -261,22 +280,27 @@ public static InstantiatingGrpcChannelProvider.Builder defaultGrpcTransportProvi
261280
Collections.singletonList(InstantiatingGrpcChannelProvider.HardBoundTokenTypes.ALTS));
262281
}
263282
}
264-
return grpcTransportProviderBuilder
265-
.setChannelPoolSettings(
266-
ChannelPoolSettings.builder()
267-
.setInitialChannelCount(10)
268-
.setMinRpcsPerChannel(1)
269-
// Keep it conservative as we scale the channel size every 1min
270-
// and delta is 2 channels.
271-
.setMaxRpcsPerChannel(25)
272-
.setPreemptiveRefreshEnabled(true)
273-
.build())
274-
.setMaxInboundMessageSize(MAX_MESSAGE_SIZE)
275-
.setKeepAliveTime(Duration.ofSeconds(30)) // sends ping in this interval
276-
.setKeepAliveTimeout(
277-
Duration.ofSeconds(10)); // wait this long before considering the connection dead
283+
return commonTraits(grpcTransportProviderBuilder);
278284
}
279285

286+
/** Applies Direct Access traits (DirectPath & ALTS) to an existing builder. */
287+
public static InstantiatingGrpcChannelProvider.Builder applyDirectAccessTraits(
288+
InstantiatingGrpcChannelProvider.Builder builder) {
289+
290+
builder
291+
.setAttemptDirectPathXds()
292+
.setAttemptDirectPath(true)
293+
.setAllowNonDefaultServiceAccount(true);
294+
295+
if (!DIRECT_PATH_BOUND_TOKEN_DISABLED) {
296+
builder.setAllowHardBoundTokenTypes(
297+
Collections.singletonList(InstantiatingGrpcChannelProvider.HardBoundTokenTypes.ALTS));
298+
}
299+
300+
return builder;
301+
}
302+
303+
280304
@SuppressWarnings("WeakerAccess")
281305
public static TransportChannelProvider defaultTransportChannelProvider() {
282306
return defaultGrpcTransportProviderBuilder().build();

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

Lines changed: 5 additions & 7 deletions
Original file line numberDiff line numberDiff line change
@@ -16,11 +16,11 @@
1616
package com.google.cloud.bigtable.data.v2.stub;
1717

1818
import com.google.api.core.ApiFuture;
19+
import com.google.api.core.ApiFutures;
1920
import com.google.api.core.InternalApi;
20-
import com.google.api.core.SettableApiFuture;
2121
import com.google.bigtable.v2.PingAndWarmResponse;
2222
import com.google.cloud.bigtable.gaxx.grpc.ChannelPrimer;
23-
import io.grpc.ManagedChannel;
23+
import io.grpc.Channel;
2424

2525
@InternalApi
2626
public class NoOpChannelPrimer implements ChannelPrimer {
@@ -31,14 +31,12 @@ static NoOpChannelPrimer create() {
3131
private NoOpChannelPrimer() {}
3232

3333
@Override
34-
public void primeChannel(ManagedChannel channel) {
34+
public void primeChannel(Channel channel) {
3535
// No op
3636
}
3737

3838
@Override
39-
public ApiFuture<PingAndWarmResponse> sendPrimeRequestsAsync(ManagedChannel channel) {
40-
SettableApiFuture future = SettableApiFuture.create();
41-
future.set(PingAndWarmResponse.getDefaultInstance());
42-
return future;
39+
public ApiFuture<PingAndWarmResponse> sendPrimeRequestsAsync(Channel channel) {
40+
return ApiFutures.immediateFuture(PingAndWarmResponse.getDefaultInstance());
4341
}
4442
}

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

Lines changed: 81 additions & 9 deletions
Original file line numberDiff line numberDiff line change
@@ -24,12 +24,16 @@
2424
import com.google.api.gax.rpc.TransportChannelProvider;
2525
import com.google.auth.Credentials;
2626
import com.google.cloud.bigtable.data.v2.internal.csm.tracers.ChannelPoolMetricsTracer;
27+
import com.google.cloud.bigtable.data.v2.stub.EnhancedBigtableStubSettings;
2728
import com.google.common.base.Preconditions;
29+
import io.grpc.Channel;
2830
import io.grpc.ManagedChannel;
2931
import java.io.IOException;
3032
import java.util.Map;
3133
import java.util.concurrent.Executor;
3234
import java.util.concurrent.ScheduledExecutorService;
35+
import java.util.logging.Level;
36+
import java.util.logging.Logger;
3337
import javax.annotation.Nullable;
3438

3539
/**
@@ -38,20 +42,25 @@
3842
*/
3943
@InternalApi
4044
public final class BigtableTransportChannelProvider implements TransportChannelProvider {
45+
private static final Logger LOG =
46+
Logger.getLogger(BigtableTransportChannelProvider.class.getName());
4147
private final InstantiatingGrpcChannelProvider delegate;
4248
private final ChannelPrimer channelPrimer;
4349
@Nullable private final ChannelPoolMetricsTracer channelPoolMetricsTracer;
4450
@Nullable private final ScheduledExecutorService backgroundExecutor;
51+
@Nullable private final Map<String, String> headers;
4552

4653
private BigtableTransportChannelProvider(
4754
InstantiatingGrpcChannelProvider instantiatingGrpcChannelProvider,
4855
ChannelPrimer channelPrimer,
4956
ChannelPoolMetricsTracer channelPoolMetricsTracer,
50-
ScheduledExecutorService backgroundExecutor) {
57+
ScheduledExecutorService backgroundExecutor,
58+
@Nullable Map<String, String> headers) {
5159
delegate = Preconditions.checkNotNull(instantiatingGrpcChannelProvider);
5260
this.channelPrimer = channelPrimer;
5361
this.channelPoolMetricsTracer = channelPoolMetricsTracer;
5462
this.backgroundExecutor = backgroundExecutor;
63+
this.headers = headers;
5564
}
5665

5766
@Override
@@ -76,7 +85,7 @@ public BigtableTransportChannelProvider withExecutor(Executor executor) {
7685
InstantiatingGrpcChannelProvider newChannelProvider =
7786
(InstantiatingGrpcChannelProvider) delegate.withExecutor(executor);
7887
return new BigtableTransportChannelProvider(
79-
newChannelProvider, channelPrimer, channelPoolMetricsTracer, backgroundExecutor);
88+
newChannelProvider, channelPrimer, channelPoolMetricsTracer, backgroundExecutor, headers);
8089
}
8190

8291
@Override
@@ -89,7 +98,7 @@ public TransportChannelProvider withBackgroundExecutor(ScheduledExecutorService
8998
InstantiatingGrpcChannelProvider newChannelProvider =
9099
(InstantiatingGrpcChannelProvider) delegate.withBackgroundExecutor(executor);
91100
return new BigtableTransportChannelProvider(
92-
newChannelProvider, channelPrimer, channelPoolMetricsTracer, executor);
101+
newChannelProvider, channelPrimer, channelPoolMetricsTracer, executor, headers);
93102
}
94103

95104
@Override
@@ -102,7 +111,7 @@ public BigtableTransportChannelProvider withHeaders(Map<String, String> headers)
102111
InstantiatingGrpcChannelProvider newChannelProvider =
103112
(InstantiatingGrpcChannelProvider) delegate.withHeaders(headers);
104113
return new BigtableTransportChannelProvider(
105-
newChannelProvider, channelPrimer, channelPoolMetricsTracer, backgroundExecutor);
114+
newChannelProvider, channelPrimer, channelPoolMetricsTracer, backgroundExecutor, headers);
106115
}
107116

108117
@Override
@@ -115,7 +124,7 @@ public TransportChannelProvider withEndpoint(String endpoint) {
115124
InstantiatingGrpcChannelProvider newChannelProvider =
116125
(InstantiatingGrpcChannelProvider) delegate.withEndpoint(endpoint);
117126
return new BigtableTransportChannelProvider(
118-
newChannelProvider, channelPrimer, channelPoolMetricsTracer, backgroundExecutor);
127+
newChannelProvider, channelPrimer, channelPoolMetricsTracer, backgroundExecutor, headers);
119128
}
120129

121130
@Deprecated
@@ -130,20 +139,83 @@ public TransportChannelProvider withPoolSize(int size) {
130139
InstantiatingGrpcChannelProvider newChannelProvider =
131140
(InstantiatingGrpcChannelProvider) delegate.withPoolSize(size);
132141
return new BigtableTransportChannelProvider(
133-
newChannelProvider, channelPrimer, channelPoolMetricsTracer, backgroundExecutor);
142+
newChannelProvider, channelPrimer, channelPoolMetricsTracer, backgroundExecutor, headers);
143+
}
144+
145+
// We need this for direct access checker.
146+
private Map<String, String> updateFeatureFlags(
147+
Map<String, String> originalHeaders, boolean isDirectAccessEligible) {
148+
if (originalHeaders == null) {
149+
return java.util.Collections.emptyMap();
150+
}
151+
java.util.Map<String, String> newHeaders = new java.util.HashMap<>(originalHeaders);
152+
String encodedFlags = newHeaders.get("bigtable-features");
153+
154+
if (encodedFlags != null) {
155+
try {
156+
byte[] decoded = java.util.Base64.getUrlDecoder().decode(encodedFlags);
157+
com.google.bigtable.v2.FeatureFlags flags =
158+
com.google.bigtable.v2.FeatureFlags.parseFrom(decoded);
159+
160+
com.google.bigtable.v2.FeatureFlags updatedFlags =
161+
flags.toBuilder()
162+
.setDirectAccessRequested(isDirectAccessEligible)
163+
.setTrafficDirectorEnabled(isDirectAccessEligible)
164+
.build();
165+
166+
newHeaders.put(
167+
"bigtable-features",
168+
java.util.Base64.getUrlEncoder().encodeToString(updatedFlags.toByteArray()));
169+
} catch (Exception e) {
170+
// use original headers
171+
}
172+
}
173+
return newHeaders;
134174
}
135175

136176
/** Expected to only be called once when BigtableClientContext is created */
137177
@Override
138178
public TransportChannel getTransportChannel() throws IOException {
179+
Map<String, String> directAccessEligibleHeaders = updateFeatureFlags(this.headers, true);
180+
181+
InstantiatingGrpcChannelProvider.Builder directAccessProvider =
182+
EnhancedBigtableStubSettings.applyDirectAccessTraits(delegate.toBuilder())
183+
.setChannelPoolSettings(ChannelPoolSettings.staticallySized(1));
184+
185+
InstantiatingGrpcChannelProvider directAccessProviderWithHeaders =
186+
(InstantiatingGrpcChannelProvider)
187+
directAccessProvider.build().withHeaders(directAccessEligibleHeaders);
188+
GrpcTransportChannel directAccessTransportChannel =
189+
(GrpcTransportChannel) directAccessProviderWithHeaders.getTransportChannel();
190+
Channel maybeDirectAccessChannel = directAccessTransportChannel.getChannel();
191+
DirectAccessChecker directAccessChecker = UnaryDirectAccessChecker.create(channelPrimer);
192+
boolean isDirectAccessEligible = false;
193+
194+
try {
195+
isDirectAccessEligible = directAccessChecker.check(maybeDirectAccessChannel);
196+
} catch (Exception e) {
197+
LOG.log(Level.INFO, "Client is not direct access eligible, using standard transport.", e);
198+
}
199+
200+
InstantiatingGrpcChannelProvider selectedProvider;
201+
202+
if (isDirectAccessEligible) {
203+
selectedProvider = directAccessProviderWithHeaders;
204+
} else {
205+
Map<String, String> fallbackHeaders = updateFeatureFlags(this.headers, false);
206+
selectedProvider = (InstantiatingGrpcChannelProvider) delegate.withHeaders(fallbackHeaders);
207+
}
208+
139209
// This provider's main purpose is to replace the default GAX ChannelPool
140210
// with a custom BigtableChannelPool, reusing the delegate's configuration.
141211

142212
// To create our pool, we need a factory for raw gRPC channels.
143213
// We achieve this by configuring our delegate to not use its own pooling
144214
// (by setting pool size to 1) and then calling getTransportChannel() on it.
145215
InstantiatingGrpcChannelProvider singleChannelProvider =
146-
delegate.toBuilder().setChannelPoolSettings(ChannelPoolSettings.staticallySized(1)).build();
216+
selectedProvider.toBuilder()
217+
.setChannelPoolSettings(ChannelPoolSettings.staticallySized(1))
218+
.build();
147219

148220
ChannelFactory channelFactory =
149221
() -> {
@@ -187,7 +259,7 @@ public TransportChannelProvider withCredentials(Credentials credentials) {
187259
InstantiatingGrpcChannelProvider newChannelProvider =
188260
(InstantiatingGrpcChannelProvider) delegate.withCredentials(credentials);
189261
return new BigtableTransportChannelProvider(
190-
newChannelProvider, channelPrimer, channelPoolMetricsTracer, backgroundExecutor);
262+
newChannelProvider, channelPrimer, channelPoolMetricsTracer, backgroundExecutor, );
191263
}
192264

193265
/** Creates a BigtableTransportChannelProvider. */
@@ -200,6 +272,6 @@ public static BigtableTransportChannelProvider create(
200272
instantiatingGrpcChannelProvider,
201273
channelPrimer,
202274
outstandingRpcsMetricTracker,
203-
backgroundExecutor);
275+
backgroundExecutor, null);
204276
}
205277
}

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

Lines changed: 12 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -18,11 +18,21 @@
1818
import com.google.api.core.ApiFuture;
1919
import com.google.api.core.InternalApi;
2020
import com.google.bigtable.v2.PingAndWarmResponse;
21+
import io.grpc.Channel;
2122
import io.grpc.ManagedChannel;
2223

2324
@InternalApi("For internal use by google-cloud-java clients only")
2425
public interface ChannelPrimer {
25-
void primeChannel(ManagedChannel channel);
26+
@Deprecated
27+
default void primeChannel(ManagedChannel channel) {
28+
primeChannel((Channel) channel);
29+
}
2630

27-
ApiFuture<PingAndWarmResponse> sendPrimeRequestsAsync(ManagedChannel channel);
31+
void primeChannel(Channel channel);
32+
33+
@Deprecated
34+
default ApiFuture<PingAndWarmResponse> sendPrimeRequestsAsync(ManagedChannel channel) {
35+
return sendPrimeRequestsAsync((Channel) channel);
36+
}
37+
ApiFuture<PingAndWarmResponse> sendPrimeRequestsAsync(Channel channel);
2838
}

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

Lines changed: 26 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,26 @@
1+
/*
2+
* Copyright 2026 Google LLC
3+
*
4+
* Licensed under the Apache License, Version 2.0 (the "License");
5+
* you may not use this file except in compliance with the License.
6+
* You may obtain a copy of the License at
7+
*
8+
* https://www.apache.org/licenses/LICENSE-2.0
9+
*
10+
* Unless required by applicable law or agreed to in writing, software
11+
* distributed under the License is distributed on an "AS IS" BASIS,
12+
* WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
13+
* See the License for the specific language governing permissions and
14+
* limitations under the License.
15+
*/
16+
package com.google.cloud.bigtable.gaxx.grpc;
17+
18+
import com.google.api.core.InternalApi;
19+
import io.grpc.Channel;
20+
21+
@InternalApi
22+
/* Evaluates whether a given channel supports Direct Access. */
23+
public interface DirectAccessChecker {
24+
/// Performs a request on the provided channel to check for Direct Access eligibility.
25+
boolean check(Channel channel );
26+
}

0 commit comments

Comments
 (0)