Add batchRequestThrottled to tracers

googleapis · Oct 29, 2021 · c0723f1 · c0723f1
1 parent 810960c
commit c0723f1
Show file tree

Hide file tree

Showing 4 changed files with 10 additions and 5 deletions.
diff --git a/...bigtable/src/main/java/com/google/cloud/bigtable/data/v2/stub/metrics/BigtableTracer.java b/...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/...-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/stub/metrics/MetricsTracer.java b/...-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/.../main/java/com/google/cloud/bigtable/data/v2/stub/metrics/TracedBatcherUnaryCallable.java b/.../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<ResponseT> 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/...table/src/test/java/com/google/cloud/bigtable/data/v2/stub/metrics/MetricsTracerTest.java b/...table/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();