feat: tracing for aggregate queries, bulkwriter, partition queries, a…

google-cloud-firestore/pom.xml

@@ -39,10 +39,6 @@
       <artifactId>grpc-google-cloud-firestore-v1</artifactId>
       <scope>test</scope>
     </dependency>
-    <dependency>
-      <groupId>io.opencensus</groupId>
-      <artifactId>opencensus-contrib-grpc-util</artifactId>
-    </dependency>
     <dependency>
       <groupId>com.google.code.findbugs</groupId>
       <artifactId>jsr305</artifactId>
@@ -91,10 +87,6 @@
       <groupId>io.grpc</groupId>
       <artifactId>grpc-stub</artifactId>
     </dependency>
-    <dependency>
-      <groupId>io.opencensus</groupId>
-      <artifactId>opencensus-api</artifactId>
-    </dependency>
     <dependency>
       <groupId>com.google.auth</groupId>
       <artifactId>google-auth-library-credentials</artifactId>
@@ -112,7 +104,23 @@
       <groupId>com.google.protobuf</groupId>
       <artifactId>protobuf-java-util</artifactId>
     </dependency>
-
+    <!-- OpenTelemetry -->
+    <dependency>
+      <groupId>io.opentelemetry</groupId>
+      <artifactId>opentelemetry-api</artifactId>
+      <version>1.29.0</version>
+    </dependency>
+    <dependency>
+      <groupId>io.opentelemetry</groupId>
+      <artifactId>opentelemetry-context</artifactId>
+      <version>1.29.0</version>
+    </dependency>
+    <dependency>
+      <groupId>io.opentelemetry.instrumentation</groupId>
+      <artifactId>opentelemetry-grpc-1.6</artifactId>
+      <version>1.29.0-alpha</version>
+    </dependency>
+    <!-- END OpenTelemetry -->
 
     <!-- Test dependencies -->
     <dependency>
@@ -173,6 +181,38 @@
       <version>3.14.0</version>
       <scope>test</scope>
     </dependency>
+    <!-- OpenTelemetry -->
+    <dependency>
+      <groupId>io.opentelemetry</groupId>
+      <artifactId>opentelemetry-sdk-testing</artifactId>
+      <version>1.32.0</version>
+      <scope>test</scope>
+    </dependency>
+    <dependency>
+      <groupId>io.opentelemetry</groupId>
+      <artifactId>opentelemetry-semconv</artifactId>
+      <version>1.29.0-alpha</version>
+      <scope>test</scope>
+    </dependency>
+    <dependency>
+      <groupId>io.opentelemetry</groupId>
+      <artifactId>opentelemetry-sdk-trace</artifactId>
+      <version>1.29.0</version>
+      <scope>test</scope>
+    </dependency>
+    <dependency>
+      <groupId>io.opentelemetry</groupId>
+      <artifactId>opentelemetry-sdk</artifactId>
+      <version>1.29.0</version>
+      <scope>test</scope>
+    </dependency>
+    <dependency>
+      <groupId>io.opentelemetry</groupId>
+      <artifactId>opentelemetry-sdk-common</artifactId>
+      <version>1.29.0</version>
+      <scope>test</scope>
+    </dependency>
+    <!-- END OpenTelemetry -->
   </dependencies>
 
   <reporting>

google-cloud-firestore/src/main/java/com/google/cloud/firestore/AggregateQuery.java

@@ -16,6 +16,8 @@
 
 package com.google.cloud.firestore;
 
+import static com.google.cloud.firestore.telemetry.TraceUtil.SPAN_NAME_RUN_AGGREGATION_QUERY;
+
 import com.google.api.core.ApiFuture;
 import com.google.api.core.InternalExtensionOnly;
 import com.google.api.core.SettableApiFuture;
@@ -24,7 +26,10 @@
 import com.google.api.gax.rpc.StatusCode;
 import com.google.api.gax.rpc.StreamController;
 import com.google.cloud.Timestamp;
+import com.google.cloud.firestore.telemetry.TraceUtil;
+import com.google.cloud.firestore.telemetry.TraceUtil.Scope;
 import com.google.cloud.firestore.v1.FirestoreSettings;
+import com.google.common.collect.ImmutableMap;
 import com.google.firestore.v1.RunAggregationQueryRequest;
 import com.google.firestore.v1.RunAggregationQueryResponse;
 import com.google.firestore.v1.RunQueryRequest;
@@ -34,6 +39,7 @@
 import com.google.firestore.v1.Value;
 import com.google.protobuf.ByteString;
 import java.util.ArrayList;
+import java.util.Collections;
 import java.util.HashMap;
 import java.util.HashSet;
 import java.util.List;
