Adding Support for TRACE_PROPAGATION_BEHAVIOR_EXTRACT

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

@@ -91,6 +91,7 @@ public final class ConfigDefaults {
 
   static final int DEFAULT_CLOCK_SYNC_PERIOD = 30; // seconds
 
+  static final String DEFAULT_TRACE_PROPAGATION_BEHAVIOR_EXTRACT = "continue";
   static final boolean DEFAULT_TRACE_PROPAGATION_EXTRACT_FIRST = false;
 
   static final boolean DEFAULT_JMX_FETCH_MULTIPLE_RUNTIME_SERVICES_ENABLED = false;

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

@@ -90,6 +90,8 @@ 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_BEHAVIOR_EXTRACT =
+      "trace.propagation.behavior.extract";
   public static final String TRACE_PROPAGATION_EXTRACT_FIRST = "trace.propagation.extract.first";
   public static final String TRACE_BAGGAGE_MAX_ITEMS = "trace.baggage.max.items";
   public static final String TRACE_BAGGAGE_MAX_BYTES = "trace.baggage.max.bytes";

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

@@ -1546,11 +1546,20 @@ private DDSpanContext buildSpanContext() {
           // Propagate external trace
           isRemote = true;
           final ExtractedContext extractedContext = (ExtractedContext) parentContext;
-          traceId = extractedContext.getTraceId();
-          parentSpanId = extractedContext.getSpanId();
-          samplingPriority = extractedContext.getSamplingPriority();
-          endToEndStartTime = extractedContext.getEndToEndStartTime();
-          propagationTags = extractedContext.getPropagationTags();
+
+          if (Config.get().getTracePropagationBehaviorExtract().equals("restart")) {
+            traceId = idGenerationStrategy.generateTraceId();
+            parentSpanId = 0;
+            samplingPriority = PrioritySampling.UNSET;
+            endToEndStartTime = 0;
+            propagationTags = propagationTagsFactory.empty();
+          } else {
+            traceId = extractedContext.getTraceId();
+            parentSpanId = extractedContext.getSpanId();
+            samplingPriority = extractedContext.getSamplingPriority();
+            endToEndStartTime = extractedContext.getEndToEndStartTime();
+            propagationTags = extractedContext.getPropagationTags();
+          }
         } else if (parentContext != null) {
           traceId =
               parentContext.getTraceId() == DDTraceId.ZERO

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

@@ -856,4 +856,8 @@ public boolean isOutbound() {
     Object spanKind = context.getTag(Tags.SPAN_KIND);
     return Tags.SPAN_KIND_CLIENT.equals(spanKind) || Tags.SPAN_KIND_PRODUCER.equals(spanKind);
   }
+
+  public List<AgentSpanLink> getLinks() {
+    return links;
+  }
 }

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

@@ -164,7 +164,10 @@ 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, config.isTracePropagationExtractFirst());
+    return new HttpCodec.CompoundExtractor(
+        extractors,
+        config.isTracePropagationExtractFirst(),
+        config.getTracePropagationBehaviorExtract());
   }
 
   public static HttpCodec.Extractor newMultiExtractor(

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

@@ -171,10 +171,11 @@ public static Extractor createExtractor(
     switch (extractors.size()) {
       case 0:
         return StubExtractor.INSTANCE;
-      case 1:
-        return extractors.get(0);
       default:
-        return new CompoundExtractor(extractors, config.isTracePropagationExtractFirst());
+        return new CompoundExtractor(
+            extractors,
+            config.isTracePropagationExtractFirst(),
+            config.getTracePropagationBehaviorExtract());
     }
   }
 
@@ -208,10 +209,13 @@ public <C> TagContext extract(C carrier, AgentPropagation.ContextVisitor<C> gett
   public static class CompoundExtractor implements Extractor {
     private final List<Extractor> extractors;
     private final boolean extractFirst;
+    private final String extractBehavior;
 
-    public CompoundExtractor(final List<Extractor> extractors, boolean extractFirst) {
+    public CompoundExtractor(
+        final List<Extractor> extractors, boolean extractFirst, String extractBehavior) {
       this.extractors = extractors;
       this.extractFirst = extractFirst;
+      this.extractBehavior = extractBehavior;
     }
 
     @Override
@@ -262,6 +266,16 @@ else if (extracted != null && partialContext == null) {
       }
 
       if (context != null) {
+        if (extractBehavior.equals("restart")) {
+          context.resetTerminatedContextLink();
+          context.addTerminatedContextLink(
+              DDSpanLink.from(
+                  context,
+                  SpanAttributes.builder()
+                      .put("reason", "propagation_behavior_extract")
+                      .put("context_headers", context.getPropagationStyle().toString())
+                      .build()));
+        }
         log.debug("Extract complete context {}", context);
         return context;
       } else if (partialContext != null) {

dd-trace-core/src/test/groovy/datadog/trace/core/CoreSpanBuilderTest.groovy

@@ -342,6 +342,33 @@ class CoreSpanBuilderTest extends DDCoreSpecification {
     new ExtractedContext(DDTraceId.from(3), 4, PrioritySampling.SAMPLER_KEEP, "some-origin", 0, ["asdf": "qwer"], [(ORIGIN_KEY): "some-origin", "zxcv": "1234"], null, PropagationTags.factory().empty(), null, DATADOG)                     | _
   }
 
+  def "build context from ExtractedContext with TRACE_PROPAGATION_BEHAVIOR_EXTRACT=restart"() {
+    setup:
+    injectSysConfig("trace.propagation.behavior.extract", "restart")
+    def thread = Thread.currentThread()
+    TagContext context = extractedContext
+    context.addTerminatedContextLink(DDSpanLink.from(extractedContext))
+    final DDSpan span = tracer.buildSpan("test", "op name")
+      .asChildOf(context).start()
+
+    expect:
+    span.traceId != extractedContext.traceId
+    span.parentId != extractedContext.spanId
+    span.samplingPriority() == PrioritySampling.UNSET
+
+    def spanLinks = span.links
+
+    assert spanLinks.size() == 1
+    def link = spanLinks[0]
+    link.traceId() == extractedContext.traceId
+    link.spanId() == extractedContext.spanId
+    link.traceState() == extractedContext.propagationTags.headerValue(PropagationTags.HeaderType.W3C)
+
+    where:
+    extractedContext                                                                                                                                                                                                                         | _
+    new ExtractedContext(DDTraceId.ONE, 2, PrioritySampling.SAMPLER_DROP, null, 0, [:], [:], null, PropagationTags.factory().fromHeaderValue(PropagationTags.HeaderType.DATADOG, "_dd.p.dm=934086a686-4,_dd.p.anytag=value"), null, DATADOG) | _
+  }
+
   def "TagContext should populate default span details"() {
     setup:
     def thread = Thread.currentThread()

dd-trace-core/src/test/groovy/datadog/trace/core/propagation/HttpExtractorTest.groovy

@@ -33,6 +33,7 @@ class HttpExtractorTest extends DDSpecification {
     Config config = Mock(Config) {
       getTracePropagationStylesToExtract() >> styles
       isTracePropagationExtractFirst() >> extractFirst
+      getTracePropagationBehaviorExtract() >> "continue"
     }
     DynamicConfig dynamicConfig = DynamicConfig.create()
       .setHeaderTags(["SOME_HEADER": "some-tag"])
@@ -119,6 +120,7 @@ class HttpExtractorTest extends DDSpecification {
     setup:
     Config config = Mock(Config) {
       getTracePropagationStylesToExtract() >> styles
+      getTracePropagationBehaviorExtract() >> "continue"
     }
     DynamicConfig dynamicConfig = DynamicConfig.create().apply()
     HttpCodec.Extractor extractor = HttpCodec.createExtractor(config, { dynamicConfig.captureTraceConfig() })
@@ -171,6 +173,7 @@ class HttpExtractorTest extends DDSpecification {
     setup:
     Config config = Mock(Config) {
       getTracePropagationStylesToExtract() >> styles
+      getTracePropagationBehaviorExtract() >> "continue"
     }
     DynamicConfig dynamicConfig = DynamicConfig.create().apply()
     HttpCodec.Extractor extractor = HttpCodec.createExtractor(config, { dynamicConfig.captureTraceConfig() })
@@ -206,4 +209,50 @@ class HttpExtractorTest extends DDSpecification {
     [TRACECONTEXT, B3MULTI, DATADOG] | "2"               | "2"               | "1"               | "b"               | W3C_TRACE_PARENT | W3C_TRACE_STATE_NO_P | [DATADOG]
     // spotless:on
   }
+
+  def "verify extract behavior with TRACE_PROPAGATION_BEHAVIOR_EXTRACT"() {
+    setup:
+    Config config = Mock(Config) {
+      getTracePropagationStylesToExtract() >> [DATADOG, TRACECONTEXT]
+      getTracePropagationBehaviorExtract() >> behaviorExtract
+      isTracePropagationExtractFirst() >> extractFirst
+    }
+    DynamicConfig dynamicConfig = DynamicConfig.create().apply()
+    HttpCodec.Extractor extractor = HttpCodec.createExtractor(config, { dynamicConfig.captureTraceConfig() })
+
+    final Map<String, String> actual = [
+      (DatadogHttpCodec.TRACE_ID_KEY.toUpperCase()): datadogTraceId,
+      (DatadogHttpCodec.SPAN_ID_KEY.toUpperCase()) : datadogSpanId,
+      (W3CHttpCodec.TRACE_PARENT_KEY.toUpperCase()): w3cTraceParent,
+      (W3CHttpCodec.TRACE_STATE_KEY.toUpperCase()) : traceState
+    ]
+
+    when:
+    final TagContext context = extractor.extract(actual, ContextVisitors.stringValuesMap())
+
+    then:
+    def links = context.getTerminatedContextLinks()
+    assert links.size() == expectedLinkStyle.size()
+    for (int i = 0; i < links.size(); i++) {
+      assert links[i].traceId().toString() == linkTraceId
+      assert links[i].spanId().toString() == linkSpanId
+      assert links[i].attributes().asMap()["reason"] == linkReason
+      assert links[i].attributes().asMap()["context_headers"] == expectedLinkStyle[i].toString()
+      if (expectedLinkStyle[i] == TRACECONTEXT || linkReason == "propagation_behavior_extract" && datadogTraceId == getW3C_TRACE_ID()) {
+        assert links[i].traceState() == traceState
+      }
+    }
+
+    // no tests for behaviorExtract=ignore since no extraction will occur in this situation
+    where:
+    // spotless:off
+    behaviorExtract | extractFirst | datadogTraceId | datadogSpanId | w3cTraceParent   | traceState           | expectedLinkStyle | linkTraceId | linkSpanId            | linkReason
+    "continue"      | false        | "1"            | "2"           | W3C_TRACE_PARENT | ""                   | []                | ""          | "1311768467463790320" | ""
+    "continue"      | false        | "2"            | "2"           | W3C_TRACE_PARENT | W3C_TRACE_STATE_NO_P | [TRACECONTEXT]    | "1"         | "1311768467463790320" | "terminated_context"
+    "restart"       | false        | "1"            | "2"           | W3C_TRACE_PARENT | "dd=foo=1"           | [DATADOG]         | "1"         | "1311768467463790320" | "propagation_behavior_extract"
+    "restart"       | false        | "2"            | "2"           | W3C_TRACE_PARENT | ""                   | [DATADOG]         | "2"         | "2"                   | "propagation_behavior_extract"
+    "restart"       | true         | "2"            | "2"           | W3C_TRACE_PARENT | ""                   | [DATADOG]         | "2"         | "2"                   | "propagation_behavior_extract"
+    "restart"       | true         | "1"            | "2"           | W3C_TRACE_PARENT | ""                   | [DATADOG]         | "1"         | "2"                   | "propagation_behavior_extract"
+    // spotless:on
+  }
 }

internal-api/src/main/java/datadog/trace/api/Config.java

@@ -189,6 +189,7 @@ public static String getHostName() {
   private final boolean tracePropagationStyleB3PaddingEnabled;
   private final Set<TracePropagationStyle> tracePropagationStylesToExtract;
   private final Set<TracePropagationStyle> tracePropagationStylesToInject;
+  private final String tracePropagationBehaviorExtract;
   private final boolean tracePropagationExtractFirst;
   private final int traceBaggageMaxItems;
   private final int traceBaggageMaxBytes;
@@ -924,6 +925,11 @@ private Config(final ConfigProvider configProvider, final InstrumenterConfig ins
 
     tracePropagationStyleB3PaddingEnabled =
         isEnabled(true, TRACE_PROPAGATION_STYLE, ".b3.padding.enabled");
+
+    tracePropagationBehaviorExtract =
+        configProvider.getString(
+            TRACE_PROPAGATION_BEHAVIOR_EXTRACT, DEFAULT_TRACE_PROPAGATION_BEHAVIOR_EXTRACT);
+
     {
       // The dd.propagation.style.(extract|inject) settings have been deprecated in
       // favor of dd.trace.propagation.style(|.extract|.inject) settings.
@@ -995,8 +1001,14 @@ private Config(final ConfigProvider configProvider, final InstrumenterConfig ins
         logOverriddenDeprecatedSettingWarning(PROPAGATION_STYLE_INJECT, injectOrigin, inject);
       }
       // Now we can check if we should pick the default injection/extraction
+
+      if (extract.isEmpty()) {
+        extract = DEFAULT_TRACE_PROPAGATION_STYLE;
+      }
+
       tracePropagationStylesToExtract =
-          extract.isEmpty() ? DEFAULT_TRACE_PROPAGATION_STYLE : extract;
+          tracePropagationBehaviorExtract.equals("ignore") ? new HashSet<>() : extract;
+
       tracePropagationStylesToInject = inject.isEmpty() ? DEFAULT_TRACE_PROPAGATION_STYLE : inject;
 
       traceBaggageMaxItems =
@@ -2260,6 +2272,10 @@ public Set<TracePropagationStyle> getTracePropagationStylesToInject() {
     return tracePropagationStylesToInject;
   }
 
+  public String getTracePropagationBehaviorExtract() {
+    return tracePropagationBehaviorExtract;
+  }
+
   public boolean isTracePropagationExtractFirst() {
     return tracePropagationExtractFirst;
   }
@@ -4441,6 +4457,8 @@ public String toString() {
         + tracePropagationStylesToExtract
         + ", tracePropagationStylesToInject="
         + tracePropagationStylesToInject
+        + ", tracePropagationBehaviorExtract="
+        + tracePropagationBehaviorExtract
         + ", tracePropagationExtractFirst="
         + tracePropagationExtractFirst
         + ", clockSyncPeriod="

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

@@ -89,6 +89,10 @@ public void addTerminatedContextLink(AgentSpanLink link) {
     this.terminatedContextLinks.add(link);
   }
 
+  public void resetTerminatedContextLink() {
+    this.terminatedContextLinks = null;
+  }
+
   @Override
   public String getForwarded() {
     return httpHeaders.forwarded;

internal-api/src/test/groovy/datadog/trace/api/ConfigTest.groovy

@@ -123,6 +123,7 @@ import static datadog.trace.api.config.TracerConfig.SPLIT_BY_TAGS
 import static datadog.trace.api.config.TracerConfig.TRACE_AGENT_PORT
 import static datadog.trace.api.config.TracerConfig.TRACE_AGENT_URL
 import static datadog.trace.api.config.TracerConfig.TRACE_PROPAGATION_EXTRACT_FIRST
+import static datadog.trace.api.config.TracerConfig.TRACE_PROPAGATION_BEHAVIOR_EXTRACT
 import static datadog.trace.api.config.TracerConfig.TRACE_RATE_LIMIT
 import static datadog.trace.api.config.TracerConfig.TRACE_REPORT_HOSTNAME
 import static datadog.trace.api.config.TracerConfig.TRACE_RESOLVER_ENABLED
@@ -205,6 +206,7 @@ class ConfigTest extends DDSpecification {
     prop.setProperty(PROPAGATION_STYLE_EXTRACT, "Datadog, B3")
     prop.setProperty(PROPAGATION_STYLE_INJECT, "B3, Datadog")
     prop.setProperty(TRACE_PROPAGATION_EXTRACT_FIRST, "false")
+    prop.setProperty(TRACE_PROPAGATION_BEHAVIOR_EXTRACT, "restart")
     prop.setProperty(JMX_FETCH_ENABLED, "false")
     prop.setProperty(JMX_FETCH_METRICS_CONFIGS, "/foo.yaml,/bar.yaml")
     prop.setProperty(JMX_FETCH_CHECK_PERIOD, "100")
@@ -295,6 +297,7 @@ class ConfigTest extends DDSpecification {
     config.tracePropagationStylesToExtract.toList() == [DATADOG, B3SINGLE, B3MULTI]
     config.tracePropagationStylesToInject.toList() == [B3SINGLE, B3MULTI, DATADOG]
     config.tracePropagationExtractFirst == false
+    config.tracePropagationBehaviorExtract == "restart"
     config.jmxFetchEnabled == false
     config.jmxFetchMetricsConfigs == ["/foo.yaml", "/bar.yaml"]
     config.jmxFetchCheckPeriod == 100
@@ -386,6 +389,7 @@ class ConfigTest extends DDSpecification {
     System.setProperty(PREFIX + PROPAGATION_STYLE_EXTRACT, "Datadog, B3")
     System.setProperty(PREFIX + PROPAGATION_STYLE_INJECT, "B3, Datadog")
     System.setProperty(PREFIX + TRACE_PROPAGATION_EXTRACT_FIRST, "false")
+    System.setProperty(PREFIX + TRACE_PROPAGATION_BEHAVIOR_EXTRACT, "restart")
     System.setProperty(PREFIX + JMX_FETCH_ENABLED, "false")
     System.setProperty(PREFIX + JMX_FETCH_METRICS_CONFIGS, "/foo.yaml,/bar.yaml")
     System.setProperty(PREFIX + JMX_FETCH_CHECK_PERIOD, "100")
@@ -475,6 +479,7 @@ class ConfigTest extends DDSpecification {
     config.tracePropagationStylesToExtract.toList() == [DATADOG, B3SINGLE, B3MULTI]
     config.tracePropagationStylesToInject.toList() == [B3SINGLE, B3MULTI, DATADOG]
     config.tracePropagationExtractFirst == false
+    config.tracePropagationBehaviorExtract == "restart"
     config.jmxFetchEnabled == false
     config.jmxFetchMetricsConfigs == ["/foo.yaml", "/bar.yaml"]
     config.jmxFetchCheckPeriod == 100
@@ -2604,4 +2609,18 @@ class ConfigTest extends DDSpecification {
     config.finalDebuggerSnapshotUrl == "http://localhost:8126/debugger/v1/input"
     config.finalDebuggerSymDBUrl == "http://localhost:8126/symdb/v1/input"
   }
+
+  def "specify overrides for PROPAGATION_STYLE_EXTRACT when TRACE_PROPAGATION_BEHAVIOR_EXTRACT=ignore"() {
+    setup:
+    def prop = new Properties()
+    prop.setProperty(PROPAGATION_STYLE_EXTRACT, "Datadog, B3")
+    prop.setProperty(TRACE_PROPAGATION_BEHAVIOR_EXTRACT, "ignore")
+
+    when:
+    Config config = Config.get(prop)
+
+    then:
+    config.tracePropagationBehaviorExtract == "ignore"
+    config.tracePropagationStylesToExtract.toList() == []
+  }
 }