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 d491f9a93..3e9081d66 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 @@ -157,7 +157,7 @@ private io.opentelemetry.api.trace.Span startSpanWithParentContext( .with( io.opentelemetry.api.trace.Span.wrap( parentSpanContext.getSpanContext()))); - return spanBuilder.startSpan(); + return otelTraceUtil.addSettingsAttributesToCurrentSpan(spanBuilder).startSpan(); } @Override @@ -540,7 +540,10 @@ com.google.datastore.v1.LookupResponse lookup( .put("Received", response.getFoundCount()) .put("Missing", response.getMissingCount()) .put("Deferred", response.getDeferredCount()) - .put("isTransactional", isTransactional) + .put("transactional", isTransactional) + .put( + "transaction_id", + isTransactional ? readOptions.getTransaction().toStringUtf8() : "") .build()); return response; }, @@ -772,6 +775,11 @@ public Void call() throws DatastoreException { retrySettings, EXCEPTION_HANDLER, getOptions().getClock()); + span.addEvent( + com.google.cloud.datastore.telemetry.TraceUtil.SPAN_NAME_ROLLBACK, + new ImmutableMap.Builder() + .put("transaction_id", requestPb.getTransaction().toStringUtf8()) + .build()); } catch (RetryHelperException e) { span.end(e); throw DatastoreException.translateAndThrow(e); diff --git a/google-cloud-datastore/src/main/java/com/google/cloud/datastore/telemetry/DisabledTraceUtil.java b/google-cloud-datastore/src/main/java/com/google/cloud/datastore/telemetry/DisabledTraceUtil.java index 6ba0fd81c..06941c721 100644 --- a/google-cloud-datastore/src/main/java/com/google/cloud/datastore/telemetry/DisabledTraceUtil.java +++ b/google-cloud-datastore/src/main/java/com/google/cloud/datastore/telemetry/DisabledTraceUtil.java @@ -22,6 +22,7 @@ import com.google.cloud.datastore.telemetry.TraceUtil.SpanContext; import io.grpc.ManagedChannelBuilder; import io.opentelemetry.api.trace.Span; +import io.opentelemetry.api.trace.SpanBuilder; import io.opentelemetry.api.trace.Tracer; import io.opentelemetry.api.trace.TracerProvider; import java.util.Map; @@ -115,6 +116,10 @@ public TraceUtil.Span startSpan(String spanName, TraceUtil.SpanContext parentSpa return new Span(); } + public SpanBuilder addSettingsAttributesToCurrentSpan(SpanBuilder spanBuilder) { + return getTracer().spanBuilder("TRACING_DISABLED_NO_OP"); + } + @Nonnull @Override public TraceUtil.Span getCurrentSpan() { diff --git a/google-cloud-datastore/src/main/java/com/google/cloud/datastore/telemetry/EnabledTraceUtil.java b/google-cloud-datastore/src/main/java/com/google/cloud/datastore/telemetry/EnabledTraceUtil.java index 438395cb1..3b962754d 100644 --- a/google-cloud-datastore/src/main/java/com/google/cloud/datastore/telemetry/EnabledTraceUtil.java +++ b/google-cloud-datastore/src/main/java/com/google/cloud/datastore/telemetry/EnabledTraceUtil.java @@ -238,7 +238,8 @@ public Scope makeCurrent() { } /** Applies the current Datastore instance settings as attributes to the current Span */ - private SpanBuilder addSettingsAttributesToCurrentSpan(SpanBuilder spanBuilder) { + @Override + public SpanBuilder addSettingsAttributesToCurrentSpan(SpanBuilder spanBuilder) { spanBuilder = spanBuilder.setAllAttributes( Attributes.builder() 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 dd1dcf29e..dce53e952 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 @@ -22,6 +22,7 @@ import com.google.cloud.datastore.DatastoreOptions; import io.grpc.ManagedChannelBuilder; import io.opentelemetry.api.trace.Span; +import io.opentelemetry.api.trace.SpanBuilder; import io.opentelemetry.api.trace.Tracer; import java.util.Map; import javax.annotation.Nonnull; @@ -144,6 +145,12 @@ interface Scope extends AutoCloseable { */ Span startSpan(String spanName, SpanContext parentSpanContext); + /** + * Adds common SpanAttributes to the current span, useful when hand-creating a new Span without + * using the TraceUtil.Span interface. + */ + SpanBuilder addSettingsAttributesToCurrentSpan(SpanBuilder spanBuilder); + /** Returns the current span. */ @Nonnull Span getCurrentSpan(); diff --git a/google-cloud-datastore/src/test/java/com/google/cloud/datastore/it/ITTracingTest.java b/google-cloud-datastore/src/test/java/com/google/cloud/datastore/it/ITTracingTest.java index 6432cbc9e..485f3272e 100644 --- a/google-cloud-datastore/src/test/java/com/google/cloud/datastore/it/ITTracingTest.java +++ b/google-cloud-datastore/src/test/java/com/google/cloud/datastore/it/ITTracingTest.java @@ -16,7 +16,9 @@ package com.google.cloud.datastore.it; +import static com.google.cloud.datastore.aggregation.Aggregation.count; import static com.google.cloud.datastore.telemetry.TraceUtil.*; +import static com.google.common.truth.Truth.assertThat; import static io.opentelemetry.semconv.resource.attributes.ResourceAttributes.SERVICE_NAME; import static org.junit.Assert.assertEquals; import static org.junit.Assert.assertFalse; @@ -24,6 +26,9 @@ import static org.junit.Assert.assertNull; import static org.junit.Assert.assertTrue; +import com.google.cloud.datastore.AggregationQuery; +import com.google.cloud.datastore.AggregationResult; +import com.google.cloud.datastore.AggregationResults; import com.google.cloud.datastore.Datastore; import com.google.cloud.datastore.DatastoreOpenTelemetryOptions; import com.google.cloud.datastore.DatastoreOptions; @@ -33,7 +38,10 @@ import com.google.cloud.datastore.KeyFactory; import com.google.cloud.datastore.Query; import com.google.cloud.datastore.QueryResults; +import com.google.cloud.datastore.ReadOption; +import com.google.cloud.datastore.StructuredQuery; import com.google.cloud.datastore.StructuredQuery.PropertyFilter; +import com.google.cloud.datastore.Transaction; import com.google.cloud.datastore.testing.RemoteDatastoreHelper; import com.google.common.base.Preconditions; import com.google.testing.junit.testparameterinjector.TestParameter; @@ -91,6 +99,10 @@ protected String datastoreNamedDatabase() { private static Key KEY2; + private static Key KEY3; + + private static Key KEY4; + private static OpenTelemetrySdk openTelemetrySdk; // We use an InMemorySpanExporter for testing which keeps all generated trace spans @@ -166,7 +178,14 @@ public void before() { Key.newBuilder(projectId, kind1, "key2", options.getDatabaseId()) .setNamespace(options.getNamespace()) .build(); - + KEY3 = + Key.newBuilder(projectId, kind1, "key3", options.getDatabaseId()) + .setNamespace(options.getNamespace()) + .build(); + KEY4 = + Key.newBuilder(projectId, kind1, "key4", options.getDatabaseId()) + .setNamespace(options.getNamespace()) + .build(); cleanupTestSpanContext(); } @@ -295,11 +314,6 @@ void assertHasExpectedAttributes(SpanData spanData, String... additionalExpected "gcp.datastore.memoryUtilization", "gcp.datastore.settings.host", "gcp.datastore.settings.databaseId", - "gcp.datastore.settings.channel.needsCredentials", - "gcp.datastore.settings.channel.needsEndpoint", - "gcp.datastore.settings.channel.needsHeaders", - "gcp.datastore.settings.channel.shouldAutoClose", - "gcp.datastore.settings.channel.transportName", "gcp.datastore.settings.retrySettings.maxRpcTimeout", "gcp.datastore.settings.retrySettings.retryDelayMultiplier", "gcp.datastore.settings.retrySettings.initialRetryDelay", @@ -380,6 +394,8 @@ public void lookupTraceTest() throws Exception { Entity entity = datastore.get(KEY1); assertNull(entity); + waitForTracesToComplete(); + List spans = prepareSpans(); assertEquals(1, spans.size()); assertSpanHierarchy(SPAN_NAME_LOOKUP); @@ -403,6 +419,8 @@ public void allocateIdsTraceTest() throws Exception { IncompleteKey pk1 = keyFactory.newKey(); Key key1 = datastore.allocateId(pk1); + waitForTracesToComplete(); + List spans = prepareSpans(); assertEquals(1, spans.size()); assertSpanHierarchy(SPAN_NAME_ALLOCATE_IDS); @@ -416,6 +434,8 @@ public void reserveIdsTraceTest() throws Exception { List keyList = datastore.reserveIds(key1, key2); assertEquals(2, keyList.size()); + waitForTracesToComplete(); + List spans = prepareSpans(); assertEquals(1, spans.size()); assertSpanHierarchy(SPAN_NAME_RESERVE_IDS); @@ -427,6 +447,8 @@ public void commitTraceTest() throws Exception { Entity response = datastore.add(entity1); assertEquals(entity1, response); + waitForTracesToComplete(); + List spans = prepareSpans(); assertEquals(1, spans.size()); assertSpanHierarchy(SPAN_NAME_COMMIT); @@ -438,6 +460,8 @@ public void putTraceTest() throws Exception { Entity response = datastore.put(entity1); assertEquals(entity1, response); + waitForTracesToComplete(); + List spans = prepareSpans(); assertEquals(1, spans.size()); assertSpanHierarchy(SPAN_NAME_COMMIT); @@ -476,6 +500,8 @@ public void updateTraceTest() throws Exception { Entity entity2_update = Entity.newBuilder(entity2).set("test_field", "new_test_value1").build(); datastore.update(entity1_update, entity2_update); + waitForTracesToComplete(); + spans = prepareSpans(); assertEquals(1, spans.size()); assertSpanHierarchy(SPAN_NAME_COMMIT); @@ -506,6 +532,9 @@ public void deleteTraceTest() throws Exception { cleanupTestSpanContext(); datastore.delete(entity1.getKey()); + + waitForTracesToComplete(); + spans = prepareSpans(); assertEquals(1, spans.size()); assertSpanHierarchy(SPAN_NAME_COMMIT); @@ -544,6 +573,8 @@ public void runQueryTraceTest() throws Exception { assertEquals(entity1, queryResults.next()); assertFalse(queryResults.hasNext()); + waitForTracesToComplete(); + List spans = prepareSpans(); assertEquals(1, spans.size()); assertSpanHierarchy(SPAN_NAME_RUN_QUERY); @@ -563,20 +594,255 @@ public void runQueryTraceTest() throws Exception { } @Test - public void runAggregationQueryTraceTest() throws Exception {} + public void runAggregationQueryTraceTest() throws Exception { + Entity entity1 = + Entity.newBuilder(KEY1) + .set("pepper_name", "jalapeno") + .set("max_scoville_level", 10000) + .build(); + Entity entity2 = + Entity.newBuilder(KEY2) + .set("pepper_name", "serrano") + .set("max_scoville_level", 25000) + .build(); + Entity entity3 = + Entity.newBuilder(KEY3) + .set("pepper_name", "habanero") + .set("max_scoville_level", 350000) + .build(); + Entity entity4 = + Entity.newBuilder(KEY4) + .set("pepper_name", "ghost") + .set("max_scoville_level", 1500000) + .build(); - @Test - public void newTransactionReadTraceTest() throws Exception {} + List entityList = new ArrayList<>(); + entityList.add(entity1); + entityList.add(entity2); + entityList.add(entity3); + entityList.add(entity4); + + List response = datastore.add(entity1, entity2, entity3, entity4); + assertEquals(entityList, response); + + // Clean Up test span context to verify RunAggregationQuery spans + cleanupTestSpanContext(); + + PropertyFilter mediumSpicyFilters = PropertyFilter.lt("max_scoville_level", 100000); + StructuredQuery mediumSpicyQuery = + Query.newEntityQueryBuilder().setKind(KEY1.getKind()).setFilter(mediumSpicyFilters).build(); + AggregationQuery countSpicyPeppers = + Query.newAggregationQueryBuilder() + .addAggregation(count().as("count")) + .over(mediumSpicyQuery) + .build(); + AggregationResults results = datastore.runAggregation(countSpicyPeppers); + assertThat(results.size()).isEqualTo(1); + AggregationResult result = results.get(0); + assertThat(result.getLong("count")).isEqualTo(2L); + + waitForTracesToComplete(); + + List spans = prepareSpans(); + assertEquals(1, spans.size()); + assertSpanHierarchy(SPAN_NAME_RUN_AGGREGATION_QUERY); + } @Test - public void newTransactionQueryTest() throws Exception {} + public void newTransactionReadWriteTraceTest() throws Exception { + // Transaction.Begin + Transaction transaction = datastore.newTransaction(); + + // Transaction.Lookup + Entity entity = datastore.get(KEY1, ReadOption.transactionId(transaction.getTransactionId())); + assertNull(entity); + + Entity updatedEntity = Entity.newBuilder(KEY1).set("test_field", "new_test_value1").build(); + transaction.put(updatedEntity); + + // Transaction.Commit + transaction.commit(); + + waitForTracesToComplete(); + + List spans = prepareSpans(); + assertEquals(3, spans.size()); + + assertSpanHierarchy(SPAN_NAME_BEGIN_TRANSACTION); + assertSpanHierarchy(SPAN_NAME_TRANSACTION_LOOKUP); + SpanData span = getSpanByName(SPAN_NAME_TRANSACTION_LOOKUP); + assertTrue( + hasEvent( + span, + SPAN_NAME_TRANSACTION_LOOKUP, + Attributes.builder() + .put("Deferred", 0) + .put("Missing", 1) + .put("Received", 0) + .put("transactional", true) + .put("transaction_id", transaction.getTransactionId().toStringUtf8()) + .build())); + + assertSpanHierarchy(SPAN_NAME_TRANSACTION_COMMIT); + span = getSpanByName(SPAN_NAME_TRANSACTION_COMMIT); + assertTrue( + hasEvent( + span, + SPAN_NAME_TRANSACTION_COMMIT, + Attributes.builder() + .put("doc_count", 1) + .put("transactional", true) + .put("transaction_id", transaction.getTransactionId().toStringUtf8()) + .build())); + } @Test - public void newTransactionReadWriteTraceTest() throws Exception {} + public void newTransactionQueryTest() throws Exception { + 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); + + // Clean Up test span context to verify Transaction RunQuery spans + cleanupTestSpanContext(); + + Transaction transaction = datastore.newTransaction(); + PropertyFilter filter = PropertyFilter.eq("test_field", entity1.getValue("test_field")); + Query query = + Query.newEntityQueryBuilder().setKind(KEY1.getKind()).setFilter(filter).build(); + QueryResults queryResults = transaction.run(query); + transaction.commit(); + assertTrue(queryResults.hasNext()); + assertEquals(entity1, queryResults.next()); + assertFalse(queryResults.hasNext()); + + waitForTracesToComplete(); + + List spans = prepareSpans(); + assertEquals(3, spans.size()); + + assertSpanHierarchy(SPAN_NAME_BEGIN_TRANSACTION); + assertSpanHierarchy(SPAN_NAME_TRANSACTION_RUN_QUERY); + assertSpanHierarchy(SPAN_NAME_TRANSACTION_COMMIT); + SpanData span = getSpanByName(SPAN_NAME_TRANSACTION_RUN_QUERY); + assertTrue( + hasEvent( + span, + SPAN_NAME_TRANSACTION_RUN_QUERY, + Attributes.builder() + .put("response_count", 1) + .put("transactional", true) + .put("read_consistency", "READ_CONSISTENCY_UNSPECIFIED") + .put("more_results", "NO_MORE_RESULTS") + .put("transaction_id", transaction.getTransactionId().toStringUtf8()) + .build())); + } @Test - public void newTransactionRollbackTest() throws Exception {} + public void newTransactionRollbackTest() throws Exception { + Entity entity1 = Entity.newBuilder(KEY1).set("pepper_type", "jalapeno").build(); + Entity entity2 = Entity.newBuilder(KEY2).set("pepper_type", "habanero").build(); + List entityList = new ArrayList<>(); + entityList.add(entity1); + entityList.add(entity2); + + List response = datastore.add(entity1, entity2); + assertEquals(entityList, response); + + // Clean Up test span context to verify Transaction Rollback spans + cleanupTestSpanContext(); + + String simplified_spice_level = "not_spicy"; + Entity entity1update = + Entity.newBuilder(entity1).set("spice_level", simplified_spice_level).build(); + Transaction transaction = datastore.newTransaction(); + entity1 = transaction.get(KEY1); + switch (entity1.getString("pepper_type")) { + case "jalapeno": + simplified_spice_level = "mild"; + break; + + case "habanero": + simplified_spice_level = "hot"; + break; + } + transaction.update(entity1update); + transaction.delete(KEY2); + transaction.rollback(); + assertFalse(transaction.isActive()); + + waitForTracesToComplete(); + + List spans = prepareSpans(); + assertEquals(3, spans.size()); + + assertSpanHierarchy(SPAN_NAME_BEGIN_TRANSACTION); + assertSpanHierarchy(SPAN_NAME_TRANSACTION_LOOKUP); + SpanData span = getSpanByName(SPAN_NAME_TRANSACTION_LOOKUP); + assertTrue( + hasEvent( + span, + SPAN_NAME_TRANSACTION_LOOKUP, + Attributes.builder() + .put("Deferred", 0) + .put("Missing", 0) + .put("Received", 1) + .put("transactional", true) + .put("transaction_id", transaction.getTransactionId().toStringUtf8()) + .build())); + + assertSpanHierarchy(SPAN_NAME_ROLLBACK); + span = getSpanByName(SPAN_NAME_ROLLBACK); + assertTrue( + hasEvent( + span, + SPAN_NAME_ROLLBACK, + Attributes.builder() + .put("transaction_id", transaction.getTransactionId().toStringUtf8()) + .build())); + } @Test - public void runInTransactionQueryTest() throws Exception {} + public void runInTransactionQueryTest() throws Exception { + // Set up + 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); + + // Clean Up test span context to verify Transaction Rollback spans + cleanupTestSpanContext(); + + PropertyFilter filter = PropertyFilter.eq("test_field", entity1.getValue("test_field")); + Query query = + Query.newEntityQueryBuilder().setKind(KEY1.getKind()).setFilter(filter).build(); + Datastore.TransactionCallable callable = + transaction -> { + QueryResults queryResults = datastore.run(query); + assertTrue(queryResults.hasNext()); + assertEquals(entity1, queryResults.next()); + assertFalse(queryResults.hasNext()); + return true; + }; + datastore.runInTransaction(callable); + + waitForTracesToComplete(); + + List spans = prepareSpans(); + assertEquals(4, spans.size()); + + // Since the runInTransaction method runs the TransactionCallable opaquely in a transaction + // there is no way for the API user to know the transaction ID, so we will not validate it here. + assertSpanHierarchy(SPAN_NAME_TRANSACTION_RUN, SPAN_NAME_BEGIN_TRANSACTION); + assertSpanHierarchy(SPAN_NAME_TRANSACTION_RUN, SPAN_NAME_RUN_QUERY); + assertSpanHierarchy(SPAN_NAME_TRANSACTION_RUN, SPAN_NAME_TRANSACTION_COMMIT); + } }