fix(bigtable): resolve session diversion through a single shared rule

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

@@ -29,6 +29,9 @@
 import com.google.bigtable.v2.GetClientConfigurationRequest;
 import com.google.bigtable.v2.PeerInfo;
 import com.google.bigtable.v2.PeerInfo.TransportType;
+import com.google.bigtable.v2.Row;
+import com.google.bigtable.v2.SessionReadModifyWriteRowResponse;
+import com.google.bigtable.v2.SessionReadRowResponse;
 import com.google.cloud.bigtable.data.v2.internal.api.ChannelProviders.ChannelProvider;
 import com.google.cloud.bigtable.data.v2.internal.api.ChannelProviders.ConfiguredChannelProvider;
 import com.google.cloud.bigtable.data.v2.internal.api.Client;
@@ -47,6 +50,7 @@
 import com.google.cloud.bigtable.data.v2.internal.csm.tracers.DirectPathCompatibleTracer;
 import com.google.cloud.bigtable.data.v2.internal.dp.DirectAccessInvestigator;
 import com.google.cloud.bigtable.data.v2.internal.util.ClientConfigurationManager;
+import com.google.cloud.bigtable.data.v2.internal.util.SessionScope;
 import com.google.cloud.bigtable.data.v2.models.ConditionalRowMutation;
 import com.google.cloud.bigtable.data.v2.models.Query;
 import com.google.cloud.bigtable.data.v2.models.ReadModifyWriteRow;
@@ -69,6 +73,7 @@
 import java.util.concurrent.Executors;
 import java.util.concurrent.ScheduledExecutorService;
 import java.util.concurrent.TimeUnit;
+import java.util.function.Function;
 import java.util.logging.Level;
 import java.util.logging.Logger;
 import javax.annotation.Nullable;
@@ -82,6 +87,19 @@ public class ShimImpl implements Shim {
 
   private static final Duration DA_CHECK_TIMEOUT = Duration.ofSeconds(5);
 
+  // A ReadRow on a missing key legitimately has no row, and the classic path surfaces that as a
+  // null Row.
+  @VisibleForTesting
+  static final Function<SessionReadRowResponse, Row> READ_ROW_EXTRACTOR =
+      r -> r.hasRow() ? r.getRow() : null;
+
+  // ReadModifyWriteRow always mutates and so always yields a row. The classic path builds a Row
+  // from the (possibly unset) response row rather than returning null, so extract unconditionally
+  // to keep both paths non-null.
+  @VisibleForTesting
+  static final Function<SessionReadModifyWriteRowResponse, Row> READ_MODIFY_WRITE_ROW_EXTRACTOR =
+      SessionReadModifyWriteRowResponse::getRow;
+
   private final ClientConfigurationManager configManager;
   private final Resource<ClientConfigurationManager> configManagerResource;
   private final Client client;
@@ -374,11 +392,10 @@ public <RowT> UnaryCallable<Query, RowT> decorateReadRow(
     return new DivertingUnaryCallable<>(
         configManager,
         classic,
-        new RowBuilderShim<>(readRowShim, rowAdapter, r -> r.hasRow() ? r.getRow() : null),
+        new RowBuilderShim<>(readRowShim, rowAdapter, READ_ROW_EXTRACTOR),
         Util.extractTimeout(settings),
         debugTagTracer,
-        "method:Bigtable.ReadRow",
-        /* fallbackToGlobalSessionLoad= */ true);
+        SessionScope.READ_ROW);
   }
 
   @Override
@@ -390,8 +407,7 @@ public UnaryCallable<RowMutation, Void> decorateMutateRow(
         mutateRowShim,
         Util.extractTimeout(settings),
         debugTagTracer,
-        "method:Bigtable.MutateRow",
-        /* fallbackToGlobalSessionLoad= */ true);
+        SessionScope.MUTATE_ROW);
   }
 
   @Override
@@ -403,8 +419,7 @@ public UnaryCallable<ConditionalRowMutation, Boolean> decorateCheckAndMutateRow(
         checkAndMutateRowShim,
         Util.extractTimeout(settings),
         debugTagTracer,
-        "method:Bigtable.CheckAndMutateRow",
-        /* fallbackToGlobalSessionLoad= */ false);
+        SessionScope.CHECK_AND_MUTATE_ROW);
   }
 
   @Override
@@ -415,11 +430,9 @@ public <RowT> UnaryCallable<ReadModifyWriteRow, RowT> decorateReadModifyWriteRow
     return new DivertingUnaryCallable<>(
         configManager,
         classic,
-        new RowBuilderShim<>(
-            readModifyWriteRowShim, rowAdapter, r -> r.hasRow() ? r.getRow() : null),
+        new RowBuilderShim<>(readModifyWriteRowShim, rowAdapter, READ_MODIFY_WRITE_ROW_EXTRACTOR),
         Util.extractTimeout(settings),
         debugTagTracer,
-        "method:Bigtable.ReadModifyWriteRow",
-        /* fallbackToGlobalSessionLoad= */ false);
+        SessionScope.READ_MODIFY_WRITE_ROW);
   }
 }

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

