Optimize KSR update path with static rate→TagValue cache

bm1549 · claude · bm1549 · commit bcc0957a743f · 2026-03-20T11:32:56.000-04:00
Two changes: 1. Benchmark (address dougqh feedback): - Switch to Throughput mode + @threads(8) to surface GC pressure - @State(Scope.Thread): each thread gets its own PTags, models real traces - Add updateRateFreshTrace: resets instance cache each iteration to model per-trace cost (the actual hot path Doug was concerned about) - Update run instructions to include -t 8 -prof gc 2. Static cache for KSR TagValue (Option A): - Add static volatile cachedKsrRate + cachedKsrTagValue to PTags - On updateKnuthSamplingRate, check static cache before formatting - Eliminates char[]+String allocation on the per-trace path in steady state (every new PTags starts with NaN; without the cache, each trace root re-formats even when the rate is constant) - Race is benign: two threads computing the same rate produce equal TagValues - gc.alloc.rate.norm: updateRateFreshTrace goes from 80 B/op → ≈ 0 B/op structurally (not JIT-dependent) Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>
diff --git a/dd-trace-core/src/jmh/java/datadog/trace/core/propagation/ptags/KnuthSamplingRateFormatBenchmark.java b/dd-trace-core/src/jmh/java/datadog/trace/core/propagation/ptags/KnuthSamplingRateFormatBenchmark.java
@@ -1,6 +1,5 @@
 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;
@@ -16,6 +15,7 @@
 import org.openjdk.jmh.annotations.Scope;
 import org.openjdk.jmh.annotations.Setup;
 import org.openjdk.jmh.annotations.State;
+import org.openjdk.jmh.annotations.Threads;
 import org.openjdk.jmh.annotations.Warmup;
 import org.openjdk.jmh.infra.Blackhole;
 
@@ -29,14 +29,20 @@
  *
  * <pre>
  *   ./gradlew :dd-trace-core:jmhJar
- *   java -jar dd-trace-core/build/libs/dd-trace-core-*-jmh.jar KnuthSamplingRateFormatBenchmark
+ *   java -jar dd-trace-core/build/libs/dd-trace-core-*-jmh.jar KnuthSamplingRateFormatBenchmark \
+ *     -t 8 -prof gc
  * </pre>
+ *
+ * <p>Use {@code -t 8} (or higher) to surface GC pressure from multi-threaded allocation. Use {@code
+ * -prof gc} to see alloc rate (bytes/op) per benchmark — that's the primary signal for whether the
+ * hot path is allocation-free.
  */
-@State(Scope.Benchmark)
+@State(Scope.Thread)
 @Warmup(iterations = 3, time = 10, timeUnit = SECONDS)
 @Measurement(iterations = 5, time = 10, timeUnit = SECONDS)
-@BenchmarkMode(Mode.AverageTime)
-@OutputTimeUnit(NANOSECONDS)
+@BenchmarkMode(Mode.Throughput)
+@OutputTimeUnit(SECONDS)
+@Threads(8)
 @Fork(value = 1)
 public class KnuthSamplingRateFormatBenchmark {
 
@@ -67,12 +73,27 @@ public void customFormat(Blackhole bh) {
     bh.consume(PTagsFactory.PTags.formatKnuthSamplingRate(rate));
   }
 
-  /** Cached TagValue: the full getKnuthSamplingRateTagValue() hot-path after caching. */
+  /**
+   * Cached TagValue: the full getKnuthSamplingRateTagValue() hot-path after caching. Should be
+   * near-zero allocation (volatile read only).
+   */
   @Benchmark
   public void cachedTagValue(Blackhole bh) {
     bh.consume(ptags.getKnuthSamplingRateTagValue());
   }
 
+  /**
+   * Models the per-trace allocation cost: resets the instance cache (simulating a new PTags), then
+   * calls updateKnuthSamplingRate. This is what every trace root pays. With the static cache
+   * applied, this should also be near-zero allocation after warmup.
+   */
+  @Benchmark
+  public void updateRateFreshTrace(Blackhole bh) {
+    ptags.updateKnuthSamplingRate(Double.NaN); // reset instance cache, like a new PTags
+    ptags.updateKnuthSamplingRate(rate);
+    bh.consume(ptags.getKnuthSamplingRateTagValue());
+  }
+
   // ---- old implementation for comparison ----
 
   static String stringFormatImpl(double rate) {
diff --git a/dd-trace-core/src/main/java/datadog/trace/core/propagation/ptags/PTagsFactory.java b/dd-trace-core/src/main/java/datadog/trace/core/propagation/ptags/PTagsFactory.java
@@ -94,6 +94,12 @@ static class PTags extends PropagationTags {
     private volatile double knuthSamplingRate = Double.NaN;
     private volatile TagValue knuthSamplingRateTagValue;
 
+    // Static cache for the most-recently-seen rate → TagValue. In steady state a service uses one
+    // rate, so this eliminates the char[] + String allocation on every new PTags instance.
+    // Writes are benign-racy: two threads computing the same rate produce equal TagValues.
+    private static volatile double cachedKsrRate = Double.NaN;
+    private static volatile TagValue cachedKsrTagValue;
+
     // xDatadogTagsSize of the tagPairs, does not include the decision maker tag
     private volatile int xDatadogTagsSize = -1;
 
@@ -275,8 +281,19 @@ public void updateKnuthSamplingRate(double rate) {
         clearCachedHeader(DATADOG);
         clearCachedHeader(W3C);
         knuthSamplingRate = rate;
-        knuthSamplingRateTagValue =
-            Double.isNaN(rate) ? null : TagValue.from(formatKnuthSamplingRate(rate));
+        if (Double.isNaN(rate)) {
+          knuthSamplingRateTagValue = null;
+        } else {
+          TagValue tv;
+          if (Double.compare(cachedKsrRate, rate) == 0) {
+            tv = cachedKsrTagValue;
+          } else {
+            tv = TagValue.from(formatKnuthSamplingRate(rate));
+            cachedKsrTagValue = tv;
+            cachedKsrRate = rate;
+          }
+          knuthSamplingRateTagValue = tv;
+        }
       }
     }
 
