Skip to content

Commit 98f5d1c

Browse files
andrewlockbouwkast
andrewlock
and
bouwkast
authored
Refactor TraceSampler to separate building from evaluation (#7311)
## Summary of changes Refactors the `TraceSampler` and `ITraceSampler` abstractions to separate the building of the rules from the evaluation ## Reason for change It's neater, and it makes some subsequent changes much simpler, easier and more performant ## Implementation details Create a nested `Builder` type in `TraceSampler` and move all the `RegisterRule()` methods on it. Call `Build` to get a final list of rules to be applied. ## Test coverage Covered by existing tests ## Other details Part of a stack of branches related to sampling improvements - #7311 👈 - #7312 --------- Co-authored-by: Steven Bouwkamp <[email protected]>
1 parent 5651b4c commit 98f5d1c

File tree

5 files changed

+90
-83
lines changed

5 files changed

+90
-83
lines changed

‎tracer/src/Datadog.Trace/Ci/Sampling/CISampler.cs

Lines changed: 0 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -16,10 +16,6 @@ public SamplingDecision MakeSamplingDecision(Span span)
1616
return new SamplingDecision(SamplingPriorityValues.UserKeep, mechanism: null, rate: null, limiterRate: null);
1717
}
1818

19-
public void RegisterRule(ISamplingRule rule)
20-
{
21-
}
22-
2319
public void SetDefaultSampleRates(IReadOnlyDictionary<string, float> sampleRates)
2420
{
2521
}

‎tracer/src/Datadog.Trace/Sampling/ITraceSampler.cs

Lines changed: 0 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -14,7 +14,5 @@ internal interface ITraceSampler
1414
void SetDefaultSampleRates(IReadOnlyDictionary<string, float> sampleRates);
1515

1616
SamplingDecision MakeSamplingDecision(Span span);
17-
18-
void RegisterRule(ISamplingRule rule);
1917
}
2018
}

‎tracer/src/Datadog.Trace/Sampling/TraceSampler.cs

Lines changed: 66 additions & 55 deletions
Original file line numberDiff line numberDiff line change
@@ -17,13 +17,14 @@ internal class TraceSampler : ITraceSampler
1717
private static readonly IDatadogLogger Log = DatadogLogging.GetLoggerFor<TraceSampler>();
1818

1919
private readonly IRateLimiter _limiter;
20-
private readonly List<ISamplingRule> _rules = [];
20+
private readonly List<ISamplingRule> _rules;
21+
private readonly AgentSamplingRule? _agentSamplingRule;
2122

22-
private AgentSamplingRule? _agentSamplingRule;
23-
24-
public TraceSampler(IRateLimiter limiter)
23+
public TraceSampler(IRateLimiter limiter, List<ISamplingRule> rules, AgentSamplingRule? agentSamplingRule)
2524
{
2625
_limiter = limiter;
26+
_rules = [..rules];
27+
_agentSamplingRule = agentSamplingRule;
2728
}
2829

2930
public void SetDefaultSampleRates(IReadOnlyDictionary<string, float> sampleRates)
@@ -48,57 +49,6 @@ public SamplingDecision MakeSamplingDecision(Span span)
4849
return SamplingDecision.Default;
4950
}
5051

