igorbernstein2 · Sep 27, 2024 · Sep 27, 2024 · Sep 27, 2024 · Sep 27, 2024 · Sep 27, 2024
@@ -39,6 +39,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.retrying.SimpleStreamResumptionStrategy;
 import com.google.api.gax.rpc.Callables;
 import com.google.api.gax.rpc.ClientContext;
 import com.google.api.gax.rpc.RequestParamsExtractor;
@@ -123,6 +124,8 @@
 import com.google.cloud.bigtable.data.v2.stub.mutaterows.MutateRowsBatchingDescriptor;
 import com.google.cloud.bigtable.data.v2.stub.mutaterows.MutateRowsPartialErrorRetryAlgorithm;
 import com.google.cloud.bigtable.data.v2.stub.mutaterows.MutateRowsRetryingCallable;
+import com.google.cloud.bigtable.data.v2.stub.opt.UnaryOverStreamingCallable;
+import com.google.cloud.bigtable.data.v2.stub.opt.UnaryOverStreamingTracerModCallable;
 import com.google.cloud.bigtable.data.v2.stub.readrows.FilterMarkerRowsCallable;
 import com.google.cloud.bigtable.data.v2.stub.readrows.ReadRowsBatchingDescriptor;
 import com.google.cloud.bigtable.data.v2.stub.readrows.ReadRowsFirstCallable;
@@ -154,6 +157,7 @@
 import java.io.IOException;
 import java.net.URI;
 import java.net.URISyntaxException;
+import java.time.Duration;
 import java.util.Collections;
 import java.util.List;
 import java.util.Map;
