fix: Fixing how span hierarchy is created across threads - using Span…

… instead of Context

google-cloud-datastore/src/main/java/com/google/cloud/datastore/DatastoreImpl.java

@@ -50,6 +50,7 @@
 import com.google.datastore.v1.RunQueryResponse;
 import com.google.datastore.v1.TransactionOptions;
 import com.google.protobuf.ByteString;
+import io.opentelemetry.context.Context;
 import java.util.ArrayList;
 import java.util.Arrays;
 import java.util.Collections;
@@ -114,18 +115,18 @@ static class ReadWriteTransactionCallable<T> implements Callable<T> {
  private volatile TransactionOptions options;
  private volatile Transaction transaction;
 
- private final TraceUtil.Context parentContext;
+ private final TraceUtil.Span parentSpan;
 
  ReadWriteTransactionCallable(
  Datastore datastore,
  TransactionCallable<T> callable,
  TransactionOptions options,
- @Nullable com.google.cloud.datastore.telemetry.TraceUtil.Context parentContext) {
+ @Nullable com.google.cloud.datastore.telemetry.TraceUtil.Span parentSpan) {
  this.datastore = datastore;
  this.callable = callable;
  this.options = options;
  this.transaction = null;
- this.parentContext = parentContext;
+ this.parentSpan = parentSpan;
  }
 
  Datastore getDatastore() {
@@ -148,27 +149,19 @@ void setPrevTransactionId(ByteString transactionId) {
 
  @Override
  public T call() throws DatastoreException {
- TraceUtil.Span span =
- datastore
- .getOptions()
- .getTraceUtil()
- .startSpan(
- com.google.cloud.datastore.telemetry.TraceUtil.SPAN_NAME_TRANSACTION_RUN,
- parentContext);
- try (TraceUtil.Scope ignored = span.makeCurrent()) {
+ 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();
- span.end(ex);
  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());
@@ -179,30 +172,40 @@ 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);
+ try (Scope ignored = span.makeCurrent()) {
  return RetryHelper.runWithRetries(
- new ReadWriteTransactionCallable<T>(
- this, callable, null, otelTraceUtil.getCurrentContext()),
+ new ReadWriteTransactionCallable<T>(this, callable, null, span),
  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);
+ try (Scope ignored = span.makeCurrent()) {
  return RetryHelper.runWithRetries(
- new ReadWriteTransactionCallable<T>(
- this, callable, transactionOptions, otelTraceUtil.getCurrentContext()),
+ new ReadWriteTransactionCallable<T>(this, callable, transactionOptions, span),
  retrySettings,
  TRANSACTION_EXCEPTION_HANDLER,
  getOptions().getClock());
  } catch (RetryHelperException e) {
+ span.end(e);
  throw DatastoreException.translateAndThrow(e);
+ } finally {
+ span.end();
  }
  }
 

google-cloud-datastore/src/main/java/com/google/cloud/datastore/telemetry/DisabledTraceUtil.java

@@ -107,6 +107,11 @@ public TraceUtil.Span startSpan(String spanName, TraceUtil.Context parentContext
  return new Span();
  }
 
+ @Override
+ public TraceUtil.Span startSpan(String spanName, TraceUtil.Span parentSpan) {
+ return new Span();
+ }
+
  public SpanBuilder addSettingsAttributesToCurrentSpan(SpanBuilder spanBuilder) {
  return getTracer().spanBuilder("TRACING_DISABLED_NO_OP");
  }

google-cloud-datastore/src/main/java/com/google/cloud/datastore/telemetry/EnabledTraceUtil.java

@@ -28,10 +28,12 @@
 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;
@@ -79,6 +81,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() {
@@ -181,10 +188,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()) {
@@ -305,6 +308,16 @@ public TraceUtil.Span startSpan(String spanName, TraceUtil.Context parentContext
  return new Span(addSettingsAttributesToCurrentSpan(spanBuilder).startSpan(), spanName);
  }
 
+ @Override
+ public TraceUtil.Span startSpan(String spanName, TraceUtil.Span parentSpan) {
+ SpanBuilder spanBuilder =
+ tracer
+ .spanBuilder(spanName)
+ .setSpanKind(SpanKind.PRODUCER)
+ .setParent(io.opentelemetry.context.Context.current().with(parentSpan.getSpan()));
+ return new Span(addSettingsAttributesToCurrentSpan(spanBuilder).startSpan(), spanName);
+ }
+
  @Nonnull
  @Override
  public TraceUtil.Span getCurrentSpan() {

google-cloud-datastore/src/main/java/com/google/cloud/datastore/telemetry/TraceUtil.java

@@ -153,6 +153,8 @@ interface Scope extends AutoCloseable {
  */
  Span startSpan(String spanName, Context parentContext);
 
+ Span startSpan(String spanName, Span parentSpan);
+
  /**
  * Adds common SpanAttributes to the current span, useful when hand-creating a new Span without
  * using the TraceUtil.Span interface.