Adding Support for TRACE_PROPAGATION_BEHAVIOR_EXTRACT (#8535)

mhlidd · web-flow · commit c67bed5168fa · 2025-03-19T10:48:03.000-04:00
* initial implementation of TRACE_PROPAGATION_BEHAVIOR_EXTRACT

* initial implementation and unit tests for TRACE_PROPAGATION_BEHAVIOR_EXTRACT

* change to sampling

* comments for future work

* adding isremote() to all AgentSpanContexts

* updating span link logic for restart

* updating enum

* cleanup

* adding override to new class

* updating enum and tests

* refactoring capturing config

* forcing enum to be built by build time

* adding enum to skip test coverae

* responding to PR comments

* spotless

* modifying enum  constructor
diff --git a/dd-java-agent/agent-otel/otel-shim/src/main/java/datadog/opentelemetry/shim/trace/OtelExtractedContext.java b/dd-java-agent/agent-otel/otel-shim/src/main/java/datadog/opentelemetry/shim/trace/OtelExtractedContext.java
@@ -78,4 +78,9 @@ public Iterable<Map.Entry<String, String>> baggageItems() {
   public PathwayContext getPathwayContext() {
     return null;
   }
+
+  @Override
+  public boolean isRemote() {
+    return true;
+  }
 }
diff --git a/dd-java-agent/instrumentation/graal/native-image/src/main/java/datadog/trace/instrumentation/graal/nativeimage/NativeImageGeneratorRunnerInstrumentation.java b/dd-java-agent/instrumentation/graal/native-image/src/main/java/datadog/trace/instrumentation/graal/nativeimage/NativeImageGeneratorRunnerInstrumentation.java
@@ -97,6 +97,7 @@ public static void onEnter(@Advice.Argument(value = 0, readOnly = false) String[
               + "datadog.trace.api.ResolverCacheConfig$4:build_time,"
               + "datadog.trace.api.ResolverCacheConfig$5:build_time,"
               + "datadog.trace.api.TracePropagationStyle:build_time,"
+              + "datadog.trace.api.TracePropagationBehaviorExtract:build_time,"
               + "datadog.trace.api.telemetry.OtelEnvMetricCollector:build_time,"
               + "datadog.trace.api.profiling.ProfilingEnablement:build_time,"
               + "datadog.trace.bootstrap.config.provider.ConfigConverter:build_time,"
diff --git a/dd-java-agent/instrumentation/spark/src/main/java/datadog/trace/instrumentation/spark/DatabricksParentContext.java b/dd-java-agent/instrumentation/spark/src/main/java/datadog/trace/instrumentation/spark/DatabricksParentContext.java
@@ -93,4 +93,9 @@ public Iterable<Map.Entry<String, String>> baggageItems() {
   public PathwayContext getPathwayContext() {
     return null;
   }
+
+  @Override
+  public boolean isRemote() {
+    return false;
+  }
 }
diff --git a/dd-java-agent/instrumentation/spark/src/main/java/datadog/trace/instrumentation/spark/OpenlineageParentContext.java b/dd-java-agent/instrumentation/spark/src/main/java/datadog/trace/instrumentation/spark/OpenlineageParentContext.java
@@ -143,6 +143,11 @@ public PathwayContext getPathwayContext() {
     return null;
   }
 
+  @Override
+  public boolean isRemote() {
+    return false;
+  }
+
   public String getParentJobNamespace() {
     return parentJobNamespace;
   }
diff --git a/dd-trace-api/build.gradle b/dd-trace-api/build.gradle
@@ -15,6 +15,7 @@ excludedClassesCoverage += [
   'datadog.trace.api.GlobalTracer*',
   'datadog.trace.api.PropagationStyle',
   'datadog.trace.api.TracePropagationStyle',
+  'datadog.trace.api.TracePropagationBehaviorExtract',
   'datadog.trace.api.SpanCorrelation*',
   'datadog.trace.api.internal.TraceSegment',
   'datadog.trace.api.internal.TraceSegment.NoOp',
diff --git a/dd-trace-api/src/main/java/datadog/trace/api/ConfigDefaults.java b/dd-trace-api/src/main/java/datadog/trace/api/ConfigDefaults.java
@@ -91,6 +91,8 @@ public final class ConfigDefaults {
 
   static final int DEFAULT_CLOCK_SYNC_PERIOD = 30; // seconds
 
+  static final TracePropagationBehaviorExtract DEFAULT_TRACE_PROPAGATION_BEHAVIOR_EXTRACT =
+      TracePropagationBehaviorExtract.CONTINUE;
   static final boolean DEFAULT_TRACE_PROPAGATION_EXTRACT_FIRST = false;
 
   static final boolean DEFAULT_JMX_FETCH_MULTIPLE_RUNTIME_SERVICES_ENABLED = false;
diff --git a/dd-trace-api/src/main/java/datadog/trace/api/TracePropagationBehaviorExtract.java b/dd-trace-api/src/main/java/datadog/trace/api/TracePropagationBehaviorExtract.java
@@ -0,0 +1,24 @@
+package datadog.trace.api;
+
+import java.util.Locale;
+
+/** Trace propagation styles for injecting and extracting trace propagation headers. */
+public enum TracePropagationBehaviorExtract {
+  CONTINUE,
+  RESTART,
+  IGNORE;
+
+  private String displayName;
+
+  TracePropagationBehaviorExtract() {
+    this.displayName = name().toLowerCase(Locale.ROOT);
+  }
+
+  @Override
+  public String toString() {
+    if (displayName == null) {
+      displayName = name().toLowerCase(Locale.ROOT);
+    }
+    return displayName;
+  }
+}
diff --git a/dd-trace-api/src/main/java/datadog/trace/api/config/TracerConfig.java b/dd-trace-api/src/main/java/datadog/trace/api/config/TracerConfig.java
@@ -93,6 +93,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";
diff --git a/dd-trace-core/src/main/java/datadog/trace/core/CoreTracer.java b/dd-trace-core/src/main/java/datadog/trace/core/CoreTracer.java
@@ -32,6 +32,7 @@
 import datadog.trace.api.InstrumenterConfig;
 import datadog.trace.api.StatsDClient;
 import datadog.trace.api.TraceConfig;
+import datadog.trace.api.TracePropagationBehaviorExtract;
 import datadog.trace.api.config.GeneralConfig;
 import datadog.trace.api.datastreams.AgentDataStreamsMonitoring;
 import datadog.trace.api.datastreams.PathwayContext;
@@ -61,6 +62,8 @@
 import datadog.trace.bootstrap.instrumentation.api.BlackHoleSpan;
 import datadog.trace.bootstrap.instrumentation.api.ProfilingContextIntegration;
 import datadog.trace.bootstrap.instrumentation.api.ScopeState;
+import datadog.trace.bootstrap.instrumentation.api.SpanAttributes;
+import datadog.trace.bootstrap.instrumentation.api.SpanLink;
 import datadog.trace.bootstrap.instrumentation.api.TagContext;
 import datadog.trace.civisibility.interceptor.CiVisibilityApmProtocolInterceptor;
 import datadog.trace.civisibility.interceptor.CiVisibilityTelemetryInterceptor;
@@ -1506,6 +1509,28 @@ private DDSpanContext buildSpanContext() {
       String parentServiceName = null;
       boolean isRemote = false;
 
+      if (parentContext != null
+          && parentContext.isRemote()
+          && Config.get().getTracePropagationBehaviorExtract()
+              == TracePropagationBehaviorExtract.RESTART) {
+        SpanLink link;
+        if (parentContext instanceof ExtractedContext) {
+          ExtractedContext pc = (ExtractedContext) parentContext;
+          link =
+              DDSpanLink.from(
+                  pc,
+                  SpanAttributes.builder()
+                      .put("reason", "propagation_behavior_extract")
+                      .put("context_headers", pc.getPropagationStyle().toString())
+                      .build());
+        } else {
+          link = SpanLink.from(parentContext);
+        }
+        // reset links that may have come terminated span links
+        links = new ArrayList<>();
+        links.add(link);
+        parentContext = null;
+      }
       // Propagate internal trace.
       // Note: if we are not in the context of distributed tracing and we are starting the first
       // root span, parentContext will be null at this point.
diff --git a/dd-trace-core/src/test/groovy/datadog/trace/core/CoreSpanBuilderTest.groovy b/dd-trace-core/src/test/groovy/datadog/trace/core/CoreSpanBuilderTest.groovy
@@ -342,6 +342,27 @@ 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 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)
+    final DDSpan span = tracer.buildSpan("test", "op name")
+      .asChildOf(extractedContext).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)
+  }
+
   def "TagContext should populate default span details"() {
     setup:
     def thread = Thread.currentThread()
diff --git a/internal-api/src/main/java/datadog/trace/api/Config.java b/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 TracePropagationBehaviorExtract tracePropagationBehaviorExtract;
   private final boolean tracePropagationExtractFirst;
   private final int traceBaggageMaxItems;
   private final int traceBaggageMaxBytes;
@@ -940,6 +941,22 @@ private Config(final ConfigProvider configProvider, final InstrumenterConfig ins
 
     tracePropagationStyleB3PaddingEnabled =
         isEnabled(true, TRACE_PROPAGATION_STYLE, ".b3.padding.enabled");
+
+    TracePropagationBehaviorExtract tmpTracePropagationBehaviorExtract;
+    try {
+      tmpTracePropagationBehaviorExtract =
+          TracePropagationBehaviorExtract.valueOf(
+              configProvider
+                  .getString(
+                      TRACE_PROPAGATION_BEHAVIOR_EXTRACT,
+                      DEFAULT_TRACE_PROPAGATION_BEHAVIOR_EXTRACT.toString())
+                  .toUpperCase(Locale.ROOT));
+    } catch (IllegalArgumentException e) {
+      tmpTracePropagationBehaviorExtract = TracePropagationBehaviorExtract.CONTINUE;
+      log.warn("Error while parsing TRACE_PROPAGATION_BEHAVIOR_EXTRACT, defaulting to `continue`");
+    }
+    tracePropagationBehaviorExtract = tmpTracePropagationBehaviorExtract;
+
     {
       // The dd.propagation.style.(extract|inject) settings have been deprecated in
       // favor of dd.trace.propagation.style(|.extract|.inject) settings.
@@ -1011,8 +1028,16 @@ 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 == TracePropagationBehaviorExtract.IGNORE
+              ? new HashSet<>()
+              : extract;
+
       tracePropagationStylesToInject = inject.isEmpty() ? DEFAULT_TRACE_PROPAGATION_STYLE : inject;
 
       traceBaggageMaxItems =
@@ -2291,6 +2316,10 @@ public Set<TracePropagationStyle> getTracePropagationStylesToInject() {
     return tracePropagationStylesToInject;
   }
 
+  public TracePropagationBehaviorExtract getTracePropagationBehaviorExtract() {
+    return tracePropagationBehaviorExtract;
+  }
+
   public boolean isTracePropagationExtractFirst() {
     return tracePropagationExtractFirst;
   }
@@ -4490,6 +4519,8 @@ public String toString() {
         + tracePropagationStylesToExtract
         + ", tracePropagationStylesToInject="
         + tracePropagationStylesToInject
+        + ", tracePropagationBehaviorExtract="
+        + tracePropagationBehaviorExtract
         + ", tracePropagationExtractFirst="
         + tracePropagationExtractFirst
         + ", clockSyncPeriod="
diff --git a/internal-api/src/main/java/datadog/trace/bootstrap/instrumentation/api/AgentSpanContext.java b/internal-api/src/main/java/datadog/trace/bootstrap/instrumentation/api/AgentSpanContext.java
@@ -14,6 +14,7 @@
  * contextualize the associated Span instance.
  */
 public interface AgentSpanContext {
+
   /**
    * Gets the TraceId of the span's trace.
    *
@@ -52,6 +53,13 @@ public interface AgentSpanContext {
 
   default void mergePathwayContext(PathwayContext pathwayContext) {}
 
+  /**
+   * Gets whether the span context used is part of the local trace or from another service
+   *
+   * @return boolean representing if the span context is part of the local trace
+   */
+  boolean isRemote();
+
   interface Extracted extends AgentSpanContext {
     /**
      * Gets the span links related to the other terminated context.
diff --git a/internal-api/src/main/java/datadog/trace/bootstrap/instrumentation/api/NoopSpanContext.java b/internal-api/src/main/java/datadog/trace/bootstrap/instrumentation/api/NoopSpanContext.java
@@ -45,6 +45,11 @@ public PathwayContext getPathwayContext() {
     return NoopPathwayContext.INSTANCE;
   }
 
+  @Override
+  public boolean isRemote() {
+    return false;
+  }
+
   @Override
   public List<AgentSpanLink> getTerminatedContextLinks() {
     return emptyList();
diff --git a/internal-api/src/main/java/datadog/trace/bootstrap/instrumentation/api/NotSampledSpanContext.java b/internal-api/src/main/java/datadog/trace/bootstrap/instrumentation/api/NotSampledSpanContext.java
@@ -42,4 +42,9 @@ public Iterable<Map.Entry<String, String>> baggageItems() {
   public PathwayContext getPathwayContext() {
     return delegate.getPathwayContext();
   }
+
+  @Override
+  public boolean isRemote() {
+    return delegate.isRemote();
+  }
 }
diff --git a/internal-api/src/main/java/datadog/trace/bootstrap/instrumentation/api/TagContext.java b/internal-api/src/main/java/datadog/trace/bootstrap/instrumentation/api/TagContext.java
@@ -236,6 +236,11 @@ public PathwayContext getPathwayContext() {
     return this.pathwayContext;
   }
 
+  @Override
+  public boolean isRemote() {
+    return true;
+  }
+
   public TagContext withPathwayContext(PathwayContext pathwayContext) {
     this.pathwayContext = pathwayContext;
     return this;
diff --git a/internal-api/src/test/groovy/datadog/trace/api/ConfigTest.groovy b/internal-api/src/test/groovy/datadog/trace/api/ConfigTest.groovy
@@ -125,6 +125,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
@@ -207,6 +208,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")
@@ -300,6 +302,7 @@ class ConfigTest extends DDSpecification {
     config.tracePropagationStylesToExtract.toList() == [DATADOG, B3SINGLE, B3MULTI]
     config.tracePropagationStylesToInject.toList() == [B3SINGLE, B3MULTI, DATADOG]
     config.tracePropagationExtractFirst == false
+    config.tracePropagationBehaviorExtract == TracePropagationBehaviorExtract.RESTART
     config.jmxFetchEnabled == false
     config.jmxFetchMetricsConfigs == ["/foo.yaml", "/bar.yaml"]
     config.jmxFetchCheckPeriod == 100
@@ -394,6 +397,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")
@@ -485,6 +489,7 @@ class ConfigTest extends DDSpecification {
     config.tracePropagationStylesToExtract.toList() == [DATADOG, B3SINGLE, B3MULTI]
     config.tracePropagationStylesToInject.toList() == [B3SINGLE, B3MULTI, DATADOG]
     config.tracePropagationExtractFirst == false
+    config.tracePropagationBehaviorExtract == TracePropagationBehaviorExtract.RESTART
     config.jmxFetchEnabled == false
     config.jmxFetchMetricsConfigs == ["/foo.yaml", "/bar.yaml"]
     config.jmxFetchCheckPeriod == 100
@@ -2658,4 +2663,30 @@ 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 == TracePropagationBehaviorExtract.IGNORE
+    config.tracePropagationStylesToExtract.toList() == []
+  }
+
+  def "verify try/catch behavior for invalid strings for TRACE_PROPAGATION_BEHAVIOR_EXTRACT"() {
+    setup:
+    def prop = new Properties()
+    prop.setProperty(TRACE_PROPAGATION_BEHAVIOR_EXTRACT, "test")
+
+    when:
+    Config config = Config.get(prop)
+
+    then:
+    config.tracePropagationBehaviorExtract == TracePropagationBehaviorExtract.CONTINUE
+  }
 }

Original file line number	Diff line number	Diff line change
`@@ -78,4 +78,9 @@ public Iterable<Map.Entry<String, String>> baggageItems() {`
`78`	`78`	`public PathwayContext getPathwayContext() {`
`79`	`79`	`return null;`
`80`	`80`	`}`
	`81`	`+`
	`82`	`+ @Override`
	`83`	`+ public boolean isRemote() {`
	`84`	`+ return true;`
	`85`	`+ }`
`81`	`86`	`}`
Original file line number	Diff line number	Diff line change
`@@ -93,4 +93,9 @@ public Iterable<Map.Entry<String, String>> baggageItems() {`
`93`	`93`	`public PathwayContext getPathwayContext() {`
`94`	`94`	`return null;`
`95`	`95`	`}`
	`96`	`+`
	`97`	`+ @Override`
	`98`	`+ public boolean isRemote() {`
	`99`	`+ return false;`
	`100`	`+ }`
`96`	`101`	`}`
Original file line number	Diff line number	Diff line change
`@@ -143,6 +143,11 @@ public PathwayContext getPathwayContext() {`
`143`	`143`	`return null;`
`144`	`144`	`}`
`145`	`145`
	`146`	`+ @Override`
	`147`	`+ public boolean isRemote() {`
	`148`	`+ return false;`
	`149`	`+ }`
	`150`	`+`
`146`	`151`	`public String getParentJobNamespace() {`
`147`	`152`	`return parentJobNamespace;`
`148`	`153`	`}`