fix(propagation): resolve otel sampling state during injection

Author: MilanGarnier

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

@@ -656,10 +656,7 @@ public DDSpan setSamplingPriority(
       context
           .getPropagationTags()
           .updateOtelTraceState(
-              getTraceId().toLong(),
-              sampleRate,
-              probabilitySamplingResult.booleanValue(),
-              samplingPriority);
+              getTraceId().toLong(), sampleRate, probabilitySamplingResult.booleanValue());
     }
     return this;
   }

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

@@ -285,10 +285,8 @@ private <C> void applyTraceContextToFirstContext(
         ExtractedContext firstContext,
         ExtractedContext traceContext,
         ExtractionCache<C> extractionCache) {
-      firstContext
-          .getPropagationTags()
-          .updateW3CTracestateFrom(
-              traceContext.getPropagationTags(), firstContext.getSamplingPriority());
+      // Propagate newly extracted W3C tracestate to first valid context
+      firstContext.getPropagationTags().updateW3CTracestateFrom(traceContext.getPropagationTags());
       // Check if parent spans differ to reconcile them
       if (firstContext.getSpanId() != traceContext.getSpanId()) {
         // Override parent span id with W3C one

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

@@ -68,7 +68,7 @@ public interface Factory {
   public abstract void forceKeep(int samplingMechanism);
 
   public abstract void updateOtelTraceState(
-      long traceIdLowOrderBits, double sampleRate, boolean sampled, int samplingPriority);
+      long traceIdLowOrderBits, double sampleRate, boolean sampled);
 
   public abstract int getSamplingPriority();
 
@@ -98,11 +98,10 @@ public abstract void updateOtelTraceState(
    */
   public abstract void updateW3CTracestate(String tracestate);
 
-  /**
-   * Enriches these tags with the W3C tracestate from {@code source} and the reconciled sampling
-   * priority.
-   */
-  public abstract void updateW3CTracestateFrom(PropagationTags source, int samplingPriority);
+  /** Updates the original W3C tracestate header from {@code source}. */
+  public void updateW3CTracestateFrom(PropagationTags source) {
+    updateW3CTracestate(source.getW3CTracestate());
+  }
 
   /**
    * Constructs a header value that includes valid propagated _dd.p.* tags and possibly a new
@@ -120,6 +119,15 @@ public abstract void updateOtelTraceState(
    */
   public abstract String headerValue(HeaderType headerType, CharSequence lastParentIdOverride);
 
+  /**
+   * Like {@link #headerValue(HeaderType, CharSequence)}, but resolves W3C sampling state using the
+   * supplied sampling priority.
+   */
+  public String headerValue(
+      HeaderType headerType, CharSequence lastParentIdOverride, int samplingPriority) {
+    return headerValue(headerType, lastParentIdOverride);
+  }
+
   /**
    * Fills a provided tagMap with valid propagated _dd.p.* tags and possibly a new sampling decision
    * tags _dd.p.dm (root span only) based on the current state, or sets only an error tag if the

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

@@ -65,25 +65,28 @@ public Injector(Map<String, String> invertedBaggageMapping) {
     @Override
     public <C> void inject(
         final DDSpanContext context, final C carrier, final CarrierSetter<C> setter) {
-      injectTraceParent(context, carrier, setter);
-      injectTraceState(context, carrier, setter);
+      int samplingPriority = context.getSamplingPriority();
+      injectTraceParent(context, carrier, setter, samplingPriority);
+      injectTraceState(context, carrier, setter, samplingPriority);
       injectBaggage(context, carrier, setter);
     }
 
-    private <C> void injectTraceParent(DDSpanContext context, C carrier, CarrierSetter<C> setter) {
+    private <C> void injectTraceParent(
+        DDSpanContext context, C carrier, CarrierSetter<C> setter, int samplingPriority) {
       String traceparent =
-          W3CTraceParent.from(
-              context.getTraceId(), context.getSpanId(), context.getSamplingPriority() > 0);
+          W3CTraceParent.from(context.getTraceId(), context.getSpanId(), samplingPriority > 0);
       setter.set(carrier, TRACE_PARENT_KEY, traceparent);
     }
 
-    private <C> void injectTraceState(DDSpanContext context, C carrier, CarrierSetter<C> setter) {
+    private <C> void injectTraceState(
+        DDSpanContext context, C carrier, CarrierSetter<C> setter, int samplingPriority) {
       PropagationTags propagationTags = context.getPropagationTags();
       // Supply the injecting span's id for the W3C `p:` as a parameter rather than mutating it into
       // the (possibly trace-level, shared) tags — keeps transient per-injection identity out of
       // shared state, so concurrent sibling injects can't race on it.
       String tracestate =
-          propagationTags.headerValue(W3C, DDSpanId.toHexStringPadded(context.getSpanId()));
+          propagationTags.headerValue(
+              W3C, DDSpanId.toHexStringPadded(context.getSpanId()), samplingPriority);
       if (tracestate != null && !tracestate.isEmpty()) {
         setter.set(carrier, TRACE_STATE_KEY, tracestate);
       }

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

@@ -141,9 +141,10 @@ static OtelTraceState updateProbability(
 
     long threshold = computeThreshold(sampleRate);
     long randomValue = computeRandomValue(traceIdLowOrderBits);
-    if (sampled && randomValue < threshold) {
+    boolean finalSampled = samplingPriority > 0;
+    if (finalSampled && randomValue < threshold) {
       randomValue = threshold;
-    } else if (!sampled && randomValue >= threshold) {
+    } else if (!finalSampled && randomValue >= threshold) {
       randomValue = threshold == 0 ? 0 : threshold - 1;
     }
 
@@ -170,6 +171,13 @@ OtelTraceState removeThresholdForLimiterDemotion() {
         hasLocallyGeneratedRandomValue());
   }
 
+  OtelTraceState reconcileSamplingDecision(boolean sampled) {
+    if (randomValue == NO_VALUE || threshold == NO_VALUE || sampled == (randomValue >= threshold)) {
+      return this;
+    }
+    return removeThresholdForLimiterDemotion();
+  }
+
   String getValue() {
     return value;
   }

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

@@ -24,19 +24,20 @@ abstract class PTagsCodec {
   protected static final String PROPAGATION_ERROR_INCONSISTENT_TID = "inconsistent_tid ";
   protected static final TagKey UPSTREAM_SERVICES_DEPRECATED_TAG = TagKey.from("upstream_services");
 
-  static String headerValue(PTagsCodec codec, PTags ptags) {
-    return headerValue(codec, ptags, null);
-  }
-
-  static String headerValue(PTagsCodec codec, PTags ptags, CharSequence lastParentIdOverride) {
-    int estimate = codec.estimateHeaderSize(ptags);
+  static String headerValue(
+      PTagsCodec codec,
+      PTags ptags,
+      CharSequence lastParentIdOverride,
+      OtelTraceState otelTraceState,
+      int samplingPriority) {
+    int estimate = codec.estimateHeaderSize(ptags, otelTraceState);
     if (estimate == 0) {
       return "";
     }
 
     // No encoding validation here because we don't allow arbitrary tag change
     StringBuilder sb = new StringBuilder(estimate);
-    int size = codec.appendPrefix(sb, ptags, lastParentIdOverride);
+    int size = codec.appendPrefix(sb, ptags, lastParentIdOverride, samplingPriority);
     if (!ptags.isPropagationTagsDisabled()) {
       if (ptags.getDecisionMakerTagValue() != null) {
         size = codec.appendTag(sb, DECISION_MAKER_TAG, ptags.getDecisionMakerTagValue(), size);
@@ -72,7 +73,7 @@ static String headerValue(PTagsCodec codec, PTags ptags, CharSequence lastParent
         size = codec.appendTag(sb, tagKey, tagValue, size);
       }
     }
-    size = codec.appendSuffix(sb, ptags, size);
+    size = codec.appendSuffix(sb, ptags, size, otelTraceState);
     if (codec.isTooLarge(sb, size)) {
       return null;
     } else {
@@ -176,6 +177,10 @@ static int calcXDatadogTagsSize(int size, TagKey tagKey, TagValue tagValue) {
 
   protected abstract int estimateHeaderSize(PTags pTags);
 
+  protected int estimateHeaderSize(PTags pTags, OtelTraceState otelTraceState) {
+    return estimateHeaderSize(pTags);
+  }
+
   protected abstract int appendPrefix(StringBuilder sb, PTags ptags);
 
   /**
@@ -186,10 +191,20 @@ protected int appendPrefix(StringBuilder sb, PTags ptags, CharSequence lastParen
     return appendPrefix(sb, ptags);
   }
 
+  protected int appendPrefix(
+      StringBuilder sb, PTags ptags, CharSequence lastParentIdOverride, int samplingPriority) {
+    return appendPrefix(sb, ptags, lastParentIdOverride);
+  }
+
   protected abstract int appendTag(StringBuilder sb, TagElement key, TagElement value, int size);
 
   protected abstract int appendSuffix(StringBuilder sb, PTags ptags, int size);
 
+  protected int appendSuffix(
+      StringBuilder sb, PTags ptags, int size, OtelTraceState otelTraceState) {
+    return appendSuffix(sb, ptags, size);
+  }
+
   protected abstract boolean isTooLarge(StringBuilder sb, int size);
 
   protected abstract boolean isEmpty(StringBuilder sb, int size);

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

@@ -113,6 +113,7 @@ static class PTags extends PropagationTags {
     private volatile TagValue orgPropagationMarkerTagValue;
 
     private volatile OtelTraceState otelTraceState;
+    private volatile OtelSamplingDecision otelSamplingDecision;
 
     // Static cache for the most-recently-seen rate → TagValue. In steady state a service uses one
     // rate, so this eliminates the char[] + String allocation on every new PTags instance.
@@ -267,17 +268,15 @@ private void doUpdateTraceSamplingPriority(int samplingPriority, int samplingMec
         }
         decisionMakerTagValue = null;
       }
-      if (removeOtelProbability && otelTraceState != null) {
-        setOtelTraceState(otelTraceState.removeForNonProbabilityDecision());
+      if (removeOtelProbability) {
+        setOtelSamplingDecision(OtelSamplingDecision.NON_PROBABILITY);
       }
     }
 
     @Override
-    public void updateOtelTraceState(
-        long traceIdLowOrderBits, double sampleRate, boolean sampled, int samplingPriority) {
-      setOtelTraceState(
-          OtelTraceState.updateProbability(
-              otelTraceState, traceIdLowOrderBits, sampleRate, sampled, samplingPriority));
+    public void updateOtelTraceState(long traceIdLowOrderBits, double sampleRate, boolean sampled) {
+      setOtelSamplingDecision(
+          OtelSamplingDecision.probability(traceIdLowOrderBits, sampleRate, sampled));
     }
 
     @Override
@@ -442,7 +441,13 @@ public CharSequence getLastParentId() {
     public String headerValue(HeaderType headerType) {
       String header = getCachedHeader(headerType);
       if (header == null) {
-        header = PTagsCodec.headerValue(factory.getDecoderEncoder(headerType), this);
+        header =
+            PTagsCodec.headerValue(
+                factory.getDecoderEncoder(headerType),
+                this,
+                null,
+                resolveOtelTraceState(headerType, samplingPriority),
+                samplingPriority);
         if (header != null) {
           setCachedHeader(headerType, header);
         } else {
@@ -464,7 +469,28 @@ public String headerValue(HeaderType headerType, CharSequence lastParentIdOverri
       // Inject-time path: encode fresh with the override; do NOT cache — the W3C `p:` is
       // per-injecting-span and these tags may be shared across sibling spans.
       String header =
-          PTagsCodec.headerValue(factory.getDecoderEncoder(headerType), this, lastParentIdOverride);
+          PTagsCodec.headerValue(
+              factory.getDecoderEncoder(headerType),
+              this,
+              lastParentIdOverride,
+              resolveOtelTraceState(headerType, samplingPriority),
+              samplingPriority);
+      return (header == null || header.isEmpty()) ? null : header;
+    }
+
+    @Override
+    public String headerValue(
+        HeaderType headerType, CharSequence lastParentIdOverride, int samplingPriority) {
+      if (lastParentIdOverride == null && samplingPriority == this.samplingPriority) {
+        return headerValue(headerType);
+      }
+      String header =
+          PTagsCodec.headerValue(
+              factory.getDecoderEncoder(headerType),
+              this,
+              lastParentIdOverride,
+              resolveOtelTraceState(headerType, samplingPriority),
+              samplingPriority);
       return (header == null || header.isEmpty()) ? null : header;
     }
 
@@ -560,16 +586,13 @@ public void updateW3CTracestate(String tracestate) {
     }
 
     @Override
-    public void updateW3CTracestateFrom(PropagationTags source, int samplingPriority) {
-      String sourceTracestate = source.getW3CTracestate();
-      OtelTraceState sourceOtelTraceState =
-          source instanceof PTags
-              ? ((PTags) source).getOtelTraceState()
-              : W3CPTagsCodec.extractOtelTraceState(sourceTracestate);
-      setW3CTracestate(sourceTracestate, sourceOtelTraceState);
-      if (samplingPriority != PrioritySampling.UNSET) {
-        this.samplingPriority = samplingPriority;
+    public void updateW3CTracestateFrom(PropagationTags source) {
+      if (!(source instanceof PTags)) {
+        super.updateW3CTracestateFrom(source);
+        return;
       }
+      PTags sourcePTags = (PTags) source;
+      setW3CTracestate(sourcePTags.tracestate, sourcePTags.getOtelTraceState());
     }
 
     private void setW3CTracestate(String tracestate, OtelTraceState otelTraceState) {
@@ -589,6 +612,26 @@ void setOtelTraceState(OtelTraceState otelTraceState) {
       }
     }
 
+    private OtelTraceState resolveOtelTraceState(HeaderType headerType, int samplingPriority) {
+      OtelTraceState current = otelTraceState;
+      if (headerType != W3C) {
+        return current;
+      }
+      OtelSamplingDecision decision = otelSamplingDecision;
+      OtelTraceState resolved =
+          decision == null ? current : decision.resolve(current, samplingPriority);
+      return resolved == null || samplingPriority == PrioritySampling.UNSET
+          ? resolved
+          : resolved.reconcileSamplingDecision(samplingPriority > 0);
+    }
+
+    private void setOtelSamplingDecision(OtelSamplingDecision decision) {
+      if (otelSamplingDecision != decision) {
+        otelSamplingDecision = decision;
+        clearCachedHeader(W3C);
+      }
+    }
+
     String getError() {
       return this.error;
     }
@@ -618,5 +661,62 @@ private static boolean isProbabilitySamplingMechanism(int samplingMechanism) {
           return false;
       }
     }
+
+    private static final class OtelSamplingDecision {
+      private static final OtelSamplingDecision NON_PROBABILITY =
+          new OtelSamplingDecision(0, 0, false, false);
+
+      private final long traceIdLowOrderBits;
+      private final double sampleRate;
+      private final boolean sampled;
+      private final boolean probability;
+      private volatile Resolution resolution;
+
+      private OtelSamplingDecision(
+          long traceIdLowOrderBits, double sampleRate, boolean sampled, boolean probability) {
+        this.traceIdLowOrderBits = traceIdLowOrderBits;
+        this.sampleRate = sampleRate;
+        this.sampled = sampled;
+        this.probability = probability;
+      }
+
+      private static OtelSamplingDecision probability(
+          long traceIdLowOrderBits, double sampleRate, boolean sampled) {
+        return new OtelSamplingDecision(traceIdLowOrderBits, sampleRate, sampled, true);
+      }
+
+      private OtelTraceState resolve(OtelTraceState current, int samplingPriority) {
+        Resolution cached = resolution;
+        if (cached != null
+            && cached.source == current
+            && cached.samplingPriority == samplingPriority) {
+          return cached.resolved;
+        }
+        OtelTraceState resolved;
+        if (!probability) {
+          resolved = current == null ? null : current.removeForNonProbabilityDecision();
+        } else if (current != null) {
+          resolved = current;
+        } else {
+          resolved =
+              OtelTraceState.updateProbability(
+                  null, traceIdLowOrderBits, sampleRate, sampled, samplingPriority);
+        }
+        resolution = new Resolution(current, samplingPriority, resolved);
+        return resolved;
+      }
+
+      private static final class Resolution {
+        private final OtelTraceState source;
+        private final int samplingPriority;
+        private final OtelTraceState resolved;
+
+        private Resolution(OtelTraceState source, int samplingPriority, OtelTraceState resolved) {
+          this.source = source;
+          this.samplingPriority = samplingPriority;
+          this.resolved = resolved;
+        }
+      }
+    }
   }
 }