fix log injection and sampling rule bug

Author: khanayan123

integration-tests/log_injection.spec.js

@@ -0,0 +1,71 @@
+'use strict'
+
+const { FakeAgent, createSandbox, spawnProc, curlAndAssertMessage } = require('./helpers')
+const { assert } = require('chai')
+const path = require('path')
+const { USER_KEEP } = require('../ext/priority')
+
+describe('Log Injection', () => {
+  let agent
+  let proc
+  let sandbox
+  let cwd
+  let app
+  let env
+
+  before(async () => {
+    sandbox = await createSandbox(['express', 'winston'])
+    cwd = sandbox.folder
+    app = path.join(cwd, 'log_injection/index.js')
+  })
+
+  after(async () => {
+    await sandbox.remove()
+  })
+
+  beforeEach(async () => {
+    agent = await new FakeAgent().start()
+  })
+
+  describe('log injection with rule based sampling', () => {
+    beforeEach(async () => {
+      agent = await new FakeAgent().start()
+    })
+
+    afterEach(async () => {
+      proc.kill()
+      await agent.stop()
+    })
+
+    it('should correctly apply rule based sampling when log injection is enabled', async () => {
+      env = {
+        AGENT_PORT: agent.port,
+        lOG_INJECTION: 'true'
+      }
+      proc = await spawnProc(app, { cwd, env, execArgv: [] })
+      const url = proc.url + '/sampled'
+      return curlAndAssertMessage(agent, url, ({ headers, payload }) => {
+        assert.isArray(payload)
+        assert.strictEqual(payload.length, 1)
+        // Bug: previously got USER_REJECT instead of USER_KEEP when log injection is enabled,
+        // meaning resource rules are not applied & instead global sampling is applied
+        // Now gets USER_KEEP because resource rules are applied
+        assert.propertyVal(payload[0][0].metrics, '_sampling_priority_v1', USER_KEEP)
+      }, 20000, 1)
+    })
+
+    it('should correctly apply rule based sampling when log injection is disabled', async () => {
+      env = {
+        AGENT_PORT: agent.port,
+        lOG_INJECTION: 'false'
+      }
+      proc = await spawnProc(app, { cwd, env, execArgv: [] })
+      const url = proc.url + '/sampled'
+      return curlAndAssertMessage(agent, url, ({ headers, payload }) => {
+        assert.isArray(payload)
+        assert.strictEqual(payload.length, 1)
+        assert.propertyVal(payload[0][0].metrics, '_sampling_priority_v1', USER_KEEP)
+      }, 20000, 1)
+    })
+  })
+})

integration-tests/log_injection/index.js

@@ -0,0 +1,50 @@
+'use strict'
+
+const options = {
+  service: 'test-service',
+  sampleRate: 0.0,
+  samplingRules: [
+    {
+      resource: 'GET /sampled',
+      sampleRate: 1.0
+    }
+  ]
+}
+
+if (process.env.AGENT_PORT) {
+  options.port = process.env.AGENT_PORT
+}
+
+if (process.env.lOG_INJECTION) {
+  options.logInjection = process.env.lOG_INJECTION
+}
+
+const tracer = require('dd-trace')
+tracer.init(options)
+
+const express = require('express')
+const winston = require('winston')
+
+const app = express()
+
+// Create winston logger
+const logger = winston.createLogger({
+  level: 'info',
+  format: winston.format.json(),
+  transports: [
+    new winston.transports.Console({ silent: true })
+  ]
+})
+
+// Route WITH logging (demonstrates the bug)
+app.get('/sampled', (req, res) => {
+  // BUG: This winston.info() triggers log injection BEFORE resource.name is set
+  // which causes sampling decision to happen too early, bypassing the resource rule
+  logger.info('Processing GET /sampled request')
+  res.json({ message: 'logged request' })
+})
+
+const server = app.listen(0, () => {
+  const port = server.address().port
+  process.send({ port })
+})

packages/dd-trace/src/priority_sampler.js

@@ -5,6 +5,7 @@ const RateLimiter = require('./rate_limiter')
 const Sampler = require('./sampler')
 const { setSamplingRules } = require('./startup-log')
 const SamplingRule = require('./sampling_rule')
+const { resourceLocator } = SamplingRule
 
 const {
   SAMPLING_MECHANISM_DEFAULT,
@@ -116,7 +117,11 @@ class PrioritySampler {
       context._sampling.priority = tag
       context._sampling.mechanism = SAMPLING_MECHANISM_MANUAL
     } else if (auto) {
-      context._sampling.priority = this._getPriorityFromAuto(root)
+      const priority = this._getPriorityFromAuto(root)
+      if (priority === undefined) {
+        return // Defer sampling until resource information is available
+      }
+      context._sampling.priority = priority
     } else {
       return
     }
@@ -200,10 +205,18 @@ class PrioritySampler {
   /**
    *
    * @param span {DatadogSpan}
-   * @returns {SamplingPriority}
+   * @returns {SamplingPriority|undefined}
    */
   _getPriorityFromAuto (span) {
     const context = this._getContext(span)
+
+    // Check if resource name is available and if we have resource-based rules
+    const { _tags: tags } = context
+    const hasResource = tags.resource || tags['resource.name']
+    if (!hasResource && this.#hasResourceBasedRules()) {
+      return // Defer sampling until resource is available
+    }
+
     const rule = this.#findRule(span)
 
     return rule
@@ -338,6 +351,16 @@ class PrioritySampler {
     }
   }
 
+  /**
+   * Check if any sampling rules depend on resource names
+   * @returns {boolean}
+   */
+  #hasResourceBasedRules () {
+    return this._rules.some(rule =>
+      rule.matchers.some(matcher => matcher.locator === resourceLocator)
+    )
+  }
+
   /**
    *
    * @param span {DatadogSpan}

packages/dd-trace/src/sampling_rule.js

@@ -143,3 +143,4 @@ class SamplingRule {
 }
 
 module.exports = SamplingRule
+module.exports.resourceLocator = resourceLocator

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

@@ -527,4 +527,49 @@ describe('PrioritySampler', () => {
       expect(setPriority).to.be.calledOnceWithExactly(span, USER_KEEP, SAMPLING_MECHANISM_APPSEC)
     })
   })
+
+  describe('resource-based rule deferral', () => {
+    beforeEach(() => {
+      span.tracer = sinon.stub().returns({ _service: 'test-service' })
+    })
+
+    it('should defer sampling when resource rules exist but resource name is not set', () => {
+      prioritySampler.configure('test', {
+        sampleRate: 0,
+        rules: [{ resource: 'GET /api/*', sampleRate: 1 }]
+      })
+
+      span.context()._tags = {}
+
+      prioritySampler.sample(span)
+
+      expect(span.context()._sampling.priority).to.be.undefined
+    })
+
+    it('should sample normally when resource rules exist and resource name is set', () => {
+      prioritySampler.configure('test', {
+        sampleRate: 0,
+        rules: [{ resource: 'GET /api/*', sampleRate: 1 }]
+      })
+
+      span.context()._tags = { resource: 'GET /api/users' }
+
+      prioritySampler.sample(span)
+
+      expect(span.context()._sampling.priority).to.equal(USER_KEEP)
+    })
+
+    it('should sample normally when no resource rules exist', () => {
+      prioritySampler.configure('test', {
+        sampleRate: 1,
+        rules: [{ service: 'my-service', sampleRate: 1 }]
+      })
+
+      span.context()._tags = {}
+
+      prioritySampler.sample(span)
+
+      expect(span.context()._sampling.priority).to.equal(USER_KEEP)
+    })
+  })
 })