From cc0ad8feb5f5cfbfb41dca843504171e5ba9f88b Mon Sep 17 00:00:00 2001 From: Mattie Fu Date: Sun, 18 Oct 2026 02:11:14 +0000 Subject: [PATCH 1/2] feat(bigtable): add per-method session diversion configuration Add a methodScopeKey to DivertingUnaryCallable so session load can be configured per-RPC-method via SessionDiversionConfiguration, in addition to the existing global session_load. Resolve load from global_diversion or per_scope_diversion based on which strategy the server returns. Also move the supports() check before resolveSessionLoad() to avoid a config lookup for requests the shim can't handle, and fix ClientConfigurationManager to preserve session config when a diversion configuration is present even if the global session_load is 0. Change-Id: I3925b50b9ce0bd4e7b91b17888f9745d07dd9919 --- .../data/v2/internal/compat/ShimImpl.java | 16 +++- .../compat/ops/DivertingUnaryCallable.java | 37 +++++++-- .../util/ClientConfigurationManager.java | 16 +++- .../ops/DivertingUnaryCallableTest.java | 81 ++++++++++++++++++- 4 files changed, 137 insertions(+), 13 deletions(-) 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..d6968ede6314 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,14 +376,20 @@ public UnaryCallable decorateReadRow( classic, new RowBuilderShim<>(readRowShim, rowAdapter, r -> r.hasRow() ? r.getRow() : null), Util.extractTimeout(settings), - debugTagTracer); + debugTagTracer, + "method:Bigtable.ReadRow"); } @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"); } @Override @@ -397,7 +403,8 @@ public UnaryCallable decorateCheckAndMutateRow( classic, checkAndMutateRowShim, Util.extractTimeout(settings), - debugTagTracer); + debugTagTracer, + "method:Bigtable.CheckAndMutateRow"); } @Override @@ -414,6 +421,7 @@ public UnaryCallable decorateReadModifyWriteRow new RowBuilderShim<>( readModifyWriteRowShim, rowAdapter, r -> r.hasRow() ? r.getRow() : null), Util.extractTimeout(settings), - debugTagTracer); + debugTagTracer, + "method:Bigtable.ReadModifyWriteRow"); } } 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..8cceec8d5015 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,22 @@ public class DivertingUnaryCallable extends UnaryCallable classic, UnaryShim experimental, Duration defaultTimeout, - DebugTagTracer debugTagTracer) { + DebugTagTracer debugTagTracer, + String methodScopeKey) { this.configurationManager = configurationManager; this.classic = classic; this.experimental = experimental; this.defaultTimeout = defaultTimeout; this.debugTagTracer = debugTagTracer; + this.methodScopeKey = methodScopeKey; } @Override @@ -109,13 +116,12 @@ public ApiFuture futureCall(ReqT request, ApiCallContext context) { } private boolean useExperimental(ReqT req) { - float ratio = - configurationManager.getClientConfiguration().getSessionConfiguration().getSessionLoad(); - if (ratio == 0) { + if (!experimental.supports(req)) { return false; } - if (!experimental.supports(req)) { + float ratio = resolveSessionLoad(); + if (ratio == 0) { return false; } @@ -125,6 +131,27 @@ 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 sessionConfig.getSessionLoad(); + } + 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..4f539f206835 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,6 +21,7 @@ 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.TelemetryConfiguration; import com.google.cloud.bigtable.data.v2.internal.api.ChannelProviders.ChannelProvider; @@ -446,8 +447,11 @@ 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) { + // the server-returned config has session_load=0. Preserve the config if a diversion + // configuration is present even when session_load is 0, since diversion config encodes its own + // per-scope session loads. + if (builder.getSessionConfiguration().getSessionLoad() == 0 + && !builder.getSessionConfiguration().hasSessionDiversionConfiguration()) { builder.clearSessionConfiguration(); return builder.build(); } @@ -456,7 +460,13 @@ private ClientConfiguration normalizeConfig(ClientConfiguration cfg) { } public boolean areSessionsRequired() { - return overrideConfig.map(c -> c.getSessionConfiguration().getSessionLoad() > 0).orElse(false); + return overrideConfig + .map( + c -> { + SessionClientConfiguration session = c.getSessionConfiguration(); + return session.getSessionLoad() > 0 || session.hasSessionDiversionConfiguration(); + }) + .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..ae3333223849 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,55 @@ 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 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 new DivertingUnaryCallable<>( - configManager, classic, experimental, Duration.ofSeconds(10), tracer); + configManager, classic, experimental, Duration.ofSeconds(10), tracer, "method:Bigtable.ReadRow"); } private void stubSessionLoad(float load) { @@ -126,6 +173,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; From 21d93420d2971f00c2e85807d52315d2e0fd1671 Mon Sep 17 00:00:00 2001 From: Mattie Fu Date: Sun, 18 Oct 2026 02:15:01 +0000 Subject: [PATCH 2/2] fix(bigtable): check ratio before supports() in DivertingUnaryCallable.useExperimental Change-Id: I0501f16a70e7de8c3d9b7df3cb8dcb95040ab07c --- .../data/v2/internal/compat/ops/DivertingUnaryCallable.java | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) 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 8cceec8d5015..ea531086a1c0 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 @@ -116,12 +116,12 @@ public ApiFuture futureCall(ReqT request, ApiCallContext context) { } private boolean useExperimental(ReqT req) { - if (!experimental.supports(req)) { + float ratio = resolveSessionLoad(); + if (ratio == 0) { return false; } - float ratio = resolveSessionLoad(); - if (ratio == 0) { + if (!experimental.supports(req)) { return false; }