Skip to content

Commit 5b6a974

Browse files
nimfnimf
authored
feat(bigtable): Option to disable direct access fallback (#14193)
1 parent bbbd18c commit 5b6a974

2 files changed

Lines changed: 85 additions & 2 deletions

File tree

‎java-bigtable/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/internal/util/ClientConfigurationManager.java‎

Lines changed: 25 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -21,6 +21,7 @@
2121
import com.google.bigtable.v2.ClientConfiguration.PollingCase;
2222
import com.google.bigtable.v2.FeatureFlags;
2323
import com.google.bigtable.v2.GetClientConfigurationRequest;
24+
import com.google.bigtable.v2.SessionClientConfiguration.ChannelPoolConfiguration.DirectAccessOnly;
2425
import com.google.bigtable.v2.TelemetryConfiguration;
2526
import com.google.cloud.bigtable.data.v2.internal.api.ChannelProviders.ChannelProvider;
2627
import com.google.cloud.bigtable.data.v2.internal.api.Util;
@@ -73,6 +74,10 @@ public class ClientConfigurationManager implements AutoCloseable {
7374
private static final Logger logger = Logger.getLogger(ClientConfigurationManager.class.getName());
7475

7576
public static final String OVERRIDE_SYS_PROP_KEY = "bigtable.internal.client-config-override";
77+
public static final String DISABLE_DIRECT_ACCESS_FALLBACK_ENV_VAR =
78+
"CBT_DISABLE_DIRECT_ACCESS_FALLBACK";
79+
public static final String DISABLE_DIRECT_ACCESS_FALLBACK_SYS_PROP_KEY =
80+
"bigtable.internal.disable-direct-access-fallback";
7681

7782
public interface ConfigListener<T> {
7883
void onChange(T newValue);
@@ -126,6 +131,7 @@ public void close() {
126131

127132
private final ClientConfiguration defaultConfig;
128133
private final Optional<ClientConfiguration> overrideConfig;
134+
private final boolean disableDirectPathFallback;
129135

130136
private final Duration defaultDeadline = Duration.ofSeconds(5);
131137

@@ -177,6 +183,11 @@ public ClientConfigurationManager(
177183
"Failed to parse bigtable.internal.client-config-override", e);
178184
}
179185
});
186+
this.disableDirectPathFallback =
187+
Optional.ofNullable(System.getenv(DISABLE_DIRECT_ACCESS_FALLBACK_ENV_VAR))
188+
.orElseGet(
189+
() -> sysProps.getProperty(DISABLE_DIRECT_ACCESS_FALLBACK_SYS_PROP_KEY, "false"))
190+
.equalsIgnoreCase("true");
180191

181192
if (overrideConfig.isPresent()) {
182193
logger.log(
@@ -188,7 +199,7 @@ public ClientConfigurationManager(
188199
.printToString(overrideConfig.get()));
189200
}
190201
featureFlags = channelProvider.updateFeatureFlags(featureFlags);
191-
this.clientConfiguration = defaultConfig;
202+
this.clientConfiguration = normalizeConfig(defaultConfig);
192203

193204
this.metadata =
194205
Util.composeMetadata(
@@ -406,6 +417,18 @@ private ClientConfiguration normalizeConfig(ClientConfiguration cfg) {
406417
// Inject overrides
407418
overrideConfig.ifPresent(builder::mergeFrom);
408419

420+
if (disableDirectPathFallback
421+
&& builder
422+
.getSessionConfiguration()
423+
.getChannelConfiguration()
424+
.hasDirectAccessWithFallback()) {
425+
builder
426+
.getSessionConfigurationBuilder()
427+
.getChannelConfigurationBuilder()
428+
.clearDirectAccessWithFallback()
429+
.setDirectAccessOnly(DirectAccessOnly.getDefaultInstance());
430+
}
431+
409432
// When sessions are disabled make sure to clear out the config. Read from the builder, not
410433
// cfg, so that a nonzero session_load supplied via the override sys-prop is honoured even when
411434
// the server-returned config has session_load=0.
@@ -453,7 +476,7 @@ private void setClientConfiguration(ClientConfiguration result) {
453476
ClientConfiguration old;
454477
synchronized (this) {
455478
old = this.clientConfiguration;
456-
clientConfiguration = result;
479+
clientConfiguration = normalizeConfig(result);
457480
if (clientConfiguration.hasPollingConfiguration()) {
458481
this.validUntil =
459482
Instant.now()

‎java-bigtable/google-cloud-bigtable/src/test/java/com/google/cloud/bigtable/data/v2/internal/util/ClientConfigurationManagerTest.java‎

Lines changed: 60 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -33,6 +33,7 @@
3333
import com.google.bigtable.v2.LoadBalancingOptions.LeastInFlight;
3434
import com.google.bigtable.v2.LoadBalancingOptions.PeakEwma;
3535
import com.google.bigtable.v2.SessionClientConfiguration;
36+
import com.google.bigtable.v2.SessionClientConfiguration.ChannelPoolConfiguration.DirectAccessWithFallback;
3637
import com.google.bigtable.v2.SessionClientConfiguration.SessionPoolConfiguration;
3738
import com.google.cloud.bigtable.data.v2.FakeServiceBuilder;
3839
import com.google.cloud.bigtable.data.v2.internal.api.ChannelProviders;
@@ -465,6 +466,59 @@ public void onChange(SessionClientConfiguration newValue) {
465466
assertThat(retrievedConfig.get()).isNotNull();
466467
}
467468

469+
@Test
470+
void disableDirectPathFallbackTest() throws Exception {
471+
Properties sysProps = new Properties();
472+
sysProps.setProperty(
473+
ClientConfigurationManager.DISABLE_DIRECT_ACCESS_FALLBACK_SYS_PROP_KEY, "true");
474+
String clientConfigOverrides =
475+
TextFormat.printer()
476+
.printToString(
477+
ClientConfiguration.newBuilder()
478+
.setSessionConfiguration(
479+
SessionClientConfiguration.newBuilder()
480+
.setSessionLoad(0.75f)
481+
.setChannelConfiguration(
482+
SessionClientConfiguration.ChannelPoolConfiguration.newBuilder()
483+
.setDirectAccessWithFallback(
484+
DirectAccessWithFallback.getDefaultInstance())))
485+
.build());
486+
sysProps.setProperty(ClientConfigurationManager.OVERRIDE_SYS_PROP_KEY, clientConfigOverrides);
487+
488+
try (ClientConfigurationManager fallbackDisabledManager =
489+
new ClientConfigurationManager(
490+
sysProps, FEATURE_FLAGS, CLIENT_INFO, channelProvider, noopDebugTracer, mockExecutor)) {
491+
492+
// Check initial default config with override has fallback disabled and direct_access_only set
493+
SessionClientConfiguration.ChannelPoolConfiguration initialChannelConfig =
494+
fallbackDisabledManager
495+
.getClientConfiguration()
496+
.getSessionConfiguration()
497+
.getChannelConfiguration();
498+
assertThat(initialChannelConfig.hasDirectAccessOnly()).isTrue();
499+
assertThat(initialChannelConfig.hasDirectAccessWithFallback()).isFalse();
500+
501+
// Start manager and fetch server config (which sends direct_access_with_fallback)
502+
ClientConfiguration fetchedConfig = fallbackDisabledManager.start().get();
503+
SessionClientConfiguration.ChannelPoolConfiguration fetchedChannelConfig =
504+
fetchedConfig.getSessionConfiguration().getChannelConfiguration();
505+
506+
// Verify that direct_access_with_fallback is converted to direct_access_only
507+
assertThat(fetchedChannelConfig.hasDirectAccessOnly()).isTrue();
508+
assertThat(fetchedChannelConfig.hasDirectAccessWithFallback()).isFalse();
509+
510+
// Verify other channel pool configuration fields from the server are preserved
511+
SessionClientConfiguration.ChannelPoolConfiguration serverChannelConfig =
512+
service.config.get().getSessionConfiguration().getChannelConfiguration();
513+
assertThat(fetchedChannelConfig.getMinServerCount())
514+
.isEqualTo(serverChannelConfig.getMinServerCount());
515+
assertThat(fetchedChannelConfig.getMaxServerCount())
516+
.isEqualTo(serverChannelConfig.getMaxServerCount());
517+
assertThat(fetchedChannelConfig.getPerServerSessionCount())
518+
.isEqualTo(serverChannelConfig.getPerServerSessionCount());
519+
}
520+
}
521+
468522
static class FakeConfigService extends BigtableGrpc.BigtableImplBase {
469523
private final AtomicReference<ClientConfiguration> config = new AtomicReference<>();
470524

@@ -475,6 +529,12 @@ static class FakeConfigService extends BigtableGrpc.BigtableImplBase {
475529
public FakeConfigService() throws IOException {
476530
ClientConfiguration.Builder builder = ClientConfigurationManager.loadDefault().toBuilder();
477531
builder.getSessionConfigurationBuilder().setSessionLoad(0.25f);
532+
builder
533+
.getSessionConfigurationBuilder()
534+
.getChannelConfigurationBuilder()
535+
.setMinServerCount(2)
536+
.setMaxServerCount(10)
537+
.setPerServerSessionCount(15);
478538
config.set(builder.build());
479539
}
480540

0 commit comments

Comments
 (0)