Start extracted Datadog propagation tags from the extracted priority

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

@@ -336,6 +336,14 @@ protected void onlyTagContext() {
     this.fullContext = false;
   }
 
+  /**
+   * The extracted sampling priority that the trace will inherit, so propagation tags can start from
+   * it, or {@link PrioritySampling#UNSET} when it may not be inherited.
+   */
+  protected int inheritedSamplingPriority() {
+    return apmTracingEnabled ? samplingPriority : PrioritySampling.UNSET;
+  }
+
   protected int defaultSamplingPriority() {
     return PrioritySampling.UNSET;
   }

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

@@ -111,6 +111,9 @@ private static class DatadogContextInterpreter extends ContextInterpreter {
 
     private final boolean isAwsPropagationEnabled;
 
+    /** The raw x-datadog-tags value, parsed in {@link #build()} once the priority is known. */
+    private String datadogTags;
+
     private DatadogContextInterpreter(Config config) {
       super(config);
       isAwsPropagationEnabled = config.isAwsPropagationEnabled();
@@ -121,6 +124,12 @@ public TracePropagationStyle style() {
       return DATADOG;
     }
 
+    @Override
+    public ContextInterpreter reset(TraceConfig traceConfig) {
+      datadogTags = null;
+      return super.reset(traceConfig);
+    }
+
     @Override
     public boolean accept(String key, String value) {
       if (null == key || key.isEmpty()) {
@@ -187,7 +196,7 @@ public boolean accept(String key, String value) {
                 endToEndStartTime = extractEndToEndStartTime(firstHeaderValue(value));
                 break;
               case DD_TAGS:
-                propagationTags = propagationTagsFactory.fromHeaderValue(HeaderType.DATADOG, value);
+                datadogTags = value;
                 break;
               case OT_BAGGAGE:
                 addBaggageItem(lowerCaseKey.substring(OT_BAGGAGE_PREFIX.length()), value);
@@ -216,6 +225,16 @@ public boolean accept(String key, String value) {
 
     @Override
     protected TagContext build() {
+      if (datadogTags != null) {
+        try {
+          propagationTags =
+              propagationTagsFactory.fromDatadogHeaderValue(
+                  datadogTags, inheritedSamplingPriority());
+        } catch (RuntimeException e) {
+          invalidateContext();
+          log.debug("Exception when extracting context", e);
+        }
+      }
       restore128bTraceId();
       return super.build();
     }

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

@@ -122,6 +122,12 @@ public interface Factory {
 
     PropagationTags fromHeaderValue(HeaderType headerType, String value);
 
+    /**
+     * Like {@link #fromHeaderValue} for {@link HeaderType#DATADOG}, starting from the trace's
+     * extracted {@code samplingPriority} so it doesn't have to be applied as a later update.
+     */
+    PropagationTags fromDatadogHeaderValue(String value, int samplingPriority);
+
     /**
      * Returns a fresh PropagationTags that re-encodes only the non-{@code dd} vendor sections of
      * the supplied W3C tracestate. All Datadog-side state (sampling priority, origin, {@code

dd-trace-core/src/main/java/datadog/trace/core/propagation/ptags/DatadogPTagsCodec.java

@@ -4,6 +4,7 @@
 
 import datadog.logging.RatelimitedLogger;
 import datadog.trace.api.ProductTraceSource;
+import datadog.trace.api.sampling.PrioritySampling;
 import datadog.trace.core.propagation.PropagationTags;
 import datadog.trace.core.propagation.PropagationTags.SamplingState;
 import datadog.trace.core.propagation.ptags.PTagsFactory.PTags;
@@ -55,6 +56,11 @@ final class DatadogPTagsCodec extends PTagsCodec {
    */
   @Override
   PropagationTags fromHeaderValue(PTagsFactory tagsFactory, String value) {
+    return fromHeaderValue(tagsFactory, value, PrioritySampling.UNSET);
+  }
+
+  /** Parses {@code value} into propagation tags that start at {@code samplingPriority}. */
+  PropagationTags fromHeaderValue(PTagsFactory tagsFactory, String value, int samplingPriority) {
     if (value == null) {
       return tagsFactory.empty();
     }
@@ -150,6 +156,7 @@ PropagationTags fromHeaderValue(PTagsFactory tagsFactory, String value) {
         decisionMakerTagValue,
         traceIdTagValue,
         traceSource,
+        samplingPriority,
         orgPropagationMarkerTagValue,
         LLMObsTagValues.of(
             llmObsTraceIdTagValue,

dd-trace-core/src/main/java/datadog/trace/core/propagation/ptags/PTagsFactory.java

@@ -51,14 +51,27 @@ PTagsCodec getDecoderEncoder(@Nonnull HeaderType headerType) {
 
   @Override
   public final PropagationTags empty() {
-    return createValid(null, null, null, ProductTraceSource.UNSET, null, LLMObsTagValues.EMPTY);
+    return createValid(
+        null,
+        null,
+        null,
+        ProductTraceSource.UNSET,
+        PrioritySampling.UNSET,
+        null,
+        LLMObsTagValues.EMPTY);
   }
 
   @Override
   public final PropagationTags fromHeaderValue(@Nonnull HeaderType headerType, String value) {
     return DEC_ENC_MAP.get(headerType).fromHeaderValue(this, value);
   }
 
+  @Override
+  public final PropagationTags fromDatadogHeaderValue(String value, int samplingPriority) {
+    return ((DatadogPTagsCodec) DEC_ENC_MAP.get(DATADOG))
+        .fromHeaderValue(this, value, samplingPriority);
+  }
+
   @Override
   public final PropagationTags emptyW3C(String originalTracestate) {
     if (originalTracestate == null || originalTracestate.isEmpty()) {
@@ -72,6 +85,7 @@ PropagationTags createValid(
       TagValue decisionMakerTagValue,
       TagValue traceIdTagValue,
       int productTraceSource,
+      int samplingPriority,
       TagValue orgPropagationMarkerTagValue,
       @Nonnull LLMObsTagValues llmObsTagValues) {
     return new PTags(
@@ -80,8 +94,13 @@ PropagationTags createValid(
         decisionMakerTagValue,
         traceIdTagValue,
         productTraceSource,
+        samplingPriority,
+        null,
+        null,
         orgPropagationMarkerTagValue,
-        llmObsTagValues);
+        llmObsTagValues,
+        null,
+        null);
   }
 
   PropagationTags createInvalid(String error) {
@@ -182,29 +201,6 @@ static class PTags extends PropagationTags {
      */
     private volatile CharSequence lastParentId;
 
-    PTags(
-        PTagsFactory factory,
-        List<TagElement> tagPairs,
-        TagValue decisionMakerTagValue,
-        TagValue traceIdTagValue,
-        int traceSource,
-        TagValue orgPropagationMarkerTagValue,
-        @Nonnull LLMObsTagValues llmObsTagValues) {
-      this(
-          factory,
-          tagPairs,
-          decisionMakerTagValue,
-          traceIdTagValue,
-          traceSource,
-          PrioritySampling.UNSET,
-          null,
-          null,
-          orgPropagationMarkerTagValue,
-          llmObsTagValues,
-          null,
-          null);
-    }
-
     // Takes tracestate/otelTraceState up-front to build SamplingState in one allocation.
     PTags(
         PTagsFactory factory,

dd-trace-core/src/test/java/datadog/trace/core/DDSpanContextPropagationTagsTest.java

@@ -9,6 +9,7 @@
 import static datadog.trace.api.sampling.SamplingMechanism.DEFAULT;
 import static datadog.trace.api.sampling.SamplingMechanism.MANUAL;
 import static org.junit.jupiter.api.Assertions.assertEquals;
+import static org.junit.jupiter.api.Assertions.assertSame;
 
 import datadog.trace.api.DDTraceId;
 import datadog.trace.bootstrap.instrumentation.api.AgentSpanContext;
@@ -18,6 +19,7 @@
 import datadog.trace.test.junit.utils.tabletest.TableTestTypeConverters;
 import java.util.Map;
 import org.junit.jupiter.api.BeforeEach;
+import org.junit.jupiter.api.Test;
 import org.tabletest.junit.TableTest;
 import org.tabletest.junit.TypeConverter;
 import org.tabletest.junit.TypeConverterSources;
@@ -178,4 +180,19 @@ public static int toInt(String value) {
         throw new IllegalArgumentException("Value is invalid for " + value);
     }
   }
+
+  @Test
+  void rootSpanKeepsExtractedSamplingState() {
+    PropagationTags propagationTags =
+        tracer.getPropagationTagsFactory().fromDatadogHeaderValue("_dd.p.dm=-1", SAMPLER_KEEP);
+    PropagationTags.SamplingState extractedState = propagationTags.samplingState();
+    AgentSpanContext extracted =
+        new ExtractedContext(
+            DDTraceId.from(123), 456, SAMPLER_KEEP, null, propagationTags, DATADOG);
+
+    DDSpan span = (DDSpan) tracer.buildSpan("datadog", "top").asChildOf(extracted).start();
+
+    assertSame(extractedState, span.spanContext().getPropagationTags().samplingState());
+    assertEquals(SAMPLER_KEEP, span.spanContext().getSamplingPriority());
+  }
 }

dd-trace-core/src/test/java/datadog/trace/core/propagation/DatadogHttpExtractorTest.java

@@ -383,6 +383,27 @@ void extractMappedBaggageIsSubjectToTheItemLimit() {
     assertEquals(singletonMap(SOME_BAGGAGE, "mappedBaggageValue"), context.getBaggage());
   }
 
+  @ParameterizedTest
+  @ValueSource(booleans = {true, false})
+  void propagationTagsStartFromExtractedPriority(boolean tagsFirst) {
+    Map<String, String> headers = new LinkedHashMap<>();
+    headers.put(TRACE_ID_KEY, "1");
+    headers.put(SPAN_ID_KEY, "2");
+    if (tagsFirst) {
+      headers.put(DATADOG_TAGS_KEY, "_dd.p.dm=-1");
+      headers.put(SAMPLING_PRIORITY_KEY, String.valueOf(SAMPLER_KEEP));
+    } else {
+      headers.put(SAMPLING_PRIORITY_KEY, String.valueOf(SAMPLER_KEEP));
+      headers.put(DATADOG_TAGS_KEY, "_dd.p.dm=-1");
+    }
+
+    ExtractedContext context = (ExtractedContext) extractor.extract(headers, stringValuesMap());
+
+    PropagationTags propagationTags = context.getPropagationTags();
+    assertEquals(SAMPLER_KEEP, propagationTags.getSamplingPriority());
+    assertEquals("_dd.p.dm=-1", propagationTags.headerValue(PropagationTags.HeaderType.DATADOG));
+  }
+
   /**
    * With APM tracing disabled, an upstream sampling decision is only kept when {@code _dd.p.ts}
    * marks the trace for ASM or AI Guard. Otherwise the local sampler sets the priority.
@@ -415,6 +436,19 @@ void upstreamPriorityKeptWhenProductMarkedTheTrace() {
       assertEquals(SAMPLER_KEEP, context.getSamplingPriority());
     }
 
+    @Test
+    void propagationTagsDontStartFromUninheritedPriority() {
+      Map<String, String> headers = new HashMap<>();
+      headers.put(TRACE_ID_KEY, "1");
+      headers.put(SPAN_ID_KEY, "2");
+      headers.put(SAMPLING_PRIORITY_KEY, String.valueOf(SAMPLER_KEEP));
+      headers.put(DATADOG_TAGS_KEY, "_dd.p.dm=-1");
+
+      ExtractedContext context = (ExtractedContext) extractor.extract(headers, stringValuesMap());
+
+      assertEquals(UNSET, context.getPropagationTags().getSamplingPriority());
+    }
+
     @Test
     void upstreamPriorityDroppedWhenOnlyDsmMarkedTheTrace() {
       Map<String, String> headers = new HashMap<>();