fix: Create Span hierarchy using parent Span

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<T> implements Callable<T> {
-
+ static class TracedReadWriteTransactionCallable<T> implements Callable<T> {
  private final Datastore datastore;
  private final TransactionCallable<T> 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<T> 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<T> implements Callable<T> {
+ private final Datastore datastore;
+ private final TransactionCallable<T> callable;
+ private volatile TransactionOptions options;
+ private volatile Transaction transaction;
+
+ ReadWriteTransactionCallable(
+ Datastore datastore, TransactionCallable<T> 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> T runInTransaction(final TransactionCallable<T> callable) {
- try {
+ TraceUtil.Span span =
+ otelTraceUtil.startSpan(
+ com.google.cloud.datastore.telemetry.TraceUtil.SPAN_NAME_TRANSACTION_RUN);
+ Callable<T> transactionCallable =
+ (getOptions().getOpenTelemetryOptions().isEnabled()
+ ? new TracedReadWriteTransactionCallable<T>(
+ this, callable, /*transactionOptions=*/ null, span)
+ : new ReadWriteTransactionCallable<T>(this, callable, /*transactionOptions=*/ null));
+ try (Scope ignored = span.makeCurrent()) {
  return RetryHelper.runWithRetries(
- new ReadWriteTransactionCallable<T>(
- 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> T runInTransaction(
  final TransactionCallable<T> callable, TransactionOptions transactionOptions) {
- try {
+ TraceUtil.Span span =
+ otelTraceUtil.startSpan(
+ com.google.cloud.datastore.telemetry.TraceUtil.SPAN_NAME_TRANSACTION_RUN);
+
+ Callable<T> transactionCallable =
+ (getOptions().getOpenTelemetryOptions().isEnabled()
+ ? new TracedReadWriteTransactionCallable<T>(this, callable, transactionOptions, span)
+ : new ReadWriteTransactionCallable<T>(this, callable, transactionOptions));
+
+ try (Scope ignored = span.makeCurrent()) {
  return RetryHelper.runWithRetries(
- new ReadWriteTransactionCallable<T>(
- 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),

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);

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<ManagedChannelBuilder, ManagedChannelBuilder> 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;