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

@@ -59,19 +59,14 @@ 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(),
+        sampled,
+        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

@@ -146,18 +146,21 @@ 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(),
+              true,
               matchedRule.getMechanism());
         } else {
           span.setSamplingPriority(
               PrioritySampling.USER_DROP,
               SAMPLING_RULE_RATE,
               matchedRule.getSampler().getSampleRate(),
+              true,
               matchedRule.getMechanism());
         }
         span.setMetric(SAMPLING_LIMIT_RATE, rateLimit);
@@ -166,6 +169,7 @@ public <T extends CoreSpan<T>> void setSamplingPriority(final T span) {
             PrioritySampling.USER_DROP,
             SAMPLING_RULE_RATE,
             matchedRule.getSampler().getSampleRate(),
+            false,
             matchedRule.getMechanism());
       }
     }

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,
+      boolean probabilitySamplingResult,
+      int samplingMechanism);
 
   T setSpanSamplingPriority(double rate, int limit);
 

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

@@ -647,15 +647,18 @@ 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)) {
+      int samplingPriority,
+      CharSequence rate,
+      double sampleRate,
+      boolean probabilitySamplingResult,
+      int samplingMechanism) {
+    if (context.setSamplingPriority(
+        samplingPriority,
+        samplingMechanism,
+        sampleRate,
+        probabilitySamplingResult,
+        getTraceId().toLong())) {
       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 boolean probabilitySamplingResult,
+      final long traceIdLowOrderBits) {
+    DDSpanContext spanContext = getRootSpanContextOrThis();
+    if (!spanContext.validateSamplingPriority(newPriority, newMechanism)) {
+      return false;
+    }
+    return spanContext.propagationTags.tryUpdateProbabilitySamplingDecision(
+        newPriority,
+        newMechanism,
+        sampleRate,
+        probabilitySamplingResult,
+        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,9 @@ 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();
+      if (currentSamplingPriority != PrioritySampling.UNSET) {
+        tags.put(SAMPLE_RATE_KEY, currentSamplingPriority);
       }
       if (httpStatusCode != 0) {
         tags.put(Tags.HTTP_STATUS, (int) httpStatusCode);
@@ -1408,7 +1416,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()) {