fix context for TracedBatchingContextCallable

googleapis · Jun 23, 2021 · c307ef8 · c307ef8
1 parent c0f679f
commit c307ef8
Showing 1 changed file with 6 additions and 8 deletions.
diff --git a/...d-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/stub/EnhancedBigtableStub.java b/...d-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/stub/EnhancedBigtableStub.java
@@ -31,6 +31,7 @@
 import com.google.api.gax.retrying.RetryAlgorithm;
 import com.google.api.gax.retrying.RetryingExecutorWithContext;
 import com.google.api.gax.retrying.ScheduledRetryingExecutor;
+import com.google.api.gax.rpc.ApiCallContext;
 import com.google.api.gax.rpc.Callables;
 import com.google.api.gax.rpc.ClientContext;
 import com.google.api.gax.rpc.RequestParamsExtractor;
@@ -512,9 +513,8 @@ private UnaryCallable<BulkMutation, Void> createBulkMutateRowsCallable() {
         new HeaderTracerUnaryCallable<>(
             userFacing, settings.getHeaderTracer(), spanName.toString());
 
-    // TracedBatchingContextCallable needs the last in the callable chain so Batcher can call
-    // TracedBatchingCOntextCallable#futureCall(request, batchingCallContext) to add batching
-    // metrics to ApiTracer.
+    // TracedBatchingContextCallable needs to be the last in the callable chain so Batcher can pass
+    // batchingCallContext to add batching metrics to ApiTracer.
     UnaryCallable<BulkMutation, Void> batchingContextCallable =
         new TracedBatchingContextCallable<>(
             withHeaderTracer,
@@ -578,15 +578,13 @@ public Batcher<ByteString, Row> newBulkReadRowsBatcher(
       @Nonnull Query query, @Nullable GrpcCallContext ctx) {
     Preconditions.checkNotNull(query, "query cannot be null");
     UnaryCallable<Query, List<Row>> callable = readRowsCallable().all();
+    ApiCallContext callContext = clientContext.getDefaultCallContext();
     if (ctx != null) {
-      callable = callable.withDefaultCallContext(ctx);
+      callContext = ctx.merge(callContext);
     }
     UnaryCallable<Query, List<Row>> batchingContextCallable =
         new TracedBatchingContextCallable(
-            callable,
-            clientContext.getDefaultCallContext(),
-            clientContext.getTracerFactory(),
-            getSpanName("ReadRows"));
+            callable, callContext, clientContext.getTracerFactory(), getSpanName("ReadRows"));
     return new BatcherImpl<>(
         settings.bulkReadRowsSettings().getBatchingDescriptor(),
         batchingContextCallable,