Ehsann/tracing merge main

google-cloud-firestore/pom.xml

@@ -16,6 +16,7 @@
   </parent>
   <properties>
     <site.installationModule>google-cloud-firestore</site.installationModule>
+    <opentelemetry.version>1.29.0</opentelemetry.version>
   </properties>
   <dependencies>
     <dependency>
@@ -39,10 +40,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 +88,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>
@@ -113,6 +106,20 @@
       <artifactId>protobuf-java-util</artifactId>
     </dependency>
 
+    <!-- OpenTelemetry -->
+    <dependency>
+      <groupId>io.opentelemetry</groupId>
+      <artifactId>opentelemetry-api</artifactId>
+    </dependency>
+    <dependency>
+      <groupId>io.opentelemetry</groupId>
+      <artifactId>opentelemetry-context</artifactId>
+    </dependency>
+    <dependency>
+      <groupId>io.opentelemetry.instrumentation</groupId>
+      <artifactId>opentelemetry-grpc-1.6</artifactId>
+    </dependency>
+    <!-- END OpenTelemetry -->
 
     <!-- Test dependencies -->
     <dependency>
@@ -173,6 +180,78 @@
       <version>3.14.0</version>
       <scope>test</scope>
     </dependency>
+    <!-- OpenTelemetry -->
+    <dependency>
+      <groupId>io.opentelemetry</groupId>
+      <artifactId>opentelemetry-sdk</artifactId>
+      <version>${opentelemetry.version}</version>
+      <scope>test</scope>
+    </dependency>
+    <dependency>
+      <groupId>io.opentelemetry</groupId>
+      <artifactId>opentelemetry-sdk-testing</artifactId>
+      <version>${opentelemetry.version}</version>
+      <scope>test</scope>
+    </dependency>
+    <dependency>
+      <groupId>io.opentelemetry</groupId>
+      <artifactId>opentelemetry-semconv</artifactId>
+      <version>${opentelemetry.version}-alpha</version>
+      <scope>test</scope>
+    </dependency>
+    <dependency>
+      <groupId>io.opentelemetry</groupId>
+      <artifactId>opentelemetry-sdk-trace</artifactId>
+      <version>${opentelemetry.version}</version>
+      <scope>test</scope>
+    </dependency>
+    <dependency>
+      <groupId>io.opentelemetry</groupId>
+      <artifactId>opentelemetry-sdk-common</artifactId>
+      <version>${opentelemetry.version}</version>
+      <scope>test</scope>
+    </dependency>
+    <dependency>
+      <groupId>com.google.cloud.opentelemetry</groupId>
+      <artifactId>exporter-trace</artifactId>
+      <version>0.15.0</version>
+      <scope>test</scope>
+    </dependency>
+    <!-- END OpenTelemetry -->
+    <!-- Cloud Ops -->
+    <dependency>
+      <groupId>com.google.api.grpc</groupId>
+      <artifactId>proto-google-cloud-trace-v1</artifactId>
+      <version>1.3.0</version>
+      <scope>test</scope>
+    </dependency>
+    <dependency>
+      <groupId>com.google.cloud.opentelemetry</groupId>
+      <artifactId>exporter-trace</artifactId>
+      <version>0.15.0</version>
+      <scope>test</scope>
+    </dependency>
+    <!-- END OpenTelemetry -->
+    <!-- Cloud Ops -->
+    <dependency>
+      <groupId>com.google.api.grpc</groupId>
+      <artifactId>proto-google-cloud-trace-v1</artifactId>
+      <version>1.3.0</version>
+      <scope>test</scope>
+    </dependency>
+    <dependency>
+      <groupId>com.google.cloud</groupId>
+      <artifactId>google-cloud-trace</artifactId>
+      <version>1.3.0</version>
+      <scope>test</scope>
+    </dependency>
+    <dependency>
+      <groupId>com.google.testparameterinjector</groupId>
+      <artifactId>test-parameter-injector</artifactId>
+      <version>1.15</version>
+      <scope>test</scope>
+    </dependency>
+    <!-- END Cloud Ops -->
   </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,6 +26,8 @@
 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;
