chore: refactor TracerFactory creation (#2102)

googleapis · Feb 16, 2024 · b8fd4d5 · b8fd4d5
1 parent beb0219
commit b8fd4d5
Show file tree

Hide file tree

Showing 4 changed files with 49 additions and 27 deletions.
diff --git a/...d-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/BigtableDataClientFactory.java b/...d-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/BigtableDataClientFactory.java
@@ -112,7 +112,7 @@ public BigtableDataClient createDefault() {
               .toBuilder()
               .setTracerFactory(
                   EnhancedBigtableStub.createBigtableTracerFactory(
-                      defaultSettings.getStubSettings()))
+                      defaultSettings.getStubSettings(), sharedClientContext))
               .build();
 
       return BigtableDataClient.createWithClientContext(defaultSettings, clientContext);
@@ -140,7 +140,8 @@ public BigtableDataClient createForAppProfile(@Nonnull String appProfileId) thro
         sharedClientContext
             .toBuilder()
             .setTracerFactory(
-                EnhancedBigtableStub.createBigtableTracerFactory(settings.getStubSettings()))
+                EnhancedBigtableStub.createBigtableTracerFactory(
+                    settings.getStubSettings(), sharedClientContext))
             .build();
     return BigtableDataClient.createWithClientContext(settings, clientContext);
   }
@@ -168,7 +169,8 @@ public BigtableDataClient createForInstance(@Nonnull String projectId, @Nonnull
         sharedClientContext
             .toBuilder()
             .setTracerFactory(
-                EnhancedBigtableStub.createBigtableTracerFactory(settings.getStubSettings()))
+                EnhancedBigtableStub.createBigtableTracerFactory(
+                    settings.getStubSettings(), sharedClientContext))
             .build();
 
     return BigtableDataClient.createWithClientContext(settings, clientContext);
@@ -197,7 +199,8 @@ public BigtableDataClient createForInstance(
         sharedClientContext
             .toBuilder()
             .setTracerFactory(
-                EnhancedBigtableStub.createBigtableTracerFactory(settings.getStubSettings()))
+                EnhancedBigtableStub.createBigtableTracerFactory(
+                    settings.getStubSettings(), sharedClientContext))
             .build();
     return BigtableDataClient.createWithClientContext(settings, clientContext);
   }

