feat: add opencensus tracing/stats support

google-cloud-datastore/pom.xml

@@ -78,6 +78,10 @@
       <groupId>com.google.auth</groupId>
       <artifactId>google-auth-library-oauth2-http</artifactId>
     </dependency>
+    <dependency>
+      <groupId>io.opencensus</groupId>
+      <artifactId>opencensus-api</artifactId>
+    </dependency>
 
     <!-- Test dependencies -->
     <dependency>

google-cloud-datastore/src/main/java/com/google/cloud/datastore/DatastoreImpl.java

@@ -34,6 +34,9 @@
 import com.google.datastore.v1.ReserveIdsRequest;
 import com.google.datastore.v1.TransactionOptions;
 import com.google.protobuf.ByteString;
+import io.opencensus.common.Scope;
+import io.opencensus.trace.Span;
+import io.opencensus.trace.Status;
 import java.util.ArrayList;
 import java.util.Arrays;
 import java.util.Collections;
@@ -51,12 +54,14 @@ final class DatastoreImpl extends BaseService<DatastoreOptions> implements Datas
   private final RetrySettings retrySettings;
   private static final ExceptionHandler TRANSACTION_EXCEPTION_HANDLER =
       TransactionExceptionHandler.build();
+  private final TraceUtil traceUtil;
 
   DatastoreImpl(DatastoreOptions options) {
     super(options);
     this.datastoreRpc = options.getDatastoreRpcV1();
     retrySettings =
         MoreObjects.firstNonNull(options.getRetrySettings(), ServiceOptions.getNoRetrySettings());
+    traceUtil = TraceUtil.getInstance();
   }
 
   @Override
