Skip to content

Commit a3ca3f2

Browse files
mhliddclaude
andcommitted
Fix agent sampling rates being ignored when sampling rules are configured
When any sampling rule or default rate is configured, Sampler.Builder.forConfig returns a RuleBasedTraceSampler, which delegates to an inner RateByServiceTraceSampler for spans that match no rule. That inner sampler was never registered to receive the rates published by the trace agent: WriterFactory registers the sampler only when it is a RemoteResponseListener, and RuleBasedTraceSampler neither implemented that interface nor forwarded onResponse to its fallback. The fallback therefore stayed at its initial rate of 1.0 for the life of the JVM, keeping every rule miss and stamping _dd.agent_psr=1.0 with the AGENT_RATE mechanism. Remote configuration compounded this: a change to the trace sample rate or sampling rules rebuilds the sampler in CoreTracer.ConfigSnapshot, but the response listener is registered once, against the initial sampler, and is never re-registered, so the rebuilt sampler was disconnected as well. Two changes: - RuleBasedTraceSampler implements RemoteResponseListener and forwards agent rates to its fallback sampler. - A single RateByServiceTraceSampler instance is shared by every sampler the tracer builds. Sampler.Builder.forConfig takes it as an argument, and CoreTracer reuses the instance from the initial sampler whenever remote configuration triggers a rebuild, so rebuilt samplers stay connected to the agent and keep the rates already learned. Go, Node, Python, Ruby and .NET all apply agent rates on a rule miss; Java was the only tracer that did not. Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>
1 parent e304fe0 commit a3ca3f2

5 files changed

Lines changed: 380 additions & 5 deletions

File tree

‎dd-trace-core/src/main/java/datadog/trace/common/sampling/RuleBasedTraceSampler.java‎

Lines changed: 53 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -3,6 +3,7 @@
33
import datadog.trace.api.config.TracerConfig;
44
import datadog.trace.api.sampling.SamplingMechanism;
55
import datadog.trace.api.sampling.SamplingRule;
6+
import datadog.trace.common.writer.RemoteResponseListener;
67
import datadog.trace.core.CoreSpan;
78
import datadog.trace.core.util.SimpleRateLimiter;
89
import java.util.ArrayList;
@@ -12,7 +13,8 @@
1213
import org.slf4j.Logger;
1314
import org.slf4j.LoggerFactory;
1415