@@ -59,6 +65,11 @@ public class AggregateQuery {
     this.aliasMap = new HashMap<>();
   }
 
+  @Nonnull
+  private TraceUtil getTraceUtil() {
+    return query.getFirestore().getOptions().getTraceUtil();
+  }
+
   /** Returns the query whose aggregations will be calculated by this object. */
   @Nonnull
   public Query getQuery() {
@@ -77,17 +88,30 @@ public ApiFuture<AggregateQuerySnapshot> get() {
 
   @Nonnull
   ApiFuture<AggregateQuerySnapshot> get(@Nullable final ByteString transactionId) {
-    AggregateQueryResponseDeliverer responseDeliverer =
-        new AggregateQueryResponseDeliverer(
-            transactionId, /* startTimeNanos= */ query.rpcContext.getClock().nanoTime());
-    runQuery(responseDeliverer);
-    return responseDeliverer.getFuture();
+    TraceUtil.Span span =
+        getTraceUtil()
+            .startSpan(
+                transactionId == null
+                    ? TraceUtil.SPAN_NAME_AGGREGATION_QUERY_GET
+                    : TraceUtil.SPAN_NAME_TRANSACTION_GET_AGGREGATION_QUERY);
+    try (Scope ignored = span.makeCurrent()) {
+      AggregateQueryResponseDeliverer responseDeliverer =
+          new AggregateQueryResponseDeliverer(
+              transactionId, /* startTimeNanos= */ query.rpcContext.getClock().nanoTime());
+      runQuery(responseDeliverer, /* attempt= */ 0);
+      ApiFuture<AggregateQuerySnapshot> result = responseDeliverer.getFuture();
+      span.endAtFuture(result);
+      return result;
+    } catch (Exception error) {
+      span.end(error);
+      throw error;
+    }
   }
 
-  private void runQuery(AggregateQueryResponseDeliverer responseDeliverer) {
+  private void runQuery(AggregateQueryResponseDeliverer responseDeliverer, int attempt) {
     RunAggregationQueryRequest request = toProto(responseDeliverer.getTransactionId());
     AggregateQueryResponseObserver responseObserver =
-        new AggregateQueryResponseObserver(responseDeliverer);
+        new AggregateQueryResponseObserver(responseDeliverer, attempt);
     ServerStreamingCallable<RunAggregationQueryRequest, RunAggregationQueryResponse> callable =
         query.rpcContext.getClient().runAggregationQueryCallable();
     query.rpcContext.streamRequest(request, responseObserver, callable);
@@ -138,18 +162,36 @@ private final class AggregateQueryResponseObserver
 
     private final AggregateQueryResponseDeliverer responseDeliverer;
     private StreamController streamController;
+    private int attempt;
 
-    AggregateQueryResponseObserver(AggregateQueryResponseDeliverer responseDeliverer) {
+    AggregateQueryResponseObserver(AggregateQueryResponseDeliverer responseDeliverer, int attempt) {
       this.responseDeliverer = responseDeliverer;
+      this.attempt = attempt;
+    }
+
+    Map<String, Object> getAttemptAttributes() {
+      ImmutableMap.Builder<String, Object> builder =
+          new ImmutableMap.Builder<String, Object>().put("isRetryAttempt", attempt > 0);
+      if (attempt > 0) {
+        builder.put("attemptNumber", attempt);
+      }
+      return builder.build();
     }
 
     @Override
     public void onStart(StreamController streamController) {
+      getTraceUtil()
+          .currentSpan()
+          .addEvent(SPAN_NAME_RUN_AGGREGATION_QUERY + " Stream started.", getAttemptAttributes());
       this.streamController = streamController;
     }
 
     @Override
     public void onResponse(RunAggregationQueryResponse response) {
+      getTraceUtil()
+          .currentSpan()
+          .addEvent(
+              SPAN_NAME_RUN_AGGREGATION_QUERY + " Response Received.", getAttemptAttributes());
       // Close the stream to avoid it dangling, since we're not expecting any more responses.
       streamController.cancel();
 
@@ -165,8 +207,19 @@ public void onResponse(RunAggregationQueryResponse response) {
     @Override
     public void onError(Throwable throwable) {
       if (shouldRetry(throwable)) {
-        runQuery(responseDeliverer);
+        getTraceUtil()
+            .currentSpan()
+            .addEvent(
+                SPAN_NAME_RUN_AGGREGATION_QUERY + ": Retryable Error",
+                Collections.singletonMap("error.message", throwable.getMessage()));
+
+        runQuery(responseDeliverer, attempt + 1);
       } else {
+        getTraceUtil()
+            .currentSpan()
+            .addEvent(
+                SPAN_NAME_RUN_AGGREGATION_QUERY + ": Error",
+                Collections.singletonMap("error.message", throwable.getMessage()));
         responseDeliverer.deliverError(throwable);
       }
     }

google-cloud-firestore/src/main/java/com/google/cloud/firestore/BulkCommitBatch.java

@@ -21,13 +21,10 @@
 import com.google.api.gax.rpc.ApiException;
 import com.google.cloud.Timestamp;
 import com.google.common.base.Preconditions;
-import com.google.common.collect.ImmutableMap;
 import com.google.common.util.concurrent.MoreExecutors;
 import com.google.firestore.v1.BatchWriteRequest;
 import com.google.firestore.v1.BatchWriteResponse;
 import io.grpc.Status;
-import io.opencensus.trace.AttributeValue;
-import io.opencensus.trace.Tracing;
 import java.util.ArrayList;
 import java.util.List;
 import java.util.Set;
@@ -69,12 +66,6 @@ ApiFuture<WriteResult> wrapResult(int writeIndex) {
    * <p>The writes in the batch are not applied atomically and can be applied out of order.
    */
   ApiFuture<Void> bulkCommit() {
-    Tracing.getTracer()
-        .getCurrentSpan()
-        .addAnnotation(
-            TraceUtil.SPAN_NAME_BATCHWRITE,
-            ImmutableMap.of("numDocuments", AttributeValue.longAttributeValue(getWrites().size())));
-
     final BatchWriteRequest.Builder request = BatchWriteRequest.newBuilder();
     request.setDatabase(firestore.getDatabaseName());
 

google-cloud-firestore/src/main/java/com/google/cloud/firestore/BulkWriter.java

@@ -26,6 +26,9 @@
 import com.google.api.core.SettableApiFuture;
 import com.google.api.gax.rpc.ApiException;
 import com.google.api.gax.rpc.StatusCode.Code;
+import com.google.cloud.firestore.telemetry.TraceUtil;
+import com.google.cloud.firestore.telemetry.TraceUtil.Context;
+import com.google.cloud.firestore.telemetry.TraceUtil.Scope;
 import com.google.cloud.firestore.v1.FirestoreSettings;
 import com.google.common.annotations.VisibleForTesting;
 import com.google.common.base.Preconditions;
@@ -221,6 +224,8 @@ enum OperationType {
   @GuardedBy("lock")
   private Executor errorExecutor;
 
+  Context traceContext;
+
   /**
    * Used to track when writes are enqueued. The user handler executors cannot be changed after a
    * write has been enqueued.
@@ -237,6 +242,7 @@ enum OperationType {
     this.successExecutor = MoreExecutors.directExecutor();
     this.errorExecutor = MoreExecutors.directExecutor();
     this.bulkCommitBatch = new BulkCommitBatch(firestore, bulkWriterExecutor, maxBatchSize);
+    this.traceContext = firestore.getOptions().getTraceUtil().currentContext();
 
     if (!options.getThrottlingEnabled()) {
       this.rateLimiter =
@@ -899,21 +905,32 @@ private void scheduleCurrentBatchLocked(final boolean flush) {
 
   /** Sends the provided batch once the rate limiter does not require any delay. */
   private void sendBatchLocked(final BulkCommitBatch batch, final boolean flush) {
-    // Send the batch if it is does not require any delay, or schedule another attempt after the
+    // Send the batch if it does not require any delay, or schedule another attempt after the
     // appropriate timeout.
     boolean underRateLimit = rateLimiter.tryMakeRequest(batch.getMutationsSize());
     if (underRateLimit) {
-      batch
-          .bulkCommit()
-          .addListener(
-              () -> {
-                if (flush) {
-                  synchronized (lock) {
-                    scheduleCurrentBatchLocked(/* flush= */ true);
-                  }
+      TraceUtil.Span span =
+          firestore
+              .getOptions()
+              .getTraceUtil()
+              .startSpan(TraceUtil.SPAN_NAME_BULK_WRITER_COMMIT, traceContext)
+              .setAttribute("numDocuments", batch.getWrites().size());
+      try (Scope ignored = span.makeCurrent()) {
+        ApiFuture<Void> result = batch.bulkCommit();
+        result.addListener(
+            () -> {
+              if (flush) {
+                synchronized (lock) {
+                  scheduleCurrentBatchLocked(/* flush= */ true);
                 }
-              },
-              bulkWriterExecutor);
+              }
+            },
+            bulkWriterExecutor);
+        span.endAtFuture(result);
+      } catch (Exception error) {
+        span.end(error);
+        throw error;
+      }
     } else {
       long delayMs = rateLimiter.getNextRequestDelayMs(batch.getMutationsSize());
       logger.log(Level.FINE, () -> String.format("Backing off for %d seconds", delayMs / 1000));