From 116e341296f23e9ba072f326f8d9fb711f1934e7 Mon Sep 17 00:00:00 2001
From: Bruce Bujon <bruce.bujon@datadoghq.com>
Date: Tue, 6 May 2025 08:23:42 +0200
Subject: [PATCH] feat(otel): Add context API support OTel propagators

---
 .../shim/context/OtelContext.java             |  9 +++++
 .../propagation/AgentTextMapPropagator.java   | 38 ++++++++-----------
 2 files changed, 25 insertions(+), 22 deletions(-)

diff --git a/dd-java-agent/agent-otel/otel-shim/src/main/java/datadog/opentelemetry/shim/context/OtelContext.java b/dd-java-agent/agent-otel/otel-shim/src/main/java/datadog/opentelemetry/shim/context/OtelContext.java
index bd3fc0ff17d..0f0bded653b 100644
--- a/dd-java-agent/agent-otel/otel-shim/src/main/java/datadog/opentelemetry/shim/context/OtelContext.java
+++ b/dd-java-agent/agent-otel/otel-shim/src/main/java/datadog/opentelemetry/shim/context/OtelContext.java
@@ -89,6 +89,15 @@ public String toString() {
     return "OtelContext{" + "delegate=" + delegate + '}';
   }
 
+  /**
+   * Returns the underlying context.
+   *
+   * @return The underlying context.
+   */
+  public datadog.context.Context asContext() {
+    return this.delegate;
+  }
+
   private static datadog.context.ContextKey delegateKey(ContextKey key) {
     return DELEGATE_KEYS.computeIfAbsent(key, OtelContext::mapByKeyName);
   }
diff --git a/dd-java-agent/agent-otel/otel-shim/src/main/java/datadog/opentelemetry/shim/context/propagation/AgentTextMapPropagator.java b/dd-java-agent/agent-otel/otel-shim/src/main/java/datadog/opentelemetry/shim/context/propagation/AgentTextMapPropagator.java
index 71feb6b9f56..ae7d708b5e4 100644
--- a/dd-java-agent/agent-otel/otel-shim/src/main/java/datadog/opentelemetry/shim/context/propagation/AgentTextMapPropagator.java
+++ b/dd-java-agent/agent-otel/otel-shim/src/main/java/datadog/opentelemetry/shim/context/propagation/AgentTextMapPropagator.java
@@ -1,9 +1,7 @@
 package datadog.opentelemetry.shim.context.propagation;
 
 import static datadog.context.propagation.Propagators.defaultPropagator;
-import static datadog.opentelemetry.shim.trace.OtelSpanContext.fromRemote;
 import static datadog.trace.api.TracePropagationStyle.TRACECONTEXT;
-import static datadog.trace.bootstrap.instrumentation.api.AgentPropagation.extractContextAndGetSpanContext;
 
 import datadog.opentelemetry.shim.context.OtelContext;
 import datadog.opentelemetry.shim.trace.OtelExtractedContext;
@@ -13,8 +11,6 @@
 import datadog.trace.bootstrap.instrumentation.api.AgentSpanContext.Extracted;
 import datadog.trace.bootstrap.instrumentation.api.TagContext;
 import datadog.trace.util.PropagationUtils;
-import io.opentelemetry.api.trace.Span;
-import io.opentelemetry.api.trace.SpanContext;
 import io.opentelemetry.api.trace.TraceState;
 import io.opentelemetry.context.Context;
 import io.opentelemetry.context.propagation.TextMapGetter;
@@ -45,27 +41,25 @@ public <C> Context extract(Context context, @Nullable C carrier, TextMapGetter<C
     if (carrier == null) {
       return context;
     }
-    Extracted extracted =
-        extractContextAndGetSpanContext(
-            carrier,
-            (carrier1, classifier) -> {
-              for (String key : getter.keys(carrier1)) {
-                classifier.accept(key, getter.get(carrier1, key));
-              }
-            });
-    if (extracted == null) {
-      return context;
-    } else {
-      TraceState traceState = extractTraceState(extracted, carrier, getter);
-      SpanContext spanContext = fromRemote(extracted, traceState);
-      return Span.wrap(spanContext).storeInContext(OtelContext.ROOT);
-    }
+    datadog.context.Context extracted =
+        defaultPropagator()
+            .extract(
+                convertContext(context),
+                carrier,
+                (carrier1, classifier) -> {
+                  for (String key : getter.keys(carrier1)) {
+                    classifier.accept(key, getter.get(carrier1, key));
+                  }
+                });
+    return new OtelContext(extracted);
   }
 
   private static datadog.context.Context convertContext(Context context) {
-    // TODO Extract baggage too
-    // TODO Create fast path from OtelSpan --> AgentSpan delegate --> with() to inflate as full
-    // context if baggage
+    // Try to get the underlying context when injecting a Datadog context
+    if (context instanceof OtelContext) {
+      return ((OtelContext) context).asContext();
+    }
+    // Otherwise, fallback to extracting limited tracing context and recreating an OTel context from
     AgentSpanContext extract = OtelExtractedContext.extract(context);
     return AgentSpan.fromSpanContext(extract);
   }