Add _dd.p.ksr propagated tag for Knuth sampling rate

packages/dd-trace/src/constants.js

@@ -22,6 +22,7 @@ module.exports = {
   SPAN_SAMPLING_MAX_PER_SECOND: '_dd.span_sampling.max_per_second',
   DATADOG_LAMBDA_EXTENSION_PATH: '/opt/extensions/datadog-agent',
   DECISION_MAKER_KEY: '_dd.p.dm',
+  SAMPLING_KNUTH_RATE: '_dd.p.ksr',
   PROCESS_ID: 'process_id',
   ERROR_TYPE: 'error.type',
   ERROR_MESSAGE: 'error.message',

packages/dd-trace/src/priority_sampler.js

@@ -31,10 +31,21 @@ const {
   SAMPLING_LIMIT_DECISION,
   SAMPLING_AGENT_DECISION,
   DECISION_MAKER_KEY,
+  SAMPLING_KNUTH_RATE,
 } = require('./constants')
 
 const DEFAULT_KEY = 'service:,env:'
 
+/**
+ * Formats a sampling rate as a string with up to 6 significant digits and no trailing zeros.
+ *
+ * @param {number} rate
+ * @returns {string}
+ */
+function formatKnuthRate (rate) {
+  return Number(rate.toPrecision(6)).toString()
+}
+
 const defaultSampler = new Sampler(AUTO_KEEP)
 
 /**
@@ -254,6 +265,7 @@ class PrioritySampler {
    */
   #getPriorityByRule (context, rule) {
     context._trace[SAMPLING_RULE_DECISION] = rule.sampleRate
+    context._trace.tags[SAMPLING_KNUTH_RATE] = formatKnuthRate(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
@@ -290,9 +302,15 @@ class PrioritySampler {
     // TODO: Change underscored properties to private ones.
     const sampler = this._samplers[key] || this._samplers[DEFAULT_KEY]
 
-    context._trace[SAMPLING_AGENT_DECISION] = sampler.rate()
+    const rate = sampler.rate()
+    context._trace[SAMPLING_AGENT_DECISION] = rate
 
-    context._sampling.mechanism = sampler === defaultSampler ? SAMPLING_MECHANISM_DEFAULT : SAMPLING_MECHANISM_AGENT
+    if (sampler === defaultSampler) {
+      context._sampling.mechanism = SAMPLING_MECHANISM_DEFAULT
+    } else {
+      context._trace.tags[SAMPLING_KNUTH_RATE] = formatKnuthRate(rate)
+      context._sampling.mechanism = SAMPLING_MECHANISM_AGENT
+    }
 
     return sampler.isSampled(context) ? AUTO_KEEP : AUTO_REJECT
   }

packages/dd-trace/test/priority_sampler.spec.js

@@ -18,6 +18,7 @@ const {
   SAMPLING_MECHANISM_REMOTE_DYNAMIC,
   DECISION_MAKER_KEY,
   SAMPLING_MECHANISM_APPSEC,
+  SAMPLING_KNUTH_RATE,
 } = require('../src/constants')
 const { ASM } = require('../src/standalone/product')
 
@@ -369,6 +370,65 @@ describe('PrioritySampler', () => {
       assert.strictEqual(context._trace['_dd.limit_psr'], 1)
     })
 
+    it('should not set _dd.p.ksr tag for default sampling mechanism', () => {
+      prioritySampler.sample(span)
+
+      assert.strictEqual(context._trace.tags[SAMPLING_KNUTH_RATE], undefined)
+    })
+
+    it('should set _dd.p.ksr tag for agent sampling with explicit rates', () => {
+      prioritySampler.update({
+        'service:test,env:test': 0.5,
+      })
+      prioritySampler.sample(span)
+
+      assert.strictEqual(context._trace.tags[SAMPLING_KNUTH_RATE], '0.5')
+    })
+
+    it('should set _dd.p.ksr tag for rule-based sampling', () => {
+      prioritySampler = new PrioritySampler('test', {
+        sampleRate: 0.5,
+      })
+      prioritySampler.sample(span)
+
+      assert.strictEqual(context._trace.tags[SAMPLING_KNUTH_RATE], '0.5')
+    })
+
+    it('should set _dd.p.ksr tag for rule-based sampling with rate 0', () => {
+      prioritySampler = new PrioritySampler('test', {
+        sampleRate: 0,
+      })
+      prioritySampler.sample(span)
+
+      assert.strictEqual(context._trace.tags[SAMPLING_KNUTH_RATE], '0')
+    })
+
+    it('should set _dd.p.ksr tag for rate 1.0 as "1"', () => {
+      prioritySampler = new PrioritySampler('test', {
+        sampleRate: 1.0,
+      })
+      prioritySampler.sample(span)
+
+      assert.strictEqual(context._trace.tags[SAMPLING_KNUTH_RATE], '1')
+    })
+
+    it('should set _dd.p.ksr tag as a string type', () => {
+      prioritySampler.update({
+        'service:test,env:test': 0.5,
+      })
+      prioritySampler.sample(span)
+
+      assert.strictEqual(typeof context._trace.tags[SAMPLING_KNUTH_RATE], 'string')
+    })
+
+    it('should not set _dd.p.ksr tag for manual sampling', () => {
+      context._tags[MANUAL_KEEP] = undefined
+
+      prioritySampler.sample(context)
+
+      assert.strictEqual(context._trace.tags[SAMPLING_KNUTH_RATE], undefined)
+    })
+
     it('should ignore empty span', () => {
       prioritySampler.sample()
       prioritySampler.sample()