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 4 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
4 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
1d37bde
feat(bigtable): replace WipFeatures gate with per-method fallbackToGl…
mutianf Oct 19, 2026
bd1d4db
fix(bigtable): require a nonzero session load to keep session config
mutianf Oct 20, 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
26 changes: 16 additions & 10 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,44 +376,50 @@ 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",
/* fallbackToGlobalSessionLoad= */ true);
}

@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",
/* fallbackToGlobalSessionLoad= */ true);
}

@Override
public UnaryCallable<ConditionalRowMutation, Boolean> decorateCheckAndMutateRow(
UnaryCallable<ConditionalRowMutation, Boolean> classic, UnaryCallSettings<?, ?> settings) {
if (!WipFeatures.CHECK_AND_MUTATE_ROW_ENABLED) {
return classic;
}
return new DivertingUnaryCallable<>(

@nimf nimf Oct 21, 2026

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

now that we don't return classic anymore, we are missing adding withDefaultCallContext in the EnhancedBigtableStub that calls this.

configManager,
classic,
checkAndMutateRowShim,
Util.extractTimeout(settings),
debugTagTracer);
debugTagTracer,
"method:Bigtable.CheckAndMutateRow",
/* fallbackToGlobalSessionLoad= */ false);
}

@Override
public <RowT> UnaryCallable<ReadModifyWriteRow, RowT> decorateReadModifyWriteRow(
UnaryCallable<ReadModifyWriteRow, RowT> classic,
RowAdapter<RowT> rowAdapter,
UnaryCallSettings<?, ?> settings) {
if (!WipFeatures.READ_MODIFY_WRITE_ROW_ENABLED) {
return classic;
}
return new DivertingUnaryCallable<>(

@nimf nimf Oct 21, 2026

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

now that we don't return classic anymore, we are missing adding withDefaultCallContext in the EnhancedBigtableStub that calls this.

configManager,
classic,
new RowBuilderShim<>(
readModifyWriteRowShim, rowAdapter, r -> r.hasRow() ? r.getRow() : null),
Util.extractTimeout(settings),
debugTagTracer);
debugTagTracer,
"method:Bigtable.ReadModifyWriteRow",
/* fallbackToGlobalSessionLoad= */ false);
}
}
27 changes: 0 additions & 27 deletions ...bigtable/src/main/java/com/google/cloud/bigtable/data/v2/internal/compat/WipFeatures.java
View file
Open in desktop

This file was deleted.

40 changes: 37 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,28 @@ 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;
// When true, falls back to the flat SessionClientConfiguration.session_load when no
// SessionDiversionConfiguration is present. Used for ReadRow/MutateRow during the server
// rollout window before the server starts sending SessionDiversionConfiguration.
private final boolean fallbackToGlobalSessionLoad;

public DivertingUnaryCallable(
ClientConfigurationManager configurationManager,
UnaryCallable<ReqT, RespT> classic,
UnaryShim<ReqT, RespT> 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
Expand Down Expand Up @@ -109,8 +122,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 +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;

@nimf nimf Oct 20, 2026

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

so here if we don't have ReadRow or MutateRow in the map we will return 0, but the comment for the old session_load in the proto says:

// Deprecated: Prefer session_diversion_configuration. If both are provided,
// the client should apply this session_load to Bigtable.ReadRow &
// Bigtable.MutateRow, then process the session_diversion_configuration,
// overwriting any behavior established by this value.

We don't apply the old global session_load for ReadRow and MutateRow here. Is the comment stale?

default:
return 0f;

@nimf nimf Oct 21, 2026

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

We can get here only if SessionDiversionConfiguration is provided but with the strategy is not set or set to some new value in the future. Should we check fallbackToGlobalSessionLoad and fallback then?

}
}

return fallbackToGlobalSessionLoad ? sessionConfig.getSessionLoad() : 0f;
}

ApiException translateException(Throwable e) {
Throwable cause = e;
while (cause instanceof CompletionException || cause instanceof ExecutionException) {
Expand Down
29 changes: 27 additions & 2 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,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;
Expand Down Expand Up @@ -447,16 +449,39 @@ 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();
}

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;

@nimf nimf Oct 20, 2026

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

why are we returning true here immediately? What if session_load for ReadRow and MutateRow (or globally) is then configured to 0 in session diversion config?

}

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) {
Expand Down
115 changes: 114 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,89 @@ 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 fallbackEnabled_noDiversionConfig_usesGlobalSessionLoad() {
stubSessionLoad(1.0f);
DivertingUnaryCallable<String, String> 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<String, String> callable =
callable(/* fallbackToGlobalSessionLoad= */ false);

ApiFuture<String> 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<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 callable(true);
}

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