update on comments

Author: mutianf

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 {