feat(sampling): Support for OpenTelemetry consistent tracestate sampling

dd-java-agent/instrumentation/opentelemetry/opentelemetry-0.3/src/test/groovy/OpenTelemetryTest.groovy

@@ -287,6 +287,7 @@ class OpenTelemetryTest extends InstrumentationSpecification {
     }
     if (contextPriority == UNSET) {
       expectedTracestate += ";t.ksr:1"
+      expectedTracestate += ",ot=${span.delegate.spanContext().propagationTags.samplingState().otelTraceState}"
     }
     if (traceId.toHighOrderLong() != 0) {
       expectedDataTags << "_dd.p.tid=" + traceId.toHexStringPadded(32).substring(0, 16)

dd-java-agent/instrumentation/opentracing/opentracing-0.31/src/test/groovy/OpenTracing31Test.groovy

@@ -295,6 +295,7 @@ class OpenTracing31Test extends InstrumentationSpecification {
     }
     if (contextPriority == UNSET) {
       expectedTracestate += ";t.ksr:1"
+      expectedTracestate += ",ot=${context.delegate.propagationTags.samplingState().otelTraceState}"
       datadogTags << "_dd.p.ksr=1"
     }
     def expectedTextMap = [

dd-java-agent/instrumentation/opentracing/opentracing-0.32/src/test/groovy/OpenTracing32Test.groovy

@@ -311,6 +311,7 @@ class OpenTracing32Test extends InstrumentationSpecification {
     }
     if (contextPriority == UNSET) {
       expectedTracestate+= ";t.ksr:1"
+      expectedTracestate+= ",ot=${context.delegate.propagationTags.samplingState().otelTraceState}"
       datadogTags << "_dd.p.ksr=1"
     }
     def expectedTextMap = [

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

@@ -64,19 +64,13 @@ 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;
+    span.setSamplingPriority(
+        samplingPriority,
+        SAMPLING_AGENT_RATE,
+        sampler.getSampleRate(),
+        SamplingMechanism.AGENT_RATE);
   }
 
   private <T extends CoreSpan<T>> String getSpanEnv(final T span) {

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

@@ -186,7 +186,8 @@ public <T extends CoreSpan<T>> void setSamplingPriority(final T span) {
               PrioritySampling.USER_DROP,
               SAMPLING_RULE_RATE,
               matchedRule.getSampler().getSampleRate(),
-              matchedRule.getMechanism());
+              matchedRule.getMechanism(),
+              true);
         }
         span.setMetric(SAMPLING_LIMIT_RATE, rateLimit);
       } else {

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

@@ -124,6 +124,15 @@ default void processTagsAndBaggageWithStructuredLinks(
   T setSamplingPriority(
       int samplingPriority, CharSequence rate, double sampleRate, int samplingMechanism);
 
+  default T setSamplingPriority(
+      int samplingPriority,
+      CharSequence rate,
+      double sampleRate,
+      int samplingMechanism,
+      boolean rateLimiterRejected) {
+    return setSamplingPriority(samplingPriority, rate, sampleRate, samplingMechanism);
+  }
+
   T setSpanSamplingPriority(double rate, int limit);
 
   T setMetric(CharSequence name, int value);

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

@@ -648,14 +648,23 @@ public final DDSpan setSamplingPriority(final int newPriority, int samplingMecha
   @Override
   public DDSpan setSamplingPriority(
       int samplingPriority, CharSequence rate, double sampleRate, int samplingMechanism) {
-    if (context.setSamplingPriority(samplingPriority, samplingMechanism)) {
+    return setSamplingPriority(samplingPriority, rate, sampleRate, samplingMechanism, false);
+  }
+
+  @Override
+  public DDSpan setSamplingPriority(
+      int samplingPriority,
+      CharSequence rate,
+      double sampleRate,
+      int samplingMechanism,
+      boolean rateLimiterRejected) {
+    if (context.setSamplingPriority(
+        samplingPriority,
+        samplingMechanism,
+        sampleRate,
+        getTraceId().toLong(),
+        rateLimiterRejected)) {
       setMetric(rate, sampleRate);
-      if (samplingMechanism == SamplingMechanism.AGENT_RATE
-          || samplingMechanism == SamplingMechanism.LOCAL_USER_RULE
-          || samplingMechanism == SamplingMechanism.REMOTE_USER_RULE
-          || samplingMechanism == SamplingMechanism.REMOTE_ADAPTIVE_RULE) {
-        context.getPropagationTags().updateKnuthSamplingRate(sampleRate);
-      }
     }
     return this;
   }

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

@@ -46,7 +46,6 @@
 import java.util.TreeMap;
 import java.util.concurrent.ConcurrentHashMap;
 import java.util.concurrent.ThreadLocalRandom;
-import java.util.concurrent.atomic.AtomicIntegerFieldUpdater;
 import java.util.function.Function;
 import javax.annotation.Nonnull;
 import org.slf4j.Logger;
@@ -161,11 +160,6 @@ public class DDSpanContext
 
   private volatile boolean topLevel;
 
-  private static final AtomicIntegerFieldUpdater<DDSpanContext> SAMPLING_PRIORITY_UPDATER =
-      AtomicIntegerFieldUpdater.newUpdater(DDSpanContext.class, "samplingPriority");
-
-  private volatile int samplingPriority = PrioritySampling.UNSET;
-
   /** The origin of the trace. (eg. Synthetics, CI App) */
   private volatile CharSequence origin;
 
@@ -667,10 +661,6 @@ public void forceKeep(byte samplingMechanism) {
   }
 
   private void forceKeepThisSpan(byte samplingMechanism) {
-    // if the user really wants to keep this trace chunk, we will let them,
-    // even if the old sampling priority and mechanism have already propagated
-    SAMPLING_PRIORITY_UPDATER.set(this, PrioritySampling.USER_KEEP);
-    // record force keep decision for future distributed trace propagation
     propagationTags.forceKeep(samplingMechanism);
   }
 
@@ -710,26 +700,43 @@ private boolean setThisSpanSamplingPriority(final int newPriority, final int new
     if (!validateSamplingPriority(newPriority, newMechanism)) {
       return false;
     }
-    if (SamplingMechanism.canAvoidSamplingPriorityLock(newPriority, newMechanism)) {
-      SAMPLING_PRIORITY_UPDATER.set(this, newPriority);
-      propagationTags.updateTraceSamplingPriority(newPriority, newMechanism);
-      return true;
-    }
-    if (!SAMPLING_PRIORITY_UPDATER.compareAndSet(this, PrioritySampling.UNSET, newPriority)) {
+    boolean updated =
+        propagationTags.tryUpdateTraceSamplingPriority(
+            newPriority,
+            newMechanism,
+            SamplingMechanism.canAvoidSamplingPriorityLock(newPriority, newMechanism));
+    if (!updated) {
       if (log.isDebugEnabled()) {
         log.debug(
             "samplingPriority locked at priority: {}. Refusing to set to priority: {} mechanism: {}",
-            samplingPriority,
+            propagationTags.getSamplingPriority(),
             newPriority,
             newMechanism);
       }
       return false;
     }
-    // set trace level sampling priority tag propagationTags
-    propagationTags.updateTraceSamplingPriority(newPriority, newMechanism);
     return true;
   }
 
+  public boolean setSamplingPriority(
+      final int newPriority,
+      final int newMechanism,
+      final double sampleRate,
+      final long traceIdLowOrderBits,
+      final boolean rateLimiterRejected) {
+    DDSpanContext spanContext = getRootSpanContextOrThis();
+    if (!spanContext.validateSamplingPriority(newPriority, newMechanism)) {
+      return false;
+    }
+    return spanContext.propagationTags.tryUpdateProbabilitySamplingDecision(
+        newPriority,
+        newMechanism,
+        sampleRate,
+        rateLimiterRejected,
+        traceIdLowOrderBits,
+        SamplingMechanism.canAvoidSamplingPriorityLock(newPriority, newMechanism));
+  }
+
   private boolean validateSamplingPriority(final int newPriority, final int newMechanism) {
     if (newPriority == PrioritySampling.UNSET) {
       log.debug("{}: Refusing to set samplingPriority to UNSET", this);
@@ -758,7 +765,7 @@ private boolean validateSamplingPriority(final int newPriority, final int newMec
 
   @Override
   public int getSamplingPriority() {
-    return getRootSpanContextOrThis().samplingPriority;
+    return getRootSpanContextOrThis().propagationTags.getSamplingPriority();
   }
 
   public void setSpanSamplingPriority(double rate, int limit) {
@@ -790,7 +797,7 @@ public boolean lockSamplingPriority() {
       return rootSpan.spanContext().lockSamplingPriority();
     }
 
-    return SAMPLING_PRIORITY_UPDATER.get(this) != PrioritySampling.UNSET;
+    return propagationTags.getSamplingPriority() != PrioritySampling.UNSET;
   }
 
   public CharSequence getOrigin() {
@@ -1258,8 +1265,11 @@ public TagMap getTags() {
       tags.put(DDTags.THREAD_ID, threadId);
       // maintain previously observable type of the thread name :|
       tags.put(DDTags.THREAD_NAME, threadName.toString());
-      if (samplingPriority != PrioritySampling.UNSET) {
-        tags.put(SAMPLE_RATE_KEY, samplingPriority);
+      int currentSamplingPriority = getSamplingPriority();
+      // add _sample_rate tag only on the root/local span owning the decision
+      if (getRootSpanContextIfDifferent() == null
+          && currentSamplingPriority != PrioritySampling.UNSET) {
+        tags.put(SAMPLE_RATE_KEY, currentSamplingPriority);
       }
       if (httpStatusCode != 0) {
         tags.put(Tags.HTTP_STATUS, (int) httpStatusCode);
@@ -1408,7 +1418,7 @@ void processTagsAndBaggage(
               threadName,
               unsafeTags,
               baggageItemsWithPropagationTags,
-              samplingPriority != PrioritySampling.UNSET ? samplingPriority : getSamplingPriority(),
+              getSamplingPriority(),
               measured,
               topLevel,
               httpStatusCode,

dd-trace-core/src/main/java/datadog/trace/core/otlp/trace/OtlpTraceJson.java

@@ -34,6 +34,7 @@
 import datadog.trace.core.MetadataConsumer;
 import datadog.trace.core.PendingTrace;
 import datadog.trace.core.propagation.PropagationTags;
+import datadog.trace.core.propagation.PropagationTags.SamplingState;
 import java.util.List;
 import java.util.Map;
 
@@ -51,13 +52,14 @@ private OtlpTraceJson() {}
   public static void writeSpan(
       JsonWriter writer, DDSpan span, MetaWriter metaWriter, List<? extends AgentSpanLink> links) {
     PropagationTags propagationTags = span.spanContext().getPropagationTags();
+    SamplingState samplingState = propagationTags.samplingState();
 
     writer.beginObject();
 
     writer.name("traceId").value(hexTraceId(span.getTraceId()));
     writer.name("spanId").value(hexSpanId(span.getSpanId()));
 
-    String tracestate = propagationTags.getW3CTracestate();
+    String tracestate = propagationTags.getW3CTracestate(samplingState);
     if (tracestate != null) {
       writer.name("traceState").value(tracestate);
     }
@@ -67,7 +69,7 @@ public static void writeSpan(
     }
 
     int traceFlags = NO_TRACE_FLAGS;
-    if (span.samplingPriority() > 0) {
+    if (samplingState.getSamplingPriority() > 0) {
       traceFlags |= SAMPLED_TRACE_FLAG;
     }
     if (span.spanContext().isRemote()) {

dd-trace-core/src/main/java/datadog/trace/core/otlp/trace/OtlpTraceProto.java

@@ -48,6 +48,7 @@
 import datadog.trace.core.PendingTrace;
 import datadog.trace.core.otlp.common.OtlpProtoBuffer;
 import datadog.trace.core.propagation.PropagationTags;
+import datadog.trace.core.propagation.PropagationTags.SamplingState;
 
 /** Provides optimized writers for OpenTelemetry's "trace.proto" wire protocol. */
 public final class OtlpTraceProto {
@@ -84,14 +85,15 @@ public static int recordSpanMessage(
       int nestedSpanLinkBytes,
       OtlpProtoBuffer protobuf) {
     PropagationTags propagationTags = span.spanContext().getPropagationTags();
+    SamplingState samplingState = propagationTags.samplingState();
 
     writeTag(buf, 1, LEN_WIRE_TYPE);
     writeTraceId(buf, span.getTraceId());
 
     writeTag(buf, 2, LEN_WIRE_TYPE);
     writeSpanId(buf, span.getSpanId());
 
-    String tracestate = propagationTags.getW3CTracestate();
+    String tracestate = propagationTags.getW3CTracestate(samplingState);
     if (tracestate != null) {
       writeTag(buf, 3, LEN_WIRE_TYPE);
       writeString(buf, tracestate);
@@ -103,7 +105,7 @@ public static int recordSpanMessage(
     }
 
     int traceFlags = NO_TRACE_FLAGS;
-    if (span.samplingPriority() > 0) {
+    if (samplingState.getSamplingPriority() > 0) {
       traceFlags |= SAMPLED_TRACE_FLAG;
     }
     if (span.spanContext().isRemote()) {

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

@@ -22,6 +22,47 @@
  */
 public abstract class PropagationTags {
 
+  public static final class SamplingState {
+    private final int samplingPriority;
+    private final String tracestate;
+    private final CharSequence otelTraceState;
+    private final CharSequence decisionMaker;
+    private final CharSequence knuthSamplingRate;
+
+    public SamplingState(
+        int samplingPriority,
+        String tracestate,
+        CharSequence otelTraceState,
+        CharSequence decisionMaker,
+        CharSequence knuthSamplingRate) {
+      this.samplingPriority = samplingPriority;
+      this.tracestate = tracestate;
+      this.otelTraceState = otelTraceState;
+      this.decisionMaker = decisionMaker;
+      this.knuthSamplingRate = knuthSamplingRate;
+    }
+
+    public int getSamplingPriority() {
+      return samplingPriority;
+    }
+
+    public String getTracestate() {
+      return tracestate;
+    }
+
+    public CharSequence getOtelTraceState() {
+      return otelTraceState;
+    }
+
+    public CharSequence getDecisionMaker() {
+      return decisionMaker;
+    }
+
+    public CharSequence getKnuthSamplingRate() {
+      return knuthSamplingRate;
+    }
+  }
+
   public static PropagationTags.Factory factory(Config config) {
     return factory(config.getxDatadogTagsMaxLength());
   }
@@ -65,10 +106,23 @@ public interface Factory {
    */
   public abstract void updateTraceSamplingPriority(int samplingPriority, int samplingMechanism);
 
+  public abstract boolean tryUpdateTraceSamplingPriority(
+      int samplingPriority, int samplingMechanism, boolean allowOverride);
+
+  public abstract boolean tryUpdateProbabilitySamplingDecision(
+      int samplingPriority,
+      int samplingMechanism,
+      double sampleRate,
+      boolean rateLimiterRejected,
+      long traceIdLowOrderBits,
+      boolean allowOverride);
+
   public abstract void forceKeep(int samplingMechanism);
 
   public abstract int getSamplingPriority();
 
+  public abstract SamplingState samplingState();
+
   public abstract void updateTraceOrigin(CharSequence origin);
 
   public abstract CharSequence getOrigin();
@@ -87,6 +141,8 @@ public interface Factory {
    */
   public abstract String getW3CTracestate();
 
+  public abstract String getW3CTracestate(SamplingState samplingState);
+
   /**
    * Stores the original <a href="https://www.w3.org/TR/trace-context/#tracestate-header">W3C
    * tracestate header</a> value.
@@ -116,6 +172,9 @@ public void updateW3CTracestateFrom(PropagationTags source) {
    */
   public abstract String headerValue(HeaderType headerType, CharSequence lastParentIdOverride);
 
+  public abstract String headerValue(
+      HeaderType headerType, CharSequence lastParentIdOverride, SamplingState samplingState);
+
   /**
    * 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