Skip to content

Commit 00f7bea

Browse files
authored
fix(bigtable): resolve session diversion through a single shared rule (#14738)
Follow-up to #14681, from reviewing it after it merged. Three independent fixes. ### 1. Two drifting resolution rules (the main one) `DivertingUnaryCallable.resolveSessionLoad` (per RPC) and `ClientConfigurationManager.divertsAnyTrafficToSessions` (client-wide) answered "does traffic go to sessions?" with separately maintained rules that had already diverged: | | per-RPC | client-wide | |---|---|---| | diversion config present | authoritative, flat load ignored | ORs with flat `session_load` | | per-scope map | only this method's key | **any** key, incl. methods the client never serves | So `per_scope_diversion { "method:Bigtable.MutateRows": {session_load: 1.0} }` made `areSessionsRequired()` true — flipping the `sessions_required` / `sessions_compatible` feature flags sent to the server and keeping the session config and channel pool alive — while `DivertingUnaryCallable` diverted exactly zero requests, because no callable asks for that scope. Same for `session_load: 1.0` plus a CheckAndMutateRow-only map. Both now go through a new `SessionScope` enum, which owns the four scopes this client serves, whether each honours the deprecated flat `session_load`, and the single resolution rule (moved verbatim out of `DivertingUnaryCallable`). `ShimImpl` passes `SessionScope.READ_ROW` instead of a `String` + `boolean` pair. ### 2. The default config shadowed the flat `session_load` `bigtable-default-client-config.textproto` carried a zero-valued `per_scope_diversion` for ReadRow and MutateRow. Since a diversion configuration is authoritative once present, those entries forced **every** scope to 0 in any config derived from the default: - the initial `normalizeConfig(defaultConfig)` before the first `GetClientConfiguration` response, - the `handleFailedFetch` → `setClientConfiguration(defaultConfig)` path when the validity window expires during a config-service outage, - anything merged from the `bigtable.internal.client-config-override` sys-prop, including the `bigtable-session-it` Maven profile. This surfaced as 10 `ClientConfigurationManagerTest` failures the moment fix 1 made the two checks agree — they had been passing only because the client-wide check disagreed and reported sessions as on. The entries were purely documentary (an absent key and an absent diversion config both resolve to 0), so they are removed. ### 3. ReadModifyWriteRow could return `null` on the session path The extractor was `r -> r.hasRow() ? r.getRow() : null`, so a `SessionReadModifyWriteRowResponse` without a row handed user code a `null` the classic path never produces (`EnhancedBigtableStub` always builds a `Row` from the possibly-unset `resp.getRow()`). Unlike ReadRow — where a missing key legitimately means `null` — ReadModifyWriteRow always mutates and always yields a row, so it now extracts unconditionally. Both extractors are named constants with tests pinning the split. Also forwards the per-method retry settings to the default call context for CheckAndMutateRow and ReadModifyWriteRow, matching the existing ReadRow and MutateRow wiring. ### Testing New `SessionScopeTest` (7 cases) covers the resolution rule and the agreement invariant; new cases in `ShimImplTest` pin the ReadRow-null / ReadModifyWriteRow-non-null split, and `ClientConfigurationManagerTest` gains one asserting the bundled default never shadows a flat `session_load`. `DivertingUnaryCallableTest`, `ClientConfigurationManagerTest`, `ShimImplTest`, `SessionDeadlineTest`, `SessionScopeTest`, `RowBuilderShimTest`, `ReadModifyWriteRowShimTest` — 57 tests, all passing.
1 parent afae40e commit 00f7bea

10 files changed

Lines changed: 326 additions & 94 deletions

File tree

‎java-bigtable/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/internal/compat/ShimImpl.java‎

Lines changed: 24 additions & 11 deletions
Original file line numberDiff line numberDiff line change
@@ -29,6 +29,9 @@
2929
import com.google.bigtable.v2.GetClientConfigurationRequest;
3030
import com.google.bigtable.v2.PeerInfo;
3131
import com.google.bigtable.v2.PeerInfo.TransportType;
32+
import com.google.bigtable.v2.Row;
33+
import com.google.bigtable.v2.SessionReadModifyWriteRowResponse;
34+
import com.google.bigtable.v2.SessionReadRowResponse;
3235
import com.google.cloud.bigtable.data.v2.internal.api.ChannelProviders.ChannelProvider;
3336
import com.google.cloud.bigtable.data.v2.internal.api.ChannelProviders.ConfiguredChannelProvider;
3437
import com.google.cloud.bigtable.data.v2.internal.api.Client;
@@ -47,6 +50,7 @@
4750
import com.google.cloud.bigtable.data.v2.internal.csm.tracers.DirectPathCompatibleTracer;
4851
import com.google.cloud.bigtable.data.v2.internal.dp.DirectAccessInvestigator;
4952
import com.google.cloud.bigtable.data.v2.internal.util.ClientConfigurationManager;
53+
import com.google.cloud.bigtable.data.v2.internal.util.SessionScope;
5054
import com.google.cloud.bigtable.data.v2.models.ConditionalRowMutation;
5155
import com.google.cloud.bigtable.data.v2.models.Query;
5256
import com.google.cloud.bigtable.data.v2.models.ReadModifyWriteRow;
@@ -69,6 +73,7 @@
6973
import java.util.concurrent.Executors;
7074
import java.util.concurrent.ScheduledExecutorService;
7175
import java.util.concurrent.TimeUnit;
76+
import java.util.function.Function;
7277
import java.util.logging.Level;
7378
import java.util.logging.Logger;
7479
import javax.annotation.Nullable;
@@ -82,6 +87,19 @@ public class ShimImpl implements Shim {
8287

8388
private static final Duration DA_CHECK_TIMEOUT = Duration.ofSeconds(5);
8489

90+
// A ReadRow on a missing key legitimately has no row, and the classic path surfaces that as a
91+
// null Row.
92+
@VisibleForTesting
93+
static final Function<SessionReadRowResponse, Row> READ_ROW_EXTRACTOR =
94+
r -> r.hasRow() ? r.getRow() : null;
95+
96+
// ReadModifyWriteRow always mutates and so always yields a row. The classic path builds a Row
97+
// from the (possibly unset) response row rather than returning null, so extract unconditionally
98+
// to keep both paths non-null.
99+
@VisibleForTesting
100+
static final Function<SessionReadModifyWriteRowResponse, Row> READ_MODIFY_WRITE_ROW_EXTRACTOR =
101+
SessionReadModifyWriteRowResponse::getRow;
102+
85103
private final ClientConfigurationManager configManager;
86104
private final Resource<ClientConfigurationManager> configManagerResource;
87105
private final Client client;
@@ -374,11 +392,10 @@ public <RowT> UnaryCallable<Query, RowT> decorateReadRow(
374392
return new DivertingUnaryCallable<>(
375393
configManager,
376394
classic,
377-
new RowBuilderShim<>(readRowShim, rowAdapter, r -> r.hasRow() ? r.getRow() : null),
395+
new RowBuilderShim<>(readRowShim, rowAdapter, READ_ROW_EXTRACTOR),
378396
Util.extractTimeout(settings),
379397
debugTagTracer,
380-
"method:Bigtable.ReadRow",
381-
/* fallbackToGlobalSessionLoad= */ true);
398+
SessionScope.READ_ROW);
382399
}
383400

384401
@Override
@@ -390,8 +407,7 @@ public UnaryCallable<RowMutation, Void> decorateMutateRow(
390407
mutateRowShim,
391408
Util.extractTimeout(settings),
392409
debugTagTracer,
393-
"method:Bigtable.MutateRow",
394-
/* fallbackToGlobalSessionLoad= */ true);
410+
SessionScope.MUTATE_ROW);
395411
}
396412

