Skip to content

Commit e559aab

Browse files
bm1549bm1549
committed
Add consistent sampling via knuth's method
1 parent 9515d54 commit e559aab

File tree

8 files changed

+148
-13
lines changed

8 files changed

+148
-13
lines changed

‎benchmark/core.js

+1-1
Original file line numberDiff line numberDiff line change
@@ -93,7 +93,7 @@ suite
9393
sampler = new Sampler(0.5)
9494
},
9595
fn () {
96-
sampler.isSampled()
96+
sampler.isSampled(span.context())
9797
}
9898
})
9999
.add('format', {

‎packages/datadog-plugin-openai/src/tracing.js

+1-1
Original file line numberDiff line numberDiff line change
@@ -329,7 +329,7 @@ class OpenAiTracingPlugin extends TracingPlugin {
329329
sendLog (methodName, span, tags, openaiStore, error) {
330330
if (!openaiStore) return
331331
if (!Object.keys(openaiStore).length) return
332-
if (!this.sampler.isSampled()) return
332+
if (!this.sampler.isSampled(span.context())) return
333333

334334
const log = {
335335
status: error ? 'error' : 'info',

‎packages/dd-trace/src/appsec/api_security_sampler.js

+7-3
Original file line numberDiff line numberDiff line change
@@ -8,12 +8,16 @@ const { AUTO_REJECT, USER_REJECT } = require('../../../../ext/priority')
88
const MAX_SIZE = 4096
99

1010
let enabled
11+
12+
/**
13+
* @type {TTLCache}
14+
*/
1115
let sampledRequests
1216

1317
class NoopTTLCache {
1418
clear () { }
15-
set (key) { return undefined }
16-
has (key) { return false }
19+
set (_key, _value) { return undefined }
20+
has (_key) { return false }
1721
}
1822

1923
function configure ({ apiSecurity }) {
@@ -48,7 +52,7 @@ function sampleRequest (req, res, force = false) {
4852
}
4953

5054
if (force) {
51-
sampledRequests.set(key)
55+
sampledRequests.set(key, true)
5256
}
5357

5458
return true

‎packages/dd-trace/src/opentracing/span.js

+3
Original file line numberDiff line numberDiff line change
@@ -143,6 +143,9 @@ class DatadogSpan {
143143
return `Span${json}`
144144
}
145145

146+
/**
147+
* @returns {DatadogSpanContext}
148+
*/
146149
context () {
147150
return this._spanContext
148151
}

‎packages/dd-trace/src/opentracing/span_context.js

+4
Original file line numberDiff line numberDiff line change
@@ -52,6 +52,10 @@ class DatadogSpanContext {
5252
return this._traceId.toString(10)
5353
}
5454

55+
toTraceIdNumber (get128BitId = false) {
56+
return parseInt(this.toTraceId(get128BitId), 16)
57+
}
58+
5559
toSpanId (get128bitId = false) {
5660
if (get128bitId) {
5761
return this._spanId.toString(16).padStart(16, '0')

‎packages/dd-trace/src/priority_sampler.js

+107-4
Original file line numberDiff line numberDiff line change
@@ -39,12 +39,28 @@ const DEFAULT_KEY = 'service:,env:'
3939

4040
const defaultSampler = new Sampler(AUTO_KEEP)
4141

42+
/**
43+
* from config.js
44+
* @typedef { sampleRate: number, provenance: string, rateLimit: number, rules: SamplingRule[] } SamplingConfig
45+
*
46+
* empirically defined
47+
* @typedef {2|-1|1|0} SamplingPriority
48+
*/
4249
class PrioritySampler {
50+
/**
51+
* @param env {string}
52+
* @param config {SamplingConfig}
53+
*/
4354
constructor (env, config) {
4455
this.configure(env, config)
4556
this.update({})
4657
}
4758

59+
/**
60+
*
61+
* @param env {string}
62+
* @param opts {SamplingConfig}
63+
*/
4864
configure (env, opts = {}) {
4965
const { sampleRate, provenance = undefined, rateLimit = 100, rules = [] } = opts
5066
this._env = env
@@ -55,12 +71,22 @@ class PrioritySampler {
5571
setSamplingRules(this._rules)
5672
}
5773

74+
/**
75+
* @param span {DatadogSpan}
76+
* @returns {boolean}
77+
*/
5878
isSampled (span) {
5979
const priority = this._getPriorityFromAuto(span)
6080
log.trace(span)
6181
return priority === USER_KEEP || priority === AUTO_KEEP
6282
}
6383

84+
/**
85+
*
86+
* @param span {DatadogSpan}
87+
* @param auto {boolean}
88+
* @returns {void}
89+
*/
6490
sample (span, auto = true) {
6591
if (!span) return
6692

@@ -73,7 +99,7 @@ class PrioritySampler {
7399

74100
log.trace(span, auto)
75101

76-
const tag = this._getPriorityFromTags(context._tags, context)
102+
const tag = this._getPriorityFromTags(context._tags)
77103

78104
if (this.validate(tag)) {
79105
context._sampling.priority = tag
@@ -87,14 +113,17 @@ class PrioritySampler {
87113
this._addDecisionMaker(root)
88114
}
89115

116+
/**
117+
*
118+
* @param rates {Object<string, number>}
119+
* @returns {void}
120+
*/
90121
update (rates) {
91122
const samplers = {}
92123

93124
for (const key in rates) {
94125
const rate = rates[key]
95-
const sampler = new Sampler(rate)
96-
97-
samplers[key] = sampler
126+
samplers[key] = new Sampler(rate)
98127
}
99128

100129
samplers[DEFAULT_KEY] = samplers[DEFAULT_KEY] || defaultSampler
@@ -104,6 +133,11 @@ class PrioritySampler {
104133
log.trace(rates)
105134
}
106135

136+
/**
137+
*
138+
* @param samplingPriority {SamplingPriority}
139+
* @returns {boolean}
140+
*/
107141
validate (samplingPriority) {
108142
switch (samplingPriority) {
109143
case USER_REJECT:
@@ -116,6 +150,12 @@ class PrioritySampler {
116150
}
117151
}
118152

153+
/**
154+
*
155+
* @param span {DatadogSpan}
156+
* @param samplingPriority {SamplingPriority}
157+
* @param product {import('./standalone/product').PRODUCTS}
158+
*/
119159
setPriority (span, samplingPriority, product) {
120160
if (!span || !this.validate(samplingPriority)) return
121161

@@ -137,10 +177,22 @@ class PrioritySampler {
137177
this._addDecisionMaker(root)
138178
}
139179

180+
/**
181+
*
182+
* @param span {DatadogSpan}
183+
* @returns {DatadogSpanContext}
184+
* @private
185+
*/
140186
_getContext (span) {
141187
return typeof span.context === 'function' ? span.context() : span
142188
}
143189

190+
/**
191+
*
192+
* @param span {DatadogSpan}
193+
* @returns {SamplingPriority}
194+
* @private
195+
*/
144196
_getPriorityFromAuto (span) {
145197
const context = this._getContext(span)
146198
const rule = this._findRule(span)
@@ -150,6 +202,12 @@ class PrioritySampler {
150202
: this._getPriorityByAgent(context)
151203
}
152204

205+
/**
206+
*
207+
* @param tags {Map<string, Object>}
208+
* @returns {SamplingPriority}
209+
* @private
210+
*/
153211
_getPriorityFromTags (tags) {
154212
if (hasOwn(tags, MANUAL_KEEP) && tags[MANUAL_KEEP] !== false) {
155213
return USER_KEEP
@@ -166,6 +224,13 @@ class PrioritySampler {
166224
}
167225
}
168226

227+
/**
228+
*
229+
* @param context {DatadogSpanContext}
230+
* @param rule {SamplingRule}
231+
* @returns {SamplingPriority}
232+
* @private
233+
*/
169234
_getPriorityByRule (context, rule) {
170235
context._trace[SAMPLING_RULE_DECISION] = rule.sampleRate
171236
context._sampling.mechanism = SAMPLING_MECHANISM_RULE
@@ -177,6 +242,12 @@ class PrioritySampler {
177242
: USER_REJECT
178243
}
179244

245+
/**
246+
*
247+
* @param context {DatadogSpanContext}
248+
* @returns {boolean}
249+
* @private
250+
*/
180251
_isSampledByRateLimit (context) {
181252
const allowed = this._limiter.isAllowed()
182253

@@ -185,6 +256,12 @@ class PrioritySampler {
185256
return allowed
186257
}
187258

259+
/**
260+
*
261+
* @param context {DatadogSpanContext}
262+
* @returns {SamplingPriority}
263+
* @private
264+
*/
188265
_getPriorityByAgent (context) {
189266
const key = `service:${context._tags[SERVICE_NAME]},env:${this._env}`
190267
const sampler = this._samplers[key] || this._samplers[DEFAULT_KEY]
@@ -200,6 +277,12 @@ class PrioritySampler {
200277
return sampler.isSampled(context) ? AUTO_KEEP : AUTO_REJECT
201278
}
202279

280+
/**
281+
*
282+
* @param span {DatadogSpan}
283+
* @private
284+
* @returns {void}
285+
*/
203286
_addDecisionMaker (span) {
204287
const context = span.context()
205288
const trace = context._trace
@@ -215,6 +298,15 @@ class PrioritySampler {
215298
}
216299
}
217300

301+
/**
302+
*
303+
* @param rules {SamplingRule[]}
304+
* @param sampleRate {number}
305+
* @param rateLimit {number}
306+
* @param provenance {string}
307+
* @returns {SamplingRule[]}
308+
* @private
309+
*/
218310
_normalizeRules (rules, sampleRate, rateLimit, provenance) {
219311
rules = [].concat(rules || [])
220312

@@ -225,12 +317,23 @@ class PrioritySampler {
225317
.map(SamplingRule.from)
226318
}
227319

320+
/**
321+
*
322+
* @param span {DatadogSpan}
323+
* @returns {SamplingRule}
324+
* @private
325+
*/
228326
_findRule (span) {
229327
for (const rule of this._rules) {
230328
if (rule.match(span)) return rule
231329
}
232330
}
233331

332+
/**
333+
*
334+
* @param span {DatadogSpan}
335+
* @param product {import('./standalone/product').PRODUCTS}
336+
*/
234337
static keepTrace (span, product) {
235338
span?._prioritySampler?.setPriority(span, USER_KEEP, product)
236339
}

‎packages/dd-trace/src/sampler.js

+19-2
Original file line numberDiff line numberDiff line change
@@ -1,16 +1,33 @@
11
'use strict'
22

3+
const KNUTH_FACTOR = 1111111111111111111
4+
const MAX_TRACE_ID = 2**64-1
5+
36
class Sampler {
7+
/**
8+
*
9+
* @param rate {number}
10+
*/
411
constructor (rate) {
512
this._rate = rate
13+
this._sampling_id_threshold = MAX_TRACE_ID * rate
614
}
715

16+
/**
17+
* @returns {number}
18+
*/
819
rate () {
920
return this._rate
1021
}
1122

12-
isSampled () {
13-
return this._rate === 1 || Math.random() < this._rate
23+
/**
24+
*
25+
* @param context {DatadogSpanContext}
26+
* @returns {boolean}
27+
*/
28+
isSampled (context) {
29+
return this._rate === 1 ||
30+
((context.toTraceIdNumber(false) * KNUTH_FACTOR) % MAX_TRACE_ID) <= this._sampling_id_threshold
1431
}
1532
}
1633

‎packages/dd-trace/src/standalone/product.js

+6-2
Original file line numberDiff line numberDiff line change
@@ -13,12 +13,16 @@ function getProductRateLimiter (config) {
1313
return dropAll
1414
}
1515

16-
module.exports = {
16+
const PRODUCTS = {
1717
APM: { id: 1 << 0 },
1818
ASM: { id: 1 << 1, mechanism: SAMPLING_MECHANISM_APPSEC },
1919
DSM: { id: 1 << 2 },
2020
DJM: { id: 1 << 3 },
21-
DBM: { id: 1 << 4 },
21+
DBM: { id: 1 << 4 }
22+
}
23+
24+
module.exports = {
25+
...PRODUCTS,
2226

2327
getProductRateLimiter
2428
}

0 commit comments

Comments
 (0)