@@ -22,14 +22,12 @@
 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;
 import com.google.cloud.bigtable.data.v2.internal.csm.tracers.DebugTagTracer;
 import com.google.cloud.bigtable.data.v2.internal.util.ClientConfigurationManager;
+import com.google.cloud.bigtable.data.v2.internal.util.SessionScope;
 import com.google.common.base.Throwables;
 import io.grpc.Context;
 import io.grpc.Context.CancellableContext;
@@ -52,28 +50,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;
-  // 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;
+  // The diversion scope this callable's method belongs to.
+  private final SessionScope scope;
 
   public DivertingUnaryCallable(
       ClientConfigurationManager configurationManager,
       UnaryCallable<ReqT, RespT> classic,
       UnaryShim<ReqT, RespT> experimental,
       Duration defaultTimeout,
       DebugTagTracer debugTagTracer,
-      String methodScopeKey,
-      boolean fallbackToGlobalSessionLoad) {
+      SessionScope scope) {
     this.configurationManager = configurationManager;
     this.classic = classic;
     this.experimental = experimental;
     this.defaultTimeout = defaultTimeout;
     this.debugTagTracer = debugTagTracer;
-    this.methodScopeKey = methodScopeKey;
-    this.fallbackToGlobalSessionLoad = fallbackToGlobalSessionLoad;
+    this.scope = scope;
   }
 
   @Override
@@ -138,25 +130,8 @@ private boolean useExperimental(ReqT req) {
   }
 
   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 fallbackToGlobalSessionLoad ? sessionConfig.getSessionLoad() : 0f;
+    return scope.resolveSessionLoad(
+        configurationManager.getClientConfiguration().getSessionConfiguration());
   }
 
   ApiException translateException(Throwable e) {

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

@@ -23,7 +23,6 @@
 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;
@@ -458,24 +457,12 @@ private ClientConfiguration normalizeConfig(ClientConfiguration cfg) {
   }
 
   /**
-   * 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.
+   * Returns true if the configuration routes a nonzero share of traffic to sessions for any scope
+   * this client serves. Shares the resolution rule with the per-RPC diversion decision so the two
+   * cannot disagree; see {@link SessionScope}.
    */
   private static boolean divertsAnyTrafficToSessions(SessionClientConfiguration session) {
-    if (session.getSessionLoad() > 0) {
-      return true;
-    }
-
-    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;
-    }
+    return SessionScope.divertsAnyTraffic(session);
   }
 
   public boolean areSessionsRequired() {

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

@@ -0,0 +1,94 @@
+/*
+ * Copyright 2026 Google LLC
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *     https://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+package com.google.cloud.bigtable.data.v2.internal.util;
+
+import com.google.bigtable.v2.SessionClientConfiguration;
+import com.google.bigtable.v2.SessionDiversionConfiguration;
+import com.google.bigtable.v2.SessionScopeDiversionConfiguration;
+
+/**
+ * The per-method diversion scopes this client implements, along with the single rule for resolving
+ * a session load for each of them.
+ *
+ * <p>Two decisions depend on this rule: the per-RPC fork in {@code DivertingUnaryCallable}, and the
+ * client wide "are sessions in use at all" decision in {@link ClientConfigurationManager}. They
+ * must agree. If the client-wide check is the more permissive of the two, the client keeps session
+ * pools alive and advertises session feature flags while never diverting a single request. Keeping
+ * the scope list and the resolution rule in one place is what keeps them in sync.
+ */
+public enum SessionScope {
+  READ_ROW("method:Bigtable.ReadRow", /* fallbackToGlobalSessionLoad= */ true),
+  MUTATE_ROW("method:Bigtable.MutateRow", /* fallbackToGlobalSessionLoad= */ true),
+  CHECK_AND_MUTATE_ROW(
+      "method:Bigtable.CheckAndMutateRow", /* fallbackToGlobalSessionLoad= */ false),
+  READ_MODIFY_WRITE_ROW(
+      "method:Bigtable.ReadModifyWriteRow", /* fallbackToGlobalSessionLoad= */ false);
+
+  private final String key;
+  private final boolean fallbackToGlobalSessionLoad;
+
+  SessionScope(String key, boolean fallbackToGlobalSessionLoad) {
+    this.key = key;
+    this.fallbackToGlobalSessionLoad = fallbackToGlobalSessionLoad;
+  }
+
+  /** The {@code method:<service>.<method>} key the server uses to address this scope. */
+  public String getKey() {
+    return key;
+  }
+
+  /**
+   * Returns the fraction of this scope's traffic that should be served by sessions.
+   *
+   * <p>A {@link SessionDiversionConfiguration} is authoritative when present: the server only sends
+   * the deprecated flat {@code session_load} to keep ReadRow and MutateRow working on older
+   * clients, so this client ignores it entirely once a diversion configuration arrives.
+   */
+  @SuppressWarnings("deprecation")
+  public float resolveSessionLoad(SessionClientConfiguration sessionConfig) {
+    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(key);
+          return scopeConfig != null ? scopeConfig.getSessionLoad() : 0f;
+        default:
+          return 0f;
+      }
+    }
+
+    return fallbackToGlobalSessionLoad ? sessionConfig.getSessionLoad() : 0f;
+  }
+
+  /**
+   * Returns true if the configuration routes a nonzero share of traffic to sessions for at least
+   * one scope this client implements. Scopes the client does not serve are ignored: a diversion
+   * entry for them cannot produce any session traffic, so it must not keep the session machinery
+   * alive either.
+   */
+  public static boolean divertsAnyTraffic(SessionClientConfiguration sessionConfig) {
+    for (SessionScope scope : values()) {
+      if (scope.resolveSessionLoad(sessionConfig) > 0) {
+        return true;
+      }
+    }
+    return false;
+  }
+}

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