397413
@Override
@@ -403,8 +419,7 @@ public UnaryCallable<ConditionalRowMutation, Boolean> decorateCheckAndMutateRow(
403419
checkAndMutateRowShim,
404420
Util.extractTimeout(settings),
405421
debugTagTracer,
406-
"method:Bigtable.CheckAndMutateRow",
407-
/* fallbackToGlobalSessionLoad= */ false);
422+
SessionScope.CHECK_AND_MUTATE_ROW);
408423
}
409424

410425
@Override
@@ -415,11 +430,9 @@ public <RowT> UnaryCallable<ReadModifyWriteRow, RowT> decorateReadModifyWriteRow
415430
return new DivertingUnaryCallable<>(
416431
configManager,
417432
classic,
418-
new RowBuilderShim<>(
419-
readModifyWriteRowShim, rowAdapter, r -> r.hasRow() ? r.getRow() : null),
433+
new RowBuilderShim<>(readModifyWriteRowShim, rowAdapter, READ_MODIFY_WRITE_ROW_EXTRACTOR),
420434
Util.extractTimeout(settings),
421435
debugTagTracer,
422-
"method:Bigtable.ReadModifyWriteRow",
423-
/* fallbackToGlobalSessionLoad= */ false);
436+
SessionScope.READ_MODIFY_WRITE_ROW);
424437
}
425438
}

