feat: Adding Commit RPC Trace Instrumentation (#1440)

jimit-j-shah · jimit-j-shah · commit 5fa9386bfa82 · 2024-06-13T10:05:00.000-07:00
- Added end-to-end test for Datastore operationsput, add, update and delete.
- Updated E2E Test to use the namespace correctly for efficient clean-up of test data
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
@@ -619,8 +619,11 @@ 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);
-    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_COMMIT);
+    span.setAttribute("isTransactional", requestPb.hasTransaction());
+
+    try (com.google.cloud.datastore.telemetry.TraceUtil.Scope ignored = span.makeCurrent()) {
       return RetryHelper.runWithRetries(
           () -> datastoreRpc.commit(requestPb),
           retrySettings,
@@ -629,10 +632,10 @@ com.google.datastore.v1.CommitResponse commit(
               : 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();
     }
   }
 
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,9 @@ public interface TraceUtil {
   static final String LIBRARY_NAME = "com.google.cloud.datastore";
 
   static final String SPAN_NAME_LOOKUP = "Lookup";
+
+  static final String SPAN_NAME_COMMIT = "Commit";
+
   /**
    * Creates and returns an instance of the TraceUtil class.
    *
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
@@ -16,6 +16,7 @@
 
 package com.google.cloud.datastore.it;
 
+import static com.google.cloud.datastore.telemetry.TraceUtil.SPAN_NAME_COMMIT;
 import static com.google.cloud.datastore.telemetry.TraceUtil.SPAN_NAME_LOOKUP;
 import static io.opentelemetry.semconv.resource.attributes.ResourceAttributes.SERVICE_NAME;
 import static org.junit.Assert.assertEquals;
@@ -207,6 +208,8 @@ private boolean dfsContainsCallStack(long spanId, List<String> expectedCallStack
 
   private static Key KEY1;
 
+  private static Key KEY2;
+
   // Random int generator for trace ID and span ID
   private static Random random;
 
@@ -233,9 +236,15 @@ private boolean dfsContainsCallStack(long spanId, List<String> expectedCallStack
 
   private static Datastore datastore;
 
+  private static RemoteDatastoreHelper remoteDatastoreHelper;
+
   @TestParameter boolean useGlobalOpenTelemetrySDK;
 
-  @TestParameter({"default", "test-db"})
+  @TestParameter({
+    /*(default)*/
+    "",
+    "test-db"
+  })
   String datastoreNamedDatabase;
 
   @BeforeClass