@@ -131,29 +136,41 @@ public T call() throws DatastoreException {
   @Override
   public <T> T runInTransaction(final TransactionCallable<T> callable) {
     final DatastoreImpl self = this;
+    Span span = traceUtil.startSpan(TraceUtil.SPAN_NAME_TRANSACTION);
+    Scope scope = traceUtil.getTracer().withSpan(span);
     try {
       return RetryHelper.runWithRetries(
           new ReadWriteTransactionCallable<T>(self, callable, null),
           retrySettings,
           TRANSACTION_EXCEPTION_HANDLER,
           getOptions().getClock());
     } catch (RetryHelperException e) {
+      span.setStatus(Status.UNKNOWN.withDescription(e.getMessage()));
       throw DatastoreException.translateAndThrow(e);
+    } finally {
+      scope.close();
+      span.end(TraceUtil.END_SPAN_OPTIONS);
     }
   }
 
   @Override
   public <T> T runInTransaction(
       final TransactionCallable<T> callable, TransactionOptions transactionOptions) {
     final DatastoreImpl self = this;
+    Span span = traceUtil.startSpan(TraceUtil.SPAN_NAME_TRANSACTION);
+    Scope scope = traceUtil.getTracer().withSpan(span);
     try {
       return RetryHelper.runWithRetries(
           new ReadWriteTransactionCallable<T>(self, callable, transactionOptions),
           retrySettings,
           TRANSACTION_EXCEPTION_HANDLER,
           getOptions().getClock());
     } catch (RetryHelperException e) {
+      span.setStatus(Status.UNKNOWN.withDescription(e.getMessage()));
       throw DatastoreException.translateAndThrow(e);
+    } finally {
+      scope.close();
+      span.end(TraceUtil.END_SPAN_OPTIONS);
     }
   }
 
@@ -173,6 +190,8 @@ <T> QueryResults<T> run(com.google.datastore.v1.ReadOptions readOptionsPb, Query
 
   com.google.datastore.v1.RunQueryResponse runQuery(
       final com.google.datastore.v1.RunQueryRequest requestPb) {
+    Span span = traceUtil.startSpan(TraceUtil.SPAN_NAME_RUNQUERY);
+    Scope scope = traceUtil.getTracer().withSpan(span);
     try {
       return RetryHelper.runWithRetries(
           new Callable<com.google.datastore.v1.RunQueryResponse>() {
@@ -185,7 +204,11 @@ public com.google.datastore.v1.RunQueryResponse call() throws DatastoreException
           EXCEPTION_HANDLER,
           getOptions().getClock());
     } catch (RetryHelperException e) {
+      span.setStatus(Status.UNKNOWN.withDescription(e.getMessage()));
       throw DatastoreException.translateAndThrow(e);
+    } finally {
+      scope.close();
+      span.end(TraceUtil.END_SPAN_OPTIONS);
     }
   }
 
@@ -225,6 +248,8 @@ public List<Key> allocateId(IncompleteKey... keys) {
 
   private com.google.datastore.v1.AllocateIdsResponse allocateIds(
       final com.google.datastore.v1.AllocateIdsRequest requestPb) {
+    Span span = traceUtil.startSpan(TraceUtil.SPAN_NAME_ALLOCATEIDS);
+    Scope scope = traceUtil.getTracer().withSpan(span);
     try {
       return RetryHelper.runWithRetries(
           new Callable<com.google.datastore.v1.AllocateIdsResponse>() {
@@ -237,7 +262,11 @@ public com.google.datastore.v1.AllocateIdsResponse call() throws DatastoreExcept
           EXCEPTION_HANDLER,
           getOptions().getClock());
     } catch (RetryHelperException e) {
+      span.setStatus(Status.UNKNOWN.withDescription(e.getMessage()));
       throw DatastoreException.translateAndThrow(e);
+    } finally {
+      scope.close();
+      span.end(TraceUtil.END_SPAN_OPTIONS);
     }
   }
 
@@ -385,6 +414,8 @@ protected Entity computeNext() {
 
   com.google.datastore.v1.LookupResponse lookup(
       final com.google.datastore.v1.LookupRequest requestPb) {
+    Span span = traceUtil.startSpan(TraceUtil.SPAN_NAME_LOOKUP);
+    Scope scope = traceUtil.getTracer().withSpan(span);
     try {
       return RetryHelper.runWithRetries(
           new Callable<com.google.datastore.v1.LookupResponse>() {
@@ -397,7 +428,11 @@ public com.google.datastore.v1.LookupResponse call() throws DatastoreException {
           EXCEPTION_HANDLER,
           getOptions().getClock());
     } catch (RetryHelperException e) {
+      span.setStatus(Status.UNKNOWN.withDescription(e.getMessage()));
       throw DatastoreException.translateAndThrow(e);
+    } finally {
+      scope.close();
+      span.end(TraceUtil.END_SPAN_OPTIONS);
     }
   }
 
@@ -419,6 +454,8 @@ public List<Key> reserveIds(Key... keys) {
 
   com.google.datastore.v1.ReserveIdsResponse reserveIds(
       final com.google.datastore.v1.ReserveIdsRequest requestPb) {
+    Span span = traceUtil.startSpan(TraceUtil.SPAN_NAME_RESERVEIDS);
+    Scope scope = traceUtil.getTracer().withSpan(span);
     try {
       return RetryHelper.runWithRetries(
           new Callable<com.google.datastore.v1.ReserveIdsResponse>() {
@@ -431,7 +468,11 @@ public com.google.datastore.v1.ReserveIdsResponse call() throws DatastoreExcepti
           EXCEPTION_HANDLER,
           getOptions().getClock());
     } catch (RetryHelperException e) {
+      span.setStatus(Status.UNKNOWN.withDescription(e.getMessage()));
       throw DatastoreException.translateAndThrow(e);
+    } finally {
+      scope.close();
+      span.end(TraceUtil.END_SPAN_OPTIONS);
     }
   }
 
@@ -523,6 +564,8 @@ private com.google.datastore.v1.CommitResponse commitMutation(
 
   com.google.datastore.v1.CommitResponse commit(
       final com.google.datastore.v1.CommitRequest requestPb) {
+    Span span = traceUtil.startSpan(TraceUtil.SPAN_NAME_COMMIT);
+    Scope scope = traceUtil.getTracer().withSpan(span);
     try {
       return RetryHelper.runWithRetries(
           new Callable<com.google.datastore.v1.CommitResponse>() {
@@ -535,7 +578,11 @@ public com.google.datastore.v1.CommitResponse call() throws DatastoreException {
           EXCEPTION_HANDLER,
           getOptions().getClock());
     } catch (RetryHelperException e) {
+      span.setStatus(Status.UNKNOWN.withDescription(e.getMessage()));
       throw DatastoreException.translateAndThrow(e);
+    } finally {
+      scope.close();
+      span.end(TraceUtil.END_SPAN_OPTIONS);
     }
   }
 
@@ -546,6 +593,8 @@ ByteString requestTransactionId(
 
   com.google.datastore.v1.BeginTransactionResponse beginTransaction(
       final com.google.datastore.v1.BeginTransactionRequest requestPb) {
+    Span span = traceUtil.startSpan(TraceUtil.SPAN_NAME_BEGINTRANSACTION);
+    Scope scope = traceUtil.getTracer().withSpan(span);
     try {
       return RetryHelper.runWithRetries(
           new Callable<com.google.datastore.v1.BeginTransactionResponse>() {
@@ -559,7 +608,11 @@ public com.google.datastore.v1.BeginTransactionResponse call()
           EXCEPTION_HANDLER,
           getOptions().getClock());
     } catch (RetryHelperException e) {
+      span.setStatus(Status.UNKNOWN.withDescription(e.getMessage()));
       throw DatastoreException.translateAndThrow(e);
+    } finally {
+      scope.close();
+      span.end(TraceUtil.END_SPAN_OPTIONS);
     }
   }
 
@@ -571,6 +624,8 @@ void rollbackTransaction(ByteString transaction) {
   }
 
   void rollback(final com.google.datastore.v1.RollbackRequest requestPb) {
+    Span span = traceUtil.startSpan(TraceUtil.SPAN_NAME_ROLLBACK);
+    Scope scope = traceUtil.getTracer().withSpan(span);
     try {
       RetryHelper.runWithRetries(
           new Callable<Void>() {
@@ -584,7 +639,11 @@ public Void call() throws DatastoreException {
           EXCEPTION_HANDLER,
           getOptions().getClock());
     } catch (RetryHelperException e) {
+      span.setStatus(Status.UNKNOWN.withDescription(e.getMessage()));
       throw DatastoreException.translateAndThrow(e);
+    } finally {
+      scope.close();
+      span.end(TraceUtil.END_SPAN_OPTIONS);
     }
   }
 }

google-cloud-datastore/src/main/java/com/google/cloud/datastore/TraceUtil.java

@@ -0,0 +1,77 @@
+/*
+ * Copyright 2020 Google LLC
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *       http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package com.google.cloud.datastore;
+
+import com.google.cloud.datastore.spi.v1.HttpDatastoreRpc;
+import io.opencensus.trace.EndSpanOptions;
+import io.opencensus.trace.Span;
+import io.opencensus.trace.Tracer;
+import io.opencensus.trace.Tracing;
+
+/**
+ * Helper class for tracing utility. It is used for instrumenting {@link HttpDatastoreRpc} with Open
+ * Census APIs.
+ *
+ * <p>TraceUtil Instance are created by the {@link TraceUtil#getInstance()} method.
+ */
+public class TraceUtil {
+  private final Tracer tracer = Tracing.getTracer();
+  private static TraceUtil traceUtil;
+  static final String SPAN_NAME_ALLOCATEIDS = "CloudDatastoreOperation.allocateIds";
+  static final String SPAN_NAME_TRANSACTION = "CloudDatastoreOperation.readWriteTransaction";
+  static final String SPAN_NAME_BEGINTRANSACTION = "CloudDatastoreOperation.beginTransaction";
+  static final String SPAN_NAME_COMMIT = "CloudDatastoreOperation.commit";
+  static final String SPAN_NAME_LOOKUP = "CloudDatastoreOperation.lookup";
+  static final String SPAN_NAME_RESERVEIDS = "CloudDatastoreOperation.reserveIds";
+  static final String SPAN_NAME_ROLLBACK = "CloudDatastoreOperation.rollback";
+  static final String SPAN_NAME_RUNQUERY = "CloudDatastoreOperation.runQuery";
+  static final EndSpanOptions END_SPAN_OPTIONS =
+      EndSpanOptions.builder().setSampleToLocalSpanStore(true).build();
+
+  /**
+   * Starts a new span.
+   *
+   * @param spanName The name of the returned Span.
+   * @return The newly created {@link Span}.
+   */
+  protected Span startSpan(String spanName) {
+    return tracer.spanBuilder(spanName).startSpan();
+  }
+
+  /**
+   * Return the global {@link Tracer}.
+   *
+   * @return The global {@link Tracer}.
+   */
+  public Tracer getTracer() {
+    return tracer;
+  }
+
+  /**
+   * Return TraceUtil Object.
+   *
+   * @return An instance of {@link TraceUtil}
+   */
+  public static TraceUtil getInstance() {
+    if (traceUtil == null) {
+      traceUtil = new TraceUtil();
+    }
+    return traceUtil;
+  }
+
+  private TraceUtil() {}
+}

google-cloud-datastore/src/main/java/com/google/cloud/datastore/spi/v1/HttpDatastoreRpc.java

@@ -21,6 +21,8 @@
 import com.google.api.client.http.HttpTransport;
 import com.google.cloud.datastore.DatastoreException;
 import com.google.cloud.datastore.DatastoreOptions;
+import com.google.cloud.datastore.TraceUtil;
+import com.google.cloud.http.CensusHttpModule;
 import com.google.cloud.http.HttpTransportOptions;
 import com.google.datastore.v1.AllocateIdsRequest;
 import com.google.datastore.v1.AllocateIdsResponse;
@@ -75,12 +77,18 @@ public HttpDatastoreRpc(DatastoreOptions options) {
 
   private HttpRequestInitializer getHttpRequestInitializer(
       final DatastoreOptions options, HttpTransportOptions httpTransportOptions) {
-    final HttpRequestInitializer delegate = httpTransportOptions.getHttpRequestInitializer(options);
+    HttpRequestInitializer delegate = httpTransportOptions.getHttpRequestInitializer(options);
+    // Open Census initialization
+    CensusHttpModule censusHttpModule =
+        new CensusHttpModule(TraceUtil.getInstance().getTracer(), true);
+    delegate = censusHttpModule.getHttpRequestInitializer(delegate);
+
     final String applicationName = options.getApplicationName();
+    final HttpRequestInitializer localDelegate = delegate;
     return new HttpRequestInitializer() {
       @Override
       public void initialize(HttpRequest httpRequest) throws IOException {
-        delegate.initialize(httpRequest);
+        localDelegate.initialize(httpRequest);
         httpRequest.getHeaders().setUserAgent(applicationName);
       }
     };

pom.xml

@@ -162,6 +162,7 @@
     <guava.version>29.0-android</guava.version>
     <threeten.version>1.4.4</threeten.version>
     <javax.annotations.version>1.3.2</javax.annotations.version>
+    <opencensus.version>0.24.0</opencensus.version>
   </properties>
 
   <dependencyManagement>
@@ -255,7 +256,11 @@
         <artifactId>threetenbp</artifactId>
         <version>${threeten.version}</version>
       </dependency>
-
+      <dependency>
+        <groupId>io.opencensus</groupId>
+        <artifactId>opencensus-api</artifactId>
+        <version>${opencensus.version}</version>
+      </dependency>
       <dependency>
         <groupId>junit</groupId>
         <artifactId>junit</artifactId>