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 be2601dbdc..a878c5df6b 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 @@ -70,7 +70,7 @@ import com.google.cloud.bigtable.data.v2.models.RowAdapter; import com.google.cloud.bigtable.data.v2.models.RowMutation; import com.google.cloud.bigtable.data.v2.models.RowMutationEntry; -import com.google.cloud.bigtable.data.v2.stub.metrics.CompositeTracerFactory; +import com.google.cloud.bigtable.data.v2.stub.metrics.BigtableTracerFactory; import com.google.cloud.bigtable.data.v2.stub.metrics.HeaderTracerStreamingCallable; import com.google.cloud.bigtable.data.v2.stub.metrics.HeaderTracerUnaryCallable; import com.google.cloud.bigtable.data.v2.stub.metrics.MetricsTracerFactory; @@ -90,7 +90,6 @@ import com.google.common.collect.ImmutableList; import com.google.common.collect.ImmutableMap; import com.google.protobuf.ByteString; -import io.grpc.Metadata; import io.opencensus.stats.Stats; import io.opencensus.stats.StatsRecorder; import io.opencensus.tags.TagKey; @@ -120,11 +119,6 @@ */ @InternalApi public class EnhancedBigtableStub implements AutoCloseable { - static final Metadata.Key ATTEMPT_HEADER_KEY = - Metadata.Key.of("attempt", Metadata.ASCII_STRING_MARSHALLER); - static final Metadata.Key TIMESTAMP_HEADER_KEY = - Metadata.Key.of("client-timing", Metadata.ASCII_STRING_MARSHALLER); - private static final String CLIENT_NAME = "Bigtable"; private static final long FLOW_CONTROL_ADJUSTING_INTERVAL_MS = TimeUnit.SECONDS.toMillis(20); @@ -197,7 +191,7 @@ public static EnhancedBigtableStubSettings finalizeSettings( .build(); // Inject Opencensus instrumentation builder.setTracerFactory( - new CompositeTracerFactory( + new BigtableTracerFactory( ImmutableList.of( // Add OpenCensus Tracing new OpencensusTracerFactory( diff --git a/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/stub/ExtraHeadersSeverStreamingCallable.java b/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/stub/ExtraHeadersSeverStreamingCallable.java index bba09a8832..5585207a1f 100644 --- a/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/stub/ExtraHeadersSeverStreamingCallable.java +++ b/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/stub/ExtraHeadersSeverStreamingCallable.java @@ -19,13 +19,12 @@ import com.google.api.gax.rpc.ApiCallContext; import com.google.api.gax.rpc.ResponseObserver; import com.google.api.gax.rpc.ServerStreamingCallable; -import com.google.cloud.bigtable.data.v2.stub.metrics.CompositeTracer; -import com.google.common.collect.ImmutableMap; -import java.util.Arrays; -import java.util.List; -import java.util.Map; +import com.google.cloud.bigtable.data.v2.stub.metrics.Util; -/** A callable that injects client timestamp and attempt count to request headers. */ +/** + * A callable that injects client timestamp and current attempt number to request headers. Attempt + * number starts from 0. + */ final class ExtraHeadersSeverStreamingCallable extends ServerStreamingCallable { private final ServerStreamingCallable innerCallable; @@ -39,17 +38,8 @@ public void call( RequestT request, ResponseObserver responseObserver, ApiCallContext apiCallContext) { - int attemptCount = -1; - if (apiCallContext.getTracer() instanceof CompositeTracer) { - attemptCount = ((CompositeTracer) apiCallContext.getTracer()).getAttempt(); - } - Map> headers = - ImmutableMap.of( - EnhancedBigtableStub.ATTEMPT_HEADER_KEY.name(), - Arrays.asList(String.valueOf(attemptCount)), - EnhancedBigtableStub.TIMESTAMP_HEADER_KEY.name(), - Arrays.asList(String.valueOf(System.currentTimeMillis()))); - ApiCallContext newCallContext = apiCallContext.withExtraHeaders(headers); + ApiCallContext newCallContext = + apiCallContext.withExtraHeaders(Util.createExtraHeaders(apiCallContext)); innerCallable.call(request, responseObserver, newCallContext); } } diff --git a/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/stub/ExtraHeadersUnaryCallable.java b/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/stub/ExtraHeadersUnaryCallable.java index 3c33ba1adc..25bd843e76 100644 --- a/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/stub/ExtraHeadersUnaryCallable.java +++ b/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/stub/ExtraHeadersUnaryCallable.java @@ -19,13 +19,12 @@ import com.google.api.core.ApiFuture; import com.google.api.gax.rpc.ApiCallContext; import com.google.api.gax.rpc.UnaryCallable; -import com.google.cloud.bigtable.data.v2.stub.metrics.CompositeTracer; -import com.google.common.collect.ImmutableMap; -import java.util.Arrays; -import java.util.List; -import java.util.Map; +import com.google.cloud.bigtable.data.v2.stub.metrics.Util; -/** A callable that injects client timestamp and attempt count to request headers. */ +/** + * A callable that injects client timestamp and current attempt number to request headers. Attempt + * number starts from 0. + */ final class ExtraHeadersUnaryCallable extends UnaryCallable { private final UnaryCallable innerCallable; @@ -35,18 +34,10 @@ public ExtraHeadersUnaryCallable(UnaryCallable innerCallable) { } @Override + @SuppressWarnings("unchecked") public ApiFuture futureCall(RequestT request, ApiCallContext apiCallContext) { - int attemptCount = -1; - if (apiCallContext.getTracer() instanceof CompositeTracer) { - attemptCount = ((CompositeTracer) apiCallContext.getTracer()).getAttempt(); - } - Map> headers = - ImmutableMap.of( - EnhancedBigtableStub.ATTEMPT_HEADER_KEY.name(), - Arrays.asList(String.valueOf(attemptCount)), - EnhancedBigtableStub.TIMESTAMP_HEADER_KEY.name(), - Arrays.asList(String.valueOf(System.currentTimeMillis()))); - ApiCallContext newCallContext = apiCallContext.withExtraHeaders(headers); + ApiCallContext newCallContext = + apiCallContext.withExtraHeaders(Util.createExtraHeaders(apiCallContext)); return innerCallable.futureCall(request, newCallContext); } } diff --git a/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/stub/metrics/CompositeTracer.java b/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/stub/metrics/BigtableTracer.java similarity index 84% rename from google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/stub/metrics/CompositeTracer.java rename to google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/stub/metrics/BigtableTracer.java index 31f4817e76..3fc1674eac 100644 --- a/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/stub/metrics/CompositeTracer.java +++ b/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/stub/metrics/BigtableTracer.java @@ -16,6 +16,7 @@ package com.google.cloud.bigtable.data.v2.stub.metrics; import com.google.api.core.InternalApi; +import com.google.api.gax.rpc.ApiCallContext; import com.google.api.gax.tracing.ApiTracer; import com.google.api.gax.tracing.BaseApiTracer; import com.google.common.collect.ImmutableList; @@ -23,13 +24,18 @@ import java.util.List; import org.threeten.bp.Duration; -/** Combines multiple {@link ApiTracer}s into a single {@link ApiTracer}. */ +/** + * A Bigtable specific {@link ApiTracer} that will be used to plumb additional context through the + * call chains as well as combines multiple user defined {@link ApiTracer}s into a single one. This + * will ensure that operation lifecycle events are plumbed through while maintaining user configured + * functionalities. + */ @InternalApi("For internal use only") -public class CompositeTracer extends BaseApiTracer { +public class BigtableTracer extends BaseApiTracer { private final List children; private int attempt = 0; - CompositeTracer(List children) { + BigtableTracer(List children) { this.children = ImmutableList.copyOf(children); } @@ -157,6 +163,11 @@ public void batchRequestSent(long elementCount, long requestSize) { } } + /** + * Get the attempt number of the current call. Attempt number for the current call is passed in + * and recorded in {@link #attemptStarted(int)}. With the getter we can access it from {@link + * ApiCallContext}. + */ public int getAttempt() { return attempt; } diff --git a/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/stub/metrics/CompositeTracerFactory.java b/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/stub/metrics/BigtableTracerFactory.java similarity index 82% rename from google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/stub/metrics/CompositeTracerFactory.java rename to google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/stub/metrics/BigtableTracerFactory.java index 2d9256a5ea..f980c4b7ce 100644 --- a/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/stub/metrics/CompositeTracerFactory.java +++ b/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/stub/metrics/BigtableTracerFactory.java @@ -24,12 +24,15 @@ import java.util.ArrayList; import java.util.List; -/** Combines multiple {@link ApiTracerFactory} into a single {@link ApiTracerFactory}. */ +/** + * A Bigtable specific {@link ApiTracerFactory} that combines multiple {@link ApiTracerFactory} into + * a single one. + */ @InternalApi("For internal use only") -public class CompositeTracerFactory extends BaseApiTracerFactory { +public class BigtableTracerFactory extends BaseApiTracerFactory { private final List apiTracerFactories; - public CompositeTracerFactory(List apiTracerFactories) { + public BigtableTracerFactory(List apiTracerFactories) { this.apiTracerFactories = ImmutableList.copyOf(apiTracerFactories); } @@ -40,6 +43,6 @@ public ApiTracer newTracer(ApiTracer parent, SpanName spanName, OperationType op for (ApiTracerFactory factory : apiTracerFactories) { children.add(factory.newTracer(parent, spanName, operationType)); } - return new CompositeTracer(children); + return new BigtableTracer(children); } } diff --git a/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/stub/metrics/Util.java b/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/stub/metrics/Util.java index ff40aca387..43d6946b6e 100644 --- a/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/stub/metrics/Util.java +++ b/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/stub/metrics/Util.java @@ -15,20 +15,33 @@ */ package com.google.cloud.bigtable.data.v2.stub.metrics; +import com.google.api.core.InternalApi; +import com.google.api.gax.rpc.ApiCallContext; import com.google.api.gax.rpc.ApiException; import com.google.api.gax.rpc.StatusCode; import com.google.api.gax.rpc.StatusCode.Code; +import com.google.common.collect.ImmutableMap; +import io.grpc.Metadata; import io.grpc.Status; import io.grpc.StatusException; import io.grpc.StatusRuntimeException; import io.opencensus.tags.TagValue; +import java.util.Arrays; +import java.util.List; +import java.util.Map; import java.util.concurrent.CancellationException; import java.util.concurrent.ExecutionException; import java.util.concurrent.Future; import javax.annotation.Nullable; /** Utilities to help integrating with OpenCensus. */ -class Util { +@InternalApi("For internal use only") +public class Util { + public static final Metadata.Key ATTEMPT_HEADER_KEY = + Metadata.Key.of("attempt", Metadata.ASCII_STRING_MARSHALLER); + public static final Metadata.Key TIMESTAMP_HEADER_KEY = + Metadata.Key.of("client-timing", Metadata.ASCII_STRING_MARSHALLER); + private static final TagValue OK_STATUS = TagValue.create(StatusCode.Code.OK.toString()); /** Convert an exception into a value that can be used as an OpenCensus tag value. */ @@ -71,4 +84,22 @@ static TagValue extractStatus(Future future) { } return extractStatus(error); } + + /** + * Create extra headers with attempt number and client timestamp from api call context. Attempt + * number starts from 0. + */ + public static Map> createExtraHeaders(ApiCallContext apiCallContext) { + int attemptCount = -1; + if (apiCallContext.getTracer() instanceof BigtableTracer) { + attemptCount = ((BigtableTracer) apiCallContext.getTracer()).getAttempt(); + } + ImmutableMap.Builder headers = ImmutableMap.>builder(); + if (attemptCount != -1) { + headers.put(ATTEMPT_HEADER_KEY.name(), Arrays.asList(String.valueOf(attemptCount))); + } + headers.put( + TIMESTAMP_HEADER_KEY.name(), Arrays.asList(String.valueOf(System.currentTimeMillis()))); + return headers.build(); + } } diff --git a/google-cloud-bigtable/src/test/java/com/google/cloud/bigtable/data/v2/stub/ExtraHeadersCallableTest.java b/google-cloud-bigtable/src/test/java/com/google/cloud/bigtable/data/v2/stub/ExtraHeadersCallableTest.java index 0c99e8787b..644d8edbba 100644 --- a/google-cloud-bigtable/src/test/java/com/google/cloud/bigtable/data/v2/stub/ExtraHeadersCallableTest.java +++ b/google-cloud-bigtable/src/test/java/com/google/cloud/bigtable/data/v2/stub/ExtraHeadersCallableTest.java @@ -40,6 +40,7 @@ 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.models.RowMutationEntry; +import com.google.cloud.bigtable.data.v2.stub.metrics.Util; import com.google.common.collect.Queues; import com.google.protobuf.ByteString; import com.google.protobuf.BytesValue; @@ -52,8 +53,6 @@ import io.grpc.Status; import io.grpc.StatusRuntimeException; import io.grpc.stub.StreamObserver; -import io.opencensus.impl.stats.StatsComponentImpl; -import io.opencensus.stats.StatsComponent; import java.util.concurrent.BlockingQueue; import java.util.concurrent.atomic.AtomicInteger; import org.junit.After; @@ -68,7 +67,6 @@ public class ExtraHeadersCallableTest { private FakeService fakeService = new FakeService(); - private final StatsComponent localStats = new StatsComponentImpl(); private EnhancedBigtableStub stub; private static final String PROJECT_ID = "fake-project"; @@ -261,11 +259,11 @@ private void verifyHeaders(int expectedAttemptCounts, long startTimestamp) throw for (int i = 0; i < expectedAttemptCounts; i++) { Metadata headers = metadataInterceptor.headers.take(); - String attemptCount = headers.get(EnhancedBigtableStub.ATTEMPT_HEADER_KEY); + String attemptCount = headers.get(Util.ATTEMPT_HEADER_KEY); assertThat(attemptCount).isNotNull(); assertThat(Integer.valueOf(attemptCount)).isEqualTo(i); - String clientTimeStr = headers.get(EnhancedBigtableStub.TIMESTAMP_HEADER_KEY); + String clientTimeStr = headers.get(Util.TIMESTAMP_HEADER_KEY); assertThat(clientTimeStr).isNotNull(); long clientTime = Long.valueOf(clientTimeStr); assertThat(clientTime).isAtLeast(timestamp); diff --git a/google-cloud-bigtable/src/test/java/com/google/cloud/bigtable/data/v2/stub/metrics/CompositeTracerTest.java b/google-cloud-bigtable/src/test/java/com/google/cloud/bigtable/data/v2/stub/metrics/BigtableTracerTest.java similarity index 83% rename from google-cloud-bigtable/src/test/java/com/google/cloud/bigtable/data/v2/stub/metrics/CompositeTracerTest.java rename to google-cloud-bigtable/src/test/java/com/google/cloud/bigtable/data/v2/stub/metrics/BigtableTracerTest.java index cedb227bad..f1b464ff4d 100644 --- a/google-cloud-bigtable/src/test/java/com/google/cloud/bigtable/data/v2/stub/metrics/CompositeTracerTest.java +++ b/google-cloud-bigtable/src/test/java/com/google/cloud/bigtable/data/v2/stub/metrics/BigtableTracerTest.java @@ -34,17 +34,17 @@ import org.threeten.bp.Duration; @RunWith(JUnit4.class) -public class CompositeTracerTest { +public class BigtableTracerTest { @Rule public final MockitoRule mockitoRule = MockitoJUnit.rule(); @Mock private ApiTracer child1; @Mock private ApiTracer child2; - private CompositeTracer compositeTracer; + private BigtableTracer bigtableTracer; @Before public void setup() { - compositeTracer = new CompositeTracer(ImmutableList.of(child1, child2)); + bigtableTracer = new BigtableTracer(ImmutableList.of(child1, child2)); } @Test @@ -55,7 +55,7 @@ public void testInScope() { Scope scope2 = mock(Scope.class); when(child2.inScope()).thenReturn(scope2); - Scope parentScope = compositeTracer.inScope(); + Scope parentScope = bigtableTracer.inScope(); parentScope.close(); verify(scope1, times(1)).close(); @@ -63,14 +63,14 @@ public void testInScope() { @Test public void testOperationSucceeded() { - compositeTracer.operationSucceeded(); + bigtableTracer.operationSucceeded(); verify(child1, times(1)).operationSucceeded(); verify(child2, times(1)).operationSucceeded(); } @Test public void testOperationCancelled() { - compositeTracer.operationCancelled(); + bigtableTracer.operationCancelled(); verify(child1, times(1)).operationCancelled(); verify(child2, times(1)).operationCancelled(); } @@ -78,35 +78,35 @@ public void testOperationCancelled() { @Test public void testOperationFailed() { RuntimeException error = new RuntimeException(); - compositeTracer.operationFailed(error); + bigtableTracer.operationFailed(error); verify(child1, times(1)).operationFailed(error); verify(child2, times(1)).operationFailed(error); } @Test public void testConnectionSelected() { - compositeTracer.connectionSelected("connection-one"); + bigtableTracer.connectionSelected("connection-one"); verify(child1, times(1)).connectionSelected("connection-one"); verify(child2, times(1)).connectionSelected("connection-one"); } @Test public void testAttemptStarted() { - compositeTracer.attemptStarted(3); + bigtableTracer.attemptStarted(3); verify(child1, times(1)).attemptStarted(3); verify(child2, times(1)).attemptStarted(3); } @Test public void testAttemptSucceeded() { - compositeTracer.attemptSucceeded(); + bigtableTracer.attemptSucceeded(); verify(child1, times(1)).attemptSucceeded(); verify(child2, times(1)).attemptSucceeded(); } @Test public void testAttemptCancelled() { - compositeTracer.attemptCancelled(); + bigtableTracer.attemptCancelled(); verify(child1, times(1)).attemptCancelled(); verify(child2, times(1)).attemptCancelled(); } @@ -115,7 +115,7 @@ public void testAttemptCancelled() { public void testAttemptFailed() { RuntimeException error = new RuntimeException(); Duration delay = Duration.ofMillis(10); - compositeTracer.attemptFailed(error, delay); + bigtableTracer.attemptFailed(error, delay); verify(child1, times(1)).attemptFailed(error, delay); verify(child2, times(1)).attemptFailed(error, delay); } @@ -123,7 +123,7 @@ public void testAttemptFailed() { @Test public void testAttemptFailedRetriesExhausted() { RuntimeException error = new RuntimeException(); - compositeTracer.attemptFailedRetriesExhausted(error); + bigtableTracer.attemptFailedRetriesExhausted(error); verify(child1, times(1)).attemptFailedRetriesExhausted(error); verify(child2, times(1)).attemptFailedRetriesExhausted(error); } @@ -131,7 +131,7 @@ public void testAttemptFailedRetriesExhausted() { @Test public void testAttemptPermanentFailure() { RuntimeException error = new RuntimeException(); - compositeTracer.attemptPermanentFailure(error); + bigtableTracer.attemptPermanentFailure(error); verify(child1, times(1)).attemptPermanentFailure(error); verify(child2, times(1)).attemptPermanentFailure(error); } @@ -139,35 +139,35 @@ public void testAttemptPermanentFailure() { @Test public void testLroStartFailed() { RuntimeException error = new RuntimeException(); - compositeTracer.lroStartFailed(error); + bigtableTracer.lroStartFailed(error); verify(child1, times(1)).lroStartFailed(error); verify(child2, times(1)).lroStartFailed(error); } @Test public void testLroStartSucceeded() { - compositeTracer.lroStartSucceeded(); + bigtableTracer.lroStartSucceeded(); verify(child1, times(1)).lroStartSucceeded(); verify(child2, times(1)).lroStartSucceeded(); } @Test public void testResponseReceived() { - compositeTracer.responseReceived(); + bigtableTracer.responseReceived(); verify(child1, times(1)).responseReceived(); verify(child2, times(1)).responseReceived(); } @Test public void testRequestSent() { - compositeTracer.requestSent(); + bigtableTracer.requestSent(); verify(child1, times(1)).requestSent(); verify(child2, times(1)).requestSent(); } @Test public void testBatchRequestSent() { - compositeTracer.batchRequestSent(2, 20); + bigtableTracer.batchRequestSent(2, 20); verify(child1, times(1)).batchRequestSent(2, 20); verify(child2, times(1)).batchRequestSent(2, 20); }