From 730aa427d8562ce3906fb80c3809cb5a27af0ac2 Mon Sep 17 00:00:00 2001 From: Jimit Shah <57637300+jimit-j-shah@users.noreply.github.com> Date: Mon, 20 May 2024 11:13:55 -0700 Subject: [PATCH] feat: Adding Commit RPC Trace Instrumentation (#1440) - 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 --- .../google/cloud/datastore/DatastoreImpl.java | 11 +- .../cloud/datastore/telemetry/TraceUtil.java | 3 + .../cloud/datastore/it/ITE2ETracingTest.java | 109 +++++++++++++++++- 3 files changed, 115 insertions(+), 8 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 401e3ed54..cb60685c7 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 @@ -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 index c867a5525..fe4effb1d 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,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 index 24a9a7149..2fb139289 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 @@ -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 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 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 entityList = new ArrayList<>(); + entityList.add(entity1); + entityList.add(entity2); + + List 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); + } }