update on comments

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

@@ -154,7 +154,7 @@ public static EnhancedBigtableStubSettings finalizeSettings(
       EnhancedBigtableStubSettings settings, Tagger tagger, StatsRecorder stats)
       throws IOException {
     EnhancedBigtableStubSettings.Builder builder = settings.toBuilder();
-    StatsWrapper statsWrapper = StatsWrapper.get();
+    StatsWrapper statsWrapper = StatsWrapper.create();
 
     // TODO: this implementation is on the cusp of unwieldy, if we end up adding more features
     // consider splitting it up by feature.

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

@@ -35,7 +35,7 @@ public void attemptStarted(int attemptNumber) {
     this.attempt = attemptNumber;
   }
 
-  /** annotate when onRequest is called */
+  /** annotate when onRequest is called. This will be called in BuiltinMetricsTracer. */
   public void onRequest() {
     // noop
   }
@@ -63,7 +63,10 @@ public void batchRequestThrottled(long throttledTimeMs) {
     // noop
   }
 
-  /** Set the Bigtable zone and cluster so metrics can be tagged with location information. */
+  /**
+   * Set the Bigtable zone and cluster so metrics can be tagged with location information. This will
+   * be called in BuiltinMetricsTracer.
+   */
   public void setLocations(String zone, String cluster) {
     // noop
   }

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

@@ -66,11 +66,12 @@ class BuiltinMetricsTracer extends BigtableTracer {
       SpanName spanName,
       Map<String, String> attributes,
       StatsWrapper statsWrapper,
-      StatsRecorderWrapper recorder) {
+      @Nullable StatsRecorderWrapper statsRecorderWrapper) {
     this.operationType = operationType;
     this.spanName = spanName;
-    if (recorder != null) {
-      this.recorder = recorder;
+    if (statsRecorderWrapper != null) {
+      // A workaround for test to pass in a mock StatsRecorderWrapper
+      this.recorder = statsRecorderWrapper;
     } else {
       this.recorder = new StatsRecorderWrapper(operationType, spanName, attributes, statsWrapper);
     }

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

@@ -41,21 +41,22 @@ public static BuiltinMetricsTracerFactory create(
     return new BuiltinMetricsTracerFactory(statsWrapper, statsAttributes, null);
   }
 
+  // A workaround for test to pass in a mock StatsRecorderWrapper
   @VisibleForTesting
-  static BuiltinMetricsTracerFactory create(
+  public static BuiltinMetricsTracerFactory createWithRecorder(
       StatsWrapper statsWrapper,
       ImmutableMap<String, String> statsAttributes,
-      StatsRecorderWrapper recorder) {
-    return new BuiltinMetricsTracerFactory(statsWrapper, statsAttributes, recorder);
+      StatsRecorderWrapper statsRecorderWrapper) {
+    return new BuiltinMetricsTracerFactory(statsWrapper, statsAttributes, statsRecorderWrapper);
   }
 
   private BuiltinMetricsTracerFactory(
       StatsWrapper statsWrapper,
       ImmutableMap<String, String> statsAttributes,
-      StatsRecorderWrapper recorder) {
+      StatsRecorderWrapper statsRecorderWrapper) {
     this.statsAttributes = statsAttributes;
     this.statsWrapper = statsWrapper;
-    this.statsRecorderWrapper = recorder;
+    this.statsRecorderWrapper = statsRecorderWrapper;
   }
 
   @Override

google-cloud-bigtable/src/test/java/com/google/cloud/bigtable/data/v2/stub/metrics/BuiltinMetricsTracerTest.java

@@ -1,5 +1,5 @@
 /*
- * Copyright 2021 Google LLC
+ * Copyright 2022 Google LLC
  *
  * Licensed under the Apache License, Version 2.0 (the "License");
  * you may not use this file except in compliance with the License.
@@ -61,7 +61,8 @@
 import org.junit.Rule;
 import org.junit.Test;
 import org.mockito.ArgumentCaptor;
-import org.mockito.Mockito;
+import org.mockito.Captor;
+import org.mockito.Mock;
 import org.mockito.junit.MockitoJUnit;
 import org.mockito.junit.MockitoRule;
 import org.threeten.bp.Duration;
@@ -118,13 +119,14 @@ public class BuiltinMetricsTracerTest {
 
   private EnhancedBigtableStub stub;
 
-  private StatsRecorderWrapper statsRecorderWrapper;
-  private ArgumentCaptor<Long> longValue;
-  private ArgumentCaptor<Integer> intValue;
-  private ArgumentCaptor<String> status;
-  private ArgumentCaptor<String> tableId;
-  private ArgumentCaptor<String> zone;
-  private ArgumentCaptor<String> cluster;
+  @Mock private StatsRecorderWrapper statsRecorderWrapper;
+
+  @Captor private ArgumentCaptor<Long> longValue;
+  @Captor private ArgumentCaptor<Integer> intValue;
+  @Captor private ArgumentCaptor<String> status;
+  @Captor private ArgumentCaptor<String> tableId;
+  @Captor private ArgumentCaptor<String> zone;
+  @Captor private ArgumentCaptor<String> cluster;
 
   private Stopwatch serverRetryDelayStopwatch;
   private AtomicLong serverTotalRetryDelay;
@@ -136,8 +138,7 @@ public void setUp() throws Exception {
     serverRetryDelayStopwatch = Stopwatch.createUnstarted();
     serverTotalRetryDelay = new AtomicLong(0);
 
-    // Add an interceptor to send location information in the trailers and add server-timing in
-    // headers
+    // Add an interceptor to add server-timing in headers
     ServerInterceptor trailersInterceptor =
         new ServerInterceptor() {
           private AtomicInteger count = new AtomicInteger(0);
@@ -156,20 +157,13 @@ public void sendHeaders(Metadata headers) {
                         String.format("gfet4t7; dur=%d", FAKE_SERVER_TIMING));
                     super.sendHeaders(headers);
                   }
-
-                  @Override
-                  public void close(Status status, Metadata trailers) {
-                    super.close(status, trailers);
-                  }
                 },
                 metadata);
           }
         };
 
     server = FakeServiceBuilder.create(mockService).intercept(trailersInterceptor).start();
 
-    statsRecorderWrapper = Mockito.mock(StatsRecorderWrapper.class);
-
     BigtableDataSettings settings =
         BigtableDataSettings.newBuilderForEmulator(server.getPort())
             .setProjectId(PROJECT_ID)
@@ -183,18 +177,11 @@ public void close(Status status, Metadata trailers) {
         .retrySettings()
         .setInitialRetryDelay(Duration.ofMillis(200));
     stubSettingsBuilder.setTracerFactory(
-        BuiltinMetricsTracerFactory.create(
-            StatsWrapper.get(), ImmutableMap.of(), statsRecorderWrapper));
+        BuiltinMetricsTracerFactory.createWithRecorder(
+            StatsWrapper.create(), ImmutableMap.of(), statsRecorderWrapper));
 
     EnhancedBigtableStubSettings stubSettings = stubSettingsBuilder.build();
     stub = new EnhancedBigtableStub(stubSettings, ClientContext.create(stubSettings));
-
-    longValue = ArgumentCaptor.forClass(long.class);
-    intValue = ArgumentCaptor.forClass(int.class);
-    status = ArgumentCaptor.forClass(String.class);
-    tableId = ArgumentCaptor.forClass(String.class);
-    zone = ArgumentCaptor.forClass(String.class);
-    cluster = ArgumentCaptor.forClass(String.class);
   }
 
   @After
@@ -291,17 +278,13 @@ public void testMutateRowAttempts() {
     stub.mutateRowCallable()
         .call(RowMutation.create(TABLE_ID, "random-row").setCell("cf", "q", "value"));
 
+    // record will get called 4 times, 3 times for attempts and 1 for recording operation level
+    // metrics.
     verify(statsRecorderWrapper, times(4))
         .record(status.capture(), tableId.capture(), zone.capture(), cluster.capture());
-    assertThat(zone.getAllValues().get(0)).isEqualTo(UNDEFINED);
-    assertThat(zone.getAllValues().get(1)).isEqualTo(UNDEFINED);
-    assertThat(zone.getAllValues().get(2)).isEqualTo(UNDEFINED);
-    assertThat(cluster.getAllValues().get(0)).isEqualTo(UNDEFINED);
-    assertThat(cluster.getAllValues().get(1)).isEqualTo(UNDEFINED);
-    assertThat(cluster.getAllValues().get(2)).isEqualTo(UNDEFINED);
-    assertThat(status.getAllValues().get(0)).isEqualTo("UNAVAILABLE");
-    assertThat(status.getAllValues().get(1)).isEqualTo("UNAVAILABLE");
-    assertThat(status.getAllValues().get(2)).isEqualTo("OK");
+    assertThat(zone.getAllValues()).containsExactly(UNDEFINED, UNDEFINED, UNDEFINED, UNDEFINED);
+    assertThat(cluster.getAllValues()).containsExactly(UNDEFINED, UNDEFINED, UNDEFINED, UNDEFINED);
+    assertThat(status.getAllValues()).containsExactly("UNAVAILABLE", "UNAVAILABLE", "OK", "OK");
   }
 
   private class FakeService extends BigtableGrpc.BigtableImplBase {