Skip to content

feat(bigtable): add per-method session diversion configuration #14681

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
mutianf wants to merge 2 commits into
base: main
Choose a base branch
from
Open

feat(bigtable): add per-method session diversion configuration #14681

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
2 commits
Select commit Hold shift + click to select a range
cc0ad8f
feat(bigtable): add per-method session diversion configuration
mutianf Oct 18, 2026
21d9342
fix(bigtable): check ratio before supports() in DivertingUnaryCallabl…
mutianf Oct 18, 2026
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
16 changes: 12 additions & 4 deletions ...ud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/internal/compat/ShimImpl.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -376,14 +376,20 @@ public <RowT> UnaryCallable<Query, RowT> decorateReadRow(
classic,
new RowBuilderShim<>(readRowShim, rowAdapter, r -> r.hasRow() ? r.getRow() : null),
Util.extractTimeout(settings),
debugTagTracer);
debugTagTracer,
"method:Bigtable.ReadRow");
Comment thread
mutianf marked this conversation as resolved.
}

@Override
public UnaryCallable<RowMutation, Void> decorateMutateRow(
UnaryCallable<RowMutation, Void> classic, UnaryCallSettings<?, ?> settings) {
return new DivertingUnaryCallable<>(
configManager, classic, mutateRowShim, Util.extractTimeout(settings), debugTagTracer);
configManager,
classic,
mutateRowShim,
Util.extractTimeout(settings),
debugTagTracer,
"method:Bigtable.MutateRow");
}

@Override
Expand All @@ -397,7 +403,8 @@ public UnaryCallable<ConditionalRowMutation, Boolean> decorateCheckAndMutateRow(
classic,
checkAndMutateRowShim,
Util.extractTimeout(settings),
debugTagTracer);
debugTagTracer,
"method:Bigtable.CheckAndMutateRow");
}

@Override
Expand All @@ -414,6 +421,7 @@ public <RowT> UnaryCallable<ReadModifyWriteRow, RowT> decorateReadModifyWriteRow
new RowBuilderShim<>(
readModifyWriteRowShim, rowAdapter, r -> r.hasRow() ? r.getRow() : null),
Util.extractTimeout(settings),
debugTagTracer);
debugTagTracer,
"method:Bigtable.ReadModifyWriteRow");
}
}
33 changes: 30 additions & 3 deletions ...in/java/com/google/cloud/bigtable/data/v2/internal/compat/ops/DivertingUnaryCallable.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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;
Expand Down Expand Up @@ -49,18 +52,22 @@ public class DivertingUnaryCallable<ReqT, RespT> extends UnaryCallable<ReqT, Res

private final Duration defaultTimeout;
private final DebugTagTracer debugTagTracer;
// Scope key for per-method diversion lookup, e.g. "method:Bigtable.ReadRow".
private final String methodScopeKey;

public DivertingUnaryCallable(
ClientConfigurationManager configurationManager,
UnaryCallable<ReqT, RespT> classic,
UnaryShim<ReqT, RespT> 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
Expand Down Expand Up @@ -109,8 +116,7 @@ public ApiFuture<RespT> futureCall(ReqT request, ApiCallContext context) {
}

private boolean useExperimental(ReqT req) {
float ratio =
configurationManager.getClientConfiguration().getSessionConfiguration().getSessionLoad();
float ratio = resolveSessionLoad();
if (ratio == 0) {
return false;
}
Expand All @@ -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) {
Expand Down
16 changes: 13 additions & 3 deletions ...main/java/com/google/cloud/bigtable/data/v2/internal/util/ClientConfigurationManager.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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;
Expand Down Expand Up @@ -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();
}
Expand All @@ -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) {
Expand Down
81 changes: 80 additions & 1 deletion ...ava/com/google/cloud/bigtable/data/v2/internal/compat/ops/DivertingUnaryCallableTest.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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;
Expand Down Expand Up @@ -111,11 +114,55 @@ void sessionLoadZero_routesToClassicWithoutCredentialsCheck() throws Exception {
assertThat(tracer.tags).isEmpty();
}

@Test
void globalDiversion_usesGlobalSessionLoad() {
stubGlobalDiversion(1.0f);
DivertingUnaryCallable<String, String> 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<String, String> callable = callable();

ApiFuture<String> 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<String, String> 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<String, String> callable = callable();

ApiFuture<String> result = callable.futureCall("req", GrpcCallContext.createDefault());

assertThat(result.get()).isEqualTo("classic");
assertThat(experimental.callCount).isEqualTo(0);
}

// --- helpers ---

private DivertingUnaryCallable<String, String> 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) {
Expand All @@ -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<ReqT, RespT> extends UnaryCallable<ReqT, RespT> {
final List<ApiCallContext> capturedContexts = new ArrayList<>();
private final ApiFuture<RespT> response;
Expand Down
Loading