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 7e0743fc6e..d5f361fd28 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 @@ -170,8 +170,12 @@ public int getAttempt() { return attempt; } + /** Adds an annotation of the total throttled time of a batch. */ public void batchRequestThrottled(long throttledTimeMs) { for (ApiTracer child : children) { + if (child instanceof MetricsTracer) { + ((MetricsTracer) child).batchRequestThrottled(throttledTimeMs); + } } } } diff --git a/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/stub/metrics/MetricsTracer.java b/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/stub/metrics/MetricsTracer.java index ed8091ee87..add80d10ff 100644 --- a/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/stub/metrics/MetricsTracer.java +++ b/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/stub/metrics/MetricsTracer.java @@ -203,7 +203,6 @@ public void batchRequestSent(long elementCount, long requestSize) { // noop } - @Override public void batchRequestThrottled(long totalThrottledMs) { MeasureMap measures = stats diff --git a/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/stub/metrics/TracedBatcherUnaryCallable.java b/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/stub/metrics/TracedBatcherUnaryCallable.java index 2629308dd8..ba1860bb26 100644 --- a/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/stub/metrics/TracedBatcherUnaryCallable.java +++ b/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/stub/metrics/TracedBatcherUnaryCallable.java @@ -15,7 +15,6 @@ */ package com.google.cloud.bigtable.data.v2.stub.metrics; -import com.google.api.gax.tracing.TracedUnaryCallable; import com.google.api.core.ApiFuture; import com.google.api.core.InternalApi; import com.google.api.gax.batching.Batcher; @@ -40,7 +39,10 @@ public TracedBatcherUnaryCallable(UnaryCallable innerCallable) { public ApiFuture futureCall(RequestT request, ApiCallContext context) { if (context.getOption(Batcher.THROTTLED_TIME_KEY) != null) { ApiTracer tracer = context.getTracer(); - tracer.batchRequestThrottled(context.getOption(Batcher.THROTTLED_TIME_KEY)); + if (tracer instanceof BigtableTracer) { + ((BigtableTracer) tracer) + .batchRequestThrottled(context.getOption(Batcher.THROTTLED_TIME_KEY)); + } } return innerCallable.futureCall(request, context); } diff --git a/google-cloud-bigtable/src/test/java/com/google/cloud/bigtable/data/v2/stub/metrics/MetricsTracerTest.java b/google-cloud-bigtable/src/test/java/com/google/cloud/bigtable/data/v2/stub/metrics/MetricsTracerTest.java index 1fb6e137ca..856627e054 100644 --- a/google-cloud-bigtable/src/test/java/com/google/cloud/bigtable/data/v2/stub/metrics/MetricsTracerTest.java +++ b/google-cloud-bigtable/src/test/java/com/google/cloud/bigtable/data/v2/stub/metrics/MetricsTracerTest.java @@ -380,7 +380,7 @@ public Object answer(InvocationOnMock invocation) { } }) .when(mockService) - .readRows(any(ReadRowsRequest.class), anyObserver(ReadRowsResponse.class)); + .readRows(any(ReadRowsRequest.class), any()); try (Batcher batcher = stub.newBulkReadRowsBatcher(Query.create(TABLE_ID), GrpcCallContext.createDefault())) { @@ -437,7 +437,7 @@ public Object answer(InvocationOnMock invocation) { } }) .when(mockService) - .mutateRows(any(MutateRowsRequest.class), anyObserver(MutateRowsResponse.class)); + .mutateRows(any(MutateRowsRequest.class), any()); ApiCallContext defaultContext = GrpcCallContext.createDefault();