feat: Add tracing for ReserveIds operation (#1490)

- added end-to-end test
googleapis · Aug 5, 2024 · 975464a · 975464a
1 parent 552d04e
commit 975464a
Show file tree

Hide file tree

Showing 3 changed files with 27 additions and 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
@@ -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
@@ -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
@@ -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);