fix sampling priority applied incorrectly (#1266)

Author: rochdev

packages/dd-trace/src/format.js

@@ -7,6 +7,9 @@ const id = require('./id')
 
 const SAMPLING_PRIORITY_KEY = constants.SAMPLING_PRIORITY_KEY
 const ANALYTICS_KEY = constants.ANALYTICS_KEY
+const SAMPLING_RULE_DECISION = constants.SAMPLING_RULE_DECISION
+const SAMPLING_LIMIT_DECISION = constants.SAMPLING_LIMIT_DECISION
+const SAMPLING_AGENT_DECISION = constants.SAMPLING_AGENT_DECISION
 const ANALYTICS = tags.ANALYTICS
 const MEASURED = tags.MEASURED
 const ORIGIN_KEY = constants.ORIGIN_KEY
@@ -22,6 +25,7 @@ function format (span) {
   const formatted = formatSpan(span)
 
   extractError(formatted, span)
+  extractRootTags(formatted, span)
   extractTags(formatted, span)
   extractAnalytics(formatted, span)
 
@@ -96,6 +100,18 @@ function extractTags (trace, span) {
   addTag(trace.meta, trace.metrics, HOSTNAME_KEY, hostname)
 }
 
+function extractRootTags (trace, span) {
+  const context = span.context()
+  const isLocalRoot = span === context._trace.started[0]
+  const parentId = context._parentId
+
+  if (!isLocalRoot || (parentId && parentId.toString(10) !== '0')) return
+
+  addTag({}, trace.metrics, SAMPLING_RULE_DECISION, context._trace[SAMPLING_RULE_DECISION])
+  addTag({}, trace.metrics, SAMPLING_LIMIT_DECISION, context._trace[SAMPLING_LIMIT_DECISION])
+  addTag({}, trace.metrics, SAMPLING_AGENT_DECISION, context._trace[SAMPLING_AGENT_DECISION])
+}
+
 function extractError (trace, span) {
   const error = span.context()._tags['error']
 

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

@@ -113,6 +113,8 @@ class DatadogSpan extends Span {
 
   _addTags (keyValuePairs) {
     tagger.add(this._spanContext._tags, keyValuePairs)
+
+    this._prioritySampler.sample(this, false)
   }
 
   _finish (finishTime) {

packages/dd-trace/src/priority_sampler.js

@@ -41,7 +41,7 @@ class PrioritySampler {
       : this._isSampledByAgent(context)
   }
 
-  sample (span) {
+  sample (span, auto = true) {
     if (!span) return
 
     const context = this._getContext(span)
@@ -55,7 +55,9 @@ class PrioritySampler {
       return
     }
 
-    context._sampling.priority = this.isSampled(span) ? AUTO_KEEP : AUTO_REJECT
+    if (auto) {
+      context._sampling.priority = this.isSampled(span) ? AUTO_KEEP : AUTO_REJECT
+    }
   }
 
   update (rates) {
@@ -106,15 +108,15 @@ class PrioritySampler {
   }
 
   _isSampledByRule (context, rule) {
-    context._tags[SAMPLING_RULE_DECISION] = rule.sampleRate
+    context._trace[SAMPLING_RULE_DECISION] = rule.sampleRate
 
     return rule.sampler.isSampled(context)
   }
 
   _isSampledByRateLimit (context) {
     const allowed = this._limiter.isAllowed()
 
-    context._tags[SAMPLING_LIMIT_DECISION] = this._limiter.effectiveRate()
+    context._trace[SAMPLING_LIMIT_DECISION] = this._limiter.effectiveRate()
 
     return allowed
   }
@@ -123,7 +125,7 @@ class PrioritySampler {
     const key = `service:${context._tags[SERVICE_NAME]},env:${this._env}`
     const sampler = this._samplers[key] || this._samplers[DEFAULT_KEY]
 
-    context._tags[SAMPLING_AGENT_DECISION] = sampler.rate()
+    context._trace[SAMPLING_AGENT_DECISION] = sampler.rate()
 
     return sampler.isSampled(context)
   }

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

@@ -10,6 +10,9 @@ const ANALYTICS = tags.ANALYTICS
 const MEASURED = tags.MEASURED
 const ORIGIN_KEY = constants.ORIGIN_KEY
 const HOSTNAME_KEY = constants.HOSTNAME_KEY
+const SAMPLING_AGENT_DECISION = constants.SAMPLING_AGENT_DECISION
+const SAMPLING_LIMIT_DECISION = constants.SAMPLING_LIMIT_DECISION
+const SAMPLING_RULE_DECISION = constants.SAMPLING_RULE_DECISION
 
 const spanId = id('0234567812345678')
 
@@ -27,7 +30,9 @@ describe('format', () => {
       _tags: {},
       _metrics: {},
       _sampling: {},
-      _trace: {},
+      _trace: {
+        started: []
+      },
       _name: 'operation'
     }
 
@@ -40,6 +45,8 @@ describe('format', () => {
       _duration: 100
     }
 
+    spanContext._trace.started.push(span)
+
     format = require('../src/format')
   })
 
@@ -84,6 +91,35 @@ describe('format', () => {
       expect(trace.resource).to.equal('resource')
     })
 
+    it('should extract Datadog specific root tags', () => {
+      spanContext._parentId = null
+      spanContext._trace[SAMPLING_AGENT_DECISION] = 0.8
+      spanContext._trace[SAMPLING_LIMIT_DECISION] = 0.2
+      spanContext._trace[SAMPLING_RULE_DECISION] = 0.5
+
+      trace = format(span)
+
+      expect(trace.metrics).to.include({
+        [SAMPLING_AGENT_DECISION]: 0.8,
+        [SAMPLING_LIMIT_DECISION]: 0.2,
+        [SAMPLING_RULE_DECISION]: 0.5
+      })
+    })
+
+    it('should not extract Datadog specific root tags from non-root spans', () => {
+      spanContext._trace[SAMPLING_AGENT_DECISION] = 0.8
+      spanContext._trace[SAMPLING_LIMIT_DECISION] = 0.2
+      spanContext._trace[SAMPLING_RULE_DECISION] = 0.5
+
+      trace = format(span)
+
+      expect(trace.metrics).to.not.have.keys(
+        SAMPLING_AGENT_DECISION,
+        SAMPLING_LIMIT_DECISION,
+        SAMPLING_RULE_DECISION
+      )
+    })
+
     it('should discard user-defined tags with name HOSTNAME_KEY by default', () => {
       spanContext._tags[HOSTNAME_KEY] = 'some_hostname'
 

packages/dd-trace/test/opentracing/span.spec.js

@@ -250,6 +250,14 @@ describe('Span', () => {
 
       expect(tagger.add).to.have.been.calledWith(span.context()._tags, tags)
     })
+
+    it('should sample based on the tags', () => {
+      const tags = { foo: 'bar' }
+
+      span.addTags(tags)
+
+      expect(prioritySampler.sample).to.have.been.calledWith(span, false)
+    })
   })
 
   describe('finish', () => {

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

@@ -22,7 +22,10 @@ describe('PrioritySampler', () => {
   beforeEach(() => {
     context = {
       _tags: {},
-      _sampling: {}
+      _sampling: {},
+      _trace: {
+        started: []
+      }
     }
 
     span = {
@@ -262,7 +265,7 @@ describe('PrioritySampler', () => {
     it('should add metrics for agent sample rate', () => {
       prioritySampler.sample(span)
 
-      expect(context._tags).to.have.property('_dd.agent_psr', 1)
+      expect(context._trace).to.have.property('_dd.agent_psr', 1)
     })
 
     it('should add metrics for rule sample rate', () => {
@@ -271,8 +274,8 @@ describe('PrioritySampler', () => {
       })
       prioritySampler.sample(span)
 
-      expect(context._tags).to.have.property('_dd.rule_psr', 0)
-      expect(context._tags).to.not.have.property('_dd.limit_psr')
+      expect(context._trace).to.have.property('_dd.rule_psr', 0)
+      expect(context._trace).to.not.have.property('_dd.limit_psr')
     })
 
     it('should add metrics for rate limiter sample rate', () => {
@@ -282,14 +285,21 @@ describe('PrioritySampler', () => {
       })
       prioritySampler.sample(span)
 
-      expect(context._tags).to.have.property('_dd.rule_psr', 0.5)
-      expect(context._tags).to.have.property('_dd.limit_psr', 1)
+      expect(context._trace).to.have.property('_dd.rule_psr', 0.5)
+      expect(context._trace).to.have.property('_dd.limit_psr', 1)
     })
 
     it('should ignore empty span', () => {
+      expect(() => {
+        prioritySampler.sample()
+      }).to.not.throw()
       prioritySampler.sample()
+    })
+
+    it('should support manual only sampling', () => {
+      prioritySampler.sample(span, false)
 
-      expect(context._sampling).to.not.have.property('priority')
+      expect(context._sampling.priority).to.be.undefined
     })
   })