‎java-bigtable/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/internal/compat/ops/DivertingUnaryCallable.java‎

Lines changed: 7 additions & 32 deletions
Original file line numberDiff line numberDiff line change
@@ -22,14 +22,12 @@
2222
import com.google.api.gax.rpc.ApiException;
2323
import com.google.api.gax.rpc.ApiExceptionFactory;
2424
import com.google.api.gax.rpc.UnaryCallable;
25-
import com.google.bigtable.v2.SessionClientConfiguration;
26-
import com.google.bigtable.v2.SessionDiversionConfiguration;
27-
import com.google.bigtable.v2.SessionScopeDiversionConfiguration;
2825
import com.google.bigtable.v2.TelemetryConfiguration;
2926
import com.google.cloud.bigtable.data.v2.internal.compat.FutureAdapter;
3027
import com.google.cloud.bigtable.data.v2.internal.compat.Util;
3128
import com.google.cloud.bigtable.data.v2.internal.csm.tracers.DebugTagTracer;
3229
import com.google.cloud.bigtable.data.v2.internal.util.ClientConfigurationManager;
30+
import com.google.cloud.bigtable.data.v2.internal.util.SessionScope;
3331
import com.google.common.base.Throwables;
3432
import io.grpc.Context;
3533
import io.grpc.Context.CancellableContext;
@@ -52,28 +50,22 @@ public class DivertingUnaryCallable<ReqT, RespT> extends UnaryCallable<ReqT, Res
5250

5351
private final Duration defaultTimeout;
5452
private final DebugTagTracer debugTagTracer;
55-
// Scope key for per-method diversion lookup, e.g. "method:Bigtable.ReadRow".
56-
private final String methodScopeKey;
57-
// When true, falls back to the flat SessionClientConfiguration.session_load when no
58-
// SessionDiversionConfiguration is present. Used for ReadRow/MutateRow during the server
59-
// rollout window before the server starts sending SessionDiversionConfiguration.
60-
private final boolean fallbackToGlobalSessionLoad;
53+
// The diversion scope this callable's method belongs to.
54+
private final SessionScope scope;
6155

6256
public DivertingUnaryCallable(
6357
ClientConfigurationManager configurationManager,
6458
UnaryCallable<ReqT, RespT> classic,
6559
UnaryShim<ReqT, RespT> experimental,
6660
Duration defaultTimeout,
6761
DebugTagTracer debugTagTracer,
68-
String methodScopeKey,
69-
boolean fallbackToGlobalSessionLoad) {
62+
SessionScope scope) {
7063
this.configurationManager = configurationManager;
7164
this.classic = classic;
7265
this.experimental = experimental;
7366
this.defaultTimeout = defaultTimeout;
7467
this.debugTagTracer = debugTagTracer;
75-
this.methodScopeKey = methodScopeKey;
76-
this.fallbackToGlobalSessionLoad = fallbackToGlobalSessionLoad;
68+
this.scope = scope;
7769
}
7870

7971
@Override
@@ -138,25 +130,8 @@ private boolean useExperimental(ReqT req) {
138130
}
139131