@@ -280,8 +289,7 @@ public void before() throws Exception {
     // but because gRPC traces need to be deterministically force-flushed for every test
     String namedDb = datastoreNamedDatabase();
     logger.log(Level.INFO, "Integration test using named database " + namedDb);
-    RemoteDatastoreHelper remoteDatastoreHelper =
-        RemoteDatastoreHelper.create(namedDb, openTelemetrySdk);
+    remoteDatastoreHelper = RemoteDatastoreHelper.create(namedDb, openTelemetrySdk);
     options = remoteDatastoreHelper.getOptions();
     datastore = options.getService();
 
@@ -292,7 +300,14 @@ public void before() throws Exception {
 
     String projectId = options.getProjectId();
     String kind1 = "kind1";
-    KEY1 = Key.newBuilder(projectId, kind1, "name", options.getDatabaseId()).build();
+    KEY1 =
+        Key.newBuilder(projectId, kind1, "name1", options.getDatabaseId())
+            .setNamespace(options.getNamespace())
+            .build();
+    KEY2 =
+        Key.newBuilder(projectId, kind1, "name2", options.getDatabaseId())
+            .setNamespace(options.getNamespace())
+            .build();
 
     // Set up the tracer for custom TraceID injection
     rootSpanName =
@@ -319,6 +334,7 @@ public void after() throws Exception {
     if (isUsingGlobalOpenTelemetrySDK()) {
       GlobalOpenTelemetry.resetForTest();
     }
+    remoteDatastoreHelper.deleteNamespace();
     rootSpanName = null;
     tracer = null;
     retrievedTrace = null;
@@ -527,4 +543,89 @@ public void lookupTraceTest() throws Exception {
 
     fetchAndValidateTrace(customSpanContext.getTraceId(), SPAN_NAME_LOOKUP);
   }
+
+  @Test
+  public void commitTraceTest() throws Exception {
+    assertNotNull(customSpanContext);
+
+    Span rootSpan = getNewRootSpanWithContext();
+
+    Entity entity1 = Entity.newBuilder(KEY1).set("test_key", "test_value").build();
+    try (Scope ignored = rootSpan.makeCurrent()) {
+      Entity response = datastore.add(entity1);
+      assertEquals(entity1, response);
+    } finally {
+      rootSpan.end();
+    }
+    waitForTracesToComplete();
+
+    fetchAndValidateTrace(customSpanContext.getTraceId(), SPAN_NAME_COMMIT);
+  }
+
+  @Test
+  public void putTraceTest() throws Exception {
+    assertNotNull(customSpanContext);
+
+    Span rootSpan = getNewRootSpanWithContext();
+
+    Entity entity1 = Entity.newBuilder(KEY1).set("test_key", "test_value").build();
+    try (Scope ignored = rootSpan.makeCurrent()) {
+      Entity response = datastore.put(entity1);
+      assertEquals(entity1, response);
+    } finally {
+      rootSpan.end();
+    }
+    waitForTracesToComplete();
+
+    fetchAndValidateTrace(customSpanContext.getTraceId(), SPAN_NAME_COMMIT);
+  }
+
+  @Test
+  public void updateTraceTest() throws Exception {
+    assertNotNull(customSpanContext);
+
+    Entity entity1 = Entity.newBuilder(KEY1).set("test_field", "test_value1").build();
+    Entity entity2 = Entity.newBuilder(KEY2).set("test_field", "test_value2").build();
+    List<Entity> entityList = new ArrayList<>();
+    entityList.add(entity1);
+    entityList.add(entity2);
+
+    List<Entity> response = datastore.add(entity1, entity2);
+    assertEquals(entityList, response);
+
+    Span rootSpan = getNewRootSpanWithContext();
+
+    try (Scope ignored = rootSpan.makeCurrent()) {
+      Entity entity1_update =
+          Entity.newBuilder(entity1).set("test_field", "new_test_value1").build();
+      Entity entity2_update =
+          Entity.newBuilder(entity2).set("test_field", "new_test_value1").build();
+      datastore.update(entity1_update, entity2_update);
+    } finally {
+      rootSpan.end();
+    }
+    waitForTracesToComplete();
+
+    fetchAndValidateTrace(customSpanContext.getTraceId(), SPAN_NAME_COMMIT);
+  }
+
+  @Test
+  public void deleteTraceTest() throws Exception {
+    assertNotNull(customSpanContext);
+
+    Entity entity1 = Entity.newBuilder(KEY1).set("test_key", "test_value").build();
+    Entity response = datastore.put(entity1);
+    assertEquals(entity1, response);
+
+    Span rootSpan = getNewRootSpanWithContext();
+
+    try (Scope ignored = rootSpan.makeCurrent()) {
+      datastore.delete(entity1.getKey());
+    } finally {
+      rootSpan.end();
+    }
+    waitForTracesToComplete();
+
+    fetchAndValidateTrace(customSpanContext.getTraceId(), SPAN_NAME_COMMIT);
+  }
 }

Original file line number	Diff line number	Diff line change
`@@ -33,6 +33,9 @@ public interface TraceUtil {`
`33`	`33`	`static final String LIBRARY_NAME = "com.google.cloud.datastore";`
`34`	`34`
`35`	`35`	`static final String SPAN_NAME_LOOKUP = "Lookup";`
	`36`	`+`
	`37`	`+ static final String SPAN_NAME_COMMIT = "Commit";`
	`38`	`+`
`36`	`39`	`/**`
`37`	`40`	`* Creates and returns an instance of the TraceUtil class.`
`38`	`41`	`*`