51-
/// <summary>
52-
/// Register a new sampling rule. To register a <see cref="AgentSamplingRule"/>,
53-
/// use <see cref="RegisterAgentSamplingRule"/> instead.
54-
/// </summary>
55-
/// <remarks>
56-
/// The order that rules are registered is important, as they are evaluated in order.
57-
/// The first rule that matches will be used to determine the sampling rate.
58-
/// </remarks>
59-
public void RegisterRule(ISamplingRule rule)
60-
{
61-
_rules.Add(rule);
62-
}
63-
64-
/// <summary>
65-
/// Register new sampling rules. To register a <see cref="AgentSamplingRule"/>,
66-
/// use <see cref="RegisterAgentSamplingRule"/> instead.
67-
/// </summary>
68-
/// <remarks>
69-
/// The order that rules are registered is important, as they are evaluated in order.
70-
/// The first rule that matches will be used to determine the sampling rate.
71-
/// </remarks>
72-
public void RegisterRules(IEnumerable<ISamplingRule> rules)
73-
{
74-
_rules.AddRange(rules);
75-
}
76-
77-
/// <summary>
78-
/// Register a new agent sampling rule. This rule should be registered last,
79-
/// after any calls to <see cref="RegisterRule"/> or <see cref="RegisterRules"/>.
80-
/// </summary>
81-
/// <remarks>
82-
/// The order that rules are registered is important, as they are evaluated in order.
83-
/// The first rule that matches will be used to determine the sampling rate.
84-
/// </remarks>
85-
public void RegisterAgentSamplingRule(AgentSamplingRule rule)
86-
{
87-
// only register the one AgentSamplingRule
88-
if (Interlocked.Exchange(ref _agentSamplingRule, rule) == null)
89-
{
90-
// keep a reference to this rule so we can call SetDefaultSampleRates() later
91-
// to update the agent sampling rates
92-
_agentSamplingRule = rule;
93-
94-
RegisterRule(rule);
95-
}
96-
else
97-
{
98-
Log.Warning("AgentSamplingRule already registered. Ignoring additional registration.");
99-
}
100-
}
101-
10252
// used for testing
10353
internal IReadOnlyList<ISamplingRule> GetRules()
10454
{
@@ -137,5 +87,66 @@ private SamplingDecision MakeSamplingDecision(Span span, float rate, string mech
13787

13888
return new SamplingDecision(priority, mechanism, rate, limiterRate);
13989
}
90+
91+
public class Builder(IRateLimiter limiter)
92+
{
93+
private readonly IRateLimiter _limiter = limiter;
94+
private readonly List<ISamplingRule> _rules = [];
95+
private AgentSamplingRule? _agentSamplingRule;
96+
97+
public TraceSampler Build()
98+
{
99+
return new TraceSampler(_limiter, _rules, _agentSamplingRule);
100+
}
101+
102+
/// <summary>
103+
/// Register a new sampling rule. To register a <see cref="AgentSamplingRule"/>,
104+
/// use <see cref="RegisterAgentSamplingRule"/> instead.
105+
/// </summary>
106+
/// <remarks>
107+
/// The order that rules are registered is important, as they are evaluated in order.
108+
/// The first rule that matches will be used to determine the sampling rate.
109+
/// </remarks>
110+
public void RegisterRule(ISamplingRule rule)
111+
{
112+
_rules.Add(rule);
113+
}
114+
115+
/// <summary>
116+
/// Register new sampling rules. To register a <see cref="AgentSamplingRule"/>,
117+
/// use <see cref="RegisterAgentSamplingRule"/> instead.
118+
/// </summary>
119+
/// <remarks>
120+
/// The order that rules are registered is important, as they are evaluated in order.
121+
/// The first rule that matches will be used to determine the sampling rate.
122+
/// </remarks>
123+
public void RegisterRules(IEnumerable<ISamplingRule> rules)
124+
{
125+
_rules.AddRange(rules);
126+
}
127+
128+
/// <summary>
129+
/// Register a new agent sampling rule. This rule should be registered last,
130+
/// after any calls to <see cref="RegisterRule"/> or <see cref="RegisterRules"/>.
131+
/// </summary>
132+
/// <remarks>
133+
/// The order that rules are registered is important, as they are evaluated in order.
134+
/// The first rule that matches will be used to determine the sampling rate.
135+
/// </remarks>
136+
public void RegisterAgentSamplingRule(AgentSamplingRule rule)
137+
{
138+
// only register the one AgentSamplingRule
139+
// keep a reference to this rule so we can call SetDefaultSampleRates() later
140+
// to update the agent sampling rates
141+
if (Interlocked.Exchange(ref _agentSamplingRule, rule) == null)
142+
{
143+
RegisterRule(rule);
144+
}
145+
else
146+
{
147+
Log.Warning("AgentSamplingRule already registered. Ignoring additional registration.");
148+
}
149+
}
150+
}
140151
}
141152
}

‎tracer/src/Datadog.Trace/TracerManagerFactory.cs

Lines changed: 4 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -275,12 +275,12 @@ protected virtual ITraceSampler GetSampler(TracerSettings settings)
275275
(Security.Instance.Settings.AppsecEnabled || Iast.Iast.Instance.Settings.Enabled))
276276
{
277277
// Custom sampler for ASM and IAST standalone billing mode
278-
var samplerStandalone = new TraceSampler(new TracerRateLimiter(maxTracesPerInterval: 1, intervalMilliseconds: 60_000));
278+
var samplerStandalone = new TraceSampler.Builder(new TracerRateLimiter(maxTracesPerInterval: 1, intervalMilliseconds: 60_000));
279279
samplerStandalone.RegisterRule(new GlobalSamplingRateRule(1.0f));
280-
return samplerStandalone;
280+
return samplerStandalone.Build();
281281
}
282282

283-
var sampler = new TraceSampler(new TracerRateLimiter(maxTracesPerInterval: settings.MaxTracesSubmittedPerSecond, intervalMilliseconds: null));
283+
var sampler = new TraceSampler.Builder(new TracerRateLimiter(maxTracesPerInterval: settings.MaxTracesSubmittedPerSecond, intervalMilliseconds: null));
284284

285285
// sampling rules (remote value overrides local value)
286286
var samplingRulesJson = settings.CustomSamplingRules;
@@ -336,7 +336,7 @@ protected virtual ITraceSampler GetSampler(TracerSettings settings)
336336
// This rule is always present, even if the agent has not yet provided any sampling rates.
337337
sampler.RegisterAgentSamplingRule(new AgentSamplingRule());
338338

339-
return sampler;
339+
return sampler.Build();
340340
}
341341

342342
protected virtual ISpanSampler GetSpanSampler(TracerSettings settings)

‎tracer/test/Datadog.Trace.Tests/Sampling/TraceSamplerTests.cs

