Skip to content

WIP: chore(sampling): track knuth sampling in distributed trace #9518

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Draft
wants to merge 5 commits into
base: master
Choose a base branch
from
Draft

WIP: chore(sampling): track knuth sampling in distributed trace #9518

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
5 commits
Select commit Hold shift + click to select a range
9f7317e
WIP: chore(sampling): track knuth sampling in distributed trace
rachelyangdog Sep 11, 2025
b6e1df9
lint build errors
rachelyangdog Sep 15, 2025
c6ef173
lint build errors
rachelyangdog Sep 15, 2025
ad8190f
Merge branch 'master' into rachel.yang/knuth-sampling-tag
rachelyangdog Sep 25, 2025
2ac621f
fix formatting
rachelyangdog Oct 1, 2025
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
1 change: 1 addition & 0 deletions dd-java-agent/testing/src/main/groovy/datadog/trace/agent/test/asserts/TagsAssert.groovy
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -86,6 +86,7 @@ class TagsAssert {
assertedTags.add(RateByServiceTraceSampler.SAMPLING_AGENT_RATE)
assertedTags.add(TraceMapper.SAMPLING_PRIORITY_KEY.toString())
assertedTags.add("_sample_rate")
assertedTags.add(RateByServiceTraceSampler.KNUTH_SAMPLING_RATE)
assertedTags.add(DDTags.PID_TAG)
assertedTags.add(DDTags.SCHEMA_VERSION_TAG_KEY)
assertedTags.add(DDTags.PROFILING_ENABLED)
Expand Down
38 changes: 38 additions & 0 deletions dd-trace-core/src/main/java/datadog/trace/common/sampling/DeterministicSampler.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -40,10 +40,12 @@ protected <T extends CoreSpan<T>> long getSamplingId(T span) {

private final float rate;
private final long threshold;
private final String knuthSampleRate;

public DeterministicSampler(final double rate) {
this.rate = (float) rate;
this.threshold = cutoff(rate);
this.knuthSampleRate = formatKnuthSamplingRate(rate);
}

@Override
Expand All @@ -59,6 +61,11 @@ public double getSampleRate() {
return rate;
}

@Override
public String getKnuthSampleRate() {
return knuthSampleRate;
}

public static long cutoff(double rate) {
if (rate < 0.5) {
return (long) (rate * MAX) + Long.MIN_VALUE;
Expand All @@ -68,4 +75,35 @@ public static long cutoff(double rate) {
}
return Long.MAX_VALUE;
}

/**
* Thread-safe custom formatter for Knuth sampling rates. Formats rates to up to 6 decimal places,
* removing trailing zeros. Assumes the value is between 0 and 1 (inclusive).
*/
private static String formatKnuthSamplingRate(double value) {
if (value <= 0) {
return "0";
} else if (value >= 1) {
return "1";
} else {
// Scale to 6 decimal places and round
long scaled = Math.round(value * 1_000_000);

// Handle rounding to 1.0 case
if (scaled >= 1_000_000) {
return "1";
}

// Convert back to string with proper decimal formatting
String result = "0." + String.format("%06d", scaled);

// Remove trailing zeros
int endIndex = result.length();
while (endIndex > 2 && result.charAt(endIndex - 1) == '0') {
endIndex--;
}

return result.substring(0, endIndex);
}
}
}
19 changes: 19 additions & 0 deletions dd-trace-core/src/main/java/datadog/trace/common/sampling/RateByServiceTraceSampler.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -6,7 +6,10 @@
import datadog.trace.api.sampling.SamplingMechanism;
import datadog.trace.common.writer.RemoteResponseListener;
import datadog.trace.core.CoreSpan;
import java.text.DecimalFormat;
import java.text.DecimalFormatSymbols;
import java.util.Collections;
import java.util.Locale;
import java.util.Map;
import java.util.TreeMap;
import java.util.function.Function;
Expand All @@ -22,8 +25,15 @@ public class RateByServiceTraceSampler implements Sampler, PrioritySampler, Remo

private static final Logger log = LoggerFactory.getLogger(RateByServiceTraceSampler.class);
public static final String SAMPLING_AGENT_RATE = "_dd.agent_psr";
public static final String KNUTH_SAMPLING_RATE = "_dd.p.ksr";

private static final double DEFAULT_RATE = 1.0;
private static final DecimalFormat DECIMAL_FORMAT;

static {
DECIMAL_FORMAT = new DecimalFormat("0", DecimalFormatSymbols.getInstance(Locale.ENGLISH));
DECIMAL_FORMAT.setMaximumFractionDigits(6);
}

private volatile RateSamplersByEnvAndService serviceRates = new RateSamplersByEnvAndService();

Expand Down Expand Up @@ -56,12 +66,21 @@ public <T extends CoreSpan<T>> void setSamplingPriority(final T span) {
sampler.getSampleRate(),
SamplingMechanism.AGENT_RATE);
}

// Set Knuth sampling rate tag
String ksrRate = formatKnuthSamplingRate(sampler.getSampleRate());
span.setTag(KNUTH_SAMPLING_RATE, ksrRate);
}

private <T extends CoreSpan<T>> String getSpanEnv(final T span) {
return span.getTag("env", "");
}

