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

## Summary

- Adds a `methodScopeKey` to `DivertingUnaryCallable` so session load
can be configured per-RPC-method via `SessionDiversionConfiguration`, in
addition to the existing global `session_load`.
- `resolveSessionLoad()` checks `global_diversion` or
`per_scope_diversion` depending on which strategy the server returns;
falls back to flat `session_load` when no diversion config is present.
- Fixes `ClientConfigurationManager` to preserve session config when a
`SessionDiversionConfiguration` is present even if the global
`session_load` is 0 (diversion config encodes its own per-scope session
loads).

## Test plan

- [ ] `DivertingUnaryCallableTest`: covers global diversion (load=1
routes to experimental, load=0 falls back to classic), per-scope match
(routes to experimental), and per-scope non-match (falls back to
classic)
- [ ] Existing session routing tests continue to pass unchanged

Author: mutianf

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

@@ -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<>(
         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<>(
         configManager,
         classic,
         new RowBuilderShim<>(
             readModifyWriteRowShim, rowAdapter, r -> r.hasRow() ? r.getRow() : null),
         Util.extractTimeout(settings),
-        debugTagTracer);
+        debugTagTracer,
+        "method:Bigtable.ReadModifyWriteRow",
+        /* fallbackToGlobalSessionLoad= */ false);
   }
 }

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

@@ -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;
@@ -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
@@ -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;
     }
@@ -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;
+        default:
+          return 0f;
+      }
+    }
+
+    return fallbackToGlobalSessionLoad ? sessionConfig.getSessionLoad() : 0f;
+  }
+
   ApiException translateException(Throwable e) {
     Throwable cause = e;
     while (cause instanceof CompletionException || cause instanceof ExecutionException) {

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

@@ -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;
@@ -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;
+    }
+
+    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) {

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

@@ -940,7 +940,8 @@ private UnaryCallable<ConditionalRowMutation, Boolean> createCheckAndMutateRowCa
 
     return bigtableClientContext
         .getSessionShim()
-        .decorateCheckAndMutateRow(classic, perOpSettings.checkAndMutateRowSettings);
+        .decorateCheckAndMutateRow(classic, perOpSettings.checkAndMutateRowSettings)
+        .withDefaultCallContext(bigtableClientContext.getClientContext().getDefaultCallContext());
   }
 
   /**
@@ -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());
   }
 
   /**

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

@@ -12,7 +12,22 @@ polling_configuration {
 }
 
 session_configuration {
-  session_load: 0
+  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