Track OTel propagated context not captured on the scope stack (#7114)

DataDog · Jun 5, 2024 · 6860183 · 6860183
1 parent aa92d53
commit 6860183
Show file tree

Hide file tree

Showing 6 changed files with 43 additions and 7 deletions.
diff --git a/...nt/agent-otel/otel-shim/src/main/java/datadog/opentelemetry/shim/context/OtelContext.java b/...nt/agent-otel/otel-shim/src/main/java/datadog/opentelemetry/shim/context/OtelContext.java
@@ -20,6 +20,9 @@ public class OtelContext implements Context {
   private static final String OTEL_CONTEXT_SPAN_KEY = "opentelemetry-trace-span-key";
   private static final String OTEL_CONTEXT_ROOT_SPAN_KEY = "opentelemetry-traces-local-root-span";
 
+  /** Keep track of propagated context that has not been captured on the scope stack. */
+  private static final ThreadLocal<OtelContext> lastPropagated = new ThreadLocal<>();
+
   private final Span currentSpan;
   private final Span rootSpan;
 
@@ -51,15 +54,31 @@ public <V> Context with(ContextKey<V> k1, V v1) {
 
   @Override
   public Scope makeCurrent() {
-    Scope scope = Context.super.makeCurrent();
+    final Scope scope = Context.super.makeCurrent();
     if (this.currentSpan instanceof OtelSpan) {
+      // only keep propagated context until next span activation
+      lastPropagated.remove();
       AgentScope agentScope = ((OtelSpan) this.currentSpan).activate();
-      scope = new OtelScope(scope, agentScope);
+      return new OtelScope(scope, agentScope);
+    } else {
+      // propagated context not on the scope stack, capture it here
+      lastPropagated.set(this);
+      return new Scope() {
+        @Override
+        public void close() {
+          lastPropagated.remove();
+          scope.close();
+        }
+      };
     }
-    return scope;
   }
 
   public static Context current() {
+    // Check for propagated context not on the scope stack
+    Context context = lastPropagated.get();
+    if (null != context) {
+      return context;
+    }
     // Check empty context
     AgentSpan agentCurrentSpan = AgentTracer.activeSpan();
     if (null == agentCurrentSpan) {
@@ -91,6 +110,12 @@ public static Context current() {
     return new OtelContext(otelCurrentSpan, otelRootSpan);
   }
 
+  /** Last propagated context not on the scope stack; {@code null} if there's no such context. */
+  @Nullable
+  public static Context lastPropagated() {
+    return lastPropagated.get();
+  }
+
   @Override
   public String toString() {
     return "OtelContext{"

diff --git a/.../agent-otel/otel-shim/src/main/java/datadog/opentelemetry/shim/trace/OtelSpanBuilder.java b/.../agent-otel/otel-shim/src/main/java/datadog/opentelemetry/shim/trace/OtelSpanBuilder.java
@@ -10,6 +10,7 @@
 import static java.lang.Boolean.parseBoolean;
 import static java.util.Locale.ROOT;
 
+import datadog.opentelemetry.shim.context.OtelContext;
 import datadog.trace.bootstrap.instrumentation.api.AgentSpan;
 import datadog.trace.bootstrap.instrumentation.api.AgentTracer;
 import io.opentelemetry.api.common.AttributeKey;
@@ -27,6 +28,7 @@
 public class OtelSpanBuilder implements SpanBuilder {
   private final AgentTracer.SpanBuilder delegate;
   private boolean spanKindSet;
+  private boolean ignoreActiveSpan;
   /**
    * Operation name overridden value by {@link OtelConventions#OPERATION_NAME_SPECIFIC_ATTRIBUTE}
    * reserved attribute ({@code null} if not set).
@@ -51,14 +53,15 @@ public SpanBuilder setParent(Context context) {
     AgentSpan.Context extractedContext = extract(context);
     if (extractedContext != null) {
       this.delegate.asChildOf(extractedContext);
+      this.ignoreActiveSpan = true;
     }
     return this;
   }
 
   @Override
   public SpanBuilder setNoParent() {
-    this.delegate.asChildOf(null);
-    this.delegate.ignoreActiveSpan();
+    this.delegate.ignoreActiveSpan().asChildOf(null);
+    this.ignoreActiveSpan = true;
     return this;
   }
 
@@ -163,6 +166,13 @@ public Span startSpan() {
     if (!this.spanKindSet) {
       setSpanKind(INTERNAL);
     }
+    if (!this.ignoreActiveSpan) {
+      // support automatic parenting from propagated context not on the scope stack
+      Context context = OtelContext.lastPropagated();
+      if (null != context) {
+        setParent(context);
+      }
+    }
     AgentSpan delegate = this.delegate.start();
     // Apply overrides
     if (this.overriddenOperationName != null) {

diff --git a/...main/java/datadog/trace/instrumentation/opentelemetry14/OpenTelemetryInstrumentation.java b/...main/java/datadog/trace/instrumentation/opentelemetry14/OpenTelemetryInstrumentation.java
@@ -58,6 +58,7 @@ public boolean onlyMatchKnownTypes() {
   public String[] helperClassNames() {
     return new String[] {
       "datadog.opentelemetry.shim.context.OtelContext",
+      "datadog.opentelemetry.shim.context.OtelContext$1",
       "datadog.opentelemetry.shim.context.OtelScope",
       "datadog.opentelemetry.shim.context.propagation.AgentTextMapPropagator",
       "datadog.opentelemetry.shim.context.propagation.OtelContextPropagators",

diff --git a/...og/trace/instrumentation/opentelemetry14/context/OpenTelemetryContextInstrumentation.java b/...og/trace/instrumentation/opentelemetry14/context/OpenTelemetryContextInstrumentation.java
@@ -55,6 +55,7 @@ public boolean onlyMatchKnownTypes() {
   public String[] helperClassNames() {
     return new String[] {
       "datadog.opentelemetry.shim.context.OtelContext",
+      "datadog.opentelemetry.shim.context.OtelContext$1",
       "datadog.opentelemetry.shim.context.OtelScope",
       "datadog.opentelemetry.shim.trace.OtelExtractedContext",
       "datadog.opentelemetry.shim.trace.OtelConventions",

diff --git a/...e/instrumentation/opentelemetry14/context/OpenTelemetryContextStorageInstrumentation.java b/...e/instrumentation/opentelemetry14/context/OpenTelemetryContextStorageInstrumentation.java
@@ -56,6 +56,7 @@ public boolean onlyMatchKnownTypes() {
   public String[] helperClassNames() {
     return new String[] {
       "datadog.opentelemetry.shim.context.OtelContext",
+      "datadog.opentelemetry.shim.context.OtelContext$1",
       "datadog.opentelemetry.shim.context.OtelScope",
       "datadog.opentelemetry.shim.trace.OtelExtractedContext",
       "datadog.opentelemetry.shim.trace.OtelConventions",

diff --git a/...nstrumentation/opentelemetry/opentelemetry-1.4/src/test/groovy/OpenTelemetry14Test.groovy b/...nstrumentation/opentelemetry/opentelemetry-1.4/src/test/groovy/OpenTelemetry14Test.groovy
@@ -12,7 +12,6 @@ import io.opentelemetry.context.Context
 import io.opentelemetry.context.ThreadLocalContextStorage
 import opentelemetry14.context.propagation.TextMap
 import org.skyscreamer.jsonassert.JSONAssert
-import spock.lang.Ignore
 import spock.lang.Subject
 
 import java.security.InvalidParameterException
@@ -100,7 +99,6 @@ class OpenTelemetry14Test extends AgentTestRunner {
     }
   }
 
-  @Ignore("Core tracer is not picking incomplete span context from context")
   def "test parent span using propagation data"() {
     setup:
     def traceId = '00000000000000001111111111111111'