15-
public class RuleBasedTraceSampler<T extends CoreSpan<T>> implements Sampler, PrioritySampler {
16+
public class RuleBasedTraceSampler<T extends CoreSpan<T>>
17+
implements Sampler, PrioritySampler, RemoteResponseListener {
1618

1719
private static final Logger log = LoggerFactory.getLogger(RuleBasedTraceSampler.class);
1820
private final List<RateSamplingRule> samplingRules;
@@ -47,6 +49,31 @@ public static RuleBasedTraceSampler build(
4749
final List<? extends SamplingRule.TraceSamplingRule> traceSamplingRules,
4850
final Double defaultRate,
4951
final int rateLimit) {
52+
return build(
53+
serviceRules,
54+
operationRules,
55+
traceSamplingRules,
56+
defaultRate,
57+
rateLimit,
58+
new RateByServiceTraceSampler());
59+
}
60+
61+
/**
62+
* Builds a rule based sampler that delegates to {@code fallbackSampler} for spans that match no
63+
* rule.
64+
*
65+
* <p>The fallback is supplied by the caller so that a single agent rate sampler instance can be
66+
* shared across sampler rebuilds triggered by remote configuration. That instance is the one
67+
* registered to receive agent sampling rates, so reusing it keeps the rates that have already
68+
* been learned and keeps the rule miss path connected to the agent.
69+
*/
70+
public static RuleBasedTraceSampler build(
71+
@Deprecated final Map<String, String> serviceRules,
72+
@Deprecated final Map<String, String> operationRules,
73+
final List<? extends SamplingRule.TraceSamplingRule> traceSamplingRules,
74+
final Double defaultRate,
75+
final int rateLimit,
76+
final PrioritySampler fallbackSampler) {
5077

5178
final List<RateSamplingRule> samplingRules = new ArrayList<>();
5279

@@ -113,7 +140,31 @@ public static RuleBasedTraceSampler build(
113140
samplingRules.add(samplingRule);
114141
}
115142

116-
return new RuleBasedTraceSampler(samplingRules, rateLimit, new RateByServiceTraceSampler());
143+
return new RuleBasedTraceSampler(samplingRules, rateLimit, fallbackSampler);
144+
}
145+
146+
/**
147+
* Returns the sampler applying agent published rates to spans that match no rule, or {@code null}
148+
* when the fallback does not use agent rates.
149+
*/
150+
public RateByServiceTraceSampler agentSampler() {
151+
return fallbackSampler instanceof RateByServiceTraceSampler
152+
? (RateByServiceTraceSampler) fallbackSampler
153+
: null;
154+
}
155+
156+
/**
157+
* Forwards agent sampling rates to the fallback sampler.
158+
*
159+
* <p>Without this the inner {@link RateByServiceTraceSampler} would never observe the rates
160+
* published by the agent, and every span matching no rule would be kept at a rate of 1.0.
161+
*/
162+
@Override
163+
public void onResponse(
164+
final String endpoint, final Map<String, Map<String, Number>> responseJson) {
165+
if (fallbackSampler instanceof RemoteResponseListener) {
166+
((RemoteResponseListener) fallbackSampler).onResponse(endpoint, responseJson);
167+
}
117168
}
118169

119170
private static byte samplingMechanism(SamplingRule.Provenance provenance) {

‎dd-trace-core/src/main/java/datadog/trace/common/sampling/Sampler.java‎

Lines changed: 18 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -34,6 +34,21 @@ final class Builder {
3434
private static final Logger log = LoggerFactory.getLogger(Builder.class);
3535

3636
public static Sampler forConfig(final Config config, final TraceConfig traceConfig) {
37+
return forConfig(config, traceConfig, new RateByServiceTraceSampler());
38+
}
39+
40+
/**
41+
* Builds the sampler described by {@code config} and {@code traceConfig}, using {@code
42+
* agentSampler} wherever agent published sampling rates apply.
43+
*
44+
* <p>Callers that rebuild the sampler when remote configuration changes must pass the same
45+
* {@code agentSampler} instance every time: it is the instance registered to receive agent
46+
* sampling rates, and a fresh one would silently start over at a rate of 1.0.
47+
*/
48+
public static Sampler forConfig(
49+
final Config config,
50+
final TraceConfig traceConfig,
51+
final RateByServiceTraceSampler agentSampler) {
3752
Sampler sampler;
3853
if (config != null) {
3954
if (!config.isApmTracingEnabled() && isAsmEnabled(config)) {
@@ -75,7 +90,8 @@ public static Sampler forConfig(final Config config, final TraceConfig traceConf
7590
operationRules,
7691
traceSamplingRules,
7792
traceSampleRate,
78-
config.getTraceRateLimit());
93+
config.getTraceRateLimit(),
94+
agentSampler);
7995
} catch (final IllegalArgumentException e) {
8096
log.error("Invalid sampler configuration. Using AllSampler", e);
8197
sampler = new AllSampler();
@@ -95,7 +111,7 @@ public static Sampler forConfig(final Config config, final TraceConfig traceConf
95111
"OTLP traces export enabled. Using ParentBasedAlwaysOnSampler instead of RateByServiceTraceSampler.");
96112
sampler = new ParentBasedAlwaysOnSampler();
97113
} else {
98-
sampler = new RateByServiceTraceSampler();
114+
sampler = agentSampler;
99115
}
100116
} else if (config.isTraceOtlpExporterEnabled()) {
101117
// AllSampler does not emit a sampling priority; OTLP export requires one.

‎dd-trace-core/src/main/java/datadog/trace/core/CoreTracer.java‎

Lines changed: 37 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -85,6 +85,8 @@
8585
import datadog.trace.common.GitMetadataTraceInterceptor;
8686
import datadog.trace.common.metrics.MetricsAggregator;
8787
import datadog.trace.common.metrics.NoOpMetricsAggregator;
88+
import datadog.trace.common.sampling.RateByServiceTraceSampler;
89+
import datadog.trace.common.sampling.RuleBasedTraceSampler;
8890
import datadog.trace.common.sampling.Sampler;
8991
import datadog.trace.common.sampling.SingleSpanSampler;
9092
import datadog.trace.common.sampling.SpanSamplingRules;
@@ -181,6 +183,17 @@ public static CoreTracerBuilder builder() {
181183
/** Sampler defines the sampling policy in order to reduce the number of traces for instance */
182184
final Sampler initialSampler;
183185

186+
/**
187+
* The sampler applying the rates published by the agent, shared by every sampler this tracer
188+
* builds.
189+
*
190+
* <p>Only the initial sampler is registered to receive agent rates, and that registration is
191+
* never repeated. Reusing this instance when the sampler is rebuilt on a remote configuration
192+
* change keeps the rebuilt sampler connected to those rates, and keeps the rates already learned.
193+
* May be {@code null} when the tracer was given a sampler that does not use agent rates.
194+
*/
195+
final RateByServiceTraceSampler agentSampler;
196+
184197
/** Scope manager is in charge of managing the scopes from which spans are created */
185198
final ContinuableScopeManager scopeManager;
186199

@@ -679,6 +692,7 @@ private CoreTracer(
679692
this.initialConfig = config;
680693
this.apmTracingEnabled = config.isApmTracingEnabled();
681694
this.initialSampler = sampler;
695+
this.agentSampler = agentSamplerOf(sampler);
682696

683697
// Get initial Trace Sampling Rules from config
684698
String traceSamplingRulesJson = config.getTraceSamplingRules();
@@ -1641,6 +1655,22 @@ private static <K, V> Map<V, K> invertMap(Map<K, V> map) {
16411655
return Collections.unmodifiableMap(inverted);
16421656
}
16431657

1658+
/**
1659+
* Extracts the sampler applying agent published rates from the sampler this tracer was built
1660+
* with, so that it can be reused whenever the sampler is rebuilt.
1661+
*
1662+
* @return the agent rate sampler in use, or {@code null} if this tracer does not use one.
1663+
*/
1664+
private static RateByServiceTraceSampler agentSamplerOf(final Sampler sampler) {
1665+
if (sampler instanceof RateByServiceTraceSampler) {
1666+
return (RateByServiceTraceSampler) sampler;
1667+
}
1668+
if (sampler instanceof RuleBasedTraceSampler) {
1669+
return ((RuleBasedTraceSampler<?>) sampler).agentSampler();
1670+
}
1671+
return null;
1672+
}
1673+
16441674
/** Spans are built using this builder */
16451675
public abstract static class CoreSpanBuilder implements AgentTracer.SpanBuilder {
16461676
protected static final boolean USE_SCOPE = false;
@@ -2547,7 +2577,13 @@ protected ConfigSnapshot(
25472577
&& Objects.equals(getTraceSamplingRules(), oldSnapshot.getTraceSamplingRules())) {
25482578
sampler = oldSnapshot.sampler;
25492579
} else {
2550-
sampler = Sampler.Builder.forConfig(CoreTracer.this.initialConfig, this);
2580+
// Reuse the agent sampler so the rebuilt sampler keeps receiving agent rates: the response
2581+
// listener is registered once, against the initial sampler, and is never re-registered.
2582+
RateByServiceTraceSampler agentSampler = CoreTracer.this.agentSampler;
2583+
sampler =
2584+
agentSampler == null
2585+
? Sampler.Builder.forConfig(CoreTracer.this.initialConfig, this)
2586+
: Sampler.Builder.forConfig(CoreTracer.this.initialConfig, this, agentSampler);
25512587
}
25522588

25532589
if (null == oldSnapshot) {
Lines changed: 180 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,180 @@
1+
package datadog.trace.common.sampling;
2+
3+
import static datadog.trace.api.config.TracerConfig.TRACE_SAMPLE_RATE;
4+
import static datadog.trace.api.config.TracerConfig.TRACE_SAMPLING_RULES;
5+
import static datadog.trace.api.sampling.PrioritySampling.SAMPLER_DROP;
6+
import static datadog.trace.api.sampling.PrioritySampling.SAMPLER_KEEP;
7+
import static datadog.trace.api.sampling.PrioritySampling.USER_KEEP;
8+
import static org.junit.jupiter.api.Assertions.assertEquals;
9+
import static org.junit.jupiter.api.Assertions.assertInstanceOf;
10+
import static org.junit.jupiter.api.Assertions.assertNull;
11+
import static org.junit.jupiter.api.Assertions.assertSame;
12+
13+
import datadog.trace.api.Config;
14+
import datadog.trace.common.writer.ListWriter;
15+
import datadog.trace.common.writer.RemoteResponseListener;
16+
import datadog.trace.core.CoreTracer;
17+
import datadog.trace.core.DDCoreJavaSpecification;
18+
import datadog.trace.core.DDSpan;
19+
import java.util.HashMap;
20+
import java.util.Map;
21+
import java.util.Properties;
22+
import org.junit.jupiter.api.Test;
23+
24+
/**
25+
* Agent published sampling rates must still apply to spans that match no sampling rule.
26+
*
27+
* <p>Only one sampler is registered to receive those rates, so a rule based sampler has to forward
28+
* them to the fallback it delegates to.
29+
*/
30+
class RuleBasedSamplerAgentRatesTest extends DDCoreJavaSpecification {
31+
32+
private static final String RULE_MISSING_SERVICE =
33+
"[{\"service\": \"other-service\", \"sample_rate\": 1}]";
34+
35+
@Test
36+
void ruleBasedSamplerForwardsAgentRatesToFallback() {
37+
Properties properties = new Properties();
38+
properties.setProperty(TRACE_SAMPLING_RULES, RULE_MISSING_SERVICE);
39+
Sampler sampler = Sampler.Builder.forConfig(properties);
40+
assertInstanceOf(RuleBasedTraceSampler.class, sampler);
41+
42+
// The agent asks for everything to be dropped for this service.
43+
assertInstanceOf(RemoteResponseListener.class, sampler);
44+
((RemoteResponseListener) sampler).onResponse("traces", rateByService("service", "bar", 0.0));
45+
46+
CoreTracer tracer = tracerBuilder().writer(new ListWriter()).build();
47+
try {
48+
DDSpan span = buildSpan(tracer, "service");
49+
((PrioritySampler) sampler).setSamplingPriority(span);
50+
51+
assertEquals(SAMPLER_DROP, (int) span.getSamplingPriority());
52+
assertEquals(0.0, span.getTag(RateByServiceTraceSampler.SAMPLING_AGENT_RATE));
53+
} finally {
54+
tracer.close();
55+
}
56+
}
57+
58+
@Test
59+
void agentRatesDoNotOverrideMatchedRules() {
60+
Properties properties = new Properties();
61+
properties.setProperty(
62+
TRACE_SAMPLING_RULES, "[{\"service\": \"service\", \"sample_rate\": 1}]");
63+
Sampler sampler = Sampler.Builder.forConfig(properties);
64+
65+
((RemoteResponseListener) sampler).onResponse("traces", rateByService("service", "bar", 0.0));
66+
67+
CoreTracer tracer = tracerBuilder().writer(new ListWriter()).build();
68+
try {
69+
DDSpan span = buildSpan(tracer, "service");
70+
((PrioritySampler) sampler).setSamplingPriority(span);
71+
72+
assertEquals(USER_KEEP, (int) span.getSamplingPriority());
73+
assertEquals(1.0, span.getTag(RuleBasedTraceSampler.SAMPLING_RULE_RATE));
74+
assertNull(span.getTag(RateByServiceTraceSampler.SAMPLING_AGENT_RATE));
75+
} finally {
76+
tracer.close();
77+
}
78+
}
79+
80+
@Test
81+
void defaultRateKeepsAgentRatesUnused() {
82+
Properties properties = new Properties();
83+
properties.setProperty(TRACE_SAMPLING_RULES, RULE_MISSING_SERVICE);
84+
// A default rate matches every span, so the fallback is never reached.
85+
properties.setProperty(TRACE_SAMPLE_RATE, "1");
86+
Sampler sampler = Sampler.Builder.forConfig(properties);
87+
88+
((RemoteResponseListener) sampler).onResponse("traces", rateByService("service", "bar", 0.0));
89+
90+
CoreTracer tracer = tracerBuilder().writer(new ListWriter()).build();
91+
try {
92+
DDSpan span = buildSpan(tracer, "service");
93+
((PrioritySampler) sampler).setSamplingPriority(span);
94+
95+
assertEquals(USER_KEEP, (int) span.getSamplingPriority());
96+
assertNull(span.getTag(RateByServiceTraceSampler.SAMPLING_AGENT_RATE));
97+
} finally {
98+
tracer.close();
99+
}
100+
}
101+
102+
@Test
103+
void rulesKeepTheSuppliedAgentSampler() {
104+
RateByServiceTraceSampler agentSampler = new RateByServiceTraceSampler();
105+
Properties properties = new Properties();
106+
properties.setProperty(TRACE_SAMPLING_RULES, RULE_MISSING_SERVICE);
107+
108+
Sampler sampler = Sampler.Builder.forConfig(Config.get(properties), null, agentSampler);
109+
110+
assertInstanceOf(RuleBasedTraceSampler.class, sampler);
111+
assertSame(agentSampler, ((RuleBasedTraceSampler<?>) sampler).agentSampler());
112+
}
113+
114+
@Test
115+
void noRulesReturnsTheSuppliedAgentSampler() {
116+
RateByServiceTraceSampler agentSampler = new RateByServiceTraceSampler();
117+
118+
Sampler sampler = Sampler.Builder.forConfig(Config.get(new Properties()), null, agentSampler);
119+
120+
assertSame(agentSampler, sampler);
121+
}
122+
123+
@Test
124+
void ratesLearnedBeforeARebuildStillApplyAfterIt() {
125+
RateByServiceTraceSampler agentSampler = new RateByServiceTraceSampler();
126+
agentSampler.onResponse("traces", rateByService("service", "bar", 0.0));
127+
128+
// Rebuilding the sampler with the same agent sampler keeps the rates it has already learned.
129+
Properties properties = new Properties();
130+
properties.setProperty(TRACE_SAMPLING_RULES, RULE_MISSING_SERVICE);
131+
Sampler sampler = Sampler.Builder.forConfig(Config.get(properties), null, agentSampler);
132+
133+
CoreTracer tracer = tracerBuilder().writer(new ListWriter()).build();
134+
try {
135+
DDSpan span = buildSpan(tracer, "service");
136+
((PrioritySampler) sampler).setSamplingPriority(span);
137+
138+
assertEquals(SAMPLER_DROP, (int) span.getSamplingPriority());
139+
} finally {
140+
tracer.close();
141+
}
142+
}
143+
144+
@Test
145+
void unknownServiceIsKeptUntilTheAgentReportsARate() {
146+
Properties properties = new Properties();
147+
properties.setProperty(TRACE_SAMPLING_RULES, RULE_MISSING_SERVICE);
148+
Sampler sampler = Sampler.Builder.forConfig(properties);
149+
150+
CoreTracer tracer = tracerBuilder().writer(new ListWriter()).build();
151+
try {
152+
DDSpan span = buildSpan(tracer, "service");
153+
((PrioritySampler) sampler).setSamplingPriority(span);
154+
155+
assertEquals(SAMPLER_KEEP, (int) span.getSamplingPriority());
156+
assertEquals(1.0, span.getTag(RateByServiceTraceSampler.SAMPLING_AGENT_RATE));
157+
} finally {
158+
tracer.close();
159+
}
160+
}
161+
162+
private static DDSpan buildSpan(CoreTracer tracer, String serviceName) {
163+
return (DDSpan)
164+
tracer
165+
.buildSpan("datadog", "operation")
166+
.withServiceName(serviceName)
167+
.withTag("env", "bar")
168+
.ignoreActiveSpan()
169+
.start();
170+
}
171+
172+
private static Map<String, Map<String, Number>> rateByService(
173+
String service, String env, double rate) {
174+
Map<String, Number> byService = new HashMap<>();
175+
byService.put("service:" + service + ",env:" + env, rate);
176+
Map<String, Map<String, Number>> response = new HashMap<>();
177+
response.put("rate_by_service", byService);
178+
return response;
179+
}
180+
}

0 commit comments

Comments
 (0)