Adding Support for TRACE_PROPAGATION_BEHAVIOR_EXTRACT

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;
+ }
 }

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,"

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;
+ }
 }

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;
  }

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',

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;

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;
+ }
+}

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";

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.

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()

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="

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.

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();

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();
+ }
 }

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;

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
+ }
 }