RC support for sampling rules patch (#4381)

DataDog · Jun 10, 2024 · b5892a2 · b5892a2
1 parent 1f19792
commit b5892a2
Show file tree

Hide file tree

Showing 6 changed files with 41 additions and 8 deletions.
diff --git a/packages/dd-trace/src/config.js b/packages/dd-trace/src/config.js
@@ -974,8 +974,8 @@ class Config {
     this._setTags(opts, 'tags', tags)
     this._setBoolean(opts, 'tracing', options.tracing_enabled)
     // ignore tags for now since rc sampling rule tags format is not supported
-    this._setSamplingRule(opts, 'sampler.rules', this._ignoreTags(options.trace_sample_rules))
-    this._remoteUnprocessed['sampler.rules'] = options.trace_sample_rules
+    this._setSamplingRule(opts, 'sampler.rules', this._ignoreTags(options.tracing_sampling_rules))
+    this._remoteUnprocessed['sampler.rules'] = options.tracing_sampling_rules
   }
 
   _ignoreTags (samplingRules) {

diff --git a/packages/dd-trace/src/constants.js b/packages/dd-trace/src/constants.js
@@ -15,6 +15,8 @@ module.exports = {
   SAMPLING_MECHANISM_MANUAL: 4,
   SAMPLING_MECHANISM_APPSEC: 5,
   SAMPLING_MECHANISM_SPAN: 8,
+  SAMPLING_MECHANISM_REMOTE_USER: 11,
+  SAMPLING_MECHANISM_REMOTE_DYNAMIC: 12,
   SPAN_SAMPLING_MECHANISM: '_dd.span_sampling.mechanism',
   SPAN_SAMPLING_RULE_RATE: '_dd.span_sampling.rule_rate',
   SPAN_SAMPLING_MAX_PER_SECOND: '_dd.span_sampling.max_per_second',

diff --git a/packages/dd-trace/src/priority_sampler.js b/packages/dd-trace/src/priority_sampler.js
@@ -10,6 +10,8 @@ const {
   SAMPLING_MECHANISM_AGENT,
   SAMPLING_MECHANISM_RULE,
   SAMPLING_MECHANISM_MANUAL,
+  SAMPLING_MECHANISM_REMOTE_USER,
+  SAMPLING_MECHANISM_REMOTE_DYNAMIC,
   SAMPLING_RULE_DECISION,
   SAMPLING_LIMIT_DECISION,
   SAMPLING_AGENT_DECISION,
@@ -41,9 +43,9 @@ class PrioritySampler {
     this.update({})
   }
 
-  configure (env, { sampleRate, rateLimit = 100, rules = [] } = {}) {
+  configure (env, { sampleRate, provenance = undefined, rateLimit = 100, rules = [] } = {}) {
     this._env = env
-    this._rules = this._normalizeRules(rules, sampleRate, rateLimit)
+    this._rules = this._normalizeRules(rules, sampleRate, rateLimit, provenance)
     this._limiter = new RateLimiter(rateLimit)
 
     setSamplingRules(this._rules)
@@ -137,6 +139,8 @@ class PrioritySampler {
   _getPriorityByRule (context, rule) {
     context._trace[SAMPLING_RULE_DECISION] = rule.sampleRate
     context._sampling.mechanism = SAMPLING_MECHANISM_RULE
+    if (rule.provenance === 'customer') context._sampling.mechanism = SAMPLING_MECHANISM_REMOTE_USER
+    if (rule.provenance === 'dynamic') context._sampling.mechanism = SAMPLING_MECHANISM_REMOTE_DYNAMIC
 
     return rule.sample() && this._isSampledByRateLimit(context)
       ? USER_KEEP
@@ -181,11 +185,11 @@ class PrioritySampler {
     }
   }
 
-  _normalizeRules (rules, sampleRate, rateLimit) {
+  _normalizeRules (rules, sampleRate, rateLimit, provenance) {
     rules = [].concat(rules || [])
 
     return rules
-      .concat({ sampleRate, maxPerSecond: rateLimit })
+      .concat({ sampleRate, maxPerSecond: rateLimit, provenance })
       .map(rule => ({ ...rule, sampleRate: parseFloat(rule.sampleRate) }))
       .filter(rule => !isNaN(rule.sampleRate))
       .map(SamplingRule.from)

diff --git a/packages/dd-trace/src/sampling_rule.js b/packages/dd-trace/src/sampling_rule.js
@@ -64,7 +64,7 @@ function serviceLocator (span) {
 }
 
 class SamplingRule {
-  constructor ({ name, service, resource, tags, sampleRate = 1.0, maxPerSecond } = {}) {
+  constructor ({ name, service, resource, tags, sampleRate = 1.0, provenance = undefined, maxPerSecond } = {}) {
     this.matchers = []
 
     if (name) {
@@ -82,6 +82,7 @@ class SamplingRule {
 
     this._sampler = new Sampler(sampleRate)
     this._limiter = undefined
+    this.provenance = provenance
 
     if (Number.isFinite(maxPerSecond)) {
       this._limiter = new RateLimiter(maxPerSecond)

diff --git a/packages/dd-trace/test/config.spec.js b/packages/dd-trace/test/config.spec.js
@@ -1474,7 +1474,7 @@ describe('Config', () => {
     const config = new Config()
 
     config.configure({
-      trace_sample_rules: [
+      tracing_sampling_rules: [
         {
           resource: '*',
           tags: [{ key: 'tag-a', value_glob: 'tag-a-val*' }],

diff --git a/packages/dd-trace/test/priority_sampler.spec.js b/packages/dd-trace/test/priority_sampler.spec.js
@@ -9,6 +9,8 @@ const {
   SAMPLING_MECHANISM_AGENT,
   SAMPLING_MECHANISM_RULE,
   SAMPLING_MECHANISM_MANUAL,
+  SAMPLING_MECHANISM_REMOTE_USER,
+  SAMPLING_MECHANISM_REMOTE_DYNAMIC,
   DECISION_MAKER_KEY
 } = require('../src/constants')
 
@@ -202,6 +204,30 @@ describe('PrioritySampler', () => {
       expect(context._sampling.mechanism).to.equal(SAMPLING_MECHANISM_RULE)
     })
 
+    it('should support a customer-defined remote configuration sampling', () => {
+      prioritySampler = new PrioritySampler('test', {
+        rules: [
+          { sampleRate: 1, service: 'test', resource: /res.*/, provenance: 'customer' }
+        ]
+      })
+      prioritySampler.sample(context)
+
+      expect(context._sampling).to.have.property('priority', USER_KEEP)
+      expect(context._sampling.mechanism).to.equal(SAMPLING_MECHANISM_REMOTE_USER)
+    })
+
+    it('should support a dynamic remote configuration sampling', () => {
+      prioritySampler = new PrioritySampler('test', {
+        rules: [
+          { sampleRate: 0, service: 'test', resource: /res.*/, provenance: 'dynamic' }
+        ]
+      })
+      prioritySampler.sample(context)
+
+      expect(context._sampling).to.have.property('priority', USER_REJECT)
+      expect(context._sampling.mechanism).to.equal(SAMPLING_MECHANISM_REMOTE_DYNAMIC)
+    })
+
     it('should validate JSON rule into an array', () => {
       context._name = 'test'
       context._tags['service.name'] = 'test'