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 11ff58bc9..79b647acf 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 @@ -18,9 +18,6 @@ import static com.google.cloud.datastore.telemetry.TraceUtil.ATTRIBUTES_KEY_DEFERRED; import static com.google.cloud.datastore.telemetry.TraceUtil.ATTRIBUTES_KEY_DOCUMENT_COUNT; -import static com.google.cloud.datastore.telemetry.TraceUtil.ATTRIBUTES_KEY_EXCEPTION_MESSAGE; -import static com.google.cloud.datastore.telemetry.TraceUtil.ATTRIBUTES_KEY_EXCEPTION_STACKTRACE; -import static com.google.cloud.datastore.telemetry.TraceUtil.ATTRIBUTES_KEY_EXCEPTION_TYPE; import static com.google.cloud.datastore.telemetry.TraceUtil.ATTRIBUTES_KEY_MISSING; import static com.google.cloud.datastore.telemetry.TraceUtil.ATTRIBUTES_KEY_MORE_RESULTS; import static com.google.cloud.datastore.telemetry.TraceUtil.ATTRIBUTES_KEY_READ_CONSISTENCY; @@ -37,10 +34,10 @@ 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; import com.google.cloud.datastore.telemetry.TraceUtil.Scope; 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; @@ -53,10 +50,6 @@ import com.google.datastore.v1.RunQueryResponse; import com.google.datastore.v1.TransactionOptions; import com.google.protobuf.ByteString; -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; @@ -115,25 +108,24 @@ public Transaction newTransaction() { return new TransactionImpl(this); } - static class ReadWriteTransactionCallable implements Callable { - + static class TracedReadWriteTransactionCallable implements Callable { private final Datastore datastore; private final TransactionCallable callable; private volatile TransactionOptions options; private volatile Transaction transaction; - private final com.google.cloud.datastore.telemetry.TraceUtil.SpanContext parentSpanContext; + private final TraceUtil.Span parentSpan; - ReadWriteTransactionCallable( + TracedReadWriteTransactionCallable( Datastore datastore, TransactionCallable callable, TransactionOptions options, - @Nullable com.google.cloud.datastore.telemetry.TraceUtil.SpanContext parentSpanContext) { + @Nullable com.google.cloud.datastore.telemetry.TraceUtil.Span parentSpan) { this.datastore = datastore; this.callable = callable; this.options = options; this.transaction = null; - this.parentSpanContext = parentSpanContext; + this.parentSpan = parentSpan; } Datastore getDatastore() { @@ -154,57 +146,75 @@ void setPrevTransactionId(ByteString transactionId) { options = options.toBuilder().setReadWrite(readWrite).build(); } - private io.opentelemetry.api.trace.Span startSpanWithParentContext( - 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 otelTraceUtil.addSettingsAttributesToCurrentSpan(spanBuilder).startSpan(); + @Override + public T call() throws DatastoreException { + try (io.opentelemetry.context.Scope ignored = + Context.current().with(parentSpan.getSpan()).makeCurrent()) { + transaction = datastore.newTransaction(options); + T value = callable.run(transaction); + transaction.commit(); + return value; + } catch (Exception ex) { + transaction.rollback(); + throw DatastoreException.propagateUserException(ex); + } finally { + if (transaction.isActive()) { + transaction.rollback(); + } + if (options != null + && options.getModeCase().equals(TransactionOptions.ModeCase.READ_WRITE)) { + setPrevTransactionId(transaction.getTransactionId()); + } + } + } + } + + static class ReadWriteTransactionCallable implements Callable { + private final Datastore datastore; + private final TransactionCallable callable; + private volatile TransactionOptions options; + private volatile Transaction transaction; + + ReadWriteTransactionCallable( + Datastore datastore, TransactionCallable callable, TransactionOptions options) { + this.datastore = datastore; + this.callable = callable; + this.options = options; + this.transaction = null; + } + + Datastore getDatastore() { + return datastore; + } + + TransactionOptions getOptions() { + return options; + } + + Transaction getTransaction() { + return transaction; + } + + void setPrevTransactionId(ByteString transactionId) { + TransactionOptions.ReadWrite readWrite = + TransactionOptions.ReadWrite.newBuilder().setPreviousTransaction(transactionId).build(); + options = options.toBuilder().setReadWrite(readWrite).build(); } @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 = - startSpanWithParentContext( - com.google.cloud.datastore.telemetry.TraceUtil.SPAN_NAME_TRANSACTION_RUN, - parentSpanContext); - try (io.opentelemetry.context.Scope ignored = span.makeCurrent()) { + try { transaction = datastore.newTransaction(options); T value = callable.run(transaction); transaction.commit(); return value; } catch (Exception ex) { transaction.rollback(); - span.setStatus(StatusCode.ERROR, ex.getMessage()); - span.recordException( - ex, - Attributes.builder() - .put(ATTRIBUTES_KEY_EXCEPTION_MESSAGE, ex.getMessage()) - .put(ATTRIBUTES_KEY_EXCEPTION_TYPE, ex.getClass().getName()) - .put(ATTRIBUTES_KEY_EXCEPTION_STACKTRACE, Throwables.getStackTraceAsString(ex)) - .build()); - span.end(); throw DatastoreException.propagateUserException(ex); } finally { if (transaction.isActive()) { transaction.rollback(); } - span.end(); if (options != null && options.getModeCase().equals(TransactionOptions.ModeCase.READ_WRITE)) { setPrevTransactionId(transaction.getTransactionId()); @@ -215,30 +225,51 @@ public T call() throws DatastoreException { @Override public T runInTransaction(final TransactionCallable callable) { - try { + TraceUtil.Span span = + otelTraceUtil.startSpan( + com.google.cloud.datastore.telemetry.TraceUtil.SPAN_NAME_TRANSACTION_RUN); + Callable transactionCallable = + (getOptions().getOpenTelemetryOptions().isEnabled() + ? new TracedReadWriteTransactionCallable( + this, callable, /*transactionOptions=*/ null, span) + : new ReadWriteTransactionCallable(this, callable, /*transactionOptions=*/ null)); + try (Scope ignored = span.makeCurrent()) { return RetryHelper.runWithRetries( - new ReadWriteTransactionCallable( - this, callable, null, otelTraceUtil.getCurrentSpanContext()), + transactionCallable, retrySettings, TRANSACTION_EXCEPTION_HANDLER, getOptions().getClock()); } catch (RetryHelperException e) { + span.end(e); throw DatastoreException.translateAndThrow(e); + } finally { + span.end(); } } @Override public T runInTransaction( final TransactionCallable callable, TransactionOptions transactionOptions) { - try { + TraceUtil.Span span = + otelTraceUtil.startSpan( + com.google.cloud.datastore.telemetry.TraceUtil.SPAN_NAME_TRANSACTION_RUN); + + Callable transactionCallable = + (getOptions().getOpenTelemetryOptions().isEnabled() + ? new TracedReadWriteTransactionCallable(this, callable, transactionOptions, span) + : new ReadWriteTransactionCallable(this, callable, transactionOptions)); + + try (Scope ignored = span.makeCurrent()) { return RetryHelper.runWithRetries( - new ReadWriteTransactionCallable( - this, callable, transactionOptions, otelTraceUtil.getCurrentSpanContext()), + transactionCallable, retrySettings, TRANSACTION_EXCEPTION_HANDLER, getOptions().getClock()); } catch (RetryHelperException e) { + span.end(e); throw DatastoreException.translateAndThrow(e); + } finally { + span.end(); } } @@ -747,8 +778,7 @@ 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, - otelTraceUtil.getCurrentSpanContext()); + com.google.cloud.datastore.telemetry.TraceUtil.SPAN_NAME_BEGIN_TRANSACTION); 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/telemetry/DisabledTraceUtil.java b/google-cloud-datastore/src/main/java/com/google/cloud/datastore/telemetry/DisabledTraceUtil.java index 06941c721..ebb630515 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 @@ -19,7 +19,7 @@ import com.google.api.core.ApiFunction; import com.google.api.core.ApiFuture; import com.google.api.core.InternalApi; -import com.google.cloud.datastore.telemetry.TraceUtil.SpanContext; +import com.google.cloud.datastore.telemetry.TraceUtil.Context; import io.grpc.ManagedChannelBuilder; import io.opentelemetry.api.trace.Span; import io.opentelemetry.api.trace.SpanBuilder; @@ -35,14 +35,6 @@ */ @InternalApi public class DisabledTraceUtil implements TraceUtil { - - static class SpanContext implements TraceUtil.SpanContext { - @Override - public io.opentelemetry.api.trace.SpanContext getSpanContext() { - return null; - } - } - static class Span implements TraceUtil.Span { @Override public void end() {} @@ -112,7 +104,7 @@ public Span startSpan(String spanName) { } @Override - public TraceUtil.Span startSpan(String spanName, TraceUtil.SpanContext parentSpanContext) { + public TraceUtil.Span startSpan(String spanName, TraceUtil.Span parentSpan) { return new Span(); } @@ -132,12 +124,6 @@ public TraceUtil.Context getCurrentContext() { return new Context(); } - @Nonnull - @Override - public TraceUtil.SpanContext getCurrentSpanContext() { - return new SpanContext(); - } - @Override public Tracer getTracer() { return TracerProvider.noop().get(LIBRARY_NAME); 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 3b962754d..40fc7308e 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 @@ -22,19 +22,19 @@ import com.google.api.core.ApiFutures; import com.google.api.core.InternalApi; import com.google.cloud.datastore.DatastoreOptions; -import com.google.cloud.datastore.telemetry.TraceUtil.SpanContext; +import com.google.cloud.datastore.telemetry.TraceUtil.Context; +import com.google.cloud.datastore.telemetry.TraceUtil.Scope; +import com.google.cloud.datastore.telemetry.TraceUtil.Span; import com.google.common.base.Throwables; import io.grpc.ManagedChannelBuilder; import io.opentelemetry.api.GlobalOpenTelemetry; import io.opentelemetry.api.OpenTelemetry; import io.opentelemetry.api.common.Attributes; import io.opentelemetry.api.common.AttributesBuilder; -import io.opentelemetry.api.trace.Span; import io.opentelemetry.api.trace.SpanBuilder; import io.opentelemetry.api.trace.SpanKind; import io.opentelemetry.api.trace.StatusCode; import io.opentelemetry.api.trace.Tracer; -import io.opentelemetry.context.Context; import java.util.Map; import javax.annotation.Nonnull; import javax.annotation.Nullable; @@ -73,19 +73,6 @@ public ApiFunction getChannelConfi return null; } - static class SpanContext implements TraceUtil.SpanContext { - private final io.opentelemetry.api.trace.SpanContext spanContext; - - public SpanContext(io.opentelemetry.api.trace.SpanContext spanContext) { - this.spanContext = spanContext; - } - - @Override - public io.opentelemetry.api.trace.SpanContext getSpanContext() { - return this.spanContext; - } - } - static class Span implements TraceUtil.Span { private final io.opentelemetry.api.trace.Span span; private final String spanName; @@ -95,6 +82,11 @@ public Span(io.opentelemetry.api.trace.Span span, String spanName) { this.spanName = spanName; } + @Override + public io.opentelemetry.api.trace.Span getSpan() { + return this.span; + } + /** Ends this span. */ @Override public void end() { @@ -197,10 +189,6 @@ public TraceUtil.Span setAttribute(String key, boolean value) { return this; } - public io.opentelemetry.api.trace.Span getSpan() { - return this.span; - } - @Override public Scope makeCurrent() { try (io.opentelemetry.context.Scope scope = span.makeCurrent()) { @@ -307,18 +295,13 @@ public Span startSpan(String spanName) { } @Override - public TraceUtil.Span startSpan(String spanName, TraceUtil.SpanContext parentSpanContext) { + public TraceUtil.Span startSpan(String spanName, TraceUtil.Span parentSpan) { SpanBuilder spanBuilder = tracer .spanBuilder(spanName) .setSpanKind(SpanKind.PRODUCER) - .setParent( - io.opentelemetry.context.Context.current() - .with( - io.opentelemetry.api.trace.Span.wrap(parentSpanContext.getSpanContext()))); - io.opentelemetry.api.trace.Span span = - addSettingsAttributesToCurrentSpan(spanBuilder).startSpan(); - return new Span(span, spanName); + .setParent(io.opentelemetry.context.Context.current().with(parentSpan.getSpan())); + return new Span(addSettingsAttributesToCurrentSpan(spanBuilder).startSpan(), spanName); } @Nonnull @@ -333,12 +316,6 @@ public TraceUtil.Context getCurrentContext() { return new Context(io.opentelemetry.context.Context.current()); } - @Nonnull - @Override - public TraceUtil.SpanContext getCurrentSpanContext() { - return new SpanContext(io.opentelemetry.api.trace.Span.current().getSpanContext()); - } - @Override public Tracer getTracer() { return this.tracer; 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 245b825e1..57b3eab9b 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.api.trace.Span; import io.opentelemetry.api.trace.SpanBuilder; import io.opentelemetry.api.trace.Tracer; import java.util.Map; @@ -95,11 +94,6 @@ static TraceUtil getInstance(@Nonnull DatastoreOptions datastoreOptions) { @Nullable ApiFunction getChannelConfigurator(); - /** Represents a trace span's context */ - interface SpanContext { - io.opentelemetry.api.trace.SpanContext getSpanContext(); - } - /** Represents a trace span. */ interface Span { /** Adds the given event to this span. */ @@ -153,10 +147,10 @@ interface Scope extends AutoCloseable { Span startSpan(String spanName); /** - * Starts a new span with the given name and the span represented by the parentSpanContext as its - * parents, sets it as the current span and returns it. + * Starts a new span with the given name and the span represented by the parentSpan as its parent, + * sets it as the current span and returns it. */ - Span startSpan(String spanName, SpanContext parentSpanContext); + Span startSpan(String spanName, Span parentSpan); /** * Adds common SpanAttributes to the current span, useful when hand-creating a new Span without @@ -172,10 +166,6 @@ interface Scope extends AutoCloseable { @Nonnull Context getCurrentContext(); - /** Returns the current SpanContext */ - @Nonnull - SpanContext getCurrentSpanContext(); - /** Returns the current OpenTelemetry Tracer when OpenTelemetry SDK is provided. */ Tracer getTracer(); } diff --git a/google-cloud-datastore/src/test/java/com/google/cloud/datastore/DatastoreTest.java b/google-cloud-datastore/src/test/java/com/google/cloud/datastore/DatastoreTest.java index cd768f986..d88e1d1ec 100644 --- a/google-cloud-datastore/src/test/java/com/google/cloud/datastore/DatastoreTest.java +++ b/google-cloud-datastore/src/test/java/com/google/cloud/datastore/DatastoreTest.java @@ -176,6 +176,8 @@ public static void beforeClass() throws IOException, InterruptedException { public void setUp() { rpcFactoryMock = EasyMock.createStrictMock(DatastoreRpcFactory.class); rpcMock = EasyMock.createStrictMock(DatastoreRpc.class); + DatastoreOpenTelemetryOptions.Builder otelOptionsBuilder = + DatastoreOpenTelemetryOptions.newBuilder(); rpcMockOptions = options .toBuilder() 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 89c91b3a7..c80ef9353 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 @@ -38,7 +38,7 @@ public void usesDisabledSpan() { assertThat(traceUtil.getCurrentSpan() instanceof DisabledTraceUtil.Span).isTrue(); assertThat(traceUtil.startSpan("foo") instanceof DisabledTraceUtil.Span).isTrue(); assertThat( - traceUtil.startSpan("foo", traceUtil.getCurrentSpanContext()) + traceUtil.startSpan("foo", traceUtil.getCurrentSpan()) instanceof DisabledTraceUtil.Span) .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 a3620bbc2..50d7b6820 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 @@ -65,7 +65,7 @@ public void usesOpenTelemetryFromOptions() { @Test public void usesGlobalOpenTelemetryIfOpenTelemetryInstanceNotProvided() { - OpenTelemetrySdk globalOpenTelemetrySdk = OpenTelemetrySdk.builder().buildAndRegisterGlobal(); + OpenTelemetrySdk ignored = OpenTelemetrySdk.builder().buildAndRegisterGlobal(); DatastoreOptions firestoreOptions = getBaseOptions() .setOpenTelemetryOptions( @@ -92,8 +92,7 @@ public void usesEnabledSpan() { assertThat(traceUtil.getCurrentSpan() instanceof EnabledTraceUtil.Span).isTrue(); assertThat(traceUtil.startSpan("foo") != null).isTrue(); assertThat( - traceUtil.startSpan("foo", traceUtil.getCurrentSpanContext()) - instanceof EnabledTraceUtil.Span) + traceUtil.startSpan("foo", traceUtil.getCurrentSpan()) instanceof EnabledTraceUtil.Span) .isTrue(); }