Fix mismatched _dd.p.ksr under concurrent sampling decisions with different rates

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

@@ -128,9 +128,9 @@ static class PTags extends PropagationTags {
 
     // 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;
+    // The rate and its value are published together, so a reader can't pair one decision's rate
+    // with another's value; racing writers only replace each other's entry.
+    private static volatile KnuthSamplingRateEntry cachedKnuthSamplingRate;
 
     // Size of the trace-level tags only; the per-injection LLMObs tags are added on top of it by
     // DatadogPTagsCodec, so they are deliberately not part of what this caches.
@@ -459,15 +459,25 @@ private static TagValue knuthSamplingRateTagValue(double rate) {
       if (Double.isNaN(rate)) {
         return null;
       }
-      if (Double.compare(cachedKsrRate, rate) == 0) {
-        return cachedKsrTagValue;
+      KnuthSamplingRateEntry cached = cachedKnuthSamplingRate;
+      if (cached != null && Double.compare(cached.rate, rate) == 0) {
+        return cached.tagValue;
       }
       TagValue value = TagValue.from(formatKnuthSamplingRate(rate));
-      cachedKsrTagValue = value;
-      cachedKsrRate = rate;
+      cachedKnuthSamplingRate = new KnuthSamplingRateEntry(rate, value);
       return value;
     }
 
+    private static final class KnuthSamplingRateEntry {
+      final double rate;
+      final TagValue tagValue;
+
+      KnuthSamplingRateEntry(double rate, TagValue tagValue) {
+        this.rate = rate;
+        this.tagValue = tagValue;
+      }
+    }
+
     /**
      * Formats a sampling rate with up to 6 decimal digits of precision and no trailing zeros.
      *

dd-trace-core/src/test/java/datadog/trace/core/propagation/ptags/KnuthSamplingRateCacheTest.java

@@ -0,0 +1,79 @@
+package datadog.trace.core.propagation.ptags;
+
+import static datadog.trace.api.sampling.PrioritySampling.SAMPLER_KEEP;
+import static datadog.trace.api.sampling.SamplingMechanism.AGENT_RATE;
+import static org.junit.jupiter.api.Assertions.assertEquals;
+import static org.junit.jupiter.api.Assertions.assertNull;
+
+import datadog.trace.core.propagation.PropagationTags;
+import java.util.concurrent.CountDownLatch;
+import java.util.concurrent.atomic.AtomicReference;
+import org.junit.jupiter.api.Test;
+
+class KnuthSamplingRateCacheTest {
+  private static final double[] RATES = {0.5, 0.25, 0.1, 0.75};
+  private static final int THREADS = 4;
+  private static final int DECISIONS_PER_THREAD = 200_000;
+
+  private final PropagationTags.Factory factory = PropagationTags.factory();
+
+  @Test
+  void eachDecisionCarriesItsOwnRate() {
+    for (double rate : RATES) {
+      assertEquals(PTagsFactory.PTags.formatKnuthSamplingRate(rate), knuthSamplingRate(rate));
+    }
+  }
+
+  @Test
+  void noRateWithoutOne() {
+    assertNull(knuthSamplingRate(Double.NaN));
+  }
+
+  /**
+   * Concurrent decisions with different rates must each get their own rate. A cache whose rate and
+   * value were published separately could pair one decision's rate with another's value.
+   */
+  @Test
+  void concurrentDecisionsWithDifferentRatesDontMix() throws InterruptedException {
+    AtomicReference<String> mismatch = new AtomicReference<>();
+    CountDownLatch start = new CountDownLatch(1);
+    Thread[] threads = new Thread[THREADS];
+    for (int t = 0; t < THREADS; ++t) {
+      int offset = t;
+      threads[t] =
+          new Thread(
+              () -> {
+                awaitQuietly(start);
+                for (int i = 0; i < DECISIONS_PER_THREAD && mismatch.get() == null; ++i) {
+                  double rate = RATES[(i + offset) % RATES.length];
+                  String expected = PTagsFactory.PTags.formatKnuthSamplingRate(rate);
+                  String actual = knuthSamplingRate(rate);
+                  if (!expected.equals(actual)) {
+                    mismatch.compareAndSet(null, "rate " + rate + " got _dd.p.ksr=" + actual);
+                  }
+                }
+              });
+      threads[t].start();
+    }
+    start.countDown();
+    for (Thread thread : threads) {
+      thread.join();
+    }
+    assertNull(mismatch.get());
+  }
+
+  private String knuthSamplingRate(double rate) {
+    PropagationTags ptags = factory.empty();
+    ptags.tryUpdateProbabilitySamplingDecision(SAMPLER_KEEP, AGENT_RATE, rate, false, 1L, false);
+    CharSequence value = ptags.samplingState().getKnuthSamplingRate();
+    return value == null ? null : value.toString();
+  }
+
+  private static void awaitQuietly(CountDownLatch latch) {
+    try {
+      latch.await();
+    } catch (InterruptedException e) {
+      Thread.currentThread().interrupt();
+    }
+  }
+}