Add multiple context extraction support

dd-java-agent/instrumentation/opentelemetry/opentelemetry-0.3/src/test/groovy/OpenTelemetryTest.groovy

@@ -4,6 +4,8 @@ import datadog.trace.api.DDTags
 import datadog.trace.api.DDTraceId
 import datadog.trace.api.interceptor.MutableSpan
 import datadog.trace.core.propagation.PropagationTags
+
+import static datadog.trace.api.TracePropagationStyle.NONE
 import static datadog.trace.api.sampling.PrioritySampling.*
 import static datadog.trace.api.sampling.SamplingMechanism.*
 import datadog.trace.context.TraceScope
@@ -136,11 +138,11 @@ class OpenTelemetryTest extends AgentTestRunner {
     setup:
     def builder = tracer.spanBuilder("some name")
     if (parentId) {
-      def ctx = new ExtractedContext(DDTraceId.ONE, parentId, SAMPLER_DROP, null, PropagationTags.factory().empty())
+      def ctx = new ExtractedContext(DDTraceId.ONE, parentId, SAMPLER_DROP, null, PropagationTags.factory().empty(), NONE)
       builder.setParent(tracer.converter.toSpanContext(ctx))
     }
     if (linkId) {
-      def ctx = new ExtractedContext(DDTraceId.ONE, linkId, SAMPLER_DROP, null, PropagationTags.factory().empty())
+      def ctx = new ExtractedContext(DDTraceId.ONE, linkId, SAMPLER_DROP, null, PropagationTags.factory().empty(), NONE)
       builder.addLink(tracer.converter.toSpanContext(ctx))
     }
     def result = builder.startSpan()

dd-java-agent/instrumentation/opentracing/api-0.31/src/test/groovy/OpenTracing31Test.groovy

@@ -10,6 +10,7 @@ import datadog.trace.instrumentation.opentracing31.OTTracer
 import datadog.trace.instrumentation.opentracing31.TypeConverter
 import spock.lang.Shared
 
+import static datadog.trace.api.TracePropagationStyle.NONE
 import static datadog.trace.api.sampling.PrioritySampling.*
 import static datadog.trace.api.sampling.SamplingMechanism.*
 import datadog.trace.context.TraceScope
@@ -45,7 +46,7 @@ class OpenTracing31Test extends AgentTestRunner {
         .withTag("boolean", true)
     }
     if (addReference) {
-      def ctx = new ExtractedContext(DDTraceId.ONE, 2, SAMPLER_DROP, null, PropagationTags.factory().empty())
+      def ctx = new ExtractedContext(DDTraceId.ONE, 2, SAMPLER_DROP, null, PropagationTags.factory().empty(), NONE)
       builder.addReference(addReference, tracer.tracer.converter.toSpanContext(ctx))
     }
     def result = builder.start()

dd-java-agent/instrumentation/opentracing/api-0.32/src/test/groovy/OpenTracing32Test.groovy

@@ -23,6 +23,7 @@ import spock.lang.Shared
 import spock.lang.Subject
 
 import static datadog.trace.agent.test.utils.TraceUtils.runUnderTrace
+import static datadog.trace.api.TracePropagationStyle.NONE
 import static datadog.trace.api.sampling.PrioritySampling.SAMPLER_DROP
 import static datadog.trace.api.sampling.PrioritySampling.SAMPLER_KEEP
 import static datadog.trace.api.sampling.PrioritySampling.UNSET
@@ -50,7 +51,7 @@ class OpenTracing32Test extends AgentTestRunner {
         .withTag("boolean", true)
     }
     if (addReference) {
-      def ctx = new ExtractedContext(DDTraceId.ONE, 2, SAMPLER_DROP, null, PropagationTags.factory().empty())
+      def ctx = new ExtractedContext(DDTraceId.ONE, 2, SAMPLER_DROP, null, PropagationTags.factory().empty(), NONE)
       builder.addReference(addReference, tracer.tracer.converter.toSpanContext(ctx))
     }
     def result = builder.start()

dd-trace-api/src/main/java/datadog/trace/api/ConfigDefaults.java

@@ -79,6 +79,8 @@ public final class ConfigDefaults {
 
   static final int DEFAULT_CLOCK_SYNC_PERIOD = 30; // seconds
 
+  static final boolean DEFAULT_TRACE_PROPAGATION_EXTRACT_FIRST = false;
+
   static final boolean DEFAULT_JMX_FETCH_MULTIPLE_RUNTIME_SERVICES_ENABLED = false;
   static final int DEFAULT_JMX_FETCH_MULTIPLE_RUNTIME_SERVICES_LIMIT = 10;
 

dd-trace-api/src/main/java/datadog/trace/api/config/TracerConfig.java

@@ -84,6 +84,7 @@ public final class TracerConfig {
   public static final String TRACE_PROPAGATION_STYLE = "trace.propagation.style";
   public static final String TRACE_PROPAGATION_STYLE_EXTRACT = "trace.propagation.style.extract";
   public static final String TRACE_PROPAGATION_STYLE_INJECT = "trace.propagation.style.inject";
+  public static final String TRACE_PROPAGATION_EXTRACT_FIRST = "trace.propagation.extract.first";
 
   public static final String ENABLE_TRACE_AGENT_V05 = "trace.agent.v0.5.enabled";
 

dd-trace-core/src/main/java/datadog/trace/core/CoreTracer.java

@@ -1191,6 +1191,7 @@ public CoreSpanBuilder ignoreActiveSpan() {
     }
 
     private DDSpan buildSpan() {
+      addTerminatedContextAsLinks();
       DDSpan span = DDSpan.create(instrumentationName, timestampMicro, buildSpanContext(), links);
       if (span.isLocalRootSpan()) {
         EndpointTracker tracker = tracer.onRootSpanStarted(span);
@@ -1199,6 +1200,19 @@ private DDSpan buildSpan() {
       return span;
     }
 
+    private void addTerminatedContextAsLinks() {
+      if (this.parent instanceof TagContext) {
+        List<AgentSpanLink> terminatedContextLinks =
+            ((TagContext) this.parent).getTerminatedContextLinks();
+        if (!terminatedContextLinks.isEmpty()) {
+          if (this.links == null) {
+            this.links = new ArrayList<>();
+          }
+          this.links.addAll(terminatedContextLinks);
+        }
+      }
+    }
+
     @Override
     public AgentSpan start() {
       return buildSpan();

dd-trace-core/src/main/java/datadog/trace/core/propagation/B3HttpCodec.java

@@ -1,12 +1,15 @@
 package datadog.trace.core.propagation;
 
+import static datadog.trace.api.TracePropagationStyle.B3MULTI;
+import static datadog.trace.api.TracePropagationStyle.B3SINGLE;
 import static datadog.trace.core.propagation.HttpCodec.firstHeaderValue;
 
 import datadog.trace.api.Config;
 import datadog.trace.api.DD128bTraceId;
 import datadog.trace.api.DDSpanId;
 import datadog.trace.api.DDTraceId;
 import datadog.trace.api.TraceConfig;
+import datadog.trace.api.TracePropagationStyle;
 import datadog.trace.api.sampling.PrioritySampling;
 import datadog.trace.bootstrap.instrumentation.api.AgentPropagation;
 import datadog.trace.core.DDSpanContext;
@@ -161,7 +164,7 @@ static HttpCodec.Extractor newExtractor(
     final List<HttpCodec.Extractor> extractors = new ArrayList<>(2);
     extractors.add(newSingleExtractor(config, traceConfigSupplier));
     extractors.add(newMultiExtractor(config, traceConfigSupplier));
-    return new HttpCodec.CompoundExtractor(extractors);
+    return new HttpCodec.CompoundExtractor(extractors, config.isTracePropagationExtractFirst());
   }
 
   public static HttpCodec.Extractor newMultiExtractor(
@@ -212,6 +215,11 @@ private B3MultiContextInterpreter(Config config) {
       super(config);
     }
 
+    @Override
+    public TracePropagationStyle style() {
+      return B3MULTI;
+    }
+
     @Override
     public boolean accept(final String key, final String value) {
       if (null == key || key.isEmpty() || null == value || value.isEmpty()) {
@@ -267,6 +275,11 @@ public B3SingleContextInterpreter(Config config) {
       super(config);
     }
 
+    @Override
+    public TracePropagationStyle style() {
+      return B3SINGLE;
+    }
+
     @Override
     public boolean accept(String key, String value) {
       try {

dd-trace-core/src/main/java/datadog/trace/core/propagation/ContextInterpreter.java

@@ -21,6 +21,7 @@
 import datadog.trace.api.DDTraceId;
 import datadog.trace.api.Functions;
 import datadog.trace.api.TraceConfig;
+import datadog.trace.api.TracePropagationStyle;
 import datadog.trace.api.cache.DDCache;
 import datadog.trace.api.cache.DDCaches;
 import datadog.trace.api.sampling.PrioritySampling;
@@ -47,6 +48,7 @@ public abstract class ContextInterpreter implements AgentPropagation.KeyClassifi
   protected long endToEndStartTime;
   protected boolean valid;
   protected boolean fullContext;
+  protected final PropagationTags.Factory propagationTagsFactory;
   protected PropagationTags propagationTags;
 
   private TagContext.HttpHeaders httpHeaders;
@@ -58,19 +60,23 @@ public abstract class ContextInterpreter implements AgentPropagation.KeyClassifi
   protected static final boolean LOG_EXTRACT_HEADER_NAMES = Config.get().isLogExtractHeaderNames();
   private static final DDCache<String, String> CACHE = DDCaches.newFixedSizeCache(64);
 
-  protected String toLowerCase(String key) {
+  protected static String toLowerCase(String key) {
     return CACHE.computeIfAbsent(key, Functions.LowerCase.INSTANCE);
   }
 
   protected ContextInterpreter(Config config) {
     this.customIpHeaderName = config.getTraceClientIpHeader();
     this.clientIpResolutionEnabled = config.isTraceClientIpResolverEnabled();
     this.clientIpWithoutAppSec = config.isClientIpEnabled();
+    this.propagationTagsFactory = PropagationTags.factory(config);
   }
 
-  public interface Factory {
-    ContextInterpreter create();
-  }
+  /**
+   * Gets the propagation style handled by the context interpreter.
+   *
+   * @return The propagation style handled.
+   */
+  public abstract TracePropagationStyle style();
 
   protected final boolean handledForwarding(String key, String value) {
     if (value == null || !collectIpHeaders) {
@@ -229,20 +235,21 @@ public ContextInterpreter reset(TraceConfig traceConfig) {
   protected TagContext build() {
     if (valid) {
       if (fullContext && !DDTraceId.ZERO.equals(traceId)) {
-        final ExtractedContext context;
-        context =
-            new ExtractedContext(
-                traceId,
-                spanId,
-                samplingPriorityOrDefault(traceId, samplingPriority),
-                origin,
-                endToEndStartTime,
-                baggage,
-                tags,
-                httpHeaders,
-                propagationTags,
-                traceConfig);
-        return context;
+        if (propagationTags == null) {
+          propagationTags = propagationTagsFactory.empty();
+        }
+        return new ExtractedContext(
+            traceId,
+            spanId,
+            samplingPriorityOrDefault(traceId, samplingPriority),
+            origin,
+            endToEndStartTime,
+            baggage,
+            tags,
+            httpHeaders,
+            propagationTags,
+            traceConfig,
+            style());
       } else if (origin != null
           || !tags.isEmpty()
           || httpHeaders != null
@@ -254,7 +261,8 @@ protected TagContext build() {
             httpHeaders,
             baggage,
             samplingPriorityOrDefault(traceId, samplingPriority),
-            traceConfig);
+            traceConfig,
+            style());
       }
     }
     return null;
@@ -284,4 +292,8 @@ private int samplingPriorityOrDefault(DDTraceId traceId, int samplingPriority) {
         ? defaultSamplingPriority()
         : samplingPriority;
   }
+
+  public interface Factory {
+    ContextInterpreter create();
+  }
 }

dd-trace-core/src/main/java/datadog/trace/core/propagation/DatadogHttpCodec.java

@@ -1,5 +1,6 @@
 package datadog.trace.core.propagation;
 
+import static datadog.trace.api.TracePropagationStyle.DATADOG;
 import static datadog.trace.core.propagation.HttpCodec.firstHeaderValue;
 import static datadog.trace.core.propagation.XRayHttpCodec.XRayContextInterpreter.handleXRayTraceHeader;
 import static datadog.trace.core.propagation.XRayHttpCodec.X_AMZN_TRACE_ID;
@@ -12,9 +13,11 @@
 import datadog.trace.api.DDTags;
 import datadog.trace.api.DDTraceId;
 import datadog.trace.api.TraceConfig;
+import datadog.trace.api.TracePropagationStyle;
 import datadog.trace.bootstrap.instrumentation.api.AgentPropagation;
 import datadog.trace.bootstrap.instrumentation.api.TagContext;
 import datadog.trace.core.DDSpanContext;
+import datadog.trace.core.propagation.PropagationTags.HeaderType;
 import java.util.Map;
 import java.util.TreeMap;
 import java.util.function.Supplier;
@@ -75,8 +78,7 @@ public <C> void inject(
       }
 
       // inject x-datadog-tags
-      String datadogTags =
-          context.getPropagationTags().headerValue(PropagationTags.HeaderType.DATADOG);
+      String datadogTags = context.getPropagationTags().headerValue(HeaderType.DATADOG);
       if (datadogTags != null) {
         setter.set(carrier, DATADOG_TAGS_KEY, datadogTags);
       }
@@ -101,12 +103,15 @@ private static class DatadogContextInterpreter extends ContextInterpreter {
     private static final int IGNORE = -1;
 
     private final boolean isAwsPropagationEnabled;
-    private final PropagationTags.Factory datadogTagsFactory;
 
     private DatadogContextInterpreter(Config config) {
       super(config);
       isAwsPropagationEnabled = config.isAwsPropagationEnabled();
-      datadogTagsFactory = PropagationTags.factory(config);
+    }
+
+    @Override
+    public TracePropagationStyle style() {
+      return DATADOG;
     }
 
     @Override
@@ -180,8 +185,7 @@ public boolean accept(String key, String value) {
                 endToEndStartTime = extractEndToEndStartTime(firstHeaderValue(value));
                 break;
               case DD_TAGS:
-                propagationTags =
-                    datadogTagsFactory.fromHeaderValue(PropagationTags.HeaderType.DATADOG, value);
+                propagationTags = propagationTagsFactory.fromHeaderValue(HeaderType.DATADOG, value);
                 break;
               case OT_BAGGAGE:
                 {

dd-trace-core/src/main/java/datadog/trace/core/propagation/ExtractedContext.java

@@ -2,6 +2,7 @@
 
 import datadog.trace.api.DDTraceId;
 import datadog.trace.api.TraceConfig;
+import datadog.trace.api.TracePropagationStyle;
 import datadog.trace.api.sampling.PrioritySampling;
 import datadog.trace.bootstrap.instrumentation.api.TagContext;
 import java.util.Map;
@@ -20,8 +21,20 @@ public ExtractedContext(
       final long spanId,
       final int samplingPriority,
       final CharSequence origin,
-      final PropagationTags propagationTags) {
-    this(traceId, spanId, samplingPriority, origin, 0, null, null, null, propagationTags, null);
+      final PropagationTags propagationTags,
+      final TracePropagationStyle propagationStyle) {
+    this(
+        traceId,
+        spanId,
+        samplingPriority,
+        origin,
+        0,
+        null,
+        null,
+        null,
+        propagationTags,
+        null,
+        propagationStyle);
   }
 
   public ExtractedContext(
@@ -34,8 +47,9 @@ public ExtractedContext(
       final Map<String, String> tags,
       final HttpHeaders httpHeaders,
       final PropagationTags propagationTags,
-      final TraceConfig traceConfig) {
-    super(origin, tags, httpHeaders, baggage, samplingPriority, traceConfig);
+      final TraceConfig traceConfig,
+      final TracePropagationStyle propagationStyle) {
+    super(origin, tags, httpHeaders, baggage, samplingPriority, traceConfig, propagationStyle);
     this.traceId = traceId;
     this.spanId = spanId;
     this.endToEndStartTime = endToEndStartTime;

dd-trace-core/src/main/java/datadog/trace/core/propagation/HaystackHttpCodec.java

@@ -1,12 +1,14 @@
 package datadog.trace.core.propagation;
 
+import static datadog.trace.api.TracePropagationStyle.HAYSTACK;
 import static datadog.trace.core.propagation.HttpCodec.firstHeaderValue;
 
 import datadog.trace.api.Config;
 import datadog.trace.api.DD64bTraceId;
 import datadog.trace.api.DDSpanId;
 import datadog.trace.api.DDTraceId;
 import datadog.trace.api.TraceConfig;
+import datadog.trace.api.TracePropagationStyle;
 import datadog.trace.api.sampling.PrioritySampling;
 import datadog.trace.bootstrap.instrumentation.api.AgentPropagation;
 import datadog.trace.core.DDSpanContext;
@@ -139,6 +141,11 @@ private HaystackContextInterpreter(Config config) {
       super(config);
     }
 
+    @Override
+    public TracePropagationStyle style() {
+      return HAYSTACK;
+    }
+
     @Override
     public boolean accept(String key, String value) {
       if (null == key || key.isEmpty()) {