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..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 @@ -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,8 +116,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 +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;