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 aed59bd62d88..8db91d0a9308 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 @@ -29,6 +29,9 @@ import com.google.bigtable.v2.GetClientConfigurationRequest; import com.google.bigtable.v2.PeerInfo; import com.google.bigtable.v2.PeerInfo.TransportType; +import com.google.bigtable.v2.Row; +import com.google.bigtable.v2.SessionReadModifyWriteRowResponse; +import com.google.bigtable.v2.SessionReadRowResponse; import com.google.cloud.bigtable.data.v2.internal.api.ChannelProviders.ChannelProvider; import com.google.cloud.bigtable.data.v2.internal.api.ChannelProviders.ConfiguredChannelProvider; import com.google.cloud.bigtable.data.v2.internal.api.Client; @@ -47,6 +50,7 @@ import com.google.cloud.bigtable.data.v2.internal.csm.tracers.DirectPathCompatibleTracer; import com.google.cloud.bigtable.data.v2.internal.dp.DirectAccessInvestigator; import com.google.cloud.bigtable.data.v2.internal.util.ClientConfigurationManager; +import com.google.cloud.bigtable.data.v2.internal.util.SessionScope; import com.google.cloud.bigtable.data.v2.models.ConditionalRowMutation; import com.google.cloud.bigtable.data.v2.models.Query; import com.google.cloud.bigtable.data.v2.models.ReadModifyWriteRow; @@ -69,6 +73,7 @@ import java.util.concurrent.Executors; import java.util.concurrent.ScheduledExecutorService; import java.util.concurrent.TimeUnit; +import java.util.function.Function; import java.util.logging.Level; import java.util.logging.Logger; import javax.annotation.Nullable; @@ -82,6 +87,19 @@ public class ShimImpl implements Shim { private static final Duration DA_CHECK_TIMEOUT = Duration.ofSeconds(5); + // A ReadRow on a missing key legitimately has no row, and the classic path surfaces that as a + // null Row. + @VisibleForTesting + static final Function READ_ROW_EXTRACTOR = + r -> r.hasRow() ? r.getRow() : null; + + // ReadModifyWriteRow always mutates and so always yields a row. The classic path builds a Row + // from the (possibly unset) response row rather than returning null, so extract unconditionally + // to keep both paths non-null. + @VisibleForTesting + static final Function READ_MODIFY_WRITE_ROW_EXTRACTOR = + SessionReadModifyWriteRowResponse::getRow; + private final ClientConfigurationManager configManager; private final Resource configManagerResource; private final Client client; @@ -374,11 +392,10 @@ public UnaryCallable decorateReadRow( return new DivertingUnaryCallable<>( configManager, classic, - new RowBuilderShim<>(readRowShim, rowAdapter, r -> r.hasRow() ? r.getRow() : null), + new RowBuilderShim<>(readRowShim, rowAdapter, READ_ROW_EXTRACTOR), Util.extractTimeout(settings), debugTagTracer, - "method:Bigtable.ReadRow", - /* fallbackToGlobalSessionLoad= */ true); + SessionScope.READ_ROW); } @Override @@ -390,8 +407,7 @@ public UnaryCallable decorateMutateRow( mutateRowShim, Util.extractTimeout(settings), debugTagTracer, - "method:Bigtable.MutateRow", - /* fallbackToGlobalSessionLoad= */ true); + SessionScope.MUTATE_ROW); } @Override @@ -403,8 +419,7 @@ public UnaryCallable decorateCheckAndMutateRow( checkAndMutateRowShim, Util.extractTimeout(settings), debugTagTracer, - "method:Bigtable.CheckAndMutateRow", - /* fallbackToGlobalSessionLoad= */ false); + SessionScope.CHECK_AND_MUTATE_ROW); } @Override @@ -415,11 +430,9 @@ public UnaryCallable decorateReadModifyWriteRow return new DivertingUnaryCallable<>( configManager, classic, - new RowBuilderShim<>( - readModifyWriteRowShim, rowAdapter, r -> r.hasRow() ? r.getRow() : null), + new RowBuilderShim<>(readModifyWriteRowShim, rowAdapter, READ_MODIFY_WRITE_ROW_EXTRACTOR), Util.extractTimeout(settings), debugTagTracer, - "method:Bigtable.ReadModifyWriteRow", - /* fallbackToGlobalSessionLoad= */ false); + SessionScope.READ_MODIFY_WRITE_ROW); } } 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 ccaa09aad2c1..6500add3dee2 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,14 +22,12 @@ 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; import com.google.cloud.bigtable.data.v2.internal.csm.tracers.DebugTagTracer; import com.google.cloud.bigtable.data.v2.internal.util.ClientConfigurationManager; +import com.google.cloud.bigtable.data.v2.internal.util.SessionScope; import com.google.common.base.Throwables; import io.grpc.Context; import io.grpc.Context.CancellableContext; @@ -52,12 +50,8 @@ public class DivertingUnaryCallable extends UnaryCallable experimental, Duration defaultTimeout, DebugTagTracer debugTagTracer, - String methodScopeKey, - boolean fallbackToGlobalSessionLoad) { + SessionScope scope) { this.configurationManager = configurationManager; this.classic = classic; this.experimental = experimental; this.defaultTimeout = defaultTimeout; this.debugTagTracer = debugTagTracer; - this.methodScopeKey = methodScopeKey; - this.fallbackToGlobalSessionLoad = fallbackToGlobalSessionLoad; + this.scope = scope; } @Override @@ -138,25 +130,8 @@ private boolean useExperimental(ReqT req) { } 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; + return scope.resolveSessionLoad( + configurationManager.getClientConfiguration().getSessionConfiguration()); } ApiException translateException(Throwable e) { 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 b69628de763e..d51e04888da6 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 @@ -23,7 +23,6 @@ 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; @@ -458,24 +457,12 @@ private ClientConfiguration normalizeConfig(ClientConfiguration cfg) { } /** - * 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. + * Returns true if the configuration routes a nonzero share of traffic to sessions for any scope + * this client serves. Shares the resolution rule with the per-RPC diversion decision so the two + * cannot disagree; see {@link SessionScope}. */ 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; - } + return SessionScope.divertsAnyTraffic(session); } public boolean areSessionsRequired() { diff --git a/java-bigtable/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/internal/util/SessionScope.java b/java-bigtable/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/internal/util/SessionScope.java new file mode 100644 index 000000000000..e89892a5eba8 --- /dev/null +++ b/java-bigtable/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/internal/util/SessionScope.java @@ -0,0 +1,94 @@ +/* + * 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.util; + +import com.google.bigtable.v2.SessionClientConfiguration; +import com.google.bigtable.v2.SessionDiversionConfiguration; +import com.google.bigtable.v2.SessionScopeDiversionConfiguration; + +/** + * The per-method diversion scopes this client implements, along with the single rule for resolving + * a session load for each of them. + * + *

Two decisions depend on this rule: the per-RPC fork in {@code DivertingUnaryCallable}, and the + * client wide "are sessions in use at all" decision in {@link ClientConfigurationManager}. They + * must agree. If the client-wide check is the more permissive of the two, the client keeps session + * pools alive and advertises session feature flags while never diverting a single request. Keeping + * the scope list and the resolution rule in one place is what keeps them in sync. + */ +public enum SessionScope { + READ_ROW("method:Bigtable.ReadRow", /* fallbackToGlobalSessionLoad= */ true), + MUTATE_ROW("method:Bigtable.MutateRow", /* fallbackToGlobalSessionLoad= */ true), + CHECK_AND_MUTATE_ROW( + "method:Bigtable.CheckAndMutateRow", /* fallbackToGlobalSessionLoad= */ false), + READ_MODIFY_WRITE_ROW( + "method:Bigtable.ReadModifyWriteRow", /* fallbackToGlobalSessionLoad= */ false); + + private final String key; + private final boolean fallbackToGlobalSessionLoad; + + SessionScope(String key, boolean fallbackToGlobalSessionLoad) { + this.key = key; + this.fallbackToGlobalSessionLoad = fallbackToGlobalSessionLoad; + } + + /** The {@code method:.} key the server uses to address this scope. */ + public String getKey() { + return key; + } + + /** + * Returns the fraction of this scope's traffic that should be served by sessions. + * + *

A {@link SessionDiversionConfiguration} is authoritative when present: the server only sends + * the deprecated flat {@code session_load} to keep ReadRow and MutateRow working on older + * clients, so this client ignores it entirely once a diversion configuration arrives. + */ + @SuppressWarnings("deprecation") + public float resolveSessionLoad(SessionClientConfiguration sessionConfig) { + 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(key); + return scopeConfig != null ? scopeConfig.getSessionLoad() : 0f; + default: + return 0f; + } + } + + return fallbackToGlobalSessionLoad ? sessionConfig.getSessionLoad() : 0f; + } + + /** + * Returns true if the configuration routes a nonzero share of traffic to sessions for at least + * one scope this client implements. Scopes the client does not serve are ignored: a diversion + * entry for them cannot produce any session traffic, so it must not keep the session machinery + * alive either. + */ + public static boolean divertsAnyTraffic(SessionClientConfiguration sessionConfig) { + for (SessionScope scope : values()) { + if (scope.resolveSessionLoad(sessionConfig) > 0) { + return true; + } + } + return false; + } +} diff --git a/java-bigtable/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/stub/EnhancedBigtableStub.java b/java-bigtable/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/stub/EnhancedBigtableStub.java index de3ef1b672e6..cafc9e78e8b4 100644 --- a/java-bigtable/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/stub/EnhancedBigtableStub.java +++ b/java-bigtable/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/stub/EnhancedBigtableStub.java @@ -941,7 +941,11 @@ private UnaryCallable createCheckAndMutateRowCa return bigtableClientContext .getSessionShim() .decorateCheckAndMutateRow(classic, perOpSettings.checkAndMutateRowSettings) - .withDefaultCallContext(bigtableClientContext.getClientContext().getDefaultCallContext()); + .withDefaultCallContext( + bigtableClientContext + .getClientContext() + .getDefaultCallContext() + .withRetrySettings(perOpSettings.checkAndMutateRowSettings.getRetrySettings())); } /** @@ -970,7 +974,11 @@ private UnaryCallable createReadModifyWriteRowCallable( return bigtableClientContext .getSessionShim() .decorateReadModifyWriteRow(classic, rowAdapter, perOpSettings.readModifyWriteRowSettings) - .withDefaultCallContext(bigtableClientContext.getClientContext().getDefaultCallContext()); + .withDefaultCallContext( + bigtableClientContext + .getClientContext() + .getDefaultCallContext() + .withRetrySettings(perOpSettings.readModifyWriteRowSettings.getRetrySettings())); } /** diff --git a/java-bigtable/google-cloud-bigtable/src/main/resources/bigtable-default-client-config.textproto b/java-bigtable/google-cloud-bigtable/src/main/resources/bigtable-default-client-config.textproto index 51cf60561393..7a0dfd73156a 100644 --- a/java-bigtable/google-cloud-bigtable/src/main/resources/bigtable-default-client-config.textproto +++ b/java-bigtable/google-cloud-bigtable/src/main/resources/bigtable-default-client-config.textproto @@ -12,23 +12,6 @@ polling_configuration { } session_configuration { - session_diversion_configuration { - per_scope_diversion { - scope_diversions { - key: "method:Bigtable.ReadRow" - value { - session_load: 0 - } - } - scope_diversions { - key: "method:Bigtable.MutateRow" - value { - session_load: 0 - } - } - } - } - channel_configuration { min_server_count: 2 max_server_count: 25 diff --git a/java-bigtable/google-cloud-bigtable/src/test/java/com/google/cloud/bigtable/data/v2/internal/compat/ShimImplTest.java b/java-bigtable/google-cloud-bigtable/src/test/java/com/google/cloud/bigtable/data/v2/internal/compat/ShimImplTest.java index 1dc064855d16..4dfab36d0694 100644 --- a/java-bigtable/google-cloud-bigtable/src/test/java/com/google/cloud/bigtable/data/v2/internal/compat/ShimImplTest.java +++ b/java-bigtable/google-cloud-bigtable/src/test/java/com/google/cloud/bigtable/data/v2/internal/compat/ShimImplTest.java @@ -22,6 +22,8 @@ import com.google.api.gax.core.FixedExecutorProvider; import com.google.api.gax.grpc.InstantiatingGrpcChannelProvider; +import com.google.bigtable.v2.SessionReadModifyWriteRowResponse; +import com.google.bigtable.v2.SessionReadRowResponse; import com.google.cloud.bigtable.data.v2.internal.api.Client.Resource; import java.util.concurrent.CountDownLatch; import java.util.concurrent.Executor; @@ -35,6 +37,26 @@ public class ShimImplTest { private static final InstantiatingGrpcChannelProvider EMPTY_TRANSPORT = InstantiatingGrpcChannelProvider.newBuilder().build(); + /** + * ReadRow and ReadModifyWriteRow disagree on what "no row in the response" means, and the session + * path must match the classic path for each. A ReadRow miss is a null Row; a ReadModifyWriteRow + * always produces one, so a response without a row must still build a Row rather than handing + * user code a null the classic path never produces. + */ + @Test + public void readRowExtractor_noRow_returnsNull() { + assertThat(ShimImpl.READ_ROW_EXTRACTOR.apply(SessionReadRowResponse.getDefaultInstance())) + .isNull(); + } + + @Test + public void readModifyWriteRowExtractor_noRow_returnsEmptyRowNotNull() { + assertThat( + ShimImpl.READ_MODIFY_WRITE_ROW_EXTRACTOR.apply( + SessionReadModifyWriteRowResponse.getDefaultInstance())) + .isNotNull(); + } + @Test public void selectUserCallbackExecutor_transportSet_returnsSharedTransportExecutor() { Executor userExec = mock(Executor.class); 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 7b769213e825..729f378b7a07 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 @@ -33,6 +33,7 @@ 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; +import com.google.cloud.bigtable.data.v2.internal.util.SessionScope; import io.grpc.Deadline; import java.io.IOException; import java.net.URI; @@ -138,7 +139,7 @@ void globalDiversionZero_routesToClassic() throws Exception { @Test void perScopeDiversion_matchingScope_routesToExperimental() { - stubPerScopeDiversion("method:Bigtable.ReadRow", 1.0f); + stubPerScopeDiversion(SessionScope.READ_ROW.getKey(), 1.0f); DivertingUnaryCallable callable = callable(); @SuppressWarnings("unused") @@ -150,8 +151,9 @@ void perScopeDiversion_matchingScope_routesToExperimental() { @Test void fallbackEnabled_noDiversionConfig_usesGlobalSessionLoad() { stubSessionLoad(1.0f); - DivertingUnaryCallable callable = - callable(/* fallbackToGlobalSessionLoad= */ true); + // READ_ROW still honours the deprecated flat session_load while the server rolls out + // SessionDiversionConfiguration. + DivertingUnaryCallable callable = callable(SessionScope.READ_ROW); @SuppressWarnings("unused") ApiFuture ignored = callable.futureCall("req", GrpcCallContext.createDefault()); @@ -162,8 +164,8 @@ void fallbackEnabled_noDiversionConfig_usesGlobalSessionLoad() { @Test void fallbackDisabled_noDiversionConfig_routesToClassic() throws Exception { stubSessionLoad(1.0f); - DivertingUnaryCallable callable = - callable(/* fallbackToGlobalSessionLoad= */ false); + // CheckAndMutateRow is only ever enabled through a diversion configuration. + DivertingUnaryCallable callable = callable(SessionScope.CHECK_AND_MUTATE_ROW); ApiFuture result = callable.futureCall("req", GrpcCallContext.createDefault()); @@ -173,7 +175,7 @@ void fallbackDisabled_noDiversionConfig_routesToClassic() throws Exception { @Test void perScopeDiversion_nonMatchingScope_routesToClassic() throws Exception { - stubPerScopeDiversion("method:Bigtable.MutateRow", 1.0f); + stubPerScopeDiversion(SessionScope.MUTATE_ROW.getKey(), 1.0f); DivertingUnaryCallable callable = callable(); ApiFuture result = callable.futureCall("req", GrpcCallContext.createDefault()); @@ -185,18 +187,12 @@ void perScopeDiversion_nonMatchingScope_routesToClassic() throws Exception { // --- helpers --- private DivertingUnaryCallable callable() { - return callable(true); + return callable(SessionScope.READ_ROW); } - private DivertingUnaryCallable callable(boolean fallbackToGlobalSessionLoad) { + private DivertingUnaryCallable callable(SessionScope scope) { return new DivertingUnaryCallable<>( - configManager, - classic, - experimental, - Duration.ofSeconds(10), - tracer, - "method:Bigtable.ReadRow", - fallbackToGlobalSessionLoad); + configManager, classic, experimental, Duration.ofSeconds(10), tracer, scope); } private void stubSessionLoad(float load) { 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 a55255ef8fa1..fc46d09d7ede 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 @@ -285,6 +285,28 @@ void testDisabledSessionsOverriddenBySysProp() assertThat(manager.areSessionsRequired()).isTrue(); } + /** + * The bundled default config must not carry a SessionDiversionConfiguration. A diversion config + * is authoritative once present, so one baked into the default would shadow the flat session_load + * for every scope -- including in the two windows where the default is the live config: before + * the first GetClientConfiguration response, and after handleFailedFetch falls back to it when + * the validity window expires. + */ + @Test + void defaultConfigDoesNotShadowFlatSessionLoad() throws IOException { + ClientConfiguration defaultConfig = ClientConfigurationManager.loadDefault(); + assertThat(defaultConfig.getSessionConfiguration().hasSessionDiversionConfiguration()) + .isFalse(); + + // A flat session_load merged onto the default -- what the sys-prop override produces -- must + // still divert traffic. + ClientConfiguration.Builder merged = defaultConfig.toBuilder(); + merged.getSessionConfigurationBuilder().setSessionLoad(1.0f); + assertThat(SessionScope.READ_ROW.resolveSessionLoad(merged.getSessionConfiguration())) + .isEqualTo(1.0f); + assertThat(SessionScope.divertsAnyTraffic(merged.getSessionConfiguration())).isTrue(); + } + @Deprecated @Test void testMigrateSessionPool() throws ExecutionException, InterruptedException { diff --git a/java-bigtable/google-cloud-bigtable/src/test/java/com/google/cloud/bigtable/data/v2/internal/util/SessionScopeTest.java b/java-bigtable/google-cloud-bigtable/src/test/java/com/google/cloud/bigtable/data/v2/internal/util/SessionScopeTest.java new file mode 100644 index 000000000000..190a03e59528 --- /dev/null +++ b/java-bigtable/google-cloud-bigtable/src/test/java/com/google/cloud/bigtable/data/v2/internal/util/SessionScopeTest.java @@ -0,0 +1,132 @@ +/* + * 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.util; + +import static com.google.common.truth.Truth.assertThat; + +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 org.junit.jupiter.api.Test; + +@SuppressWarnings("deprecation") // exercises the deprecated flat session_load on purpose +class SessionScopeTest { + + @Test + void noDiversionConfig_fallbackScopesUseFlatSessionLoad() { + SessionClientConfiguration config = sessionLoad(0.75f).build(); + + assertThat(SessionScope.READ_ROW.resolveSessionLoad(config)).isEqualTo(0.75f); + assertThat(SessionScope.MUTATE_ROW.resolveSessionLoad(config)).isEqualTo(0.75f); + } + + @Test + void noDiversionConfig_nonFallbackScopesIgnoreFlatSessionLoad() { + SessionClientConfiguration config = sessionLoad(0.75f).build(); + + assertThat(SessionScope.CHECK_AND_MUTATE_ROW.resolveSessionLoad(config)).isEqualTo(0f); + assertThat(SessionScope.READ_MODIFY_WRITE_ROW.resolveSessionLoad(config)).isEqualTo(0f); + } + + @Test + void diversionConfigPresent_flatSessionLoadIsIgnoredEntirely() { + // The server only sends the flat load for older clients. Once a diversion config arrives it is + // authoritative, including for scopes it does not mention. + SessionClientConfiguration config = + sessionLoad(1.0f) + .setSessionDiversionConfiguration(perScope("method:Bigtable.CheckAndMutateRow", 0.1f)) + .build(); + + assertThat(SessionScope.CHECK_AND_MUTATE_ROW.resolveSessionLoad(config)).isEqualTo(0.1f); + assertThat(SessionScope.READ_ROW.resolveSessionLoad(config)).isEqualTo(0f); + assertThat(SessionScope.MUTATE_ROW.resolveSessionLoad(config)).isEqualTo(0f); + } + + @Test + void emptyDiversionConfig_resolvesToZeroForEveryScope() { + // An unset oneof, e.g. a future strategy this client cannot parse. + SessionClientConfiguration config = + sessionLoad(1.0f) + .setSessionDiversionConfiguration(SessionDiversionConfiguration.getDefaultInstance()) + .build(); + + for (SessionScope scope : SessionScope.values()) { + assertThat(scope.resolveSessionLoad(config)).isEqualTo(0f); + } + assertThat(SessionScope.divertsAnyTraffic(config)).isFalse(); + } + + @Test + void globalDiversionAppliesToEveryScope() { + SessionClientConfiguration config = + SessionClientConfiguration.newBuilder() + .setSessionDiversionConfiguration( + SessionDiversionConfiguration.newBuilder() + .setGlobalDiversion( + SessionScopeDiversionConfiguration.newBuilder().setSessionLoad(0.5f))) + .build(); + + for (SessionScope scope : SessionScope.values()) { + assertThat(scope.resolveSessionLoad(config)).isEqualTo(0.5f); + } + } + + /** + * The client-wide "are sessions in use" check must not be more permissive than the per-RPC fork. + * A diversion entry for a method this client never routes through a session cannot produce + * session traffic, so it must not keep session pools and feature flags alive either. + */ + @Test + void unservedScope_doesNotCountAsDivertingTraffic() { + SessionClientConfiguration config = + SessionClientConfiguration.newBuilder() + .setSessionDiversionConfiguration(perScope("method:Bigtable.MutateRows", 1.0f)) + .build(); + + for (SessionScope scope : SessionScope.values()) { + assertThat(scope.resolveSessionLoad(config)).isEqualTo(0f); + } + assertThat(SessionScope.divertsAnyTraffic(config)).isFalse(); + } + + @Test + void divertsAnyTraffic_trueWhenAnyServedScopeDiverts() { + assertThat( + SessionScope.divertsAnyTraffic( + SessionClientConfiguration.newBuilder() + .setSessionDiversionConfiguration( + perScope("method:Bigtable.ReadModifyWriteRow", 0.01f)) + .build())) + .isTrue(); + assertThat(SessionScope.divertsAnyTraffic(sessionLoad(0.25f).build())).isTrue(); + assertThat(SessionScope.divertsAnyTraffic(sessionLoad(0f).build())).isFalse(); + } + + private static SessionClientConfiguration.Builder sessionLoad(float load) { + return SessionClientConfiguration.newBuilder().setSessionLoad(load); + } + + private static SessionDiversionConfiguration perScope(String key, float load) { + return SessionDiversionConfiguration.newBuilder() + .setPerScopeDiversion( + PerScopeDiversion.newBuilder() + .putScopeDiversions( + key, + SessionScopeDiversionConfiguration.newBuilder().setSessionLoad(load).build())) + .build(); + } +}