From 6ce6698e9105a87e8435ce174830cc897b063c0b Mon Sep 17 00:00:00 2001 From: Douglas Q Hawkins Date: Wed, 21 Oct 2026 17:25:51 -0400 Subject: [PATCH] Fix mismatched _dd.p.ksr under concurrent decisions with different rates The single-entry Knuth sampling rate cache stored the rate and its formatted value in two separate volatile fields. A decision racing with another at a different rate could read one rate's entry with the other's value and propagate the wrong _dd.p.ksr. Publish the rate and value together in one immutable entry. Co-Authored-By: Claude Opus 5.5 --- .../core/propagation/ptags/PTagsFactory.java | 24 ++++-- .../ptags/KnuthSamplingRateCacheTest.java | 79 +++++++++++++++++++ 2 files changed, 96 insertions(+), 7 deletions(-) create mode 100644 dd-trace-core/src/test/java/datadog/trace/core/propagation/ptags/KnuthSamplingRateCacheTest.java 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 index 066727c3baf..71b59bf2105 100644 --- 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 @@ -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. * diff --git a/dd-trace-core/src/test/java/datadog/trace/core/propagation/ptags/KnuthSamplingRateCacheTest.java b/dd-trace-core/src/test/java/datadog/trace/core/propagation/ptags/KnuthSamplingRateCacheTest.java new file mode 100644 index 00000000000..c150e574648 --- /dev/null +++ b/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 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(); + } + } +}