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>
@@ -136,7 +148,6 @@
       <type>test-jar</type>
       <scope>test</scope>
     </dependency>
-
     <dependency>
       <groupId>com.google.guava</groupId>
       <artifactId>guava-testlib</artifactId>
@@ -178,12 +189,56 @@
       <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>
       <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>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();