feat(bigtable): BigtableDataClientFactory session support

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

@@ -74,10 +74,6 @@ public final class BigtableDataClientFactory implements AutoCloseable {
    */
   public static BigtableDataClientFactory create(BigtableDataSettings defaultSettings)
       throws IOException {
-    BigtableDataSettings.Builder builder = defaultSettings.toBuilder();
-    builder.stubSettings().setSessionsEnabled(false);
-    defaultSettings = builder.build();
-
     BigtableClientContext sharedClientContext =
         BigtableClientContext.create(defaultSettings.getStubSettings());
     ClientOperationSettings perOpSettings = defaultSettings.getStubSettings().getPerOpSettings();

java-bigtable/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/internal/api/Client.java

@@ -100,7 +100,7 @@ public class Client implements AutoCloseable {
   private final BigtableTimer sessionTimer;
 
   private final CallOptions defaultCallOptions;
-  private final ChannelPool channelPool;
+  private final Resource<ChannelPool> channelPool;
   private final Resource<Metrics> metrics;
   private final Resource<ClientConfigurationManager> configManager;
 
@@ -185,22 +185,26 @@ public static Client create(ClientSettings settings) throws IOException {
     return new Client(
         featureFlags,
         clientInfo,
-        settings.getChannelProvider(),
         Resource.createOwned(metrics, metrics::close),
         Resource.createOwned(configManager, configManager::close),
         Resource.createOwned(backgroundExecutor, backgroundExecutor::shutdown),
         Resource.<Executor>createOwned(
-            userCallbackExecutor, () -> shutdownAndAwait(userCallbackExecutor)));
+          userCallbackExecutor, () -> shutdownAndAwait(userCallbackExecutor)),
+        settings.getChannelProvider());
   }
 
+  /**
+   * Standard constructor used by non-factory clients. Builds an owned {@link SwitchingChannelPool}
+   * from the provided {@link ChannelProvider}.
+   */
   public Client(
       FeatureFlags featureFlags,
       ClientInfo clientInfo,
-      ChannelProvider channelProvider,
       Resource<Metrics> metrics,
       Resource<ClientConfigurationManager> configManager,
       Resource<ScheduledExecutorService> bgExecutor,
-      Resource<Executor> userCallbackExecutor)
+      Resource<Executor> userCallbackExecutor,
+      ChannelProvider channelProvider)
       throws IOException {
     this.featureFlags = featureFlags;
     this.clientInfo = clientInfo;
@@ -224,13 +228,38 @@ public Client(
                     // TODO: consider localizing this for large reads
                     .maxInboundMessageSize(256 * 1024 * 1024));
 
-    channelPool =
+    SwitchingChannelPool switchingPool =
         new SwitchingChannelPool(
             configuredChannelProvider,
             configManager.get(),
             metrics.get(),
             backgroundExecutor.get());
-    channelPool.start();
+    switchingPool.start();
+    this.channelPool = Resource.createOwned(switchingPool, switchingPool::close);
+  }
+
+  /**
+   * Factory-child constructor. Uses a pre-built, shared {@link ChannelPool} and {@link
+   * ClientConfigurationManager}. The pool and other resources should be created with
+   * Resource.createShared() and closed when the factory closes.
+   */
+  public Client(
+      FeatureFlags featureFlags,
+      ClientInfo clientInfo,
+      Resource<Metrics> metrics,
+      Resource<ClientConfigurationManager> configManager,
+      Resource<ScheduledExecutorService> bgExecutor,
+      Resource<Executor> userCallbackExecutor,
+      Resource<ChannelPool> sharedChannelPool) {
+    this.featureFlags = featureFlags;
+    this.clientInfo = clientInfo;
+    this.metrics = metrics;
+    this.configManager = configManager;
+    this.backgroundExecutor = bgExecutor;
+    this.userCallbackExecutor = userCallbackExecutor;
+    this.channelPool = sharedChannelPool;
+    this.sessionTimer = new HashedWheelTimer("bigtable-session-timer");
+    defaultCallOptions = CallOptions.DEFAULT;
   }
 
   @Override
@@ -329,7 +358,7 @@ public TableAsync openTableAsync(String tableId, Permission permission) {
               featureFlags,
               clientInfo,
               configManager.get(),
-              channelPool,
+              channelPool.get(),
               defaultCallOptions,
               tableId,
               permission,
@@ -353,7 +382,7 @@ public AuthorizedViewAsync openAuthorizedViewAsync(
               featureFlags,
               clientInfo,
               configManager.get(),
-              channelPool,
+              channelPool.get(),
               defaultCallOptions,
               tableId,
               viewId,
@@ -378,7 +407,7 @@ public MaterializedViewAsync openMaterializedViewAsync(
               featureFlags,
               clientInfo,
               configManager.get(),
-              channelPool,
+              channelPool.get(),
               defaultCallOptions,
               viewId,
               permission,
@@ -391,6 +420,21 @@ public MaterializedViewAsync openMaterializedViewAsync(
     }
   }
 
+  /** Returns the underlying channel pool (e.g. for sharing with factory children). */
+  public ChannelPool getChannelPool() {
+    return channelPool.get();
+  }
+
+  /** Returns the feature flags (e.g. for sharing with factory children). */
+  public FeatureFlags getFeatureFlags() {
+    return featureFlags;
+  }
+
+  /** Returns the underlying user callback executor (e.g. for sharing with factory children). */
+  public Executor getUserCallbackExecutor() {
+    return userCallbackExecutor.get();
+  }
+
   public static class Resource<T> {
     private final T value;
     private final Runnable closer;

java-bigtable/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/internal/api/TableBase.java

@@ -22,6 +22,8 @@
 import com.google.bigtable.v2.SessionReadRowRequest;
 import com.google.bigtable.v2.SessionReadRowResponse;
 import com.google.cloud.bigtable.data.v2.internal.channels.ChannelPool;
+import com.google.cloud.bigtable.data.v2.internal.channels.ChannelPoolOptions;
+import com.google.cloud.bigtable.data.v2.internal.channels.TenantKey;
 import com.google.cloud.bigtable.data.v2.internal.csm.Metrics;
 import com.google.cloud.bigtable.data.v2.internal.csm.attributes.ClientInfo;
 import com.google.cloud.bigtable.data.v2.internal.csm.tracers.VRpcTracer;
@@ -66,14 +68,20 @@ static <ReqT extends Message> TableBase createAndStart(
       Executor backgroundExecutor,
       Executor userCallbackExecutor) {
 
+    // Stamp the tenant key so ChannelPoolDpImpl can make tenant-aware placement decisions.
+    CallOptions stamped =
+        callOptions.withOption(
+            ChannelPoolOptions.TENANT_KEY_OPTION,
+            new TenantKey(clientInfo.getInstanceName(), clientInfo.getAppProfileId()));
+
     SessionPool<ReqT> sessionPool =
         new SessionPoolImpl<>(
             metrics,
             featureFlags,
             clientInfo,
             configManager,
             channelPool,
-            callOptions,
+            stamped,
             sessionDescriptor,
             sessionPoolName,
             timer,