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 5 commits into
base: main
Choose a base branch
from
+313 −45
Open

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

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
5 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
7678fac
add client context
mutianf Oct 21, 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<>(
Comment thread
mutianf marked this conversation as resolved.
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<>(
Comment thread
mutianf marked this conversation as resolved.
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;
Comment thread
mutianf marked this conversation as resolved.
default:
return 0f;
Comment thread
mutianf marked this conversation as resolved.
}
}

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?

@mutianf mutianf Oct 21, 2026

Copy link
Copy Markdown
Contributor Author

Choose a reason for hiding this comment

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

it just short circut the check in nomarlize 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()))

@nimf nimf Oct 22, 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.

divertsAnyTrafficToSessions is also used here, and areSessionsRequired is used in ShimImpl which would set setSessionRequired on feature flags. Wouldn't that be an issue?

@mutianf mutianf Oct 22, 2026

Copy link
Copy Markdown
Contributor Author

Choose a reason for hiding this comment

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

I think this is fine, because this uses override config to bypass session compatibility check. And it seems to be redundant because base feature flag already sets it to true. We should probably clean it up in a follow up PR.

.orElse(false);
}

private long getRetryDelay(int attempt) {
Expand Down
6 changes: 4 additions & 2 deletions ...d-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/stub/EnhancedBigtableStub.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -940,7 +940,8 @@ private UnaryCallable<ConditionalRowMutation, Boolean> createCheckAndMutateRowCa

return bigtableClientContext
.getSessionShim()
.decorateCheckAndMutateRow(classic, perOpSettings.checkAndMutateRowSettings);
.decorateCheckAndMutateRow(classic, perOpSettings.checkAndMutateRowSettings)
.withDefaultCallContext(bigtableClientContext.getClientContext().getDefaultCallContext());
}

/**
Expand Down Expand Up @@ -968,7 +969,8 @@ private UnaryCallable<ReadModifyWriteRow, Row> createReadModifyWriteRowCallable(

return bigtableClientContext
.getSessionShim()
.decorateReadModifyWriteRow(classic, rowAdapter, perOpSettings.readModifyWriteRowSettings);
.decorateReadModifyWriteRow(classic, rowAdapter, perOpSettings.readModifyWriteRowSettings)
.withDefaultCallContext(bigtableClientContext.getClientContext().getDefaultCallContext());
}

/**
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