feat: Adding Lookup RPC OpenTelemetry Tracing

google-cloud-datastore/pom.xml

@@ -18,6 +18,18 @@
     <site.installationModule>google-cloud-datastore</site.installationModule>
     <opentelemetry.version>1.37.0</opentelemetry.version>
   </properties>
+  <dependencyManagement>
+    <dependencies>
+      <dependency>
+        <!-- Artifacts from google-cloud-java monorepo -->
+        <groupId>com.google.cloud</groupId>
+        <artifactId>gapic-libraries-bom</artifactId>
+        <version>1.37.0</version>
+        <type>pom</type>
+        <scope>import</scope>
+      </dependency>
+    </dependencies>
+  </dependencyManagement>
   <dependencies>
     <dependency>
       <groupId>com.google.api.grpc</groupId>
@@ -117,6 +129,11 @@
     </dependency>
 
     <!-- OpenTelemetry -->
+    <dependency>
+      <groupId>io.opentelemetry</groupId>
+      <artifactId>opentelemetry-sdk</artifactId>
+      <version>${opentelemetry.version}</version>
+    </dependency>
     <dependency>
       <groupId>io.opentelemetry</groupId>
       <artifactId>opentelemetry-api</artifactId>
@@ -136,7 +153,6 @@
       <type>test-jar</type>
       <scope>test</scope>
     </dependency>
-
     <dependency>
       <groupId>com.google.guava</groupId>
       <artifactId>guava-testlib</artifactId>
@@ -178,12 +194,50 @@
       <version>1.4.2</version>
       <scope>test</scope>
     </dependency>
+    <dependency>
+      <groupId>com.google.testparameterinjector</groupId>
+      <artifactId>test-parameter-injector</artifactId>
+      <version>1.16</version>
+      <scope>test</scope>
+    </dependency>
+    <!-- OpenTelemetry -->
     <dependency>
       <groupId>io.opentelemetry</groupId>
-      <artifactId>opentelemetry-sdk</artifactId>
+      <artifactId>opentelemetry-sdk-common</artifactId>
+      <version>${opentelemetry.version}</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-semconv</artifactId>
+      <version>1.1.0-alpha</version>
+      <scope>test</scope>
+    </dependency>
+    <!-- END OpenTelemetry -->
+    <!-- Cloud Ops -->
+    <dependency>
+      <groupId>com.google.cloud.opentelemetry</groupId>
+      <artifactId>exporter-trace</artifactId>
+      <version>0.15.0</version>
+      <scope>test</scope>
+    </dependency>
+    <dependency>
+      <groupId>com.google.cloud</groupId>
+      <artifactId>google-cloud-trace</artifactId>
+      <scope>test</scope>
+    </dependency>
+    <dependency>
+      <groupId>com.google.api.grpc</groupId>
+      <artifactId>proto-google-cloud-trace-v1</artifactId>
+      <scope>test</scope>
+    </dependency>
+    <!-- END Cloud Ops -->
   </dependencies>
 
   <build>

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

@@ -29,6 +29,7 @@
 import com.google.common.base.Preconditions;
 import com.google.common.collect.AbstractIterator;
 import com.google.common.collect.ImmutableList;
+import com.google.common.collect.ImmutableMap;
 import com.google.common.collect.Iterables;
 import com.google.common.collect.Sets;
 import com.google.datastore.v1.ExplainOptions;
@@ -61,6 +62,9 @@ final class DatastoreImpl extends BaseService<DatastoreOptions> implements Datas
       TransactionOperationExceptionHandler.build();
   private final TraceUtil traceUtil = TraceUtil.getInstance();
 
+  private final com.google.cloud.datastore.telemetry.TraceUtil otelTraceUtil =
+      getOptions().getTraceUtil();
+
   private final ReadOptionProtoPreparer readOptionProtoPreparer;
   private final AggregationQueryExecutor aggregationQueryExecutor;
 