Lines changed: 20 additions & 18 deletions
Original file line numberDiff line numberDiff line change
@@ -30,9 +30,9 @@ public class TraceSamplerTests
3030
[Fact]
3131
public async Task RateLimiter_Never_Applied_For_DefaultRule()
3232
{
33-
var sampler = new TraceSampler(new DenyAll());
33+
var builder = new TraceSampler.Builder(new DenyAll());
3434
await RunSamplerTest(
35-
sampler,
35+
builder.Build(),
3636
iterations: 500,
3737
expectedAutoKeepRate: 1,
3838
expectedUserKeepRate: 0,
@@ -42,9 +42,9 @@ await RunSamplerTest(
4242
[Fact]
4343
public async Task RateLimiter_Denies_All_Traces()
4444
{
45-
var sampler = new TraceSampler(new DenyAll());
45+
var builder = new TraceSampler.Builder(new DenyAll());
4646

47-
sampler.RegisterRule(
47+
builder.RegisterRule(
4848
new LocalCustomSamplingRule(
4949
rate: 1,
5050
patternFormat: SamplingRulesFormat.Regex,
@@ -55,7 +55,7 @@ public async Task RateLimiter_Denies_All_Traces()
5555
timeout: Timeout));
5656

5757
await RunSamplerTest(
58-
sampler,
58+
builder.Build(),
5959
iterations: 500,
6060
expectedAutoKeepRate: 0,
6161
expectedUserKeepRate: 0,
@@ -65,9 +65,9 @@ await RunSamplerTest(
6565
[Fact]
6666
public async Task Keep_Everything_Rule()
6767
{
68-
var sampler = new TraceSampler(new NoLimits());
68+
var builder = new TraceSampler.Builder(new NoLimits());
6969

70-
sampler.RegisterRule(
70+
builder.RegisterRule(
7171
new LocalCustomSamplingRule(
7272
rate: 1,
7373
patternFormat: SamplingRulesFormat.Regex,
@@ -78,7 +78,7 @@ public async Task Keep_Everything_Rule()
7878
timeout: Timeout));
7979

8080
await RunSamplerTest(
81-
sampler,
81+
builder.Build(),
8282
iterations: 500,
8383
expectedAutoKeepRate: 0,
8484
expectedUserKeepRate: 1,
@@ -88,9 +88,9 @@ await RunSamplerTest(
8888
[Fact]
8989
public async Task Keep_Nothing_Rule()
9090
{
91-
var sampler = new TraceSampler(new NoLimits());
91+
var builder = new TraceSampler.Builder(new NoLimits());
9292

93-
sampler.RegisterRule(
93+
builder.RegisterRule(
9494
new LocalCustomSamplingRule(
9595
rate: 0,
9696
patternFormat: SamplingRulesFormat.Regex,
@@ -101,7 +101,7 @@ public async Task Keep_Nothing_Rule()
101101
timeout: Timeout));
102102

103103
await RunSamplerTest(
104-
sampler,
104+
builder.Build(),
105105
iterations: 500,
106106
expectedAutoKeepRate: 0,
107107
expectedUserKeepRate: 0,
@@ -111,9 +111,9 @@ await RunSamplerTest(
111111
[Fact]
112112
public async Task Keep_Half_Rule()
113113
{
114-
var sampler = new TraceSampler(new NoLimits());
114+
var builder = new TraceSampler.Builder(new NoLimits());
115115

116-
sampler.RegisterRule(
116+
builder.RegisterRule(
117117
new LocalCustomSamplingRule(
118118
rate: 0.5f,
119119
patternFormat: SamplingRulesFormat.Regex,
@@ -124,7 +124,7 @@ public async Task Keep_Half_Rule()
124124
timeout: Timeout));
125125

126126
await RunSamplerTest(
127-
sampler,
127+
builder.Build(),
128128
iterations: 50_000, // Higher number for lower variance
129129
expectedAutoKeepRate: 0,
130130
expectedUserKeepRate: 0.5f,
@@ -134,8 +134,9 @@ await RunSamplerTest(
134134
[Fact]
135135
public async Task No_Registered_Rules_Uses_Legacy_Rates()
136136
{
137-
var sampler = new TraceSampler(new NoLimits());
138-
sampler.RegisterAgentSamplingRule(new AgentSamplingRule());
137+
var builder = new TraceSampler.Builder(new NoLimits());
138+
builder.RegisterAgentSamplingRule(new AgentSamplingRule());
139+
var sampler = builder.Build();
139140
sampler.SetDefaultSampleRates(MockAgentRates);
140141

141142
await RunSamplerTest(
@@ -156,8 +157,9 @@ public async Task Choose_Between_Sampling_Mechanisms()
156157
scope.Span.Context.TraceContext.Environment = Env;
157158

158159
var span = scope.Span;
159-
var sampler = new TraceSampler(new NoLimits());
160-
sampler.RegisterAgentSamplingRule(new AgentSamplingRule());
160+
var builder = new TraceSampler.Builder(new NoLimits());
161+
builder.RegisterAgentSamplingRule(new AgentSamplingRule());
162+
var sampler = builder.Build();
161163

162164
// if there are no other rules, and before we have agent rates, mechanism is "Default"
163165
var (_, mechanism1, _, _) = sampler.MakeSamplingDecision(span);

0 commit comments

Comments
 (0)