private String formatKnuthSamplingRate(double rate) {
Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Should centralize this logic. Currently we define this function in two places

// Format to up to 6 decimal places, removing trailing zeros
return DECIMAL_FORMAT.format(rate);
Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Quick heads-up that java.text.DecimalFormat.format(...) is not thread-safe: https://docs.oracle.com/javase/8/docs/api/java/text/DecimalFormat.html#synchronization

But adding synchronization here would kill performance....

One option could be to write our own simple formatter, since we just want to truncate the value - for example, if we assume the value is between 0 and 1, a naive solution might be:

  private static String example(double value) {
    if (value <= 0) {
      return "0";
    } else if (value >= 1) {
      return "1";
    } else {
      return "0." + (int) (value * 1_000_000 + 0.5);
    }
  }

( you'd need to write tests to verify this behaves as expected - but you'd need to do that anyway :)

Also note you don't need to do the formatting every time the sampling rate is applied because each RateSampler is deterministic: https://github.com/DataDog/dd-trace-java/blob/master/dd-trace-core/src/main/java/datadog/trace/common/sampling/DeterministicSampler.java#L41

In other words, the sample rate for each individual RateSampler is fixed at construct time - so you could add a String getKnuthSampleRate() or similar method to RateSampler and then do the formatting once in the DeterministicSampler constructor and save it as a String field alongside the rate, at which point setting the tag just involves calling getKnuthSampleRate() to get the pre-formatted string.

}

@Override
public void onResponse(
final String endpoint, final Map<String, Map<String, Number>> responseJson) {
Expand Down
7 changes: 7 additions & 0 deletions dd-trace-core/src/main/java/datadog/trace/common/sampling/RateSampler.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -2,4 +2,11 @@

public interface RateSampler extends Sampler {
double getSampleRate();

/**
* Returns the pre-formatted Knuth sampling rate as a string. This is computed once at
* construction time to avoid thread-safety issues and performance overhead of formatting on each
* sampling operation.
*/
String getKnuthSampleRate();
}
5 changes: 5 additions & 0 deletions dd-trace-core/src/main/java/datadog/trace/common/sampling/RuleBasedTraceSampler.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -15,13 +15,15 @@
public class RuleBasedTraceSampler<T extends CoreSpan<T>> implements Sampler, PrioritySampler {

private static final Logger log = LoggerFactory.getLogger(RuleBasedTraceSampler.class);

private final List<RateSamplingRule> samplingRules;
private final PrioritySampler fallbackSampler;
private final SimpleRateLimiter rateLimiter;
private final long rateLimit;

public static final String SAMPLING_RULE_RATE = "_dd.rule_psr";
public static final String SAMPLING_LIMIT_RATE = "_dd.limit_psr";
public static final String KNUTH_SAMPLING_RATE = "_dd.p.ksr";

public RuleBasedTraceSampler(
final List<RateSamplingRule> samplingRules,
Expand Down Expand Up @@ -168,6 +170,9 @@ public <T extends CoreSpan<T>> void setSamplingPriority(final T span) {
matchedRule.getSampler().getSampleRate(),
matchedRule.getMechanism());
}

// Set Knuth sampling rate tag
span.setTag(KNUTH_SAMPLING_RATE, matchedRule.getSampler().getKnuthSampleRate());
}
}
}
58 changes: 58 additions & 0 deletions ...e-core/src/test/groovy/datadog/trace/common/sampling/DeterministicTraceSamplerTest.groovy
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -376,4 +376,62 @@ class DeterministicTraceSamplerTest extends DDSpecification {
where:
rate << (0..100)
}

def "test getKnuthSampleRate formatting: #rate -> #expected"() {
given:
DeterministicSampler sampler = new DeterministicSampler.TraceSampler(rate)

when:
String formatted = sampler.getKnuthSampleRate()

then:
formatted == expected

where:
rate | expected
0.0 | "0"
0.000001 | "0.000001"
0.000010 | "0.00001"
0.000100 | "0.0001"
0.001000 | "0.001"
0.010000 | "0.01"
0.100000 | "0.1"
0.123456 | "0.123456"
0.123450 | "0.12345"
0.123400 | "0.1234"
0.123000 | "0.123"
0.120000 | "0.12"
0.100000 | "0.1"
0.500000 | "0.5"
0.999999 | "0.999999"
1.0 | "1"
1.5 | "1" // Values > 1 are clamped to 1
-0.5 | "0" // Values < 0 are clamped to 0
}

def "test getKnuthSampleRate precision and rounding"() {
given:
// Test edge cases around rounding
DeterministicSampler sampler1 = new DeterministicSampler.TraceSampler(0.1234564) // Should round down
DeterministicSampler sampler2 = new DeterministicSampler.TraceSampler(0.1234565) // Should round up
DeterministicSampler sampler3 = new DeterministicSampler.TraceSampler(0.9999995) // Should round up to 1

expect:
sampler1.getKnuthSampleRate() == "0.123456"
sampler2.getKnuthSampleRate() == "0.123457"
sampler3.getKnuthSampleRate() == "1"
}

def "test getKnuthSampleRate is consistent"() {
given:
DeterministicSampler sampler = new DeterministicSampler.TraceSampler(0.123456)

when:
String rate1 = sampler.getKnuthSampleRate()
String rate2 = sampler.getKnuthSampleRate()

then:
rate1 == rate2
rate1 == "0.123456"
}
}