@@ -450,13 +454,26 @@ protected Entity computeNext() {
 
   com.google.datastore.v1.LookupResponse lookup(
       final com.google.datastore.v1.LookupRequest requestPb) {
-    Span span = traceUtil.startSpan(TraceUtil.SPAN_NAME_LOOKUP);
-    try (Scope scope = traceUtil.getTracer().withSpan(span)) {
+    com.google.cloud.datastore.telemetry.TraceUtil.Span span =
+        otelTraceUtil.startSpan(com.google.cloud.datastore.telemetry.TraceUtil.SPAN_NAME_LOOKUP);
+    final ReadOptions readOptions = requestPb.getReadOptions();
+    span.setAttribute(
+        "isTransactional", (readOptions.hasTransaction() || readOptions.hasNewTransaction()));
+
+    try (com.google.cloud.datastore.telemetry.TraceUtil.Scope ignored = span.makeCurrent()) {
       return RetryHelper.runWithRetries(
           new Callable<com.google.datastore.v1.LookupResponse>() {
             @Override
             public com.google.datastore.v1.LookupResponse call() throws DatastoreException {
-              return datastoreRpc.lookup(requestPb);
+              com.google.datastore.v1.LookupResponse response = datastoreRpc.lookup(requestPb);
+              span.addEvent(
+                  com.google.cloud.datastore.telemetry.TraceUtil.SPAN_NAME_LOOKUP + ": Completed",
+                  new ImmutableMap.Builder<String, Object>()
+                      .put("Received", response.getFoundCount())
+                      .put("Missing", response.getMissingCount())
+                      .put("Deferred", response.getDeferredCount())
+                      .build());
+              return response;
             }
           },
           retrySettings,
@@ -465,10 +482,10 @@ public com.google.datastore.v1.LookupResponse call() throws DatastoreException {
               : TRANSACTION_OPERATION_EXCEPTION_HANDLER,
           getOptions().getClock());
     } catch (RetryHelperException e) {
-      span.setStatus(Status.UNKNOWN.withDescription(e.getMessage()));
+      span.end(e);
       throw DatastoreException.translateAndThrow(e);
     } finally {
-      span.end(TraceUtil.END_SPAN_OPTIONS);
+      span.end();
     }
   }
 

google-cloud-datastore/src/main/java/com/google/cloud/datastore/telemetry/DisabledTraceUtil.java

@@ -61,6 +61,11 @@ public TraceUtil.Span setAttribute(String key, String value) {
       return this;
     }
 
+    @Override
+    public TraceUtil.Span setAttribute(String key, boolean value) {
+      return this;
+    }
+
     @Override
     public Scope makeCurrent() {
       return new Scope();

google-cloud-datastore/src/main/java/com/google/cloud/datastore/telemetry/EnabledTraceUtil.java

@@ -175,6 +175,12 @@ public TraceUtil.Span setAttribute(String key, String value) {
       return this;
     }
 
+    @Override
+    public TraceUtil.Span setAttribute(String key, boolean value) {
+      span.setAttribute(ATTRIBUTE_SERVICE_PREFIX + key, value);
+      return this;
+    }
+
     @Override
     public Scope makeCurrent() {
       try (io.opentelemetry.context.Scope scope = span.makeCurrent()) {

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

@@ -32,6 +32,7 @@ public interface TraceUtil {
   static final String ENABLE_TRACING_ENV_VAR = "DATASTORE_ENABLE_TRACING";
   static final String LIBRARY_NAME = "com.google.cloud.datastore";
 
+  static final String SPAN_NAME_LOOKUP = "Lookup";
   /**
    * Creates and returns an instance of the TraceUtil class.
    *
@@ -80,6 +81,9 @@ interface Span {
     /** Adds the given attribute to this span. */
     Span setAttribute(String key, String value);
 
+    /** Adds the given attribute to this span. */
+    Span setAttribute(String key, boolean value);
+
     /** Marks this span as the current span. */
     Scope makeCurrent();
 

google-cloud-datastore/src/main/java/com/google/cloud/datastore/testing/RemoteDatastoreHelper.java

@@ -19,13 +19,16 @@
 import com.google.api.core.InternalApi;
 import com.google.api.gax.retrying.RetrySettings;
 import com.google.cloud.datastore.Datastore;
+import com.google.cloud.datastore.DatastoreOpenTelemetryOptions;
 import com.google.cloud.datastore.DatastoreOptions;
 import com.google.cloud.datastore.Key;
 import com.google.cloud.datastore.Query;
 import com.google.cloud.datastore.QueryResults;
 import com.google.cloud.datastore.StructuredQuery;
 import com.google.cloud.http.HttpTransportOptions;
+import io.opentelemetry.sdk.OpenTelemetrySdk;
 import java.util.UUID;
+import javax.annotation.Nullable;
 import org.threeten.bp.Duration;
 
 /**
@@ -38,13 +41,13 @@
  * RetrySettings#getTotalTimeout()} is {@code 120000} and {@link
  * RetrySettings#getInitialRetryDelay()} is {@code 250}. {@link
  * HttpTransportOptions#getConnectTimeout()} and {@link HttpTransportOptions#getReadTimeout()} are
- * both both set to {@code 60000}.
+ * both both set to {@code 60000}. If an OpenTelemetrySdk object is passed in, OpenTelemetry Trace
+ * collection will be enabled for the Client application.
  *
  * <p>Internal testing use only
  */
 @InternalApi
 public class RemoteDatastoreHelper {
-
   private final DatastoreOptions options;
   private final Datastore datastore;
   private final String namespace;
@@ -78,18 +81,30 @@ public static RemoteDatastoreHelper create() {
   }
 
   /** Creates a {@code RemoteStorageHelper} object. */
-  public static RemoteDatastoreHelper create(String databaseId) {
+  public static RemoteDatastoreHelper create(
+      String databaseId, @Nullable OpenTelemetrySdk openTelemetrySdk) {
     HttpTransportOptions transportOptions = DatastoreOptions.getDefaultHttpTransportOptions();
     transportOptions =
         transportOptions.toBuilder().setConnectTimeout(60000).setReadTimeout(60000).build();
-    DatastoreOptions datastoreOption =
+    DatastoreOptions.Builder datastoreOptionBuilder =
         DatastoreOptions.newBuilder()
             .setDatabaseId(databaseId)
             .setNamespace(UUID.randomUUID().toString())
             .setRetrySettings(retrySettings())
-            .setTransportOptions(transportOptions)
-            .build();
-    return new RemoteDatastoreHelper(datastoreOption);
+            .setTransportOptions(transportOptions);
+
+    if (openTelemetrySdk != null) {
+      datastoreOptionBuilder.setOpenTelemetryOptions(
+          DatastoreOpenTelemetryOptions.newBuilder()
+              .setOpenTelemetry(openTelemetrySdk)
+              .setTracingEnabled(true)
+              .build());
+    }
+    return new RemoteDatastoreHelper(datastoreOptionBuilder.build());
+  }
+
+  public static RemoteDatastoreHelper create(String databaseId) {
+    return create(databaseId, /*openTelemetrySdk=*/ null);
   }
 
   private static RetrySettings retrySettings() {