@@ -563,11 +567,13 @@ public <RowT> UnaryCallable<Query, RowT> createReadRowCallable(RowAdapter<RowT>
 
     ReadRowsFirstCallable<RowT> firstRow = new ReadRowsFirstCallable<>(readRowCallable);
 
-    UnaryCallable<Query, RowT> traced =
-        new TracedUnaryCallable<>(
-            firstRow, clientContext.getTracerFactory(), getSpanName("ReadRow"));
-
-    return traced.withDefaultCallContext(clientContext.getDefaultCallContext());
+    UnaryOverStreamingTracerModCallable modCallable =
+        new UnaryOverStreamingTracerModCallable(firstRow);
+    TracedServerStreamingCallable<Query, RowT> traced =
+        new TracedServerStreamingCallable<>(
+            modCallable, clientContext.getTracerFactory(), getSpanName("ReadRow"));
+    UnaryOverStreamingCallable<Query, RowT> unaryAdapter = new UnaryOverStreamingCallable<>(traced);
+    return unaryAdapter.withDefaultCallContext(clientContext.getDefaultCallContext());
   }
 
   /**
@@ -794,8 +800,8 @@ private UnaryCallable<String, List<KeyOffset>> createSampleRowKeysCallable() {
    */
   private UnaryCallable<RowMutation, Void> createMutateRowCallable() {
     String methodName = "MutateRow";
-    UnaryCallable<MutateRowRequest, MutateRowResponse> base =
-        GrpcRawCallableFactory.createUnaryCallable(
+    ServerStreamingCallable<MutateRowRequest, MutateRowResponse> base =
+        GrpcRawCallableFactory.createServerStreamingCallable(
             GrpcCallSettings.<MutateRowRequest, MutateRowResponse>newBuilder()
                 .setMethodDescriptor(BigtableGrpc.getMutateRowMethod())
                 .setParamsExtractor(
@@ -818,17 +824,26 @@ public Map<String, String> extract(MutateRowRequest mutateRowRequest) {
                 .build(),
             settings.mutateRowSettings().getRetryableCodes());
 
-    UnaryCallable<MutateRowRequest, MutateRowResponse> withStatsHeaders =
-        new StatsHeadersUnaryCallable<>(base);
+    ServerStreamingCallable<MutateRowRequest, MutateRowResponse> withStatsHeaders =
+        new StatsHeadersServerStreamingCallable<>(base);
 
-    UnaryCallable<MutateRowRequest, MutateRowResponse> withBigtableTracer =
-        new BigtableTracerUnaryCallable<>(withStatsHeaders);
+    ServerStreamingCallable<MutateRowRequest, MutateRowResponse> withBigtableTracer =
+        new BigtableTracerStreamingCallable<>(withStatsHeaders);
 
-    UnaryCallable<MutateRowRequest, MutateRowResponse> retrying =
-        withRetries(withBigtableTracer, settings.mutateRowSettings());
+    ServerStreamingCallable<MutateRowRequest, MutateRowResponse> retrying =
+        withRetries(
+            withBigtableTracer,
+            convertUnaryToServerStreamingSettings(settings.mutateRowSettings()));
 
-    return createUserFacingUnaryCallable(
-        methodName, new MutateRowCallable(retrying, requestContext));
+    UnaryOverStreamingTracerModCallable<MutateRowRequest, MutateRowResponse> modCallable =
+        new UnaryOverStreamingTracerModCallable<>(retrying);
+    TracedServerStreamingCallable<MutateRowRequest, MutateRowResponse> traced =
+        new TracedServerStreamingCallable<>(
+            modCallable, clientContext.getTracerFactory(), getSpanName(methodName));
+    UnaryOverStreamingCallable<MutateRowRequest, MutateRowResponse> callable =
+        new UnaryOverStreamingCallable<>(traced);
+    MutateRowCallable mutateRowCallable = new MutateRowCallable(callable, requestContext);
+    return mutateRowCallable.withDefaultCallContext(clientContext.getDefaultCallContext());
   }
 
   /**
@@ -1512,6 +1527,18 @@ private SpanName getSpanName(String methodName) {
     return SpanName.of(CLIENT_NAME, methodName);
   }
 
+  private <ReqT, RespT>
+      ServerStreamingCallSettings<ReqT, RespT> convertUnaryToServerStreamingSettings(
+          UnaryCallSettings<?, ?> unarySettings) {
+    return ServerStreamingCallSettings.<ReqT, RespT>newBuilder()
+        .setResumptionStrategy(new SimpleStreamResumptionStrategy<>())
+        .setRetryableCodes(unarySettings.getRetryableCodes())
+        .setRetrySettings(unarySettings.getRetrySettings())
+        .setIdleTimeoutDuration(Duration.ZERO)
+        .setWaitTimeoutDuration(Duration.ZERO)
+        .build();
+  }
+
   @Override
   public void close() {
     if (closeClientContext) {

@@ -18,6 +18,7 @@
 import com.google.api.core.BetaApi;
 import com.google.api.gax.rpc.ApiCallContext;
 import com.google.api.gax.tracing.ApiTracer;
+import com.google.api.gax.tracing.ApiTracerFactory;
 import com.google.api.gax.tracing.BaseApiTracer;
 import javax.annotation.Nullable;
 
@@ -86,4 +87,12 @@ public void setLocations(String zone, String cluster) {
   public void grpcChannelQueuedLatencies(long queuedTimeMs) {
     // noop
   }
+
+  public void overrideOperationType(ApiTracerFactory.OperationType operationType) {
+    // noop
+  }
+
+  public void operationFinishedEarly() {
+    // noop
+  }
 }
@@ -47,7 +47,7 @@
 class BuiltinMetricsTracer extends BigtableTracer {
 
   private static final String NAME = "java-bigtable/" + Version.VERSION;
-  private final OperationType operationType;
+  private volatile OperationType operationType;
   private final SpanName spanName;
 
   // Operation level metrics
@@ -132,6 +132,21 @@ public void close() {}
     };
   }
 
+  @Override
+  public void overrideOperationType(OperationType operationType) {
+    this.operationType = operationType;
+  }
+
+  @Override
+  public void operationFinishedEarly() {
+    if (attemptTimer.isRunning()) {
+      attemptTimer.stop();
+    }
+    if (operationTimer.isRunning()) {
+      operationTimer.stop();
+    }
+  }
+
   @Override
   public void operationSucceeded() {
     recordOperationCompletion(null);
@@ -276,7 +291,11 @@ private void recordOperationCompletion(@Nullable Throwable status) {
     if (!opFinished.compareAndSet(false, true)) {
       return;
     }
-    operationTimer.stop();
+
+    // timer might've been stopped early because due to #operationFinishedEarly()
+    if (operationTimer.isRunning()) {
+      operationTimer.stop();
+    }
 
     boolean isStreaming = operationType == OperationType.ServerStreaming;
     String statusStr = Util.extractStatus(status);

@@ -16,6 +16,7 @@
 package com.google.cloud.bigtable.data.v2.stub.metrics;
 
 import com.google.api.gax.tracing.ApiTracer;
+import com.google.api.gax.tracing.ApiTracerFactory;
 import com.google.common.collect.ImmutableList;
 import java.util.ArrayList;
 import java.util.List;
@@ -62,6 +63,24 @@ public void close() {
     };
   }
 
+  @Override
+  public void overrideOperationType(ApiTracerFactory.OperationType operationType) {
+    for (ApiTracer child : children) {
+      if (child instanceof BigtableTracer) {
+        ((BigtableTracer) child).overrideOperationType(operationType);
+      }
+    }
+  }
+
+  @Override
+  public void operationFinishedEarly() {
+    for (ApiTracer child : children) {
+      if (child instanceof BigtableTracer) {
+        ((BigtableTracer) child).operationFinishedEarly();
+      }
+    }
+  }
+
   @Override
   public void operationSucceeded() {
     for (ApiTracer child : children) {

@@ -36,7 +36,7 @@
 
 class MetricsTracer extends BigtableTracer {
 
-  private final OperationType operationType;
+  private volatile OperationType operationType;
 
   private final Tagger tagger;
   private final StatsRecorder stats;
@@ -84,6 +84,17 @@ public void close() {}
     };
   }
 
+  @Override
+  public void overrideOperationType(OperationType operationType) {
+    this.operationType = operationType;
+  }
+
+  @Override
+  public void operationFinishedEarly() {
+    attemptTimer.stop();
+    operationTimer.stop();
+  }
+
   @Override
   public void operationSucceeded() {
     recordOperationCompletion(null);
@@ -103,7 +114,9 @@ private void recordOperationCompletion(@Nullable Throwable throwable) {
     if (!opFinished.compareAndSet(false, true)) {
       return;
     }
-    operationTimer.stop();
+    if (operationTimer.isRunning()) {
+      operationTimer.stop();
+    }
 
     long elapsed = operationTimer.elapsed(TimeUnit.MILLISECONDS);
 

@@ -0,0 +1,56 @@
+ package com.google.cloud.bigtable.data.v2.stub.opt;
+
+ import com.google.api.core.AbstractApiFuture;
+ import com.google.api.core.ApiFuture;
+ import com.google.api.gax.rpc.ApiCallContext;
+ import com.google.api.gax.rpc.ResponseObserver;
+ import com.google.api.gax.rpc.ServerStreamingCallable;
+ import com.google.api.gax.rpc.StreamController;
+ import com.google.api.gax.rpc.UnaryCallable;
+
+ public class UnaryOverStreamingCallable<ReqT, RespT> extends UnaryCallable<ReqT, RespT> {
+  private final ServerStreamingCallable<ReqT, RespT> inner;
+
+  public UnaryOverStreamingCallable(ServerStreamingCallable<ReqT, RespT> inner) {
+    this.inner = inner;
+  }
+
+  @Override
+  public ApiFuture<RespT> futureCall(ReqT request, ApiCallContext context) {
+    UnaryFuture<RespT> future = new UnaryFuture<>();
+    inner.call(request, future.observer, context);
+    return future;
+  }
+
+  private static class UnaryFuture<RespT> extends AbstractApiFuture<RespT> {
+    private StreamController controller;
+
+    private final ResponseObserver<RespT> observer =
+        new ResponseObserver<RespT>() {
+          @Override
+          public void onStart(StreamController streamController) {
+            UnaryFuture.this.controller = streamController;
+          }
+
+          @Override
+          public void onResponse(RespT response) {
+            UnaryFuture.this.set(response);
+          }
+
+          @Override
+          public void onError(Throwable throwable) {
+            UnaryFuture.this.setException(throwable);
+          }
+
+          @Override
+          public void onComplete() {
+            UnaryFuture.this.set(null);
+          }
+        };
+
+    @Override
+    protected void interruptTask() {
+      controller.cancel();
+    }
+  }
+ }
@@ -0,0 +1,56 @@
+ package com.google.cloud.bigtable.data.v2.stub.opt;
+
+ import com.google.api.gax.rpc.ApiCallContext;
+ import com.google.api.gax.rpc.ResponseObserver;
+ import com.google.api.gax.rpc.ServerStreamingCallable;
+ import com.google.api.gax.rpc.StreamController;
+ import com.google.api.gax.tracing.ApiTracerFactory;
+ import com.google.cloud.bigtable.data.v2.stub.metrics.BigtableTracer;
+
+ public class UnaryOverStreamingTracerModCallable<ReqT, RespT>
+    extends ServerStreamingCallable<ReqT, RespT> {
+  private final ServerStreamingCallable<ReqT, RespT> inner;
+
+  public UnaryOverStreamingTracerModCallable(ServerStreamingCallable<ReqT, RespT> inner) {
+    this.inner = inner;
+  }
+
+  @Override
+  public void call(ReqT request, ResponseObserver<RespT> responseObserver, ApiCallContext context)
+ {
+    BigtableTracer tracer = (BigtableTracer) context.getTracer();
+    tracer.overrideOperationType(ApiTracerFactory.OperationType.Unary);
+    inner.call(request, new TracerModObserver<>(responseObserver, tracer), context);
+  }
+
+  static class TracerModObserver<RespT> implements ResponseObserver<RespT> {
+    private final ResponseObserver<RespT> outer;
+    private final BigtableTracer tracer;
+
+    public TracerModObserver(ResponseObserver<RespT> outer, BigtableTracer tracer) {
+      this.outer = outer;
+      this.tracer = tracer;
+    }
+
+    @Override
+    public void onStart(StreamController controller) {
+      outer.onStart(controller);
+    }
+
+    @Override
+    public void onResponse(RespT response) {
+      outer.onResponse(response);
+      tracer.operationFinishedEarly();
+    }
+
+    @Override
+    public void onError(Throwable t) {
+      outer.onError(t);
+    }
+
+    @Override
+    public void onComplete() {
+      outer.onComplete();
+    }
+  }
+ }