use priority simpler to get span priority

DataDog · Oct 3, 2024 · ca8bb68 · ca8bb68
1 parent 6e4e4b1
commit ca8bb68
Show file tree

Hide file tree

Showing 3 changed files with 59 additions and 59 deletions.
diff --git a/packages/dd-trace/src/appsec/api_security_sampler.js b/packages/dd-trace/src/appsec/api_security_sampler.js
@@ -1,11 +1,12 @@
 'use strict'
 
 const ApiSecuritySamplerCache = require('./api_security_sampler_cache')
+const PrioritySampler = require('../priority_sampler')
 const web = require('../plugins/util/web')
-const { USER_KEEP, AUTO_KEEP } = require('../../../../ext/priority')
 
 let enabled
 let sampledRequests
+const prioritySampler = new PrioritySampler()
 
 function configure ({ apiSecurity }) {
   enabled = apiSecurity.enabled
@@ -23,27 +24,22 @@ function sampleRequest (req, res) {
   const rootSpan = web.root(req)
   if (!rootSpan) return false
 
-  const priority = getSpanPriority(rootSpan)
+  const isSampled = prioritySampler.isSampled(rootSpan)
 
-  if (priority !== AUTO_KEEP && priority !== USER_KEEP) {
+  if (!isSampled) {
     return false
   }
 
   const key = sampledRequests.computeKey(req, res)
-  const isSampled = sampledRequests.isSampled(key)
+  const alreadySampled = sampledRequests.isSampled(key)
 
-  if (isSampled) return false
+  if (alreadySampled) return false
 
   sampledRequests.set(key)
 
   return true
 }
 
-function getSpanPriority (span) {
-  const spanContext = span.context?.()
-  return spanContext._sampling?.priority // default ??
-}
-
 module.exports = {
   configure,
   disable,

diff --git a/packages/dd-trace/test/appsec/api_security_sampler.spec.js b/packages/dd-trace/test/appsec/api_security_sampler.spec.js
@@ -3,73 +3,68 @@
 const proxyquire = require('proxyquire')
 
 describe('API Security Sampler', () => {
-  let apiSecuritySampler, webStub, clock
+  let apiSecuritySampler, webStub, clock, sampler, span
 
   beforeEach(() => {
     webStub = { root: sinon.stub() }
     clock = sinon.useFakeTimers(Date.now())
 
+    sampler = sinon.stub().returns({
+      isSampled: sinon.stub()
+    })
+
     apiSecuritySampler = proxyquire('../../src/appsec/api_security_sampler', {
-      '../plugins/util/web': webStub
+      '../plugins/util/web': webStub,
+      '../priority_sampler': sampler
     })
+
+    apiSecuritySampler.configure({ apiSecurity: { enabled: true, sampleDelay: 30 } })
+
+    span = {
+      context: sinon.stub().returns({})
+    }
   })
 
   afterEach(() => {
     clock.restore()
   })
 
-  describe('sampleRequest', () => {
-    beforeEach(() => {
-      apiSecuritySampler.configure({ apiSecurity: { enabled: true, sampleDelay: 30 } })
-    })
-
-    it('should return false if not enabled', () => {
-      apiSecuritySampler.disable()
-      expect(apiSecuritySampler.sampleRequest({}, {})).to.be.false
-    })
-
-    it('should return false if no root span', () => {
-      webStub.root.returns(null)
-      expect(apiSecuritySampler.sampleRequest({}, {})).to.be.false
-    })
+  it('should return false if not enabled', () => {
+    apiSecuritySampler.disable()
+    expect(apiSecuritySampler.sampleRequest({}, {})).to.be.false
+  })
 
-    it('should return true and put request in cache if priority is AUTO_KEEP', () => {
-      const rootSpan = { context: () => ({ _sampling: { priority: 2 } }) }
-      webStub.root.returns(rootSpan)
-      const req = { url: '/test', method: 'GET' }
-      const res = { statusCode: 200 }
-      expect(apiSecuritySampler.sampleRequest(req, res)).to.be.true
-    })
+  it('should return false if no root span', () => {
+    webStub.root.returns(null)
+    expect(apiSecuritySampler.sampleRequest({}, {})).to.be.false
+  })
 
-    it('should return true and put request in cache if priority is USER_KEEP', () => {
-      const rootSpan = { context: () => ({ _sampling: { priority: 1 } }) }
-      webStub.root.returns(rootSpan)
-      const req = { url: '/test', method: 'GET' }
-      const res = { statusCode: 200 }
-      expect(apiSecuritySampler.sampleRequest(req, res)).to.be.true
-    })
+  it('should return true and put request in cache if priority is AUTO_KEEP or USER_KEEP', () => {
+    webStub.root.returns(span)
+    sampler().isSampled.returns(true)
+    const req = { url: '/test', method: 'GET' }
+    const res = { statusCode: 200 }
+    expect(apiSecuritySampler.sampleRequest(req, res)).to.be.true
+  })
 
-    it('should return false if priority is neither AUTO_KEEP nor USER_KEEP', () => {
-      const rootSpan = { context: () => ({ _sampling: { priority: 0 } }) }
-      webStub.root.returns(rootSpan)
-      expect(apiSecuritySampler.sampleRequest({}, {})).to.be.false
-    })
+  it('should return false if priority is neither AUTO_KEEP nor USER_KEEP', () => {
+    webStub.root.returns(span)
+    sampler().isSampled.returns(false)
+    expect(apiSecuritySampler.sampleRequest({}, {})).to.be.false
   })
 
-  describe('disable', () => {
-    it('should set enabled to false and clear the cache', () => {
-      const req = { url: '/test', method: 'GET' }
-      const res = { statusCode: 200 }
+  it('should set enabled to false and clear the cache', () => {
+    const req = { url: '/test', method: 'GET' }
+    const res = { statusCode: 200 }
 
-      const rootSpan = { context: () => ({ _sampling: { priority: 2 } }) }
-      webStub.root.returns(rootSpan)
+    webStub.root.returns(span)
+    sampler().isSampled.returns(true)
 
-      apiSecuritySampler.configure({ apiSecurity: { enabled: true, sampleDelay: 30 } })
-      expect(apiSecuritySampler.sampleRequest(req, res)).to.be.true
+    apiSecuritySampler.configure({ apiSecurity: { enabled: true, sampleDelay: 30 } })
+    expect(apiSecuritySampler.sampleRequest(req, res)).to.be.true
 
-      apiSecuritySampler.disable()
+    apiSecuritySampler.disable()
 
-      expect(apiSecuritySampler.sampleRequest({}, {})).to.be.false
-    })
+    expect(apiSecuritySampler.sampleRequest({}, {})).to.be.false
   })
 })
diff --git a/packages/dd-trace/test/appsec/index.spec.js b/packages/dd-trace/test/appsec/index.spec.js
@@ -46,6 +46,7 @@ describe('AppSec Index', function () {
   let graphql
   let apiSecuritySampler
   let rasp
+  let sampler
 
   const RULES = { rules: [{ a: 1 }] }
 
@@ -78,6 +79,10 @@ describe('AppSec Index', function () {
       root: sinon.stub()
     }
 
+    sampler = sinon.stub().returns({
+      isSampled: sinon.stub()
+    })
+
     blocking = {
       setTemplates: sinon.stub()
     }
@@ -103,7 +108,8 @@ describe('AppSec Index', function () {
     }
 
     apiSecuritySampler = proxyquire('../../src/appsec/api_security_sampler', {
-      '../plugins/util/web': web
+      '../plugins/util/web': web,
+      '../priority_sampler': sampler
     })
     sinon.spy(apiSecuritySampler, 'sampleRequest')
 
@@ -549,9 +555,12 @@ describe('AppSec Index', function () {
         statusCode: 201
       }
 
-      const rootSpan = { context: () => ({ _sampling: { priority: 2 } }) }
+      const span = {
+        context: sinon.stub().returns({})
+      }
 
-      web.root.returns(rootSpan)
+      web.root.returns(span)
+      sampler().isSampled.returns(true)
 
       AppSec.incomingHttpEndTranslator({ req, res })