diff --git a/java-bigtable/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/internal/compat/ShimImpl.java b/java-bigtable/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/internal/compat/ShimImpl.java index 200c96ec0409..aed59bd62d88 100644 --- a/java-bigtable/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/internal/compat/ShimImpl.java +++ b/java-bigtable/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/internal/compat/ShimImpl.java @@ -376,28 +376,35 @@ public UnaryCallable decorateReadRow( classic, new RowBuilderShim<>(readRowShim, rowAdapter, r -> r.hasRow() ? r.getRow() : null), Util.extractTimeout(settings), - debugTagTracer); + debugTagTracer, + "method:Bigtable.ReadRow", + /* fallbackToGlobalSessionLoad= */ true); } @Override public UnaryCallable decorateMutateRow( UnaryCallable classic, UnaryCallSettings settings) { return new DivertingUnaryCallable<>( - configManager, classic, mutateRowShim, Util.extractTimeout(settings), debugTagTracer); + configManager, + classic, + mutateRowShim, + Util.extractTimeout(settings), + debugTagTracer, + "method:Bigtable.MutateRow", + /* fallbackToGlobalSessionLoad= */ true); } @Override public UnaryCallable decorateCheckAndMutateRow( UnaryCallable classic, UnaryCallSettings settings) { - if (!WipFeatures.CHECK_AND_MUTATE_ROW_ENABLED) { - return classic; - } return new DivertingUnaryCallable<>( configManager, classic, checkAndMutateRowShim, Util.extractTimeout(settings), - debugTagTracer); + debugTagTracer, + "method:Bigtable.CheckAndMutateRow", + /* fallbackToGlobalSessionLoad= */ false); } @Override @@ -405,15 +412,14 @@ public UnaryCallable decorateReadModifyWriteRow UnaryCallable classic, RowAdapter rowAdapter, UnaryCallSettings settings) { - if (!WipFeatures.READ_MODIFY_WRITE_ROW_ENABLED) { - return classic; - } return new DivertingUnaryCallable<>( configManager, classic, new RowBuilderShim<>( readModifyWriteRowShim, rowAdapter, r -> r.hasRow() ? r.getRow() : null), Util.extractTimeout(settings), - debugTagTracer); + debugTagTracer, + "method:Bigtable.ReadModifyWriteRow", + /* fallbackToGlobalSessionLoad= */ false); } } diff --git a/java-bigtable/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/internal/compat/WipFeatures.java b/java-bigtable/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/internal/compat/WipFeatures.java deleted file mode 100644 index 114da0d4b542..000000000000 --- a/java-bigtable/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/internal/compat/WipFeatures.java +++ /dev/null @@ -1,27 +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.data.v2.internal.compat; - -/** Tracks work-in-progress session-path features that are not yet enabled in production. */ -public final class WipFeatures { - // Enable session-path diversion for CheckAndMutateRow once per-method diversion is supported. - public static final boolean CHECK_AND_MUTATE_ROW_ENABLED = false; - - // Enable session-path diversion for ReadModifyWriteRow once per-method diversion is supported. - public static final boolean READ_MODIFY_WRITE_ROW_ENABLED = false; - - private WipFeatures() {} -} diff --git a/java-bigtable/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/internal/compat/ops/DivertingUnaryCallable.java b/java-bigtable/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/internal/compat/ops/DivertingUnaryCallable.java index 13111478db1d..ccaa09aad2c1 100644 --- a/java-bigtable/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/internal/compat/ops/DivertingUnaryCallable.java +++ b/java-bigtable/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/internal/compat/ops/DivertingUnaryCallable.java @@ -22,6 +22,9 @@ import com.google.api.gax.rpc.ApiException; import com.google.api.gax.rpc.ApiExceptionFactory; import com.google.api.gax.rpc.UnaryCallable; +import com.google.bigtable.v2.SessionClientConfiguration; +import com.google.bigtable.v2.SessionDiversionConfiguration; +import com.google.bigtable.v2.SessionScopeDiversionConfiguration; import com.google.bigtable.v2.TelemetryConfiguration; import com.google.cloud.bigtable.data.v2.internal.compat.FutureAdapter; import com.google.cloud.bigtable.data.v2.internal.compat.Util; @@ -49,18 +52,28 @@ public class DivertingUnaryCallable extends UnaryCallable classic, UnaryShim experimental, Duration defaultTimeout, - DebugTagTracer debugTagTracer) { + DebugTagTracer debugTagTracer, + String methodScopeKey, + boolean fallbackToGlobalSessionLoad) { this.configurationManager = configurationManager; this.classic = classic; this.experimental = experimental; this.defaultTimeout = defaultTimeout; this.debugTagTracer = debugTagTracer; + this.methodScopeKey = methodScopeKey; + this.fallbackToGlobalSessionLoad = fallbackToGlobalSessionLoad; } @Override @@ -109,8 +122,7 @@ public ApiFuture futureCall(ReqT request, ApiCallContext context) { } private boolean useExperimental(ReqT req) { - float ratio = - configurationManager.getClientConfiguration().getSessionConfiguration().getSessionLoad(); + float ratio = resolveSessionLoad(); if (ratio == 0) { return false; } @@ -125,6 +137,28 @@ private boolean useExperimental(ReqT req) { return true; } + private float resolveSessionLoad() { + SessionClientConfiguration sessionConfig = + configurationManager.getClientConfiguration().getSessionConfiguration(); + + if (sessionConfig.hasSessionDiversionConfiguration()) { + SessionDiversionConfiguration diversionConfig = + sessionConfig.getSessionDiversionConfiguration(); + switch (diversionConfig.getDiversionStrategyCase()) { + case GLOBAL_DIVERSION: + return diversionConfig.getGlobalDiversion().getSessionLoad(); + case PER_SCOPE_DIVERSION: + SessionScopeDiversionConfiguration scopeConfig = + diversionConfig.getPerScopeDiversion().getScopeDiversionsMap().get(methodScopeKey); + return scopeConfig != null ? scopeConfig.getSessionLoad() : 0f; + default: + return 0f; + } + } + + return fallbackToGlobalSessionLoad ? sessionConfig.getSessionLoad() : 0f; + } + ApiException translateException(Throwable e) { Throwable cause = e; while (cause instanceof CompletionException || cause instanceof ExecutionException) { diff --git a/java-bigtable/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/internal/util/ClientConfigurationManager.java b/java-bigtable/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/internal/util/ClientConfigurationManager.java index 59c8b50bff68..b69628de763e 100644 --- a/java-bigtable/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/internal/util/ClientConfigurationManager.java +++ b/java-bigtable/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/internal/util/ClientConfigurationManager.java @@ -21,7 +21,9 @@ import com.google.bigtable.v2.ClientConfiguration.PollingCase; import com.google.bigtable.v2.FeatureFlags; import com.google.bigtable.v2.GetClientConfigurationRequest; +import com.google.bigtable.v2.SessionClientConfiguration; import com.google.bigtable.v2.SessionClientConfiguration.ChannelPoolConfiguration.DirectAccessOnly; +import com.google.bigtable.v2.SessionDiversionConfiguration; import com.google.bigtable.v2.TelemetryConfiguration; import com.google.cloud.bigtable.data.v2.internal.api.ChannelProviders.ChannelProvider; import com.google.cloud.bigtable.data.v2.internal.api.Util; @@ -447,7 +449,7 @@ private ClientConfiguration normalizeConfig(ClientConfiguration cfg) { // When sessions are disabled make sure to clear out the config. Read from the builder, not // cfg, so that a nonzero session_load supplied via the override sys-prop is honoured even when // the server-returned config has session_load=0. - if (builder.getSessionConfiguration().getSessionLoad() == 0) { + if (!divertsAnyTrafficToSessions(builder.getSessionConfiguration())) { builder.clearSessionConfiguration(); return builder.build(); } @@ -455,8 +457,31 @@ private ClientConfiguration normalizeConfig(ClientConfiguration cfg) { return builder.build(); } + /** + * Returns true if the configuration routes a nonzero share of traffic to sessions, either via the + * flat {@code session_load} or via a diversion configuration that carries a nonzero load. + */ + private static boolean divertsAnyTrafficToSessions(SessionClientConfiguration session) { + if (session.getSessionLoad() > 0) { + return true; + } + + SessionDiversionConfiguration diversion = session.getSessionDiversionConfiguration(); + switch (diversion.getDiversionStrategyCase()) { + case GLOBAL_DIVERSION: + return diversion.getGlobalDiversion().getSessionLoad() > 0; + case PER_SCOPE_DIVERSION: + return diversion.getPerScopeDiversion().getScopeDiversionsMap().values().stream() + .anyMatch(scope -> scope.getSessionLoad() > 0); + default: + return false; + } + } + public boolean areSessionsRequired() { - return overrideConfig.map(c -> c.getSessionConfiguration().getSessionLoad() > 0).orElse(false); + return overrideConfig + .map(c -> divertsAnyTrafficToSessions(c.getSessionConfiguration())) + .orElse(false); } private long getRetryDelay(int attempt) { diff --git a/java-bigtable/google-cloud-bigtable/src/test/java/com/google/cloud/bigtable/data/v2/internal/compat/ops/DivertingUnaryCallableTest.java b/java-bigtable/google-cloud-bigtable/src/test/java/com/google/cloud/bigtable/data/v2/internal/compat/ops/DivertingUnaryCallableTest.java index 1ae3c0694e76..7b769213e825 100644 --- a/java-bigtable/google-cloud-bigtable/src/test/java/com/google/cloud/bigtable/data/v2/internal/compat/ops/DivertingUnaryCallableTest.java +++ b/java-bigtable/google-cloud-bigtable/src/test/java/com/google/cloud/bigtable/data/v2/internal/compat/ops/DivertingUnaryCallableTest.java @@ -27,6 +27,9 @@ import com.google.auth.Credentials; import com.google.bigtable.v2.ClientConfiguration; import com.google.bigtable.v2.SessionClientConfiguration; +import com.google.bigtable.v2.SessionDiversionConfiguration; +import com.google.bigtable.v2.SessionDiversionConfiguration.PerScopeDiversion; +import com.google.bigtable.v2.SessionScopeDiversionConfiguration; import com.google.bigtable.v2.TelemetryConfiguration; import com.google.cloud.bigtable.data.v2.internal.csm.tracers.DebugTagTracer; import com.google.cloud.bigtable.data.v2.internal.util.ClientConfigurationManager; @@ -111,11 +114,89 @@ void sessionLoadZero_routesToClassicWithoutCredentialsCheck() throws Exception { assertThat(tracer.tags).isEmpty(); } + @Test + void globalDiversion_usesGlobalSessionLoad() { + stubGlobalDiversion(1.0f); + DivertingUnaryCallable callable = callable(); + + @SuppressWarnings("unused") + ApiFuture ignored = callable.futureCall("req", GrpcCallContext.createDefault()); + + assertThat(experimental.callCount).isEqualTo(1); + } + + @Test + void globalDiversionZero_routesToClassic() throws Exception { + stubGlobalDiversion(0.0f); + DivertingUnaryCallable callable = callable(); + + ApiFuture result = callable.futureCall("req", GrpcCallContext.createDefault()); + + assertThat(result.get()).isEqualTo("classic"); + assertThat(experimental.callCount).isEqualTo(0); + } + + @Test + void perScopeDiversion_matchingScope_routesToExperimental() { + stubPerScopeDiversion("method:Bigtable.ReadRow", 1.0f); + DivertingUnaryCallable callable = callable(); + + @SuppressWarnings("unused") + ApiFuture ignored = callable.futureCall("req", GrpcCallContext.createDefault()); + + assertThat(experimental.callCount).isEqualTo(1); + } + + @Test + void fallbackEnabled_noDiversionConfig_usesGlobalSessionLoad() { + stubSessionLoad(1.0f); + DivertingUnaryCallable callable = + callable(/* fallbackToGlobalSessionLoad= */ true); + + @SuppressWarnings("unused") + ApiFuture ignored = callable.futureCall("req", GrpcCallContext.createDefault()); + + assertThat(experimental.callCount).isEqualTo(1); + } + + @Test + void fallbackDisabled_noDiversionConfig_routesToClassic() throws Exception { + stubSessionLoad(1.0f); + DivertingUnaryCallable callable = + callable(/* fallbackToGlobalSessionLoad= */ false); + + ApiFuture result = callable.futureCall("req", GrpcCallContext.createDefault()); + + assertThat(result.get()).isEqualTo("classic"); + assertThat(experimental.callCount).isEqualTo(0); + } + + @Test + void perScopeDiversion_nonMatchingScope_routesToClassic() throws Exception { + stubPerScopeDiversion("method:Bigtable.MutateRow", 1.0f); + DivertingUnaryCallable callable = callable(); + + ApiFuture result = callable.futureCall("req", GrpcCallContext.createDefault()); + + assertThat(result.get()).isEqualTo("classic"); + assertThat(experimental.callCount).isEqualTo(0); + } + // --- helpers --- private DivertingUnaryCallable callable() { + return callable(true); + } + + private DivertingUnaryCallable callable(boolean fallbackToGlobalSessionLoad) { return new DivertingUnaryCallable<>( - configManager, classic, experimental, Duration.ofSeconds(10), tracer); + configManager, + classic, + experimental, + Duration.ofSeconds(10), + tracer, + "method:Bigtable.ReadRow", + fallbackToGlobalSessionLoad); } private void stubSessionLoad(float load) { @@ -126,6 +207,38 @@ private void stubSessionLoad(float load) { when(configManager.getClientConfiguration()).thenReturn(config); } + private void stubGlobalDiversion(float load) { + ClientConfiguration config = + ClientConfiguration.newBuilder() + .setSessionConfiguration( + SessionClientConfiguration.newBuilder() + .setSessionDiversionConfiguration( + SessionDiversionConfiguration.newBuilder() + .setGlobalDiversion( + SessionScopeDiversionConfiguration.newBuilder() + .setSessionLoad(load)))) + .build(); + when(configManager.getClientConfiguration()).thenReturn(config); + } + + private void stubPerScopeDiversion(String scopeKey, float load) { + ClientConfiguration config = + ClientConfiguration.newBuilder() + .setSessionConfiguration( + SessionClientConfiguration.newBuilder() + .setSessionDiversionConfiguration( + SessionDiversionConfiguration.newBuilder() + .setPerScopeDiversion( + PerScopeDiversion.newBuilder() + .putScopeDiversions( + scopeKey, + SessionScopeDiversionConfiguration.newBuilder() + .setSessionLoad(load) + .build())))) + .build(); + when(configManager.getClientConfiguration()).thenReturn(config); + } + private static class RecordingCallable extends UnaryCallable { final List capturedContexts = new ArrayList<>(); private final ApiFuture response; diff --git a/java-bigtable/google-cloud-bigtable/src/test/java/com/google/cloud/bigtable/data/v2/internal/util/ClientConfigurationManagerTest.java b/java-bigtable/google-cloud-bigtable/src/test/java/com/google/cloud/bigtable/data/v2/internal/util/ClientConfigurationManagerTest.java index 8a73a09037f2..a55255ef8fa1 100644 --- a/java-bigtable/google-cloud-bigtable/src/test/java/com/google/cloud/bigtable/data/v2/internal/util/ClientConfigurationManagerTest.java +++ b/java-bigtable/google-cloud-bigtable/src/test/java/com/google/cloud/bigtable/data/v2/internal/util/ClientConfigurationManagerTest.java @@ -35,6 +35,9 @@ import com.google.bigtable.v2.SessionClientConfiguration; import com.google.bigtable.v2.SessionClientConfiguration.ChannelPoolConfiguration.DirectAccessWithFallback; import com.google.bigtable.v2.SessionClientConfiguration.SessionPoolConfiguration; +import com.google.bigtable.v2.SessionDiversionConfiguration; +import com.google.bigtable.v2.SessionDiversionConfiguration.PerScopeDiversion; +import com.google.bigtable.v2.SessionScopeDiversionConfiguration; import com.google.cloud.bigtable.data.v2.FakeServiceBuilder; import com.google.cloud.bigtable.data.v2.internal.api.ChannelProviders; import com.google.cloud.bigtable.data.v2.internal.api.ChannelProviders.ForwardingChannelProvider; @@ -387,6 +390,118 @@ void testSysPropOverrideSessionLoadZero() throws Exception { assertThat(initialConfig).isEqualTo(service.config.get()); } + @Test + void testPerScopeDiversionWithNonZeroLoadEnablesSessions() throws Exception { + ClientConfiguration resolvedConfig = + startWithSessionOverride( + SessionClientConfiguration.newBuilder() + .setSessionLoad(0) + .setSessionDiversionConfiguration( + SessionDiversionConfiguration.newBuilder() + .setPerScopeDiversion( + PerScopeDiversion.newBuilder() + .putScopeDiversions( + "method:Bigtable.ReadRow", + SessionScopeDiversionConfiguration.newBuilder() + .setSessionLoad(0.5f) + .build())))); + + // The flat session_load is 0, but a scope diverts traffic, so the config must survive + // normalization and the client must advertise session support. + assertThat(resolvedConfig.hasSessionConfiguration()).isTrue(); + assertThat( + resolvedConfig + .getSessionConfiguration() + .getSessionDiversionConfiguration() + .getPerScopeDiversion() + .getScopeDiversionsMap()) + .containsKey("method:Bigtable.ReadRow"); + assertThat(manager.areSessionsRequired()).isTrue(); + } + + @Test + void testPerScopeDiversionWithAllZeroLoadsDisablesSessions() throws Exception { + ClientConfiguration resolvedConfig = + startWithSessionOverride( + SessionClientConfiguration.newBuilder() + .setSessionLoad(0) + .setSessionDiversionConfiguration( + SessionDiversionConfiguration.newBuilder() + .setPerScopeDiversion( + PerScopeDiversion.newBuilder() + .putScopeDiversions( + "method:Bigtable.ReadRow", + SessionScopeDiversionConfiguration.newBuilder() + .setSessionLoad(0) + .build())))); + + // A diversion config that routes nothing must not keep sessions alive. + assertThat(resolvedConfig.hasSessionConfiguration()).isFalse(); + assertThat(manager.areSessionsRequired()).isFalse(); + } + + @Test + void testGlobalDiversionWithZeroLoadDisablesSessions() throws Exception { + ClientConfiguration resolvedConfig = + startWithSessionOverride( + SessionClientConfiguration.newBuilder() + .setSessionLoad(0) + .setSessionDiversionConfiguration( + SessionDiversionConfiguration.newBuilder() + .setGlobalDiversion( + SessionScopeDiversionConfiguration.newBuilder().setSessionLoad(0)))); + + assertThat(resolvedConfig.hasSessionConfiguration()).isFalse(); + assertThat(manager.areSessionsRequired()).isFalse(); + } + + @Test + void testGlobalDiversionWithNonZeroLoadEnablesSessions() throws Exception { + ClientConfiguration resolvedConfig = + startWithSessionOverride( + SessionClientConfiguration.newBuilder() + .setSessionLoad(0) + .setSessionDiversionConfiguration( + SessionDiversionConfiguration.newBuilder() + .setGlobalDiversion( + SessionScopeDiversionConfiguration.newBuilder() + .setSessionLoad(0.25f)))); + + assertThat(resolvedConfig.hasSessionConfiguration()).isTrue(); + assertThat( + resolvedConfig + .getSessionConfiguration() + .getSessionDiversionConfiguration() + .getGlobalDiversion() + .getSessionLoad()) + .isEqualTo(0.25f); + assertThat(manager.areSessionsRequired()).isTrue(); + } + + /** + * Restarts the manager against a server that reports session_load=0, with the supplied session + * configuration supplied as a sys-prop override, and returns the resolved config. + */ + private ClientConfiguration startWithSessionOverride( + SessionClientConfiguration.Builder sessionOverride) throws Exception { + ClientConfiguration.Builder serverConfig = manager.getDefaultConfig().toBuilder(); + serverConfig.getSessionConfigurationBuilder().setSessionLoad(0); + service.config.set(serverConfig.build()); + + manager.close(); + Properties sysProps = new Properties(); + sysProps.setProperty( + ClientConfigurationManager.OVERRIDE_SYS_PROP_KEY, + TextFormat.printer() + .printToString( + ClientConfiguration.newBuilder().setSessionConfiguration(sessionOverride).build())); + manager = + new ClientConfigurationManager( + sysProps, FEATURE_FLAGS, CLIENT_INFO, channelProvider, noopDebugTracer, mockExecutor); + + return manager.start().get(); + } + /** * The manager must never invoke a listener while holding its own monitor, no matter which thread * ends up running the config-fetch continuation. See {@link #testDeadlockPrevention()} for the