feat(core): support OTel tracestate sampling

Author: MilanGarnier

dd-trace-core/src/main/java/datadog/trace/common/sampling/RateByServiceTraceSampler.java

@@ -58,16 +58,19 @@ public <T extends CoreSpan<T>> void setSamplingPriority(final T span) {
 
     final RateSamplersByEnvAndService rates = serviceRates;
     RateSampler sampler = rates.getSampler(env, serviceName);
+    boolean sampled = sampler.sample(span);
+    int samplingPriority = sampled ? PrioritySampling.SAMPLER_KEEP : PrioritySampling.SAMPLER_DROP;
 
-    if (sampler.sample(span)) {
+    if (rates.hasAgentRates()) {
       span.setSamplingPriority(
-          PrioritySampling.SAMPLER_KEEP,
+          samplingPriority,
           SAMPLING_AGENT_RATE,
           sampler.getSampleRate(),
-          SamplingMechanism.AGENT_RATE);
+          SamplingMechanism.AGENT_RATE,
+          sampled);
     } else {
       span.setSamplingPriority(
-          PrioritySampling.SAMPLER_DROP,
+          samplingPriority,
           SAMPLING_AGENT_RATE,
           sampler.getSampleRate(),
           SamplingMechanism.AGENT_RATE);
@@ -161,7 +164,7 @@ public void onResponse(
     if (canIncrease && anyCapped) {
       lastCappedNanos = now;
     }
-    serviceRates = new RateSamplersByEnvAndService(updatedEnvServiceRates, fallbackSampler);
+    serviceRates = new RateSamplersByEnvAndService(updatedEnvServiceRates, fallbackSampler, true);
   }
 
   private static RateSampler createRateSampler(final double sampleRate) {
@@ -183,15 +186,23 @@ private static final class RateSamplersByEnvAndService {
 
     private final Map<String, TreeMap<String, RateSampler>> envServiceRates;
     private final RateSampler fallbackSampler;
+    private final boolean hasAgentRates;
 
     RateSamplersByEnvAndService() {
-      this(Collections.emptyMap(), DEFAULT_SAMPLER);
+      this(Collections.emptyMap(), DEFAULT_SAMPLER, false);
     }
 
     RateSamplersByEnvAndService(
-        Map<String, TreeMap<String, RateSampler>> envServiceRates, RateSampler fallbackSampler) {
+        Map<String, TreeMap<String, RateSampler>> envServiceRates,
+        RateSampler fallbackSampler,
+        boolean hasAgentRates) {
       this.envServiceRates = envServiceRates;
       this.fallbackSampler = fallbackSampler;
+      this.hasAgentRates = hasAgentRates;
+    }
+
+    boolean hasAgentRates() {
+      return hasAgentRates;
     }
 
     RateSampler getFallbackSampler() {

dd-trace-core/src/main/java/datadog/trace/common/sampling/RuleBasedTraceSampler.java

@@ -146,27 +146,31 @@ public <T extends CoreSpan<T>> void setSamplingPriority(final T span) {
     if (matchedRule == null) {
       fallbackSampler.setSamplingPriority(span);
     } else {
-      if (matchedRule.sample(span)) {
+      boolean sampled = matchedRule.sample(span);
+      if (sampled) {
         if (rateLimiter.tryAcquire()) {
           span.setSamplingPriority(
               PrioritySampling.USER_KEEP,
               SAMPLING_RULE_RATE,
               matchedRule.getSampler().getSampleRate(),
-              matchedRule.getMechanism());
+              matchedRule.getMechanism(),
+              true);
         } else {
           span.setSamplingPriority(
               PrioritySampling.USER_DROP,
               SAMPLING_RULE_RATE,
               matchedRule.getSampler().getSampleRate(),
-              matchedRule.getMechanism());
+              matchedRule.getMechanism(),
+              true);
         }
         span.setMetric(SAMPLING_LIMIT_RATE, rateLimit);
       } else {
         span.setSamplingPriority(
             PrioritySampling.USER_DROP,
             SAMPLING_RULE_RATE,
             matchedRule.getSampler().getSampleRate(),
-            matchedRule.getMechanism());
+            matchedRule.getMechanism(),
+            false);
       }
     }
   }

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

@@ -124,6 +124,13 @@ default void processTagsAndBaggageWithStructuredLinks(
   T setSamplingPriority(
       int samplingPriority, CharSequence rate, double sampleRate, int samplingMechanism);
 
+  T setSamplingPriority(
+      int samplingPriority,
+      CharSequence rate,
+      double sampleRate,
+      int samplingMechanism,
+      boolean sampled);
+
   T setSpanSamplingPriority(double rate, int limit);
 
   T setMetric(CharSequence name, int value);

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

@@ -647,6 +647,27 @@ public final DDSpan setSamplingPriority(final int newPriority, int samplingMecha
   @Override
   public DDSpan setSamplingPriority(
       int samplingPriority, CharSequence rate, double sampleRate, int samplingMechanism) {
+    setSamplingPriorityWithRate(samplingPriority, rate, sampleRate, samplingMechanism);
+    return this;
+  }
+
+  @Override
+  public DDSpan setSamplingPriority(
+      int samplingPriority,
+      CharSequence rate,
+      double sampleRate,
+      int samplingMechanism,
+      boolean sampled) {
+    if (setSamplingPriorityWithRate(samplingPriority, rate, sampleRate, samplingMechanism)) {
+      context
+          .getPropagationTags()
+          .updateOtelTraceState(getTraceId().toLong(), sampleRate, sampled, samplingPriority);
+    }
+    return this;
+  }
+
+  private boolean setSamplingPriorityWithRate(
+      int samplingPriority, CharSequence rate, double sampleRate, int samplingMechanism) {
     if (context.setSamplingPriority(samplingPriority, samplingMechanism)) {
       setMetric(rate, sampleRate);
       if (samplingMechanism == SamplingMechanism.AGENT_RATE
@@ -655,8 +676,9 @@ public DDSpan setSamplingPriority(
           || samplingMechanism == SamplingMechanism.REMOTE_ADAPTIVE_RULE) {
         context.getPropagationTags().updateKnuthSamplingRate(sampleRate);
       }
+      return true;
     }
-    return this;
+    return false;
   }
 
   @Override

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

@@ -445,7 +445,13 @@ public DDSpanContext(
       setOrigin(origin);
     }
     if (samplingPriority != PrioritySampling.UNSET) {
-      setSamplingPriority(samplingPriority, SamplingMechanism.UNKNOWN);
+      if (this.propagationTags.getSamplingPriority() == samplingPriority) {
+        // Extractors already applied this priority to the propagation tags. Initialize the local
+        // field without reapplying the decision as an unknown local override.
+        SAMPLING_PRIORITY_UPDATER.set(this, samplingPriority);
+      } else {
+        setSamplingPriority(samplingPriority, SamplingMechanism.UNKNOWN);
+      }
     }
     setTag(PARENT_ID, this.propagationTags.getLastParentId());
   }

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

@@ -67,6 +67,9 @@ public interface Factory {
 
   public abstract void forceKeep(int samplingMechanism);
 
+  public abstract void updateOtelTraceState(
+      long traceIdLowOrderBits, double sampleRate, boolean sampled, int samplingPriority);
+
   public abstract int getSamplingPriority();
 
   public abstract void updateTraceOrigin(CharSequence origin);