Add _dd.p.ksr propagated tag for Knuth sampling rate

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

@@ -271,24 +271,38 @@ class OpenTelemetryTest extends InstrumentationSpecification {
     httpPropagator.inject(context, textMap, new TextMapSetter())
 
     then:
-    def expectedTraceparent = "00-${span.delegate.traceId.toHexStringPadded(32)}" +
-      "-${DDSpanId.toHexStringPadded(span.delegate.spanId)}" +
+    def traceId = span.delegate.traceId as DDTraceId
+    def spanId = span.delegate.spanId
+    def expectedTraceparent = "00-${traceId.toHexStringPadded(32)}" +
+      "-${DDSpanId.toHexStringPadded(spanId)}" +
       "-" + (propagatedPriority > 0 ? "01" : "00")
-    def expectedTracestate = "dd=s:${propagatedPriority};p:${DDSpanId.toHexStringPadded(span.delegate.spanId)}"
-    def expectedDatadogTags = null
+    def expectedTracestate = "dd=s:${propagatedPriority};p:${DDSpanId.toHexStringPadded(spanId)}"
+    def datadogTags = []
     if (propagatedMechanism != UNKNOWN) {
-      expectedDatadogTags = "_dd.p.dm=-" + propagatedMechanism
-      expectedTracestate+= ";t.dm:-" + propagatedMechanism
+      datadogTags << "_dd.p.dm=-" + propagatedMechanism
+      expectedTracestate += ";t.dm:-" + propagatedMechanism
+    }
+    if (traceId.toHighOrderLong() != 0) {
+      expectedTracestate += ";t.tid:${traceId.toHexStringPadded(32).substring(0, 16)}"
+    }
+    if (contextPriority == UNSET) {
+      expectedTracestate += ";t.ksr:1"
+    }
+    if (traceId.toHighOrderLong() != 0) {
+      datadogTags << "_dd.p.tid=" + traceId.toHexStringPadded(32).substring(0, 16)
+    }
+    if (contextPriority == UNSET) {
+      datadogTags << "_dd.p.ksr=1"
     }
     def expectedTextMap = [
-      "x-datadog-trace-id"         : "$span.delegate.traceId",
-      "x-datadog-parent-id"        : "$span.delegate.spanId",
+      "x-datadog-trace-id"         : "$traceId",
+      "x-datadog-parent-id"        : "$spanId",
       "x-datadog-sampling-priority": propagatedPriority.toString(),
       "traceparent"                : expectedTraceparent,
       "tracestate"                 : expectedTracestate,
     ]
-    if (expectedDatadogTags != null) {
-      expectedTextMap.put("x-datadog-tags", expectedDatadogTags)
+    if (!datadogTags.empty) {
+      expectedTextMap.put("x-datadog-tags", datadogTags.join(','))
     }
     textMap == expectedTextMap
 

dd-java-agent/instrumentation/opentelemetry/opentelemetry-1.4/src/test/groovy/opentelemetry14/context/propagation/DatadogPropagatorTest.groovy

@@ -37,6 +37,9 @@ class DatadogPropagatorTest extends AgentPropagatorTest {
     if (traceId.length() == 32) {
       tags+= '_dd.p.tid='+ traceId.substring(0, 16)
     }
+    if (sampling == UNSET) {
+      tags+= '_dd.p.ksr=1'
+    }
     assert headers['x-datadog-tags'] == tags.join(',')
     assert headers['x-datadog-sampling-priority'] == samplingPriority
   }

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

@@ -284,22 +284,30 @@ class OpenTracing31Test extends InstrumentationSpecification {
       "-${DDSpanId.toHexStringPadded(context.delegate.spanId)}" +
       "-" + (propagatedPriority > 0 ? "01" : "00")
     def expectedTracestate = "dd=s:${propagatedPriority};p:${DDSpanId.toHexStringPadded(context.delegate.spanId)}"
-    def expectedDatadogTags = null
+    def datadogTags = []
     if (propagatedPriority > 0) {
       def effectiveSamplingMechanism = contextPriority == UNSET ? AGENT_RATE : samplingMechanism
-      expectedDatadogTags = "_dd.p.dm=-" + effectiveSamplingMechanism
-      expectedTracestate+= ";t.dm:-" + effectiveSamplingMechanism
+      datadogTags << "_dd.p.dm=-" + effectiveSamplingMechanism
+      expectedTracestate += ";t.dm:-" + effectiveSamplingMechanism
+    }
+    def traceId = context.delegate.traceId as DDTraceId
+    if (traceId.toHighOrderLong() != 0) {
+      expectedTracestate += ";t.tid:${traceId.toHexStringPadded(32).substring(0, 16)}"
+      datadogTags << "_dd.p.tid=" + traceId.toHexStringPadded(32).substring(0, 16)
+    }
+    if (contextPriority == UNSET) {
+      expectedTracestate += ";t.ksr:1"
+      datadogTags << "_dd.p.ksr=1"
     }
-
     def expectedTextMap = [
       "x-datadog-trace-id"         : "$context.delegate.traceId",
       "x-datadog-parent-id"        : "$context.delegate.spanId",
       "x-datadog-sampling-priority": propagatedPriority.toString(),
       "traceparent"                : expectedTraceparent,
       "tracestate"                 : expectedTracestate,
     ]
-    if (expectedDatadogTags != null) {
-      expectedTextMap.put("x-datadog-tags", expectedDatadogTags)
+    if (!datadogTags.empty) {
+      expectedTextMap.put("x-datadog-tags", datadogTags.join(','))
     }
     textMap == expectedTextMap
 

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

@@ -2,6 +2,7 @@ import datadog.trace.agent.test.InstrumentationSpecification
 import datadog.trace.api.DDSpanId
 import datadog.trace.api.DDTags
 import datadog.trace.api.DDTraceId
+import datadog.trace.api.internal.util.LongStringUtils
 import datadog.trace.api.interceptor.MutableSpan
 import datadog.trace.bootstrap.instrumentation.api.AgentTracer
 import datadog.trace.bootstrap.instrumentation.api.ResourceNamePriorities
@@ -299,21 +300,30 @@ class OpenTracing32Test extends InstrumentationSpecification {
       "-${DDSpanId.toHexStringPadded(context.delegate.spanId)}" +
       "-" + (propagatedPriority > 0 ? "01" : "00")
     def expectedTracestate = "dd=s:${propagatedPriority};p:${DDSpanId.toHexStringPadded(context.delegate.spanId)}"
-    def expectedDatadogTags = null
+    def datadogTags = []
     if (propagatedPriority > 0) {
       def effectiveSamplingMechanism = contextPriority == UNSET ? AGENT_RATE : samplingMechanism
-      expectedDatadogTags = "_dd.p.dm=-" + effectiveSamplingMechanism
+      datadogTags << "_dd.p.dm=-" + effectiveSamplingMechanism
       expectedTracestate+= ";t.dm:-" + effectiveSamplingMechanism
     }
+    def traceId = context.delegate.traceId as DDTraceId
+    if (traceId.toHighOrderLong() != 0) {
+      expectedTracestate+= ";t.tid:${traceId.toHexStringPadded(32).substring(0, 16)}"
+      datadogTags << "_dd.p.tid=" + LongStringUtils.toHexStringPadded(traceId.toHighOrderLong(), 16)
+    }
+    if (contextPriority == UNSET) {
+      expectedTracestate+= ";t.ksr:1"
+      datadogTags << "_dd.p.ksr=1"
+    }
     def expectedTextMap = [
       "x-datadog-trace-id"         : "$context.delegate.traceId",
       "x-datadog-parent-id"        : "$context.delegate.spanId",
       "x-datadog-sampling-priority": propagatedPriority.toString(),
       "traceparent"                : expectedTraceparent,
       "tracestate"                 : expectedTracestate
     ]
-    if (expectedDatadogTags != null) {
-      expectedTextMap.put("x-datadog-tags", expectedDatadogTags)
+    if (!datadogTags.empty) {
+      expectedTextMap.put("x-datadog-tags", datadogTags.join(','))
     }
     textMap == expectedTextMap
 

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

@@ -646,6 +646,12 @@ public DDSpan setSamplingPriority(
       int samplingPriority, CharSequence rate, double sampleRate, int samplingMechanism) {
     if (context.setSamplingPriority(samplingPriority, samplingMechanism)) {
       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/propagation/PropagationTags.java

@@ -130,6 +130,16 @@ public interface Factory {
 
   public abstract String getDebugPropagation();
 
+  /**
+   * Updates the Knuth sampling rate (_dd.p.ksr) propagated tag. This records the sampling rate that
+   * was applied when making an agent-based or rule-based sampling decision. The rate is formatted
+   * with up to 6 significant digits and no trailing zeros, matching the Go/Python reference
+   * implementations (%.6g format).
+   *
+   * @param rate the sampling rate value
+   */
+  public abstract void updateKnuthSamplingRate(double rate);
+
   public HashMap<String, String> createTagMap() {
     HashMap<String, String> result = new HashMap<>();
     fillTagMap(result);

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

@@ -18,6 +18,7 @@ abstract class PTagsCodec {
   protected static final TagKey TRACE_ID_TAG = TagKey.from("tid");
   protected static final TagKey TRACE_SOURCE_TAG = TagKey.from("ts");
   protected static final TagKey DEBUG_TAG = TagKey.from("debug");
+  protected static final TagKey KNUTH_SAMPLING_RATE_TAG = TagKey.from("ksr");
   protected static final String PROPAGATION_ERROR_MALFORMED_TID = "malformed_tid ";
   protected static final String PROPAGATION_ERROR_INCONSISTENT_TID = "inconsistent_tid ";
   protected static final TagKey UPSTREAM_SERVICES_DEPRECATED_TAG = TagKey.from("upstream_services");
@@ -49,6 +50,11 @@ static String headerValue(PTagsCodec codec, PTags ptags) {
       if (ptags.getDebugPropagation() != null) {
         size = codec.appendTag(sb, DEBUG_TAG, TagValue.from(ptags.getDebugPropagation()), size);
       }
+      if (ptags.getKnuthSamplingRateTagValue() != null) {
+        size =
+            codec.appendTag(
+                sb, KNUTH_SAMPLING_RATE_TAG, ptags.getKnuthSamplingRateTagValue(), size);
+      }
       Iterator<TagElement> it = ptags.getTagPairs().iterator();
       while (it.hasNext() && !codec.isTooLarge(sb, size)) {
         TagElement tagKey = it.next();
@@ -103,6 +109,11 @@ static void fillTagMap(PTags propagationTags, Map<String, String> tagMap) {
       tagMap.put(
           DEBUG_TAG.forType(Encoding.DATADOG).toString(), propagationTags.getDebugPropagation());
     }
+    if (propagationTags.getKnuthSamplingRateTagValue() != null) {
+      tagMap.put(
+          KNUTH_SAMPLING_RATE_TAG.forType(Encoding.DATADOG).toString(),
+          propagationTags.getKnuthSamplingRateTagValue().forType(Encoding.DATADOG).toString());
+    }
     if (propagationTags.getTraceIdHighOrderBitsHexTagValue() != null) {
       tagMap.put(
           TRACE_ID_TAG.forType(Encoding.DATADOG).toString(),

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

@@ -3,6 +3,7 @@
 import static datadog.trace.core.propagation.PropagationTags.HeaderType.DATADOG;
 import static datadog.trace.core.propagation.PropagationTags.HeaderType.W3C;
 import static datadog.trace.core.propagation.ptags.PTagsCodec.DECISION_MAKER_TAG;
+import static datadog.trace.core.propagation.ptags.PTagsCodec.KNUTH_SAMPLING_RATE_TAG;
 import static datadog.trace.core.propagation.ptags.PTagsCodec.TRACE_ID_TAG;
 import static datadog.trace.core.propagation.ptags.PTagsCodec.TRACE_SOURCE_TAG;
 
@@ -16,6 +17,7 @@
 import java.util.Collections;
 import java.util.EnumMap;
 import java.util.List;
+import java.util.Locale;
 import java.util.Map;
 import java.util.Objects;
 import java.util.concurrent.atomic.AtomicIntegerFieldUpdater;
@@ -90,6 +92,9 @@ static class PTags extends PropagationTags {
     private volatile int traceSource;
     private volatile String debugPropagation;
 
+    // extracted Knuth sampling rate tag for easier updates
+    private volatile TagValue knuthSamplingRateTagValue;
+
     // xDatadogTagsSize of the tagPairs, does not include the decision maker tag
     private volatile int xDatadogTagsSize = -1;
 
@@ -265,6 +270,37 @@ public String getDebugPropagation() {
       return debugPropagation;
     }
 
+    @Override
+    public void updateKnuthSamplingRate(double rate) {
+      TagValue newValue = TagValue.from(formatKnuthSamplingRate(rate));
+      if (!Objects.equals(knuthSamplingRateTagValue, newValue)) {
+        clearCachedHeader(DATADOG);
+        clearCachedHeader(W3C);
+      }
+      knuthSamplingRateTagValue = newValue;
+    }
+
+    /** Formats a sampling rate with up to 6 significant digits and no trailing zeros. */
+    static String formatKnuthSamplingRate(double rate) {
+      String formatted = String.format(Locale.ROOT, "%.6g", rate);
+      int dotIndex = formatted.indexOf('.');
+      if (dotIndex >= 0) {
+        int end = formatted.length();
+        while (end > dotIndex + 1 && formatted.charAt(end - 1) == '0') {
+          end--;
+        }
+        if (formatted.charAt(end - 1) == '.') {
+          end--;
+        }
+        formatted = formatted.substring(0, end);
+      }
+      return formatted;
+    }
+
+    TagValue getKnuthSamplingRateTagValue() {
+      return knuthSamplingRateTagValue;
+    }
+
     @Override
     public int getSamplingPriority() {
       return samplingPriority;
@@ -390,6 +426,9 @@ int getXDatadogTagsSize() {
         size = PTagsCodec.calcXDatadogTagsSize(getTagPairs());
         size = PTagsCodec.calcXDatadogTagsSize(size, DECISION_MAKER_TAG, decisionMakerTagValue);
         size = PTagsCodec.calcXDatadogTagsSize(size, TRACE_ID_TAG, traceIdHighOrderBitsHexTagValue);
+        size =
+            PTagsCodec.calcXDatadogTagsSize(
+                size, KNUTH_SAMPLING_RATE_TAG, knuthSamplingRateTagValue);
         int currentProductTraceSource = traceSource;
         if (currentProductTraceSource != ProductTraceSource.UNSET) {
           size =

dd-trace-core/src/test/groovy/datadog/trace/common/writer/DDAgentApiTest.groovy

@@ -164,7 +164,7 @@ class DDAgentApiTest extends DDCoreSpecification {
     [[buildSpan(1L, "service.name", "my-service", PropagationTags.factory().fromHeaderValue(PropagationTags.HeaderType.DATADOG, "_dd.p.usr=123"))]] | [[new TreeMap<>([
       "duration" : 10,
       "error"    : 0,
-      "meta"     : ["thread.name": Thread.currentThread().getName(), "_dd.p.usr": "123", "_dd.p.dm": "-1", "_dd.svc_src" : "m"] +
+      "meta"     : ["thread.name": Thread.currentThread().getName(), "_dd.p.usr": "123", "_dd.p.dm": "-1", "_dd.p.ksr": "1", "_dd.svc_src" : "m"] +
         (Config.get().isExperimentalPropagateProcessTagsEnabled() ? ["_dd.tags.process" : ProcessTags.getTagsForSerialization().toString()] : []),
       "metrics"  : [
         (DDSpanContext.PRIORITY_SAMPLING_KEY)          : 1,
@@ -185,7 +185,7 @@ class DDAgentApiTest extends DDCoreSpecification {
     [[buildSpan(100L, "resource.name", "my-resource", PropagationTags.factory().fromHeaderValue(PropagationTags.HeaderType.DATADOG, "_dd.p.usr=123"))]] | [[new TreeMap<>([
       "duration" : 10,
       "error"    : 0,
-      "meta"     : ["thread.name": Thread.currentThread().getName(), "_dd.p.usr": "123", "_dd.p.dm": "-1"] +
+      "meta"     : ["thread.name": Thread.currentThread().getName(), "_dd.p.usr": "123", "_dd.p.dm": "-1", "_dd.p.ksr": "1"] +
         (Config.get().isExperimentalPropagateProcessTagsEnabled() ? ["_dd.tags.process" : ProcessTags.getTagsForSerialization().toString()] : []),
       "metrics"  : [
         (DDSpanContext.PRIORITY_SAMPLING_KEY)          : 1,