feat(sampling): Support for OpenTelemetry consistent tracestate sampling

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

@@ -58,20 +58,16 @@ public <T extends CoreSpan<T>> void setSamplingPriority(final T span) {
 
     final RateSamplersByEnvAndService rates = serviceRates;
     RateSampler sampler = rates.getSampler(env, serviceName);
-
-    if (sampler.sample(span)) {
-      span.setSamplingPriority(
-          PrioritySampling.SAMPLER_KEEP,
-          SAMPLING_AGENT_RATE,
-          sampler.getSampleRate(),
-          SamplingMechanism.AGENT_RATE);
-    } else {
-      span.setSamplingPriority(
-          PrioritySampling.SAMPLER_DROP,
-          SAMPLING_AGENT_RATE,
-          sampler.getSampleRate(),
-          SamplingMechanism.AGENT_RATE);
-    }
+    boolean sampled = sampler.sample(span);
+    int samplingPriority = sampled ? PrioritySampling.SAMPLER_KEEP : PrioritySampling.SAMPLER_DROP;
+
+    Boolean probabilitySamplingResult = rates.hasAgentRates() ? sampled : null;
+    span.setSamplingPriority(
+        samplingPriority,
+        SAMPLING_AGENT_RATE,
+        sampler.getSampleRate(),
+        SamplingMechanism.AGENT_RATE,
+        probabilitySamplingResult);
   }
 
   private <T extends CoreSpan<T>> String getSpanEnv(final T span) {
@@ -161,7 +157,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 +179,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,28 +146,21 @@ public <T extends CoreSpan<T>> void setSamplingPriority(final T span) {
     if (matchedRule == null) {
       fallbackSampler.setSamplingPriority(span);
     } else {
-      if (matchedRule.sample(span)) {
-        if (rateLimiter.tryAcquire()) {
-          span.setSamplingPriority(
-              PrioritySampling.USER_KEEP,
-              SAMPLING_RULE_RATE,
-              matchedRule.getSampler().getSampleRate(),
-              matchedRule.getMechanism());
-        } else {
-          span.setSamplingPriority(
-              PrioritySampling.USER_DROP,
-              SAMPLING_RULE_RATE,
-              matchedRule.getSampler().getSampleRate(),
-              matchedRule.getMechanism());
-        }
+      boolean sampled = matchedRule.sample(span);
+      int samplingPriority;
+      if (sampled) {
+        samplingPriority =
+            rateLimiter.tryAcquire() ? PrioritySampling.USER_KEEP : PrioritySampling.USER_DROP;
         span.setMetric(SAMPLING_LIMIT_RATE, rateLimit);
       } else {
-        span.setSamplingPriority(
-            PrioritySampling.USER_DROP,
-            SAMPLING_RULE_RATE,
-            matchedRule.getSampler().getSampleRate(),
-            matchedRule.getMechanism());
+        samplingPriority = PrioritySampling.USER_DROP;
       }
+      span.setSamplingPriority(
+          samplingPriority,
+          SAMPLING_RULE_RATE,
+          matchedRule.getSampler().getSampleRate(),
+          matchedRule.getMechanism(),
+          Boolean.valueOf(sampled));
     }
   }
 }

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

@@ -122,7 +122,11 @@ default void processTagsAndBaggageWithStructuredLinks(
   T setSamplingPriority(int samplingPriority, int samplingMechanism);
 
   T setSamplingPriority(
-      int samplingPriority, CharSequence rate, double sampleRate, int samplingMechanism);
+      int samplingPriority,
+      CharSequence rate,
+      double sampleRate,
+      int samplingMechanism,
+      Boolean probabilitySamplingResult);
 
   T setSpanSamplingPriority(double rate, int limit);
 

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

@@ -646,6 +646,25 @@ public final DDSpan setSamplingPriority(final int newPriority, int samplingMecha
 
   @Override
   public DDSpan setSamplingPriority(
+      int samplingPriority,
+      CharSequence rate,
+      double sampleRate,
+      int samplingMechanism,
+      Boolean probabilitySamplingResult) {
+    if (setSamplingPriorityWithRate(samplingPriority, rate, sampleRate, samplingMechanism)
+        && probabilitySamplingResult != null) {
+      context
+          .getPropagationTags()
+          .updateOtelTraceState(
+              getTraceId().toLong(),
+              sampleRate,
+              probabilitySamplingResult.booleanValue(),
+              samplingPriority);
+    }
+    return this;
+  }
+
+  private boolean setSamplingPriorityWithRate(
       int samplingPriority, CharSequence rate, double sampleRate, int samplingMechanism) {
     if (context.setSamplingPriority(samplingPriority, samplingMechanism)) {
       setMetric(rate, sampleRate);
@@ -655,8 +674,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);