initial otel dropin support

Author: mhlidd

dd-java-agent/agent-otel/otel-shim/src/main/java/datadog/opentelemetry/shim/context/OtelContext.java

@@ -5,6 +5,7 @@
 import datadog.trace.bootstrap.instrumentation.api.AgentSpan;
 import datadog.trace.bootstrap.instrumentation.api.AgentTracer;
 import datadog.trace.bootstrap.instrumentation.api.AttachableWrapper;
+import datadog.trace.bootstrap.instrumentation.api.BaggageContext;
 import io.opentelemetry.api.trace.Span;
 import io.opentelemetry.context.Context;
 import io.opentelemetry.context.ContextKey;
@@ -22,22 +23,29 @@ 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";
+  private static final String OTEL_CONTEXT_BAGGAGE_KEY = "opentelemetry-baggage";
 
   /** 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;
+  private final BaggageContext baggageContext;
 
   private final Object[] entries;
 
   public OtelContext(Span currentSpan, Span rootSpan) {
-    this(currentSpan, rootSpan, NO_ENTRIES);
+    this(currentSpan, rootSpan, null, NO_ENTRIES);
   }
 
-  public OtelContext(Span currentSpan, Span rootSpan, Object[] entries) {
+  public OtelContext(Span currentSpan, Span rootSpan, BaggageContext baggageContext) {
+    this(currentSpan, rootSpan, baggageContext, NO_ENTRIES);
+  }
+
+  public OtelContext(Span currentSpan, Span rootSpan, BaggageContext baggageContext, Object[] entries) {
     this.currentSpan = currentSpan;
     this.rootSpan = rootSpan;
+    this.baggageContext = baggageContext;
     this.entries = entries;
   }
 
@@ -49,6 +57,8 @@ public <V> V get(ContextKey<V> key) {
       return (V) this.currentSpan;
     } else if (OTEL_CONTEXT_ROOT_SPAN_KEY.equals(key.toString())) {
       return (V) this.rootSpan;
+    } else if (OTEL_CONTEXT_BAGGAGE_KEY.equals(key.toString())){
+      return (V) this.baggageContext;
     }
     for (int i = 0; i < this.entries.length; i += 2) {
       if (this.entries[i] == key) {
@@ -61,9 +71,11 @@ public <V> V get(ContextKey<V> key) {
   @Override
   public <V> Context with(ContextKey<V> key, V value) {
     if (OTEL_CONTEXT_SPAN_KEY.equals(key.toString())) {
-      return new OtelContext((Span) value, this.rootSpan, this.entries);
+      return new OtelContext((Span) value, this.rootSpan, null, this.entries);
     } else if (OTEL_CONTEXT_ROOT_SPAN_KEY.equals(key.toString())) {
-      return new OtelContext(this.currentSpan, (Span) value, this.entries);
+      return new OtelContext(this.currentSpan, (Span) value, null, this.entries);
+    } else if (OTEL_CONTEXT_BAGGAGE_KEY.equals(key.toString())){
+      return new OtelContext(null, null, this.baggageContext, this.entries);
     }
     Object[] newEntries = null;
     int oldEntriesLength = this.entries.length;
@@ -82,7 +94,7 @@ public <V> Context with(ContextKey<V> key, V value) {
       newEntries[oldEntriesLength] = key;
       newEntries[oldEntriesLength + 1] = value;
     }
-    return new OtelContext(this.currentSpan, this.rootSpan, newEntries);
+    return new OtelContext(this.currentSpan, this.rootSpan, null, newEntries);
   }
 
   @Override
@@ -149,7 +161,7 @@ public static Context current() {
         contextEntries = ((OtelScope) wrapper).contextEntries();
       }
     }
-    return new OtelContext(otelCurrentSpan, otelRootSpan, contextEntries);
+    return new OtelContext(otelCurrentSpan, otelRootSpan, null, contextEntries);
   }
 
   /** Last propagated context not on the scope stack; {@code null} if there's no such context. */
@@ -158,6 +170,10 @@ public static Context lastPropagated() {
     return lastPropagated.get();
   }
 
+  public static String getOtelContextBaggageKey(){
+    return OTEL_CONTEXT_BAGGAGE_KEY;
+  }
+
   @Override
   public String toString() {
     return "OtelContext{"

dd-java-agent/agent-otel/otel-shim/src/main/java/datadog/opentelemetry/shim/context/propagation/AgentTextMapPropagator.java

@@ -5,19 +5,23 @@
 import static datadog.trace.api.TracePropagationStyle.TRACECONTEXT;
 import static datadog.trace.bootstrap.instrumentation.api.AgentPropagation.extractContextAndGetSpanContext;
 
+import datadog.context.propagation.Propagators;
 import datadog.opentelemetry.shim.context.OtelContext;
 import datadog.opentelemetry.shim.trace.OtelExtractedContext;
 import datadog.opentelemetry.shim.trace.OtelSpan;
 import datadog.trace.api.TracePropagationStyle;
 import datadog.trace.bootstrap.instrumentation.api.AgentSpan;
 import datadog.trace.bootstrap.instrumentation.api.AgentSpanContext;
 import datadog.trace.bootstrap.instrumentation.api.AgentSpanContext.Extracted;
+import datadog.trace.bootstrap.instrumentation.api.BaggageContext;
 import datadog.trace.bootstrap.instrumentation.api.TagContext;
 import datadog.trace.util.PropagationUtils;
+import io.opentelemetry.api.baggage.Baggage;
 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.ContextKey;
 import io.opentelemetry.context.propagation.TextMapGetter;
 import io.opentelemetry.context.propagation.TextMapPropagator;
 import io.opentelemetry.context.propagation.TextMapSetter;
@@ -46,20 +50,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));
-              }
-            });
+
+    datadog.context.Context extracted = Propagators.defaultPropagator().extract(datadog.context.Context.root(), 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 new OtelContext(Span.wrap(spanContext), OtelSpan.invalid());
+    }else{
+      AgentSpan extractedSpan = AgentSpan.fromContext(extracted);
+      Extracted extractedSpanContext = extractedSpan == null ? null : (AgentSpanContext.Extracted) extractedSpan.context();
+
+      SpanContext spanContext = null;
+      if (extractedSpanContext != null){
+        TraceState traceState = extractTraceState(extractedSpanContext, carrier, getter);
+        spanContext = fromRemote(extractedSpanContext, traceState);
+      }
+      BaggageContext baggageContext = BaggageContext.fromContext(extracted);
+      return new OtelContext(Span.wrap(spanContext), OtelSpan.invalid(), baggageContext);
     }
   }
 
@@ -68,7 +77,8 @@ private static datadog.context.Context convertContext(Context context) {
     // TODO Create fast path from OtelSpan --> AgentSpan delegate --> with() to inflate as full
     // context if baggage
     AgentSpanContext extract = OtelExtractedContext.extract(context);
-    return AgentSpan.fromSpanContext(extract);
+    AgentSpan agentSpan = AgentSpan.fromSpanContext(extract);
+    return agentSpan.with(BaggageContext.getContextKey(), context.get(ContextKey.named(OtelContext.getOtelContextBaggageKey())));
   }
 
   /**

internal-api/src/main/java/datadog/trace/bootstrap/instrumentation/api/BaggageContext.java

@@ -63,6 +63,10 @@ public Map<String, String> asMap() {
     return new HashMap<>(baggage);
   }
 
+  public static ContextKey<BaggageContext> getContextKey(){
+    return CONTEXT_KEY;
+  }
+
   public static BaggageContext fromContext(Context context) {
     return context.get(CONTEXT_KEY);
   }