From 2c12c2f4d823741db19463a079d2a4f323567886 Mon Sep 17 00:00:00 2001
From: jimit-j-shah <57637300+jimit-j-shah@users.noreply.github.com>
Date: Thu, 20 Jun 2024 14:17:13 -0700
Subject: [PATCH] feat: Add tracing for ReserveIds operation - added end-to-end
 test

---
 .../google/cloud/datastore/DatastoreImpl.java | 10 ++++++----
 .../cloud/datastore/telemetry/TraceUtil.java  |  1 +
 .../cloud/datastore/it/ITE2ETracingTest.java  | 20 +++++++++++++++++++
 3 files changed, 27 insertions(+), 4 deletions(-)

diff --git a/google-cloud-datastore/src/main/java/com/google/cloud/datastore/DatastoreImpl.java b/google-cloud-datastore/src/main/java/com/google/cloud/datastore/DatastoreImpl.java
index f726f0bb2..52acb77d8 100644
--- a/google-cloud-datastore/src/main/java/com/google/cloud/datastore/DatastoreImpl.java
+++ b/google-cloud-datastore/src/main/java/com/google/cloud/datastore/DatastoreImpl.java
@@ -543,8 +543,10 @@ 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);
-    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_RESERVE_IDS);
+    try (com.google.cloud.datastore.telemetry.TraceUtil.Scope ignored = span.makeCurrent()) {
       return RetryHelper.runWithRetries(
           new Callable<com.google.datastore.v1.ReserveIdsResponse>() {
             @Override
@@ -556,10 +558,10 @@ public com.google.datastore.v1.ReserveIdsResponse call() throws DatastoreExcepti
           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();
     }
   }
 
diff --git a/google-cloud-datastore/src/main/java/com/google/cloud/datastore/telemetry/TraceUtil.java b/google-cloud-datastore/src/main/java/com/google/cloud/datastore/telemetry/TraceUtil.java
index 1198460a4..1e5226126 100644
--- a/google-cloud-datastore/src/main/java/com/google/cloud/datastore/telemetry/TraceUtil.java
+++ b/google-cloud-datastore/src/main/java/com/google/cloud/datastore/telemetry/TraceUtil.java
@@ -33,6 +33,7 @@ public interface TraceUtil {
   static final String LIBRARY_NAME = "com.google.cloud.datastore";
   static final String SPAN_NAME_LOOKUP = "Lookup";
   static final String SPAN_NAME_ALLOCATE_IDS = "AllocateIds";
+  static final String SPAN_NAME_RESERVE_IDS = "ReserveIds";
   static final String SPAN_NAME_COMMIT = "Commit";
   static final String SPAN_NAME_RUN_QUERY = "RunQuery";
   static final String SPAN_NAME_RUN_AGGREGATION_QUERY = "RunAggregationQuery";
diff --git a/google-cloud-datastore/src/test/java/com/google/cloud/datastore/it/ITE2ETracingTest.java b/google-cloud-datastore/src/test/java/com/google/cloud/datastore/it/ITE2ETracingTest.java
index a032762c8..bf7635266 100644
--- a/google-cloud-datastore/src/test/java/com/google/cloud/datastore/it/ITE2ETracingTest.java
+++ b/google-cloud-datastore/src/test/java/com/google/cloud/datastore/it/ITE2ETracingTest.java
@@ -21,6 +21,7 @@
 import static com.google.cloud.datastore.telemetry.TraceUtil.SPAN_NAME_BEGIN_TRANSACTION;
 import static com.google.cloud.datastore.telemetry.TraceUtil.SPAN_NAME_COMMIT;
 import static com.google.cloud.datastore.telemetry.TraceUtil.SPAN_NAME_LOOKUP;
+import static com.google.cloud.datastore.telemetry.TraceUtil.SPAN_NAME_RESERVE_IDS;
 import static com.google.cloud.datastore.telemetry.TraceUtil.SPAN_NAME_RUN_AGGREGATION_QUERY;
 import static com.google.cloud.datastore.telemetry.TraceUtil.SPAN_NAME_RUN_QUERY;
 import static com.google.cloud.datastore.telemetry.TraceUtil.SPAN_NAME_TRANSACTION_COMMIT;
@@ -602,6 +603,25 @@ public void allocateIdsTraceTest() throws Exception {
     fetchAndValidateTrace(customSpanContext.getTraceId(), SPAN_NAME_ALLOCATE_IDS);
   }
 
+  @Test
+  public void reserveIdsTraceTest() throws Exception {
+    assertNotNull(customSpanContext);
+
+    Span rootSpan = getNewRootSpanWithContext();
+    try (Scope ignored = rootSpan.makeCurrent()) {
+      KeyFactory keyFactory = datastore.newKeyFactory().setKind("MyKind");
+      Key key1 = keyFactory.newKey(10);
+      Key key2 = keyFactory.newKey("name");
+      List<Key> keyList = datastore.reserveIds(key1, key2);
+      assertEquals(2, keyList.size());
+    } finally {
+      rootSpan.end();
+    }
+    waitForTracesToComplete();
+
+    fetchAndValidateTrace(customSpanContext.getTraceId(), SPAN_NAME_RESERVE_IDS);
+  }
+
   @Test
   public void commitTraceTest() throws Exception {
     assertNotNull(customSpanContext);