diff --git a/...d-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/stub/EnhancedBigtableStub.java b/...d-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/stub/EnhancedBigtableStub.java
@@ -177,10 +177,14 @@ public class EnhancedBigtableStub implements AutoCloseable {
   public static EnhancedBigtableStub create(EnhancedBigtableStubSettings settings)
       throws IOException {
 
-    settings = settings.toBuilder().setTracerFactory(createBigtableTracerFactory(settings)).build();
     ClientContext clientContext = createClientContext(settings);
+    ClientContext contextWithTracer =
+        clientContext
+            .toBuilder()
+            .setTracerFactory(createBigtableTracerFactory(settings, clientContext))
+            .build();
 
-    return new EnhancedBigtableStub(settings, clientContext);
+    return new EnhancedBigtableStub(settings, contextWithTracer);
   }
 
   public static EnhancedBigtableStub createWithClientContext(
@@ -237,13 +241,17 @@ public static ClientContext createClientContext(EnhancedBigtableStubSettings set
   }
 
   public static ApiTracerFactory createBigtableTracerFactory(
-      EnhancedBigtableStubSettings settings) {
-    return createBigtableTracerFactory(settings, Tags.getTagger(), Stats.getStatsRecorder());
+      EnhancedBigtableStubSettings settings, ClientContext clientContext) {
+    return createBigtableTracerFactory(
+        settings, Tags.getTagger(), Stats.getStatsRecorder(), clientContext);
   }
 
   @VisibleForTesting
   public static ApiTracerFactory createBigtableTracerFactory(
-      EnhancedBigtableStubSettings settings, Tagger tagger, StatsRecorder stats) {
+      EnhancedBigtableStubSettings settings,
+      Tagger tagger,
+      StatsRecorder stats,
+      ClientContext clientContext) {
     String projectId = settings.getProjectId();
     String instanceId = settings.getInstanceId();
     String appProfileId = settings.getAppProfileId();

diff --git a/.../test/java/com/google/cloud/bigtable/data/v2/stub/metrics/BigtableTracerCallableTest.java b/.../test/java/com/google/cloud/bigtable/data/v2/stub/metrics/BigtableTracerCallableTest.java
@@ -44,7 +44,6 @@
 import com.google.cloud.bigtable.data.v2.models.ReadModifyWriteRow;
 import com.google.cloud.bigtable.data.v2.models.RowMutation;
 import com.google.cloud.bigtable.data.v2.stub.EnhancedBigtableStub;
-import com.google.cloud.bigtable.data.v2.stub.EnhancedBigtableStubSettings;
 import com.google.common.collect.ImmutableMap;
 import io.grpc.ForwardingServerCall.SimpleForwardingServerCall;
 import io.grpc.Metadata;
@@ -125,16 +124,21 @@ public void sendHeaders(Metadata headers) {
             .setInstanceId(INSTANCE_ID)
             .setAppProfileId(APP_PROFILE_ID)
             .build();
-    EnhancedBigtableStubSettings stubSettings =
-        settings
-            .getStubSettings()
+
+    ClientContext clientContext =
+        EnhancedBigtableStub.createClientContext(settings.getStubSettings());
+    clientContext =
+        clientContext
             .toBuilder()
             .setTracerFactory(
                 EnhancedBigtableStub.createBigtableTracerFactory(
-                    settings.getStubSettings(), Tags.getTagger(), localStats.getStatsRecorder()))
+                    settings.getStubSettings(),
+                    Tags.getTagger(),
+                    localStats.getStatsRecorder(),
+                    clientContext))
             .build();
-    attempts = stubSettings.readRowsSettings().getRetrySettings().getMaxAttempts();
-    stub = new EnhancedBigtableStub(stubSettings, ClientContext.create(stubSettings));
+    attempts = settings.getStubSettings().readRowsSettings().getRetrySettings().getMaxAttempts();
+    stub = new EnhancedBigtableStub(settings.getStubSettings(), clientContext);
 
     // Create another server without injecting the server-timing header and another stub that
     // connects to it.
@@ -146,18 +150,21 @@ public void sendHeaders(Metadata headers) {
             .setInstanceId(INSTANCE_ID)
             .setAppProfileId(APP_PROFILE_ID)
             .build();
-    EnhancedBigtableStubSettings noHeaderStubSettings =
-        noHeaderSettings
-            .getStubSettings()
+
+    ClientContext noHeaderClientContext =
+        EnhancedBigtableStub.createClientContext(noHeaderSettings.getStubSettings());
+    noHeaderClientContext =
+        noHeaderClientContext
             .toBuilder()
             .setTracerFactory(
                 EnhancedBigtableStub.createBigtableTracerFactory(
                     noHeaderSettings.getStubSettings(),
                     Tags.getTagger(),
-                    localStats.getStatsRecorder()))
+                    localStats.getStatsRecorder(),
+                    noHeaderClientContext))
             .build();
     noHeaderStub =
-        new EnhancedBigtableStub(noHeaderStubSettings, ClientContext.create(noHeaderStubSettings));
+        new EnhancedBigtableStub(noHeaderSettings.getStubSettings(), noHeaderClientContext);
   }
 
   @After

diff --git a/...table/src/test/java/com/google/cloud/bigtable/data/v2/stub/metrics/MetricsTracerTest.java b/...table/src/test/java/com/google/cloud/bigtable/data/v2/stub/metrics/MetricsTracerTest.java
@@ -39,7 +39,6 @@
 import com.google.cloud.bigtable.data.v2.models.Row;
 import com.google.cloud.bigtable.data.v2.models.RowMutationEntry;
 import com.google.cloud.bigtable.data.v2.stub.EnhancedBigtableStub;
-import com.google.cloud.bigtable.data.v2.stub.EnhancedBigtableStubSettings;
 import com.google.cloud.bigtable.data.v2.stub.mutaterows.MutateRowsBatchingDescriptor;
 import com.google.common.base.Stopwatch;
 import com.google.common.collect.ImmutableMap;
@@ -120,15 +119,20 @@ public void setUp() throws Exception {
             .setInstanceId(INSTANCE_ID)
             .setAppProfileId(APP_PROFILE_ID)
             .build();
-    EnhancedBigtableStubSettings stubSettings =
-        settings
-            .getStubSettings()
+
+    ClientContext clientContext =
+        EnhancedBigtableStub.createClientContext(settings.getStubSettings());
+    clientContext =
+        clientContext
             .toBuilder()
             .setTracerFactory(
                 EnhancedBigtableStub.createBigtableTracerFactory(
-                    settings.getStubSettings(), Tags.getTagger(), localStats.getStatsRecorder()))
+                    settings.getStubSettings(),
+                    Tags.getTagger(),
+                    localStats.getStatsRecorder(),
+                    clientContext))
             .build();
-    stub = new EnhancedBigtableStub(stubSettings, ClientContext.create(stubSettings));
+    stub = new EnhancedBigtableStub(settings.getStubSettings(), clientContext);
   }
 
   @After