@@ -941,7 +941,11 @@ private UnaryCallable<ConditionalRowMutation, Boolean> createCheckAndMutateRowCa
     return bigtableClientContext
         .getSessionShim()
         .decorateCheckAndMutateRow(classic, perOpSettings.checkAndMutateRowSettings)
-        .withDefaultCallContext(bigtableClientContext.getClientContext().getDefaultCallContext());
+        .withDefaultCallContext(
+            bigtableClientContext
+                .getClientContext()
+                .getDefaultCallContext()
+                .withRetrySettings(perOpSettings.checkAndMutateRowSettings.getRetrySettings()));
   }
 
   /**
@@ -970,7 +974,11 @@ private UnaryCallable<ReadModifyWriteRow, Row> createReadModifyWriteRowCallable(
     return bigtableClientContext
         .getSessionShim()
         .decorateReadModifyWriteRow(classic, rowAdapter, perOpSettings.readModifyWriteRowSettings)
-        .withDefaultCallContext(bigtableClientContext.getClientContext().getDefaultCallContext());
+        .withDefaultCallContext(
+            bigtableClientContext
+                .getClientContext()
+                .getDefaultCallContext()
+                .withRetrySettings(perOpSettings.readModifyWriteRowSettings.getRetrySettings()));
   }
 
   /**

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

@@ -12,23 +12,6 @@ polling_configuration {
 }
 
 session_configuration {
-  session_diversion_configuration {
-    per_scope_diversion {
-      scope_diversions {
-        key: "method:Bigtable.ReadRow"
-        value {
-          session_load: 0
-        }
-      }
-      scope_diversions {
-        key: "method:Bigtable.MutateRow"
-        value {
-          session_load: 0
-        }
-      }
-    }
-  }
-
   channel_configuration {
     min_server_count: 2
     max_server_count: 25

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

@@ -22,6 +22,8 @@
 
 import com.google.api.gax.core.FixedExecutorProvider;
 import com.google.api.gax.grpc.InstantiatingGrpcChannelProvider;
+import com.google.bigtable.v2.SessionReadModifyWriteRowResponse;
+import com.google.bigtable.v2.SessionReadRowResponse;
 import com.google.cloud.bigtable.data.v2.internal.api.Client.Resource;
 import java.util.concurrent.CountDownLatch;
 import java.util.concurrent.Executor;
@@ -35,6 +37,26 @@ public class ShimImplTest {
   private static final InstantiatingGrpcChannelProvider EMPTY_TRANSPORT =
       InstantiatingGrpcChannelProvider.newBuilder().build();
 
+  /**
+   * ReadRow and ReadModifyWriteRow disagree on what "no row in the response" means, and the session
+   * path must match the classic path for each. A ReadRow miss is a null Row; a ReadModifyWriteRow
+   * always produces one, so a response without a row must still build a Row rather than handing
+   * user code a null the classic path never produces.
+   */
+  @Test
+  public void readRowExtractor_noRow_returnsNull() {
+    assertThat(ShimImpl.READ_ROW_EXTRACTOR.apply(SessionReadRowResponse.getDefaultInstance()))
+        .isNull();
+  }
+
+  @Test
+  public void readModifyWriteRowExtractor_noRow_returnsEmptyRowNotNull() {
+    assertThat(
+            ShimImpl.READ_MODIFY_WRITE_ROW_EXTRACTOR.apply(
+                SessionReadModifyWriteRowResponse.getDefaultInstance()))
+        .isNotNull();
+  }
+
   @Test
   public void selectUserCallbackExecutor_transportSet_returnsSharedTransportExecutor() {
     Executor userExec = mock(Executor.class);