140132
private float resolveSessionLoad() {
141-
SessionClientConfiguration sessionConfig =
142-
configurationManager.getClientConfiguration().getSessionConfiguration();
143-
144-
if (sessionConfig.hasSessionDiversionConfiguration()) {
145-
SessionDiversionConfiguration diversionConfig =
146-
sessionConfig.getSessionDiversionConfiguration();
147-
switch (diversionConfig.getDiversionStrategyCase()) {
148-
case GLOBAL_DIVERSION:
149-
return diversionConfig.getGlobalDiversion().getSessionLoad();
150-
case PER_SCOPE_DIVERSION:
151-
SessionScopeDiversionConfiguration scopeConfig =
152-
diversionConfig.getPerScopeDiversion().getScopeDiversionsMap().get(methodScopeKey);
153-
return scopeConfig != null ? scopeConfig.getSessionLoad() : 0f;
154-
default:
155-
return 0f;
156-
}
157-
}
158-
159-
return fallbackToGlobalSessionLoad ? sessionConfig.getSessionLoad() : 0f;
133+
return scope.resolveSessionLoad(
134+
configurationManager.getClientConfiguration().getSessionConfiguration());
160135
}
161136

162137
ApiException translateException(Throwable e) {

‎java-bigtable/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/internal/util/ClientConfigurationManager.java‎

Lines changed: 4 additions & 17 deletions
Original file line numberDiff line numberDiff line change
@@ -23,7 +23,6 @@
2323
import com.google.bigtable.v2.GetClientConfigurationRequest;
2424
import com.google.bigtable.v2.SessionClientConfiguration;
2525
import com.google.bigtable.v2.SessionClientConfiguration.ChannelPoolConfiguration.DirectAccessOnly;
26-
import com.google.bigtable.v2.SessionDiversionConfiguration;
2726
import com.google.bigtable.v2.TelemetryConfiguration;
2827
import com.google.cloud.bigtable.data.v2.internal.api.ChannelProviders.ChannelProvider;
2928
import com.google.cloud.bigtable.data.v2.internal.api.Util;
@@ -458,24 +457,12 @@ private ClientConfiguration normalizeConfig(ClientConfiguration cfg) {
458457
}
459458

460459
/**
461-
* Returns true if the configuration routes a nonzero share of traffic to sessions, either via the
462-
* flat {@code session_load} or via a diversion configuration that carries a nonzero load.
460+
* Returns true if the configuration routes a nonzero share of traffic to sessions for any scope
461+
* this client serves. Shares the resolution rule with the per-RPC diversion decision so the two
462+
* cannot disagree; see {@link SessionScope}.
463463
*/
464464
private static boolean divertsAnyTrafficToSessions(SessionClientConfiguration session) {
465-
if (session.getSessionLoad() > 0) {
466-
return true;
467-
}
468-
469-
SessionDiversionConfiguration diversion = session.getSessionDiversionConfiguration();
470-
switch (diversion.getDiversionStrategyCase()) {
471-
case GLOBAL_DIVERSION:
472-
return diversion.getGlobalDiversion().getSessionLoad() > 0;
473-
case PER_SCOPE_DIVERSION:
474-
return diversion.getPerScopeDiversion().getScopeDiversionsMap().values().stream()
475-
.anyMatch(scope -> scope.getSessionLoad() > 0);
476-
default:
477-
return false;
478-
}
465+
return SessionScope.divertsAnyTraffic(session);
479466
}
480467

481468
public boolean areSessionsRequired() {
Lines changed: 94 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,94 @@
1+
/*
2+
* Copyright 2026 Google LLC
3+
*
4+
* Licensed under the Apache License, Version 2.0 (the "License");
5+
* you may not use this file except in compliance with the License.
6+
* You may obtain a copy of the License at
7+
*
8+
* https://www.apache.org/licenses/LICENSE-2.0
9+
*
10+
* Unless required by applicable law or agreed to in writing, software
11+
* distributed under the License is distributed on an "AS IS" BASIS,
12+
* WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
13+
* See the License for the specific language governing permissions and
14+
* limitations under the License.
15+
*/
16+
package com.google.cloud.bigtable.data.v2.internal.util;
17+
18+
import com.google.bigtable.v2.SessionClientConfiguration;
19+
import com.google.bigtable.v2.SessionDiversionConfiguration;
20+
import com.google.bigtable.v2.SessionScopeDiversionConfiguration;
21+
22+
/**
23+
* The per-method diversion scopes this client implements, along with the single rule for resolving
24+
* a session load for each of them.
25+
*
26+
* <p>Two decisions depend on this rule: the per-RPC fork in {@code DivertingUnaryCallable}, and the
27+
* client wide "are sessions in use at all" decision in {@link ClientConfigurationManager}. They
28+
* must agree. If the client-wide check is the more permissive of the two, the client keeps session
29+
* pools alive and advertises session feature flags while never diverting a single request. Keeping
30+
* the scope list and the resolution rule in one place is what keeps them in sync.
31+
*/
32+
public enum SessionScope {
33+
READ_ROW("method:Bigtable.ReadRow", /* fallbackToGlobalSessionLoad= */ true),
34+
MUTATE_ROW("method:Bigtable.MutateRow", /* fallbackToGlobalSessionLoad= */ true),
35+
CHECK_AND_MUTATE_ROW(
36+
"method:Bigtable.CheckAndMutateRow", /* fallbackToGlobalSessionLoad= */ false),
37+
READ_MODIFY_WRITE_ROW(
38+
"method:Bigtable.ReadModifyWriteRow", /* fallbackToGlobalSessionLoad= */ false);
39+
40+
private final String key;
41+
private final boolean fallbackToGlobalSessionLoad;
42+
43+
SessionScope(String key, boolean fallbackToGlobalSessionLoad) {
44+
this.key = key;
45+
this.fallbackToGlobalSessionLoad = fallbackToGlobalSessionLoad;
46+
}
47+
48+
/** The {@code method:<service>.<method>} key the server uses to address this scope. */
49+
public String getKey() {
50+
return key;
51+
}
52+
53+
/**
54+
* Returns the fraction of this scope's traffic that should be served by sessions.
55+
*
56+
* <p>A {@link SessionDiversionConfiguration} is authoritative when present: the server only sends
57+
* the deprecated flat {@code session_load} to keep ReadRow and MutateRow working on older
58+
* clients, so this client ignores it entirely once a diversion configuration arrives.
59+
*/
60+
@SuppressWarnings("deprecation")
61+
public float resolveSessionLoad(SessionClientConfiguration sessionConfig) {
62+
if (sessionConfig.hasSessionDiversionConfiguration()) {
63+
SessionDiversionConfiguration diversionConfig =
64+
sessionConfig.getSessionDiversionConfiguration();
65+
switch (diversionConfig.getDiversionStrategyCase()) {
66+
case GLOBAL_DIVERSION:
67+
return diversionConfig.getGlobalDiversion().getSessionLoad();
68+
case PER_SCOPE_DIVERSION:
69+
SessionScopeDiversionConfiguration scopeConfig =
70+
diversionConfig.getPerScopeDiversion().getScopeDiversionsMap().get(key);
71+
return scopeConfig != null ? scopeConfig.getSessionLoad() : 0f;
72+
default:
73+
return 0f;
74+
}
75+
}
76+
77+
return fallbackToGlobalSessionLoad ? sessionConfig.getSessionLoad() : 0f;
78+
}
79+
80+
/**
81+
* Returns true if the configuration routes a nonzero share of traffic to sessions for at least
82+
* one scope this client implements. Scopes the client does not serve are ignored: a diversion
83+
* entry for them cannot produce any session traffic, so it must not keep the session machinery
84+
* alive either.
85+
*/
86+
public static boolean divertsAnyTraffic(SessionClientConfiguration sessionConfig) {
87+
for (SessionScope scope : values()) {
88+
if (scope.resolveSessionLoad(sessionConfig) > 0) {
89+
return true;
90+
}
91+
}
92+
return false;
93+
}
94+
}

‎java-bigtable/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/stub/EnhancedBigtableStub.java‎

Lines changed: 10 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -941,7 +941,11 @@ private UnaryCallable<ConditionalRowMutation, Boolean> createCheckAndMutateRowCa
941941
return bigtableClientContext
942942
.getSessionShim()
943943
.decorateCheckAndMutateRow(classic, perOpSettings.checkAndMutateRowSettings)
944-
.withDefaultCallContext(bigtableClientContext.getClientContext().getDefaultCallContext());
944+
.withDefaultCallContext(
945+
bigtableClientContext
946+
.getClientContext()
947+
.getDefaultCallContext()
948+
.withRetrySettings(perOpSettings.checkAndMutateRowSettings.getRetrySettings()));
945949
}
946950

947951
/**
@@ -970,7 +974,11 @@ private UnaryCallable<ReadModifyWriteRow, Row> createReadModifyWriteRowCallable(
970974
return bigtableClientContext
971975
.getSessionShim()
972976
.decorateReadModifyWriteRow(classic, rowAdapter, perOpSettings.readModifyWriteRowSettings)
973-
.withDefaultCallContext(bigtableClientContext.getClientContext().getDefaultCallContext());
977+
.withDefaultCallContext(
978+
bigtableClientContext
979+
.getClientContext()
980+
.getDefaultCallContext()
981+
.withRetrySettings(perOpSettings.readModifyWriteRowSettings.getRetrySettings()));
974982
}
975983

976984
/**

‎java-bigtable/google-cloud-bigtable/src/main/resources/bigtable-default-client-config.textproto‎

Lines changed: 0 additions & 17 deletions
Original file line numberDiff line numberDiff line change
@@ -12,23 +12,6 @@ polling_configuration {
1212
}
1313

1414
session_configuration {
15-
session_diversion_configuration {
16-
per_scope_diversion {
17-
scope_diversions {
18-
key: "method:Bigtable.ReadRow"
19-
value {
20-
session_load: 0
21-
}
22-
}
23-
scope_diversions {
24-
key: "method:Bigtable.MutateRow"
25-
value {
26-
session_load: 0
27-
}
28-
}
29-
}
30-
}
31-
3215
channel_configuration {
3316
min_server_count: 2
3417
max_server_count: 25

‎java-bigtable/google-cloud-bigtable/src/test/java/com/google/cloud/bigtable/data/v2/internal/compat/ShimImplTest.java‎

Lines changed: 22 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -22,6 +22,8 @@
2222

2323
import com.google.api.gax.core.FixedExecutorProvider;
2424
import com.google.api.gax.grpc.InstantiatingGrpcChannelProvider;
25+
import com.google.bigtable.v2.SessionReadModifyWriteRowResponse;
26+
import com.google.bigtable.v2.SessionReadRowResponse;
2527
import com.google.cloud.bigtable.data.v2.internal.api.Client.Resource;
2628
import java.util.concurrent.CountDownLatch;
2729
import java.util.concurrent.Executor;
@@ -35,6 +37,26 @@ public class ShimImplTest {
3537
private static final InstantiatingGrpcChannelProvider EMPTY_TRANSPORT =
3638
InstantiatingGrpcChannelProvider.newBuilder().build();
3739

40+
/**
41+
* ReadRow and ReadModifyWriteRow disagree on what "no row in the response" means, and the session
42+
* path must match the classic path for each. A ReadRow miss is a null Row; a ReadModifyWriteRow
43+
* always produces one, so a response without a row must still build a Row rather than handing
44+
* user code a null the classic path never produces.
45+
*/
46+
@Test
47+
public void readRowExtractor_noRow_returnsNull() {
48+
assertThat(ShimImpl.READ_ROW_EXTRACTOR.apply(SessionReadRowResponse.getDefaultInstance()))
49+
.isNull();
50+
}
51+
52+
@Test
53+
public void readModifyWriteRowExtractor_noRow_returnsEmptyRowNotNull() {
54+
assertThat(
55+
ShimImpl.READ_MODIFY_WRITE_ROW_EXTRACTOR.apply(
56+
SessionReadModifyWriteRowResponse.getDefaultInstance()))
57+
.isNotNull();
58+
}
59+
3860
@Test
3961
public void selectUserCallbackExecutor_transportSet_returnsSharedTransportExecutor() {
4062
Executor userExec = mock(Executor.class);

0 commit comments

Comments
 (0)