@@ -35,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 +64,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() {
@@ -84,35 +94,58 @@ public ApiFuture<AggregateQuerySnapshot> get() {
    *     query execution (if any), and the query results (if any).
    */
   @Nonnull
-  public ApiFuture<ExplainResults<AggregateQuerySnapshot>> explain(ExplainOptions options) {
-    AggregateQueryExplainResponseDeliverer responseDeliverer =
-        new AggregateQueryExplainResponseDeliverer(
-            /* transactionId= */ null,
-            /* readTime= */ null,
-            /* startTimeNanos= */ query.rpcContext.getClock().nanoTime(),
-            /* explainOptions= */ options);
-    runQuery(responseDeliverer);
-    return responseDeliverer.getFuture();
+  ApiFuture<AggregateQuerySnapshot> get(
+      @Nullable final ByteString transactionId, @Nullable com.google.protobuf.Timestamp readTime) {
+    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,
+              readTime,
+              /* 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;
+    }
   }
 
   @Nonnull
-  ApiFuture<AggregateQuerySnapshot> get(
-      @Nullable final ByteString transactionId, @Nullable com.google.protobuf.Timestamp readTime) {
-    AggregateQueryResponseDeliverer responseDeliverer =
-        new AggregateQueryResponseDeliverer(
-            transactionId, readTime, /* startTimeNanos= */ query.rpcContext.getClock().nanoTime());
-    runQuery(responseDeliverer);
-    return responseDeliverer.getFuture();
+  public ApiFuture<ExplainResults<AggregateQuerySnapshot>> explain(ExplainOptions options) {
+    TraceUtil.Span span = getTraceUtil().startSpan(TraceUtil.SPAN_NAME_AGGREGATION_QUERY_GET);
+    try (Scope ignored = span.makeCurrent()) {
+      AggregateQueryExplainResponseDeliverer responseDeliverer =
+          new AggregateQueryExplainResponseDeliverer(
+              /* transactionId= */ null,
+              /* readTime= */ null,
+              /* startTimeNanos= */ query.rpcContext.getClock().nanoTime(),
+              /* explainOptions= */ options);
+      runQuery(responseDeliverer, /* attempt */ 0);
+      ApiFuture<ExplainResults<AggregateQuerySnapshot>> result = responseDeliverer.getFuture();
+      span.endAtFuture(result);
+      return result;
+    } catch (Exception error) {
+      span.end(error);
+      throw error;
+    }
   }
 
-  private <T> void runQuery(ResponseDeliverer<T> responseDeliverer) {
+  private <T> void runQuery(ResponseDeliverer<T> responseDeliverer, int attempt) {
     RunAggregationQueryRequest request =
         toProto(
             responseDeliverer.getTransactionId(),
             responseDeliverer.getReadTime(),
             responseDeliverer.getExplainOptions());
     AggregateQueryResponseObserver<T> responseObserver =
-        new AggregateQueryResponseObserver<T>(responseDeliverer);
+        new AggregateQueryResponseObserver<T>(responseDeliverer, attempt);
     ServerStreamingCallable<RunAggregationQueryRequest, RunAggregationQueryResponse> callable =
         query.rpcContext.getClient().runAggregationQueryCallable();
     query.rpcContext.streamRequest(request, responseObserver, callable);
@@ -249,20 +282,39 @@ private final class AggregateQueryResponseObserver<T>
     private Timestamp readTime = Timestamp.MAX_VALUE;
     @Nullable private Map<String, Value> aggregateFieldsMap = null;
     @Nullable private ExplainMetrics metrics = null;
+    private int attempt;
 
-    AggregateQueryResponseObserver(ResponseDeliverer<T> responseDeliverer) {
+    AggregateQueryResponseObserver(ResponseDeliverer<T> 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();
     }
 
     private boolean isExplainQuery() {
       return this.responseDeliverer.getExplainOptions() != null;
     }
 
     @Override
-    public void onStart(StreamController streamController) {}
+    public void onStart(StreamController streamController) {
+      getTraceUtil()
+          .currentSpan()
+          .addEvent(SPAN_NAME_RUN_AGGREGATION_QUERY + " Stream started.", getAttemptAttributes());
+    }
 
     @Override
     public void onResponse(RunAggregationQueryResponse response) {
+      getTraceUtil()
+          .currentSpan()
+          .addEvent(
+              SPAN_NAME_RUN_AGGREGATION_QUERY + " Response Received.", getAttemptAttributes());
       if (response.hasReadTime()) {
         readTime = Timestamp.fromProto(response.getReadTime());
       }
@@ -288,8 +340,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,18 +66,10 @@ 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() {
-
     // Follows same thread safety logic as `UpdateBuilder::commit`.
     committed = true;
     BatchWriteRequest request = buildBatchWriteRequest();
 
-    Tracing.getTracer()
-        .getCurrentSpan()
-        .addAnnotation(
-            TraceUtil.SPAN_NAME_BATCHWRITE,
-            ImmutableMap.of(
-                "numDocuments", AttributeValue.longAttributeValue(request.getWritesCount())));
-
     ApiFuture<BatchWriteResponse> response =
         processExceptions(
             firestore.sendRequest(request, firestore.getClient().batchWriteCallable()));