From e180a290bee7d78a8030b1d03386c4ac2bf2618c Mon Sep 17 00:00:00 2001 From: jimit-j-shah <57637300+jimit-j-shah@users.noreply.github.com> Date: Wed, 29 May 2024 17:24:03 -0700 Subject: [PATCH 1/3] feat: support for transactional operations - tested using newTransaction() and runInTransaction() --- .../google/cloud/datastore/DatastoreImpl.java | 49 +++++++++++------ .../cloud/datastore/TransactionImpl.java | 10 +++- .../cloud/datastore/telemetry/TraceUtil.java | 2 + .../cloud/datastore/it/ITE2ETracingTest.java | 54 ++++++++++++++++--- 4 files changed, 91 insertions(+), 24 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 66bc6c0ba..753dab596 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 @@ -25,6 +25,7 @@ import com.google.cloud.ServiceOptions; import com.google.cloud.datastore.execution.AggregationQueryExecutor; import com.google.cloud.datastore.spi.v1.DatastoreRpc; +import com.google.cloud.datastore.telemetry.TraceUtil.Context; import com.google.common.base.MoreObjects; import com.google.common.base.Preconditions; import com.google.common.collect.AbstractIterator; @@ -52,6 +53,7 @@ import java.util.Optional; import java.util.Set; import java.util.concurrent.Callable; +import javax.annotation.Nonnull; final class DatastoreImpl extends BaseService implements Datastore { @@ -103,13 +105,18 @@ static class ReadWriteTransactionCallable implements Callable { private final TransactionCallable callable; private volatile TransactionOptions options; private volatile Transaction transaction; + @Nonnull private final Context transactionTraceContext; ReadWriteTransactionCallable( - Datastore datastore, TransactionCallable callable, TransactionOptions options) { + Datastore datastore, + TransactionCallable callable, + TransactionOptions options, + @Nonnull Context parentTraceContext) { this.datastore = datastore; this.callable = callable; this.options = options; this.transaction = null; + this.transactionTraceContext = parentTraceContext; } Datastore getDatastore() { @@ -132,8 +139,9 @@ void setPrevTransactionId(ByteString transactionId) { @Override public T call() throws DatastoreException { - transaction = datastore.newTransaction(options); - try { + try (com.google.cloud.datastore.telemetry.TraceUtil.Scope ignored = + transactionTraceContext.makeCurrent()) { + transaction = datastore.newTransaction(options); T value = callable.run(transaction); transaction.commit(); return value; @@ -154,36 +162,41 @@ public T call() throws DatastoreException { @Override public T runInTransaction(final TransactionCallable callable) { - Span span = traceUtil.startSpan(TraceUtil.SPAN_NAME_TRANSACTION); - 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_TRANSACTION_RUN); + try (com.google.cloud.datastore.telemetry.TraceUtil.Scope ignored = span.makeCurrent()) { return RetryHelper.runWithRetries( - new ReadWriteTransactionCallable(this, callable, null), + new ReadWriteTransactionCallable(this, callable, null, otelTraceUtil.currentContext()), retrySettings, TRANSACTION_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(); } } @Override public T runInTransaction( final TransactionCallable callable, TransactionOptions transactionOptions) { - Span span = traceUtil.startSpan(TraceUtil.SPAN_NAME_TRANSACTION); - 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_TRANSACTION_RUN); + try (com.google.cloud.datastore.telemetry.TraceUtil.Scope ignored = span.makeCurrent()) { return RetryHelper.runWithRetries( - new ReadWriteTransactionCallable(this, callable, transactionOptions), + new ReadWriteTransactionCallable( + this, callable, transactionOptions, otelTraceUtil.currentContext()), retrySettings, TRANSACTION_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(); } } @@ -634,10 +647,14 @@ private com.google.datastore.v1.CommitResponse commitMutation( com.google.datastore.v1.CommitResponse commit( final com.google.datastore.v1.CommitRequest requestPb) { - com.google.cloud.datastore.telemetry.TraceUtil.Span span = - otelTraceUtil.startSpan(com.google.cloud.datastore.telemetry.TraceUtil.SPAN_NAME_COMMIT); + final boolean isTransactional = + requestPb.hasTransaction() || requestPb.hasSingleUseTransaction(); + final String spanName = + isTransactional + ? com.google.cloud.datastore.telemetry.TraceUtil.SPAN_NAME_TRANSACTION_COMMIT + : com.google.cloud.datastore.telemetry.TraceUtil.SPAN_NAME_COMMIT; + com.google.cloud.datastore.telemetry.TraceUtil.Span span = otelTraceUtil.startSpan(spanName); span.setAttribute("isTransactional", requestPb.hasTransaction()); - try (com.google.cloud.datastore.telemetry.TraceUtil.Scope ignored = span.makeCurrent()) { return RetryHelper.runWithRetries( () -> datastoreRpc.commit(requestPb), diff --git a/google-cloud-datastore/src/main/java/com/google/cloud/datastore/TransactionImpl.java b/google-cloud-datastore/src/main/java/com/google/cloud/datastore/TransactionImpl.java index f08a908ec..b87d175a0 100644 --- a/google-cloud-datastore/src/main/java/com/google/cloud/datastore/TransactionImpl.java +++ b/google-cloud-datastore/src/main/java/com/google/cloud/datastore/TransactionImpl.java @@ -20,6 +20,8 @@ import com.google.api.core.BetaApi; import com.google.cloud.datastore.models.ExplainOptions; +import com.google.cloud.datastore.telemetry.TraceUtil.Context; +import com.google.cloud.datastore.telemetry.TraceUtil.Scope; import com.google.common.collect.ImmutableList; import com.google.datastore.v1.ReadOptions; import com.google.datastore.v1.TransactionOptions; @@ -28,6 +30,7 @@ import java.util.Iterator; import java.util.List; import java.util.Optional; +import javax.annotation.Nonnull; final class TransactionImpl extends BaseDatastoreBatchWriter implements Transaction { @@ -37,6 +40,8 @@ final class TransactionImpl extends BaseDatastoreBatchWriter implements Transact private final ReadOptionProtoPreparer readOptionProtoPreparer; + @Nonnull private final Context transactionTraceContext; + static class ResponseImpl implements Transaction.Response { private final com.google.datastore.v1.CommitResponse response; @@ -78,6 +83,7 @@ public List getGeneratedKeys() { transactionId = datastore.requestTransactionId(requestPb); this.readOptionProtoPreparer = new ReadOptionProtoPreparer(); + this.transactionTraceContext = datastore.getOptions().getTraceUtil().currentContext(); } @Override @@ -96,7 +102,9 @@ public Iterator get(Key... keys) { @Override public List fetch(Key... keys) { validateActive(); - return DatastoreHelper.fetch(this, keys); + try (Scope ignored = transactionTraceContext.makeCurrent()) { + return DatastoreHelper.fetch(this, keys); + } } @Override 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 d43c36a07..fe0e92161 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 @@ -21,6 +21,7 @@ import com.google.api.core.InternalExtensionOnly; import com.google.cloud.datastore.DatastoreOptions; import io.grpc.ManagedChannelBuilder; +import io.opentelemetry.context.Context; import java.util.Map; import javax.annotation.Nonnull; import javax.annotation.Nullable; @@ -35,6 +36,7 @@ public interface TraceUtil { static final String SPAN_NAME_COMMIT = "Commit"; static final String SPAN_NAME_RUN_QUERY = "RunQuery"; static final String SPAN_NAME_RUN_AGGREGATION_QUERY = "RunAggregationQuery"; + static final String SPAN_NAME_TRANSACTION_RUN = "Transaction.run"; 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"; 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 67df8c44b..d2cffc5db 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 @@ -22,7 +22,9 @@ import static com.google.cloud.datastore.telemetry.TraceUtil.SPAN_NAME_LOOKUP; 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; 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.common.truth.Truth.assertThat; import static io.opentelemetry.semconv.resource.attributes.ResourceAttributes.SERVICE_NAME; import static org.junit.Assert.assertEquals; @@ -749,6 +751,7 @@ public void transactionalLookupTest() throws Exception { try (Scope ignored = rootSpan.makeCurrent()) { Transaction transaction = datastore.newTransaction(); Entity entity = datastore.get(KEY1, ReadOption.transactionId(transaction.getTransactionId())); + transaction.commit(); assertNull(entity); } finally { rootSpan.end(); @@ -757,18 +760,55 @@ public void transactionalLookupTest() throws Exception { fetchAndValidateTrace( customSpanContext.getTraceId(), - /*numExpectedSpans=*/ 2, - Collections.singletonList(Collections.singletonList(SPAN_NAME_BEGIN_TRANSACTION))); + /*numExpectedSpans=*/ 3, + Arrays.asList( + Collections.singletonList(SPAN_NAME_BEGIN_TRANSACTION), + Collections.singletonList(SPAN_NAME_TRANSACTION_LOOKUP), + Collections.singletonList(SPAN_NAME_TRANSACTION_COMMIT))); + } + + @Test + public void runInTransactionQueryTest() 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); + + assertNotNull(customSpanContext); + + Span rootSpan = getNewRootSpanWithContext(); + try (Scope ignored = rootSpan.makeCurrent()) { + 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); + } finally { + rootSpan.end(); + } + waitForTracesToComplete(); fetchAndValidateTrace( customSpanContext.getTraceId(), - /*numExpectedSpans=*/ 2, - Collections.singletonList(Collections.singletonList(SPAN_NAME_TRANSACTION_LOOKUP))); + /*numExpectedSpans=*/ 4, + Arrays.asList( + Collections.singletonList(SPAN_NAME_TRANSACTION_RUN), + Collections.singletonList(SPAN_NAME_BEGIN_TRANSACTION), + Collections.singletonList(SPAN_NAME_RUN_QUERY), + Collections.singletonList(SPAN_NAME_TRANSACTION_COMMIT))); } - @Test - public void runInTransactionQueryTest() throws Exception {} - @Test public void runInTransactionAggregationQueryTest() throws Exception {} From 96bf35c90119b7847ad2b72d1a85e38b49057346 Mon Sep 17 00:00:00 2001 From: jimit-j-shah <57637300+jimit-j-shah@users.noreply.github.com> Date: Wed, 5 Jun 2024 11:50:24 -0700 Subject: [PATCH 2/3] fix: code review comments --- .../com/google/cloud/datastore/DatastoreImpl.java | 12 ++++++------ .../com/google/cloud/datastore/TransactionImpl.java | 2 +- .../cloud/datastore/telemetry/DisabledTraceUtil.java | 4 ++-- .../cloud/datastore/telemetry/EnabledTraceUtil.java | 4 ++-- .../google/cloud/datastore/telemetry/TraceUtil.java | 5 ++--- .../datastore/telemetry/DisabledTraceUtilTest.java | 10 +++++----- .../datastore/telemetry/EnabledTraceUtilTest.java | 10 +++++----- 7 files changed, 23 insertions(+), 24 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 753dab596..a44635bfc 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 @@ -105,7 +105,7 @@ static class ReadWriteTransactionCallable implements Callable { private final TransactionCallable callable; private volatile TransactionOptions options; private volatile Transaction transaction; - @Nonnull private final Context transactionTraceContext; + @Nonnull private final Context parentTraceContext; ReadWriteTransactionCallable( Datastore datastore, @@ -116,7 +116,7 @@ static class ReadWriteTransactionCallable implements Callable { this.callable = callable; this.options = options; this.transaction = null; - this.transactionTraceContext = parentTraceContext; + this.parentTraceContext = parentTraceContext; } Datastore getDatastore() { @@ -140,7 +140,7 @@ void setPrevTransactionId(ByteString transactionId) { @Override public T call() throws DatastoreException { try (com.google.cloud.datastore.telemetry.TraceUtil.Scope ignored = - transactionTraceContext.makeCurrent()) { + parentTraceContext.makeCurrent()) { transaction = datastore.newTransaction(options); T value = callable.run(transaction); transaction.commit(); @@ -167,7 +167,7 @@ public T runInTransaction(final TransactionCallable callable) { com.google.cloud.datastore.telemetry.TraceUtil.SPAN_NAME_TRANSACTION_RUN); try (com.google.cloud.datastore.telemetry.TraceUtil.Scope ignored = span.makeCurrent()) { return RetryHelper.runWithRetries( - new ReadWriteTransactionCallable(this, callable, null, otelTraceUtil.currentContext()), + new ReadWriteTransactionCallable(this, callable, null, otelTraceUtil.getCurrentContext()), retrySettings, TRANSACTION_EXCEPTION_HANDLER, getOptions().getClock()); @@ -188,7 +188,7 @@ public T runInTransaction( try (com.google.cloud.datastore.telemetry.TraceUtil.Scope ignored = span.makeCurrent()) { return RetryHelper.runWithRetries( new ReadWriteTransactionCallable( - this, callable, transactionOptions, otelTraceUtil.currentContext()), + this, callable, transactionOptions, otelTraceUtil.getCurrentContext()), retrySettings, TRANSACTION_EXCEPTION_HANDLER, getOptions().getClock()); @@ -654,7 +654,7 @@ com.google.datastore.v1.CommitResponse commit( ? com.google.cloud.datastore.telemetry.TraceUtil.SPAN_NAME_TRANSACTION_COMMIT : com.google.cloud.datastore.telemetry.TraceUtil.SPAN_NAME_COMMIT; com.google.cloud.datastore.telemetry.TraceUtil.Span span = otelTraceUtil.startSpan(spanName); - span.setAttribute("isTransactional", requestPb.hasTransaction()); + span.setAttribute("isTransactional", isTransactional); try (com.google.cloud.datastore.telemetry.TraceUtil.Scope ignored = span.makeCurrent()) { return RetryHelper.runWithRetries( () -> datastoreRpc.commit(requestPb), diff --git a/google-cloud-datastore/src/main/java/com/google/cloud/datastore/TransactionImpl.java b/google-cloud-datastore/src/main/java/com/google/cloud/datastore/TransactionImpl.java index b87d175a0..50838f0af 100644 --- a/google-cloud-datastore/src/main/java/com/google/cloud/datastore/TransactionImpl.java +++ b/google-cloud-datastore/src/main/java/com/google/cloud/datastore/TransactionImpl.java @@ -83,7 +83,7 @@ public List getGeneratedKeys() { transactionId = datastore.requestTransactionId(requestPb); this.readOptionProtoPreparer = new ReadOptionProtoPreparer(); - this.transactionTraceContext = datastore.getOptions().getTraceUtil().currentContext(); + this.transactionTraceContext = datastore.getOptions().getTraceUtil().getCurrentContext(); } @Override 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 a4f25813a..2a42081a1 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 @@ -102,13 +102,13 @@ public TraceUtil.Span startSpan(String spanName, TraceUtil.Context parent) { @Nonnull @Override - public TraceUtil.Span currentSpan() { + public TraceUtil.Span getCurrentSpan() { return new Span(); } @Nonnull @Override - public TraceUtil.Context currentContext() { + public TraceUtil.Context getCurrentContext() { return new Context(); } } 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 3bf6a7466..50f89369a 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 @@ -300,13 +300,13 @@ public TraceUtil.Span startSpan(String spanName, TraceUtil.Context parent) { @Nonnull @Override - public TraceUtil.Span currentSpan() { + public TraceUtil.Span getCurrentSpan() { return new Span(io.opentelemetry.api.trace.Span.current(), ""); } @Nonnull @Override - public TraceUtil.Context currentContext() { + public TraceUtil.Context getCurrentContext() { return new Context(io.opentelemetry.context.Context.current()); } } 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 fe0e92161..8f45e2b62 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 @@ -21,7 +21,6 @@ import com.google.api.core.InternalExtensionOnly; import com.google.cloud.datastore.DatastoreOptions; import io.grpc.ManagedChannelBuilder; -import io.opentelemetry.context.Context; import java.util.Map; import javax.annotation.Nonnull; import javax.annotation.Nullable; @@ -135,9 +134,9 @@ interface Scope extends AutoCloseable { /** Returns the current span. */ @Nonnull - Span currentSpan(); + Span getCurrentSpan(); /** Returns the current Context. */ @Nonnull - Context currentContext(); + Context getCurrentContext(); } diff --git a/google-cloud-datastore/src/test/java/com/google/cloud/datastore/telemetry/DisabledTraceUtilTest.java b/google-cloud-datastore/src/test/java/com/google/cloud/datastore/telemetry/DisabledTraceUtilTest.java index 0f3f183cd..144fc7be8 100644 --- a/google-cloud-datastore/src/test/java/com/google/cloud/datastore/telemetry/DisabledTraceUtilTest.java +++ b/google-cloud-datastore/src/test/java/com/google/cloud/datastore/telemetry/DisabledTraceUtilTest.java @@ -29,16 +29,16 @@ public void disabledTraceUtilDoesNotProvideChannelConfigurator() { @Test public void usesDisabledContext() { DisabledTraceUtil traceUtil = new DisabledTraceUtil(); - assertThat(traceUtil.currentContext() instanceof DisabledTraceUtil.Context).isTrue(); + assertThat(traceUtil.getCurrentContext() instanceof DisabledTraceUtil.Context).isTrue(); } @Test public void usesDisabledSpan() { DisabledTraceUtil traceUtil = new DisabledTraceUtil(); - assertThat(traceUtil.currentSpan() instanceof DisabledTraceUtil.Span).isTrue(); + assertThat(traceUtil.getCurrentSpan() instanceof DisabledTraceUtil.Span).isTrue(); assertThat(traceUtil.startSpan("foo") instanceof DisabledTraceUtil.Span).isTrue(); assertThat( - traceUtil.startSpan("foo", traceUtil.currentContext()) + traceUtil.startSpan("foo", traceUtil.getCurrentContext()) instanceof DisabledTraceUtil.Span) .isTrue(); } @@ -46,8 +46,8 @@ public void usesDisabledSpan() { @Test public void usesDisabledScope() { DisabledTraceUtil traceUtil = new DisabledTraceUtil(); - assertThat(traceUtil.currentContext().makeCurrent() instanceof DisabledTraceUtil.Scope) + assertThat(traceUtil.getCurrentContext().makeCurrent() instanceof DisabledTraceUtil.Scope) .isTrue(); - assertThat(traceUtil.currentSpan().makeCurrent() instanceof DisabledTraceUtil.Scope).isTrue(); + assertThat(traceUtil.getCurrentSpan().makeCurrent() instanceof DisabledTraceUtil.Scope).isTrue(); } } diff --git a/google-cloud-datastore/src/test/java/com/google/cloud/datastore/telemetry/EnabledTraceUtilTest.java b/google-cloud-datastore/src/test/java/com/google/cloud/datastore/telemetry/EnabledTraceUtilTest.java index e88e1a849..47950ed15 100644 --- a/google-cloud-datastore/src/test/java/com/google/cloud/datastore/telemetry/EnabledTraceUtilTest.java +++ b/google-cloud-datastore/src/test/java/com/google/cloud/datastore/telemetry/EnabledTraceUtilTest.java @@ -82,23 +82,23 @@ public void enabledTraceUtilProvidesChannelConfigurator() { @Test public void usesEnabledContext() { - assertThat(newEnabledTraceUtil().currentContext() instanceof EnabledTraceUtil.Context).isTrue(); + assertThat(newEnabledTraceUtil().getCurrentContext() instanceof EnabledTraceUtil.Context).isTrue(); } @Test public void usesEnabledSpan() { EnabledTraceUtil traceUtil = newEnabledTraceUtil(); - assertThat(traceUtil.currentSpan() instanceof EnabledTraceUtil.Span).isTrue(); + assertThat(traceUtil.getCurrentSpan() instanceof EnabledTraceUtil.Span).isTrue(); assertThat(traceUtil.startSpan("foo") != null).isTrue(); assertThat( - traceUtil.startSpan("foo", traceUtil.currentContext()) instanceof EnabledTraceUtil.Span) + traceUtil.startSpan("foo", traceUtil.getCurrentContext()) instanceof EnabledTraceUtil.Span) .isTrue(); } @Test public void usesEnabledScope() { EnabledTraceUtil traceUtil = newEnabledTraceUtil(); - assertThat(traceUtil.currentContext().makeCurrent() instanceof EnabledTraceUtil.Scope).isTrue(); - assertThat(traceUtil.currentSpan().makeCurrent() instanceof EnabledTraceUtil.Scope).isTrue(); + assertThat(traceUtil.getCurrentContext().makeCurrent() instanceof EnabledTraceUtil.Scope).isTrue(); + assertThat(traceUtil.getCurrentSpan().makeCurrent() instanceof EnabledTraceUtil.Scope).isTrue(); } } From 946a80514c2fc309058e150c9f115800e656a63a Mon Sep 17 00:00:00 2001 From: jimit-j-shah <57637300+jimit-j-shah@users.noreply.github.com> Date: Wed, 5 Jun 2024 15:38:58 -0700 Subject: [PATCH 3/3] fix: lint errors --- .../google/cloud/datastore/DatastoreImpl.java | 17 +++++++++++------ .../google/cloud/datastore/TransactionImpl.java | 11 ++++------- .../telemetry/DisabledTraceUtilTest.java | 3 ++- .../telemetry/EnabledTraceUtilTest.java | 9 ++++++--- 4 files changed, 23 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 a44635bfc..eea94cc60 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 @@ -105,7 +105,6 @@ static class ReadWriteTransactionCallable implements Callable { private final TransactionCallable callable; private volatile TransactionOptions options; private volatile Transaction transaction; - @Nonnull private final Context parentTraceContext; ReadWriteTransactionCallable( Datastore datastore, @@ -116,7 +115,6 @@ static class ReadWriteTransactionCallable implements Callable { this.callable = callable; this.options = options; this.transaction = null; - this.parentTraceContext = parentTraceContext; } Datastore getDatastore() { @@ -139,8 +137,13 @@ void setPrevTransactionId(ByteString transactionId) { @Override public T call() throws DatastoreException { - try (com.google.cloud.datastore.telemetry.TraceUtil.Scope ignored = - parentTraceContext.makeCurrent()) { + com.google.cloud.datastore.telemetry.TraceUtil traceUtil = + datastore.getOptions().getTraceUtil(); + com.google.cloud.datastore.telemetry.TraceUtil.Span span = + traceUtil.startSpan( + com.google.cloud.datastore.telemetry.TraceUtil.SPAN_NAME_TRANSACTION_RUN, + datastore.getOptions().getTraceUtil().getCurrentContext()); + try (com.google.cloud.datastore.telemetry.TraceUtil.Scope ignored = span.makeCurrent()) { transaction = datastore.newTransaction(options); T value = callable.run(transaction); transaction.commit(); @@ -167,7 +170,8 @@ public T runInTransaction(final TransactionCallable callable) { com.google.cloud.datastore.telemetry.TraceUtil.SPAN_NAME_TRANSACTION_RUN); try (com.google.cloud.datastore.telemetry.TraceUtil.Scope ignored = span.makeCurrent()) { return RetryHelper.runWithRetries( - new ReadWriteTransactionCallable(this, callable, null, otelTraceUtil.getCurrentContext()), + new ReadWriteTransactionCallable( + this, callable, null, otelTraceUtil.getCurrentContext()), retrySettings, TRANSACTION_EXCEPTION_HANDLER, getOptions().getClock()); @@ -680,7 +684,8 @@ com.google.datastore.v1.BeginTransactionResponse beginTransaction( final com.google.datastore.v1.BeginTransactionRequest requestPb) { com.google.cloud.datastore.telemetry.TraceUtil.Span span = otelTraceUtil.startSpan( - com.google.cloud.datastore.telemetry.TraceUtil.SPAN_NAME_BEGIN_TRANSACTION); + com.google.cloud.datastore.telemetry.TraceUtil.SPAN_NAME_BEGIN_TRANSACTION, + otelTraceUtil.getCurrentContext()); try (com.google.cloud.datastore.telemetry.TraceUtil.Scope scope = span.makeCurrent()) { return RetryHelper.runWithRetries( () -> datastoreRpc.beginTransaction(requestPb), diff --git a/google-cloud-datastore/src/main/java/com/google/cloud/datastore/TransactionImpl.java b/google-cloud-datastore/src/main/java/com/google/cloud/datastore/TransactionImpl.java index 50838f0af..e730db81f 100644 --- a/google-cloud-datastore/src/main/java/com/google/cloud/datastore/TransactionImpl.java +++ b/google-cloud-datastore/src/main/java/com/google/cloud/datastore/TransactionImpl.java @@ -20,8 +20,7 @@ import com.google.api.core.BetaApi; import com.google.cloud.datastore.models.ExplainOptions; -import com.google.cloud.datastore.telemetry.TraceUtil.Context; -import com.google.cloud.datastore.telemetry.TraceUtil.Scope; +import com.google.cloud.datastore.telemetry.TraceUtil; import com.google.common.collect.ImmutableList; import com.google.datastore.v1.ReadOptions; import com.google.datastore.v1.TransactionOptions; @@ -40,7 +39,7 @@ final class TransactionImpl extends BaseDatastoreBatchWriter implements Transact private final ReadOptionProtoPreparer readOptionProtoPreparer; - @Nonnull private final Context transactionTraceContext; + @Nonnull private final TraceUtil traceUtil; static class ResponseImpl implements Transaction.Response { @@ -83,7 +82,7 @@ public List getGeneratedKeys() { transactionId = datastore.requestTransactionId(requestPb); this.readOptionProtoPreparer = new ReadOptionProtoPreparer(); - this.transactionTraceContext = datastore.getOptions().getTraceUtil().getCurrentContext(); + this.traceUtil = datastore.getOptions().getTraceUtil(); } @Override @@ -102,9 +101,7 @@ public Iterator get(Key... keys) { @Override public List fetch(Key... keys) { validateActive(); - try (Scope ignored = transactionTraceContext.makeCurrent()) { - return DatastoreHelper.fetch(this, keys); - } + return DatastoreHelper.fetch(this, keys); } @Override diff --git a/google-cloud-datastore/src/test/java/com/google/cloud/datastore/telemetry/DisabledTraceUtilTest.java b/google-cloud-datastore/src/test/java/com/google/cloud/datastore/telemetry/DisabledTraceUtilTest.java index 144fc7be8..a24f55597 100644 --- a/google-cloud-datastore/src/test/java/com/google/cloud/datastore/telemetry/DisabledTraceUtilTest.java +++ b/google-cloud-datastore/src/test/java/com/google/cloud/datastore/telemetry/DisabledTraceUtilTest.java @@ -48,6 +48,7 @@ public void usesDisabledScope() { DisabledTraceUtil traceUtil = new DisabledTraceUtil(); assertThat(traceUtil.getCurrentContext().makeCurrent() instanceof DisabledTraceUtil.Scope) .isTrue(); - assertThat(traceUtil.getCurrentSpan().makeCurrent() instanceof DisabledTraceUtil.Scope).isTrue(); + assertThat(traceUtil.getCurrentSpan().makeCurrent() instanceof DisabledTraceUtil.Scope) + .isTrue(); } } diff --git a/google-cloud-datastore/src/test/java/com/google/cloud/datastore/telemetry/EnabledTraceUtilTest.java b/google-cloud-datastore/src/test/java/com/google/cloud/datastore/telemetry/EnabledTraceUtilTest.java index 47950ed15..2497672d9 100644 --- a/google-cloud-datastore/src/test/java/com/google/cloud/datastore/telemetry/EnabledTraceUtilTest.java +++ b/google-cloud-datastore/src/test/java/com/google/cloud/datastore/telemetry/EnabledTraceUtilTest.java @@ -82,7 +82,8 @@ public void enabledTraceUtilProvidesChannelConfigurator() { @Test public void usesEnabledContext() { - assertThat(newEnabledTraceUtil().getCurrentContext() instanceof EnabledTraceUtil.Context).isTrue(); + assertThat(newEnabledTraceUtil().getCurrentContext() instanceof EnabledTraceUtil.Context) + .isTrue(); } @Test @@ -91,14 +92,16 @@ public void usesEnabledSpan() { assertThat(traceUtil.getCurrentSpan() instanceof EnabledTraceUtil.Span).isTrue(); assertThat(traceUtil.startSpan("foo") != null).isTrue(); assertThat( - traceUtil.startSpan("foo", traceUtil.getCurrentContext()) instanceof EnabledTraceUtil.Span) + traceUtil.startSpan("foo", traceUtil.getCurrentContext()) + instanceof EnabledTraceUtil.Span) .isTrue(); } @Test public void usesEnabledScope() { EnabledTraceUtil traceUtil = newEnabledTraceUtil(); - assertThat(traceUtil.getCurrentContext().makeCurrent() instanceof EnabledTraceUtil.Scope).isTrue(); + assertThat(traceUtil.getCurrentContext().makeCurrent() instanceof EnabledTraceUtil.Scope) + .isTrue(); assertThat(traceUtil.getCurrentSpan().makeCurrent() instanceof EnabledTraceUtil.Scope).isTrue(); } }