diff --git a/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/stub/EnhancedBigtableStub.java b/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/stub/EnhancedBigtableStub.java index 4087dc59b4..6df57061dc 100644 --- a/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/stub/EnhancedBigtableStub.java +++ b/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. diff --git a/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/stub/metrics/BigtableTracer.java b/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/stub/metrics/BigtableTracer.java index 421c5b4724..6b199e4223 100644 --- a/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/stub/metrics/BigtableTracer.java +++ b/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 } diff --git a/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/stub/metrics/BuiltinMetricsTracer.java b/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/stub/metrics/BuiltinMetricsTracer.java index d369cbfcb6..48fc7dcce8 100644 --- a/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/stub/metrics/BuiltinMetricsTracer.java +++ b/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 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); } diff --git a/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/stub/metrics/BuiltinMetricsTracerFactory.java b/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/stub/metrics/BuiltinMetricsTracerFactory.java index 3583733a45..6e0839719b 100644 --- a/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/stub/metrics/BuiltinMetricsTracerFactory.java +++ b/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 statsAttributes, - StatsRecorderWrapper recorder) { - return new BuiltinMetricsTracerFactory(statsWrapper, statsAttributes, recorder); + StatsRecorderWrapper statsRecorderWrapper) { + return new BuiltinMetricsTracerFactory(statsWrapper, statsAttributes, statsRecorderWrapper); } private BuiltinMetricsTracerFactory( StatsWrapper statsWrapper, ImmutableMap statsAttributes, - StatsRecorderWrapper recorder) { + StatsRecorderWrapper statsRecorderWrapper) { this.statsAttributes = statsAttributes; this.statsWrapper = statsWrapper; - this.statsRecorderWrapper = recorder; + this.statsRecorderWrapper = statsRecorderWrapper; } @Override diff --git a/google-cloud-bigtable/src/test/java/com/google/cloud/bigtable/data/v2/stub/metrics/BuiltinMetricsTracerTest.java b/google-cloud-bigtable/src/test/java/com/google/cloud/bigtable/data/v2/stub/metrics/BuiltinMetricsTracerTest.java index 9cc30df23b..21d4c626f9 100644 --- a/google-cloud-bigtable/src/test/java/com/google/cloud/bigtable/data/v2/stub/metrics/BuiltinMetricsTracerTest.java +++ b/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 longValue; - private ArgumentCaptor intValue; - private ArgumentCaptor status; - private ArgumentCaptor tableId; - private ArgumentCaptor zone; - private ArgumentCaptor cluster; + @Mock private StatsRecorderWrapper statsRecorderWrapper; + + @Captor private ArgumentCaptor longValue; + @Captor private ArgumentCaptor intValue; + @Captor private ArgumentCaptor status; + @Captor private ArgumentCaptor tableId; + @Captor private ArgumentCaptor zone; + @Captor private ArgumentCaptor 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,11 +157,6 @@ 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); } @@ -168,8 +164,6 @@ public void close(Status status, Metadata trailers) { 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 {