chore: move tracer creation to bigtable client context (#2797)

- BigtableClientContext stores the resource anchor (instance name/app profile)
- BigtableClientContext stores the opencensus deps
- This allows it to create all of the tracers during initialization
- BigtableClientContext can create lightweight copies of itself for BigtableDataClientFactory
- close logic of shared contexts is now stored in BigtableClientContext via isChild

Author: igorbernstein2

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

@@ -187,8 +187,7 @@ public static BigtableDataClient create(BigtableDataSettings settings) throws IO
    */
   static BigtableDataClient createWithClientContext(
       BigtableDataSettings settings, BigtableClientContext context) throws IOException {
-    EnhancedBigtableStub stub =
-        EnhancedBigtableStub.createWithClientContext(settings.getStubSettings(), context);
+    EnhancedBigtableStub stub = new EnhancedBigtableStub(settings.getStubSettings(), context);
     return new BigtableDataClient(stub);
   }
 

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

@@ -16,9 +16,8 @@
 package com.google.cloud.bigtable.data.v2;
 
 import com.google.api.core.BetaApi;
-import com.google.api.gax.rpc.ClientContext;
+import com.google.bigtable.v2.InstanceName;
 import com.google.cloud.bigtable.data.v2.stub.BigtableClientContext;
-import com.google.cloud.bigtable.data.v2.stub.EnhancedBigtableStub;
 import java.io.IOException;
 import javax.annotation.Nonnull;
 
@@ -75,8 +74,7 @@ public final class BigtableDataClientFactory implements AutoCloseable {
   public static BigtableDataClientFactory create(BigtableDataSettings defaultSettings)
       throws IOException {
     BigtableClientContext sharedClientContext =
-        EnhancedBigtableStub.createBigtableClientContext(defaultSettings.getStubSettings());
-
+        BigtableClientContext.create(defaultSettings.getStubSettings());
     return new BigtableDataClientFactory(sharedClientContext, defaultSettings);
   }
 
@@ -107,17 +105,11 @@ public void close() throws Exception {
    */
   public BigtableDataClient createDefault() {
     try {
-      ClientContext clientContext =
-          sharedClientContext.getClientContext().toBuilder()
-              .setTracerFactory(
-                  EnhancedBigtableStub.createBigtableTracerFactory(
-                      defaultSettings.getStubSettings(),
-                      sharedClientContext.getBuiltinOpenTelemetry(),
-                      sharedClientContext.getUserOpenTelemetry()))
-              .build();
+      BigtableClientContext ctx =
+          sharedClientContext.createChild(
+              sharedClientContext.getInstanceName(), sharedClientContext.getAppProfileId());
 
-      return BigtableDataClient.createWithClientContext(
-          defaultSettings, sharedClientContext.withClientContext(clientContext));
+      return BigtableDataClient.createWithClientContext(defaultSettings, ctx);
     } catch (IOException e) {
       // Should never happen because the connection has been established already
       throw new RuntimeException(
@@ -137,17 +129,12 @@ public BigtableDataClient createDefault() {
   public BigtableDataClient createForAppProfile(@Nonnull String appProfileId) throws IOException {
     BigtableDataSettings settings =
         defaultSettings.toBuilder().setAppProfileId(appProfileId).build();
+    BigtableClientContext ctx =
+        sharedClientContext.createChild(
+            InstanceName.of(settings.getProjectId(), settings.getInstanceId()),
+            settings.getAppProfileId());
 
-    ClientContext clientContext =
-        sharedClientContext.getClientContext().toBuilder()
-            .setTracerFactory(
-                EnhancedBigtableStub.createBigtableTracerFactory(
-                    settings.getStubSettings(),
-                    sharedClientContext.getBuiltinOpenTelemetry(),
-                    sharedClientContext.getUserOpenTelemetry()))
-            .build();
-    return BigtableDataClient.createWithClientContext(
-        settings, sharedClientContext.withClientContext(clientContext));
+    return BigtableDataClient.createWithClientContext(settings, ctx);
   }
 
   /**
@@ -167,18 +154,12 @@ public BigtableDataClient createForInstance(@Nonnull String projectId, @Nonnull
             .setInstanceId(instanceId)
             .setDefaultAppProfileId()
             .build();
+    BigtableClientContext ctx =
+        sharedClientContext.createChild(
+            InstanceName.of(settings.getProjectId(), settings.getInstanceId()),
+            settings.getAppProfileId());
 
-    ClientContext clientContext =
-        sharedClientContext.getClientContext().toBuilder()
-            .setTracerFactory(
-                EnhancedBigtableStub.createBigtableTracerFactory(
-                    settings.getStubSettings(),
-                    sharedClientContext.getBuiltinOpenTelemetry(),
-                    sharedClientContext.getUserOpenTelemetry()))
-            .build();
-
-    return BigtableDataClient.createWithClientContext(
-        settings, sharedClientContext.withClientContext(clientContext));
+    return BigtableDataClient.createWithClientContext(settings, ctx);
   }
 
   /**
@@ -199,15 +180,11 @@ public BigtableDataClient createForInstance(
             .setInstanceId(instanceId)
             .setAppProfileId(appProfileId)
             .build();
-    ClientContext clientContext =
-        sharedClientContext.getClientContext().toBuilder()
-            .setTracerFactory(
-                EnhancedBigtableStub.createBigtableTracerFactory(
-                    settings.getStubSettings(),
-                    sharedClientContext.getBuiltinOpenTelemetry(),
-                    sharedClientContext.getUserOpenTelemetry()))
-            .build();
-    return BigtableDataClient.createWithClientContext(
-        settings, sharedClientContext.withClientContext(clientContext));
+    BigtableClientContext ctx =
+        sharedClientContext.createChild(
+            InstanceName.of(settings.getProjectId(), settings.getInstanceId()),
+            settings.getAppProfileId());
+
+    return BigtableDataClient.createWithClientContext(settings, ctx);
   }
 }

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

@@ -24,18 +24,25 @@
 import com.google.api.gax.core.FixedExecutorProvider;
 import com.google.api.gax.grpc.InstantiatingGrpcChannelProvider;
 import com.google.api.gax.rpc.ClientContext;
+import com.google.api.gax.tracing.ApiTracerFactory;
 import com.google.auth.Credentials;
 import com.google.auth.oauth2.ServiceAccountJwtAccessCredentials;
 import com.google.bigtable.v2.InstanceName;
 import com.google.cloud.bigtable.data.v2.internal.JwtCredentialsWithAudience;
 import com.google.cloud.bigtable.data.v2.stub.metrics.BuiltinMetricsConstants;
 import com.google.cloud.bigtable.data.v2.stub.metrics.ChannelPoolMetricsTracer;
+import com.google.cloud.bigtable.data.v2.stub.metrics.CompositeTracerFactory;
 import com.google.cloud.bigtable.data.v2.stub.metrics.CustomOpenTelemetryMetricsProvider;
 import com.google.cloud.bigtable.data.v2.stub.metrics.Util;
 import com.google.cloud.bigtable.gaxx.grpc.BigtableTransportChannelProvider;
 import com.google.cloud.bigtable.gaxx.grpc.ChannelPrimer;
+import com.google.common.collect.ImmutableList;
 import io.grpc.ManagedChannelBuilder;
 import io.grpc.opentelemetry.GrpcOpenTelemetry;
+import io.opencensus.stats.Stats;
+import io.opencensus.stats.StatsRecorder;
+import io.opencensus.tags.Tagger;
+import io.opencensus.tags.Tags;
 import io.opentelemetry.api.OpenTelemetry;
 import io.opentelemetry.sdk.OpenTelemetrySdk;
 import java.io.IOException;
@@ -55,15 +62,30 @@ public class BigtableClientContext {
 
   private static final Logger logger = Logger.getLogger(BigtableClientContext.class.getName());
 
+  private final boolean isChild;
+  private final InstanceName instanceName;
+  private final String appProfileId;
+  private final ApiTracerFactory userTracerFactory;
   @Nullable private final OpenTelemetrySdk builtinOpenTelemetry;
   @Nullable private final OpenTelemetry userOpenTelemetry;
   private final ClientContext clientContext;
   // the background executor shared for OTEL instances and monitoring client and all other
   // background tasks
   private final ExecutorProvider backgroundExecutorProvider;
+  private final Tagger ocTagger;
+  private final StatsRecorder ocRecorder;
 
   public static BigtableClientContext create(EnhancedBigtableStubSettings settings)
       throws IOException {
+    return create(settings, Tags.getTagger(), Stats.getStatsRecorder());
+  }
+
+  public static BigtableClientContext create(
+      EnhancedBigtableStubSettings settings, Tagger ocTagger, StatsRecorder ocRecorder)
+      throws IOException {
+    InstanceName instanceName = InstanceName.of(settings.getProjectId(), settings.getInstanceId());
+    String appProfileId = settings.getAppProfileId();
+
     EnhancedBigtableStubSettings.Builder builder = settings.toBuilder();
 
     // Set up credentials
@@ -85,6 +107,8 @@ public static BigtableClientContext create(EnhancedBigtableStubSettings settings
         FixedExecutorProvider.create(backgroundExecutor, shouldAutoClose);
     builder.setBackgroundExecutorProvider(executorProvider);
 
+    ApiTracerFactory userTracerFactory = settings.getTracerFactory();
+
     // Set up OpenTelemetry
     @Nullable OpenTelemetry userOtel = null;
     if (settings.getMetricsProvider() instanceof CustomOpenTelemetryMetricsProvider) {
@@ -158,7 +182,17 @@ public static BigtableClientContext create(EnhancedBigtableStubSettings settings
       channelPoolMetricsTracer.start(clientContext.getExecutor());
     }
 
-    return new BigtableClientContext(clientContext, builtinOtel, userOtel, executorProvider);
+    return new BigtableClientContext(
+        false,
+        instanceName,
+        appProfileId,
+        clientContext,
+        userTracerFactory,
+        builtinOtel,
+        userOtel,
+        ocTagger,
+        ocRecorder,
+        executorProvider);
   }
 
   private static void configureGrpcOtel(
@@ -189,14 +223,53 @@ private static void configureGrpcOtel(
   }
 
   private BigtableClientContext(
+      boolean isChild,
+      InstanceName instanceName,
+      String appProfileId,
       ClientContext clientContext,
-      @Nullable OpenTelemetrySdk internalOtel,
-      @Nullable OpenTelemetry userOpenTelemetry,
-      ExecutorProvider backgroundExecutorProvider) {
-    this.clientContext = clientContext;
-    this.userOpenTelemetry = userOpenTelemetry;
-    this.builtinOpenTelemetry = internalOtel;
+      ApiTracerFactory userTracerFactory,
+      @Nullable OpenTelemetrySdk builtinOtel,
+      @Nullable OpenTelemetry userOtel,
+      Tagger ocTagger,
+      StatsRecorder ocRecorder,
+      ExecutorProvider backgroundExecutorProvider)
+      throws IOException {
+    this.isChild = isChild;
+    this.instanceName = instanceName;
+    this.appProfileId = appProfileId;
+
+    this.userTracerFactory = userTracerFactory;
+    this.builtinOpenTelemetry = builtinOtel;
+    this.userOpenTelemetry = userOtel;
+    this.ocTagger = ocTagger;
+    this.ocRecorder = ocRecorder;
     this.backgroundExecutorProvider = backgroundExecutorProvider;
+
+    ImmutableList.Builder<ApiTracerFactory> tracerFactories = ImmutableList.builder();
+    tracerFactories
+        .add(Util.createOCTracingFactory(instanceName, appProfileId))
+        .add(Util.createOCMetricsFactory(instanceName, appProfileId, ocTagger, ocRecorder))
+        .add(userTracerFactory);
+
+    if (builtinOtel != null) {
+      tracerFactories.add(Util.createOtelMetricsFactory(builtinOtel, instanceName, appProfileId));
+    }
+    if (userOtel != null) {
+      tracerFactories.add(Util.createOtelMetricsFactory(userOtel, instanceName, appProfileId));
+    }
+
+    this.clientContext =
+        clientContext.toBuilder()
+            .setTracerFactory(new CompositeTracerFactory(tracerFactories.build()))
+            .build();
+  }
+
+  public InstanceName getInstanceName() {
+    return instanceName;
+  }
+
+  public String getAppProfileId() {
+    return appProfileId;
   }
 
   @Nullable
@@ -213,12 +286,26 @@ public ClientContext getClientContext() {
     return this.clientContext;
   }
 
-  public BigtableClientContext withClientContext(ClientContext clientContext) {
+  public BigtableClientContext createChild(InstanceName instanceName, String appProfileId)
+      throws IOException {
     return new BigtableClientContext(
-        clientContext, builtinOpenTelemetry, userOpenTelemetry, backgroundExecutorProvider);
+        true,
+        instanceName,
+        appProfileId,
+        clientContext,
+        userTracerFactory,
+        builtinOpenTelemetry,
+        userOpenTelemetry,
+        ocTagger,
+        ocRecorder,
+        backgroundExecutorProvider);
   }
 
   public void close() throws Exception {
+    if (isChild) {
+      return;
+    }
+
     for (BackgroundResource resource : clientContext.getBackgroundResources()) {
       resource.close();
     }