From e0028bebbf475765493df228648a0b826f04ebb1 Mon Sep 17 00:00:00 2001 From: jimit-j-shah <57637300+jimit-j-shah@users.noreply.github.com> Date: Mon, 24 Jun 2024 10:36:41 -0700 Subject: [PATCH 1/4] feat: Added span for Transactional RunQuery - tested --- .../google/cloud/datastore/DatastoreImpl.java | 72 +++++++------------ .../cloud/datastore/telemetry/TraceUtil.java | 1 + .../cloud/datastore/it/ITE2ETracingTest.java | 44 +++++++++++- 3 files changed, 70 insertions(+), 47 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 804673b3a..ac2db3cd1 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 @@ -27,7 +27,6 @@ import com.google.cloud.datastore.spi.v1.DatastoreRpc; import com.google.common.base.MoreObjects; import com.google.common.base.Preconditions; -import com.google.common.base.Throwables; import com.google.common.collect.AbstractIterator; import com.google.common.collect.ImmutableList; import com.google.common.collect.ImmutableMap; @@ -42,10 +41,7 @@ import io.opencensus.common.Scope; import io.opencensus.trace.Span; import io.opencensus.trace.Status; -import io.opentelemetry.api.common.Attributes; -import io.opentelemetry.api.trace.SpanBuilder; -import io.opentelemetry.api.trace.SpanKind; -import io.opentelemetry.api.trace.StatusCode; +import io.opentelemetry.api.trace.SpanContext; import io.opentelemetry.context.Context; import java.util.ArrayList; import java.util.Arrays; @@ -55,6 +51,7 @@ import java.util.LinkedHashSet; import java.util.List; import java.util.Map; +import java.util.Objects; import java.util.Optional; import java.util.Set; import java.util.concurrent.Callable; @@ -111,13 +108,13 @@ static class ReadWriteTransactionCallable implements Callable { private volatile TransactionOptions options; private volatile Transaction transaction; - private final com.google.cloud.datastore.telemetry.TraceUtil.SpanContext parentSpanContext; + private final SpanContext parentSpanContext; ReadWriteTransactionCallable( Datastore datastore, TransactionCallable callable, TransactionOptions options, - @Nullable com.google.cloud.datastore.telemetry.TraceUtil.SpanContext parentSpanContext) { + @Nullable SpanContext parentSpanContext) { this.datastore = datastore; this.callable = callable; this.options = options; @@ -143,35 +140,19 @@ void setPrevTransactionId(ByteString transactionId) { options = options.toBuilder().setReadWrite(readWrite).build(); } - private io.opentelemetry.api.trace.Span startSpanInternal( - String spanName, - com.google.cloud.datastore.telemetry.TraceUtil.SpanContext parentSpanContext) { - com.google.cloud.datastore.telemetry.TraceUtil otelTraceUtil = - datastore.getOptions().getTraceUtil(); - SpanBuilder spanBuilder = - otelTraceUtil - .getTracer() - .spanBuilder(com.google.cloud.datastore.telemetry.TraceUtil.SPAN_NAME_TRANSACTION_RUN) - .setSpanKind(SpanKind.PRODUCER) - .setParent( - Context.current() - .with( - io.opentelemetry.api.trace.Span.wrap( - parentSpanContext.getSpanContext()))); - return spanBuilder.startSpan(); - } - @Override public T call() throws DatastoreException { - // TODO Instead of using OTel Spans directly, TraceUtil.Span should be used here. However, - // the same code in startSpanInternal doesn't work when EnabledTraceUtil.StartSpan is called - // probably because of some thread-local caching that is getting lost. This needs more - // debugging. The code below works and is idiomatic but could be prettier and more consistent - // with the use of TraceUtil-provided framework. io.opentelemetry.api.trace.Span span = - startSpanInternal( - com.google.cloud.datastore.telemetry.TraceUtil.SPAN_NAME_TRANSACTION_RUN, - parentSpanContext); + Objects.requireNonNull(datastore + .getOptions() + .getOpenTelemetryOptions() + .getOpenTelemetry()) + .getTracer(com.google.cloud.datastore.telemetry.TraceUtil.SPAN_NAME_TRANSACTION_RUN) + .spanBuilder(com.google.cloud.datastore.telemetry.TraceUtil.SPAN_NAME_TRANSACTION_RUN) + .setParent( + Context.current().with(io.opentelemetry.api.trace.Span.wrap(parentSpanContext))) + .startSpan(); + try (io.opentelemetry.context.Scope ignored = span.makeCurrent()) { transaction = datastore.newTransaction(options); T value = callable.run(transaction); @@ -179,14 +160,6 @@ public T call() throws DatastoreException { return value; } catch (Exception ex) { transaction.rollback(); - span.setStatus(StatusCode.ERROR, ex.getMessage()); - span.recordException( - ex, - Attributes.builder() - .put("exception.message", ex.getMessage()) - .put("exception.type", ex.getClass().getName()) - .put("exception.stacktrace", Throwables.getStackTraceAsString(ex)) - .build()); span.end(); throw DatastoreException.propagateUserException(ex); } finally { @@ -207,7 +180,7 @@ public T runInTransaction(final TransactionCallable callable) { try { return RetryHelper.runWithRetries( new ReadWriteTransactionCallable( - this, callable, null, otelTraceUtil.getCurrentSpanContext()), + this, callable, null, io.opentelemetry.api.trace.Span.current().getSpanContext()), retrySettings, TRANSACTION_EXCEPTION_HANDLER, getOptions().getClock()); @@ -222,7 +195,10 @@ public T runInTransaction( try { return RetryHelper.runWithRetries( new ReadWriteTransactionCallable( - this, callable, transactionOptions, otelTraceUtil.getCurrentSpanContext()), + this, + callable, + transactionOptions, + io.opentelemetry.api.trace.Span.current().getSpanContext()), retrySettings, TRANSACTION_EXCEPTION_HANDLER, getOptions().getClock()); @@ -285,9 +261,13 @@ public AggregationResults runAggregation( com.google.datastore.v1.RunQueryResponse runQuery( final com.google.datastore.v1.RunQueryRequest requestPb) { - com.google.cloud.datastore.telemetry.TraceUtil.Span span = - otelTraceUtil.startSpan(com.google.cloud.datastore.telemetry.TraceUtil.SPAN_NAME_RUN_QUERY); ReadOptions readOptions = requestPb.getReadOptions(); + boolean isTransactional = readOptions.hasTransaction() || readOptions.hasNewTransaction(); + String spanName = + (isTransactional + ? com.google.cloud.datastore.telemetry.TraceUtil.SPAN_NAME_TRANSACTION_RUN_QUERY + : com.google.cloud.datastore.telemetry.TraceUtil.SPAN_NAME_RUN_QUERY); + com.google.cloud.datastore.telemetry.TraceUtil.Span span = otelTraceUtil.startSpan(spanName); span.setAttribute( "isTransactional", readOptions.hasTransaction() || readOptions.hasNewTransaction()); span.setAttribute("readConsistency", readOptions.getReadConsistency().toString()); @@ -302,7 +282,7 @@ com.google.datastore.v1.RunQueryResponse runQuery( : TRANSACTION_OPERATION_EXCEPTION_HANDLER, getOptions().getClock()); span.addEvent( - com.google.cloud.datastore.telemetry.TraceUtil.SPAN_NAME_RUN_QUERY + ": Completed", + spanName + ": Completed", new ImmutableMap.Builder() .put("Received", response.getBatch().getEntityResultsCount()) .put("More results", response.getBatch().getMoreResults().toString()) 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 a1901fe0d..d63cdda34 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 @@ -42,6 +42,7 @@ public interface TraceUtil { static final String SPAN_NAME_BEGIN_TRANSACTION = "Transaction.Begin"; static final String SPAN_NAME_TRANSACTION_LOOKUP = "Transaction.Lookup"; static final String SPAN_NAME_TRANSACTION_COMMIT = "Transaction.Commit"; + static final String SPAN_NAME_TRANSACTION_RUN_QUERY = "Transaction.RunQuery"; static final String SPAN_NAME_ROLLBACK = "Transaction.Rollback"; static final String SPAN_NAME_TRANSACTION_RUN_AGGREGATION_QUERY = "Transaction.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 57f4d5b5c..161476af1 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 @@ -26,6 +26,7 @@ import static com.google.cloud.datastore.telemetry.TraceUtil.SPAN_NAME_TRANSACTION_COMMIT; import static com.google.cloud.datastore.telemetry.TraceUtil.SPAN_NAME_TRANSACTION_LOOKUP; import static com.google.cloud.datastore.telemetry.TraceUtil.SPAN_NAME_TRANSACTION_RUN; +import static com.google.cloud.datastore.telemetry.TraceUtil.SPAN_NAME_TRANSACTION_RUN_QUERY; import static com.google.common.truth.Truth.assertThat; import static io.opentelemetry.semconv.resource.attributes.ResourceAttributes.SERVICE_NAME; import static org.junit.Assert.assertEquals; @@ -795,8 +796,49 @@ public void transactionalLookupTest() throws Exception { Collections.singletonList(SPAN_NAME_TRANSACTION_COMMIT))); } + @Test + public void transactionQueryTest() 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); + + assertNotNull(customSpanContext); + + // Test + Span rootSpan = getNewRootSpanWithContext(); + try (Scope ignored = rootSpan.makeCurrent()) { + 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()); + } finally { + rootSpan.end(); + } + waitForTracesToComplete(); + + fetchAndValidateTrace( + customSpanContext.getTraceId(), + /*numExpectedSpans=*/ 3, + Arrays.asList( + Collections.singletonList(SPAN_NAME_BEGIN_TRANSACTION), + Collections.singletonList(SPAN_NAME_TRANSACTION_RUN_QUERY), + Collections.singletonList(SPAN_NAME_TRANSACTION_COMMIT))); + } + @Test 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<>(); @@ -837,7 +879,7 @@ public void runInTransactionQueryTest() throws Exception { } @Test - public void runInTransactionAggregationQueryTest() throws Exception {} + public void transactionRunQueryTest() throws Exception {} @Test public void readWriteTransactionTraceTest() throws Exception {} From a6c5bb232be5338484a7e92cae566e841eac6cc0 Mon Sep 17 00:00:00 2001 From: jimit-j-shah <57637300+jimit-j-shah@users.noreply.github.com> Date: Tue, 25 Jun 2024 16:29:25 -0700 Subject: [PATCH 2/4] fix: lint --- .../main/java/com/google/cloud/datastore/DatastoreImpl.java | 6 ++---- 1 file changed, 2 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 ac2db3cd1..c9ff74f30 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 @@ -143,10 +143,8 @@ void setPrevTransactionId(ByteString transactionId) { @Override public T call() throws DatastoreException { io.opentelemetry.api.trace.Span span = - Objects.requireNonNull(datastore - .getOptions() - .getOpenTelemetryOptions() - .getOpenTelemetry()) + Objects.requireNonNull( + datastore.getOptions().getOpenTelemetryOptions().getOpenTelemetry()) .getTracer(com.google.cloud.datastore.telemetry.TraceUtil.SPAN_NAME_TRANSACTION_RUN) .spanBuilder(com.google.cloud.datastore.telemetry.TraceUtil.SPAN_NAME_TRANSACTION_RUN) .setParent( From 9a78170f8317d2f410124d57b734b12008059d8a Mon Sep 17 00:00:00 2001 From: jimit-j-shah <57637300+jimit-j-shah@users.noreply.github.com> Date: Tue, 25 Jun 2024 16:40:41 -0700 Subject: [PATCH 3/4] fix: patch apply issues --- .../google/cloud/datastore/DatastoreImpl.java | 60 +++++++++++++------ 1 file changed, 43 insertions(+), 17 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 c9ff74f30..224aa8b79 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 @@ -27,6 +27,7 @@ import com.google.cloud.datastore.spi.v1.DatastoreRpc; import com.google.common.base.MoreObjects; import com.google.common.base.Preconditions; +import com.google.common.base.Throwables; import com.google.common.collect.AbstractIterator; import com.google.common.collect.ImmutableList; import com.google.common.collect.ImmutableMap; @@ -41,7 +42,10 @@ import io.opencensus.common.Scope; import io.opencensus.trace.Span; import io.opencensus.trace.Status; -import io.opentelemetry.api.trace.SpanContext; +import io.opentelemetry.api.common.Attributes; +import io.opentelemetry.api.trace.SpanBuilder; +import io.opentelemetry.api.trace.SpanKind; +import io.opentelemetry.api.trace.StatusCode; import io.opentelemetry.context.Context; import java.util.ArrayList; import java.util.Arrays; @@ -51,7 +55,6 @@ import java.util.LinkedHashSet; import java.util.List; import java.util.Map; -import java.util.Objects; import java.util.Optional; import java.util.Set; import java.util.concurrent.Callable; @@ -108,13 +111,13 @@ static class ReadWriteTransactionCallable implements Callable { private volatile TransactionOptions options; private volatile Transaction transaction; - private final SpanContext parentSpanContext; + private final com.google.cloud.datastore.telemetry.TraceUtil.SpanContext parentSpanContext; ReadWriteTransactionCallable( Datastore datastore, TransactionCallable callable, TransactionOptions options, - @Nullable SpanContext parentSpanContext) { + @Nullable com.google.cloud.datastore.telemetry.TraceUtil.SpanContext parentSpanContext) { this.datastore = datastore; this.callable = callable; this.options = options; @@ -140,17 +143,35 @@ void setPrevTransactionId(ByteString transactionId) { options = options.toBuilder().setReadWrite(readWrite).build(); } - @Override - public T call() throws DatastoreException { - io.opentelemetry.api.trace.Span span = - Objects.requireNonNull( - datastore.getOptions().getOpenTelemetryOptions().getOpenTelemetry()) - .getTracer(com.google.cloud.datastore.telemetry.TraceUtil.SPAN_NAME_TRANSACTION_RUN) + private io.opentelemetry.api.trace.Span startSpanInternal( + String spanName, + com.google.cloud.datastore.telemetry.TraceUtil.SpanContext parentSpanContext) { + com.google.cloud.datastore.telemetry.TraceUtil otelTraceUtil = + datastore.getOptions().getTraceUtil(); + SpanBuilder spanBuilder = + otelTraceUtil + .getTracer() .spanBuilder(com.google.cloud.datastore.telemetry.TraceUtil.SPAN_NAME_TRANSACTION_RUN) + .setSpanKind(SpanKind.PRODUCER) .setParent( - Context.current().with(io.opentelemetry.api.trace.Span.wrap(parentSpanContext))) - .startSpan(); + Context.current() + .with( + io.opentelemetry.api.trace.Span.wrap( + parentSpanContext.getSpanContext()))); + return spanBuilder.startSpan(); + } + @Override + public T call() throws DatastoreException { + // TODO Instead of using OTel Spans directly, TraceUtil.Span should be used here. However, + // the same code in startSpanInternal doesn't work when EnabledTraceUtil.StartSpan is called + // probably because of some thread-local caching that is getting lost. This needs more + // debugging. The code below works and is idiomatic but could be prettier and more consistent + // with the use of TraceUtil-provided framework. + io.opentelemetry.api.trace.Span span = + startSpanInternal( + com.google.cloud.datastore.telemetry.TraceUtil.SPAN_NAME_TRANSACTION_RUN, + parentSpanContext); try (io.opentelemetry.context.Scope ignored = span.makeCurrent()) { transaction = datastore.newTransaction(options); T value = callable.run(transaction); @@ -158,6 +179,14 @@ public T call() throws DatastoreException { return value; } catch (Exception ex) { transaction.rollback(); + span.setStatus(StatusCode.ERROR, ex.getMessage()); + span.recordException( + ex, + Attributes.builder() + .put("exception.message", ex.getMessage()) + .put("exception.type", ex.getClass().getName()) + .put("exception.stacktrace", Throwables.getStackTraceAsString(ex)) + .build()); span.end(); throw DatastoreException.propagateUserException(ex); } finally { @@ -178,7 +207,7 @@ public T runInTransaction(final TransactionCallable callable) { try { return RetryHelper.runWithRetries( new ReadWriteTransactionCallable( - this, callable, null, io.opentelemetry.api.trace.Span.current().getSpanContext()), + this, callable, null, otelTraceUtil.getCurrentSpanContext()), retrySettings, TRANSACTION_EXCEPTION_HANDLER, getOptions().getClock()); @@ -193,10 +222,7 @@ public T runInTransaction( try { return RetryHelper.runWithRetries( new ReadWriteTransactionCallable( - this, - callable, - transactionOptions, - io.opentelemetry.api.trace.Span.current().getSpanContext()), + this, callable, transactionOptions, otelTraceUtil.getCurrentSpanContext()), retrySettings, TRANSACTION_EXCEPTION_HANDLER, getOptions().getClock()); From 2e992513ad4f9c75b4f19819d614c92a4f9e609c Mon Sep 17 00:00:00 2001 From: jimit-j-shah <57637300+jimit-j-shah@users.noreply.github.com> Date: Wed, 26 Jun 2024 11:50:38 -0700 Subject: [PATCH 4/4] fix: refactor using boolean flag --- .../main/java/com/google/cloud/datastore/DatastoreImpl.java | 3 +-- 1 file changed, 1 insertion(+), 2 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 224aa8b79..7a81a469d 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 @@ -292,8 +292,7 @@ com.google.datastore.v1.RunQueryResponse runQuery( ? com.google.cloud.datastore.telemetry.TraceUtil.SPAN_NAME_TRANSACTION_RUN_QUERY : com.google.cloud.datastore.telemetry.TraceUtil.SPAN_NAME_RUN_QUERY); com.google.cloud.datastore.telemetry.TraceUtil.Span span = otelTraceUtil.startSpan(spanName); - span.setAttribute( - "isTransactional", readOptions.hasTransaction() || readOptions.hasNewTransaction()); + span.setAttribute("isTransactional", isTransactional); span.setAttribute("readConsistency", readOptions.getReadConsistency().toString()); try (com.google.cloud.datastore.telemetry.TraceUtil.Scope ignored = span.makeCurrent()) {