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 expectedDataTags = []
     if (propagatedMechanism != UNKNOWN) {
-      expectedDatadogTags = "_dd.p.dm=-" + propagatedMechanism
-      expectedTracestate+= ";t.dm:-" + propagatedMechanism
+      expectedDataTags << "_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) {
+      expectedDataTags << "_dd.p.tid=" + traceId.toHexStringPadded(32).substring(0, 16)
+    }
+    if (contextPriority == UNSET) {
+      expectedDataTags << "_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 (!expectedDataTags.empty) {
+      expectedTextMap.put("x-datadog-tags", expectedDataTags.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/jmh/java/datadog/trace/core/propagation/ptags/KnuthSamplingRateFormatBenchmark.java

@@ -0,0 +1,93 @@
+package datadog.trace.core.propagation.ptags;
+
+import static java.util.concurrent.TimeUnit.NANOSECONDS;
+import static java.util.concurrent.TimeUnit.SECONDS;
+
+import datadog.trace.core.propagation.PropagationTags;
+import java.util.Locale;
+import org.openjdk.jmh.annotations.Benchmark;
+import org.openjdk.jmh.annotations.BenchmarkMode;
+import org.openjdk.jmh.annotations.Fork;
+import org.openjdk.jmh.annotations.Level;
+import org.openjdk.jmh.annotations.Measurement;
+import org.openjdk.jmh.annotations.Mode;
+import org.openjdk.jmh.annotations.OutputTimeUnit;
+import org.openjdk.jmh.annotations.Param;
+import org.openjdk.jmh.annotations.Scope;
+import org.openjdk.jmh.annotations.Setup;
+import org.openjdk.jmh.annotations.State;
+import org.openjdk.jmh.annotations.Warmup;
+import org.openjdk.jmh.infra.Blackhole;
+
+/**
+ * Benchmarks for formatting the Knuth sampling rate (_dd.p.ksr tag value).
+ *
+ * <p>The format requirement is %.6g semantics: 6 significant figures, no trailing zeros, using
+ * fixed notation for values in [1e-4, 1] and scientific notation for smaller values.
+ *
+ * <p>Run with:
+ *
+ * <pre>
+ *   ./gradlew :dd-trace-core:jmhJar
+ *   java -jar dd-trace-core/build/libs/dd-trace-core-*-jmh.jar KnuthSamplingRateFormatBenchmark
+ * </pre>
+ */
+@State(Scope.Benchmark)
+@Warmup(iterations = 3, time = 10, timeUnit = SECONDS)
+@Measurement(iterations = 5, time = 10, timeUnit = SECONDS)
+@BenchmarkMode(Mode.AverageTime)
+@OutputTimeUnit(NANOSECONDS)
+@Fork(value = 1)
+public class KnuthSamplingRateFormatBenchmark {
+
+  /**
+   * Representative sampling rates. Most real-world rates are in [0.001, 1.0]. The 0.0001 value
+   * exercises the edge of the fixed-notation range.
+   */
+  @Param({"0.5", "0.1", "0.01", "0.001", "0.0001", "0.123456789", "0.999999"})
+  double rate;
+
+  PTagsFactory.PTags ptags;
+
+  @Setup(Level.Trial)
+  public void setUp() {
+    ptags = (PTagsFactory.PTags) PropagationTags.factory().empty();
+    ptags.updateKnuthSamplingRate(rate);
+  }
+
+  /** Baseline: old implementation using String.format + substring trimming. */
+  @Benchmark
+  public void stringFormat(Blackhole bh) {
+    bh.consume(stringFormatImpl(rate));
+  }
+
+  /** Custom formatter: char-array arithmetic, no Formatter allocation. */
+  @Benchmark
+  public void customFormat(Blackhole bh) {
+    bh.consume(PTagsFactory.PTags.formatKnuthSamplingRate(rate));
+  }
+
+  /** Cached TagValue: the full getKnuthSamplingRateTagValue() hot-path after caching. */
+  @Benchmark
+  public void cachedTagValue(Blackhole bh) {
+    bh.consume(ptags.getKnuthSamplingRateTagValue());
+  }
+
+  // ---- old implementation for comparison ----
+
+  static String stringFormatImpl(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;
+  }
+}

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(),