Skip to content

Commit 4dccc16

Browse files
na-jina-ji
authored
[DB plugins] trigger sampling of span before injecting the trace parent comment (#6653)
* feat(database): trigger sampling of span before injecting the trace parent commment The database plugins supported by the Database Monitoring require to inject the trace parent as a comment into the query statement. The trace parent to be fully valid need to contains the correct value for the "sampled" flag. However, we used to be injecting the trace parent before the decision to sample a span or not has been taken. Because of that, all the trace parents injected had the sampled flag to 00. This commits fix this issue by running the sampling decision before injecting the comment when needed. * test(database): use the correct value for trace parents in DBM comments * test(span): test that the sample method is correctly calling the sampler * refactor: update the dd-trace-api definition * test(database): add test cases for when the span was not sampled * refactor(span): remove public API to force span sampling Instead, the span processor is called directly to trigger the sampling decision * test(database): configure tracer instead of mocking priority sampler * test(database): reset tracer configuration in a afterEach clause
1 parent f87e8d6 commit 4dccc16

File tree

8 files changed

+211
-10
lines changed

8 files changed

+211
-10
lines changed

‎packages/datadog-plugin-mongodb-core/test/core.spec.js‎

Lines changed: 56 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -598,7 +598,7 @@ describe('Plugin', () => {
598598
}, () => {})
599599
})
600600

601-
it('DBM propagation should inject service mode after eixsting array comment', done => {
601+
it('DBM propagation should inject service mode after existing array comment', done => {
602602
agent
603603
.assertSomeTraces(traces => {
604604
const span = traces[0][0]
@@ -676,7 +676,7 @@ describe('Plugin', () => {
676676
`ddps='${encodeURIComponent(span.meta.service)}',` +
677677
`ddpv='${ddpv}',` +
678678
`ddprs='${encodeURIComponent(span.meta['peer.service'])}',` +
679-
`traceparent='00-${traceId}-${spanId}-00'`
679+
`traceparent='00-${traceId}-${spanId}-01'`
680680
)
681681
})
682682
.then(done)
@@ -685,6 +685,60 @@ describe('Plugin', () => {
685685
server.insert(`test.${collection}`, [{ a: 1 }], () => {})
686686
})
687687
})
688+
689+
describe('with dbmPropagationMode full but sampling disabled', () => {
690+
before(() => {
691+
tracer._tracer.configure({ env: 'tester', sampler: { sampleRate: 0 } })
692+
693+
return agent.load('mongodb-core', { dbmPropagationMode: 'full' })
694+
})
695+
696+
after(() => {
697+
tracer._tracer.configure({ env: 'tester', sampler: { sampleRate: 1 } })
698+
699+
return agent.close({ ritmReset: false })
700+
})
701+
702+
beforeEach(done => {
703+
const Server = getServer()
704+
705+
server = new Server({
706+
host: '127.0.0.1',
707+
port: 27017,
708+
reconnect: false
709+
})
710+
711+
server.on('connect', () => done())
712+
server.on('error', done)
713+
714+
server.connect()
715+
716+
startSpy = sinon.spy(MongodbCorePlugin.prototype, 'start')
717+
})
718+
719+
afterEach(() => {
720+
startSpy?.restore()
721+
})
722+
723+
it(
724+
'DBM propagation should inject full mode with traceparent as comment and the rejected sampling decision',
725+
done => {
726+
agent
727+
.assertSomeTraces(traces => {
728+
const span = traces[0][0]
729+
const traceId = span.meta['_dd.p.tid'] + span.trace_id.toString(16).padStart(16, '0')
730+
const spanId = span.span_id.toString(16).padStart(16, '0')
731+
732+
expect(startSpy.called).to.be.true
733+
const { comment } = startSpy.getCall(0).args[0].ops
734+
expect(comment).to.include(`traceparent='00-${traceId}-${spanId}-00'`)
735+
})
736+
.then(done)
737+
.catch(done)
738+
739+
server.insert(`test.${collection}`, [{ a: 1 }], () => {})
740+
})
741+
})
688742
})
689743
})
690744
})

‎packages/datadog-plugin-mongodb-core/test/mongodb.spec.js‎

Lines changed: 50 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -602,7 +602,7 @@ describe('Plugin', () => {
602602
`ddps='${encodeURIComponent(span.meta.service)}',` +
603603
`ddpv='${ddpv}',` +
604604
`ddprs='${encodeURIComponent(span.meta['peer.service'])}',` +
605-
`traceparent='00-${traceId}-${spanId}-00'`
605+
`traceparent='00-${traceId}-${spanId}-01'`
606606
)
607607
})
608608
.then(done)
@@ -614,6 +614,55 @@ describe('Plugin', () => {
614614
})
615615
})
616616

617+
describe('with dbmPropagationMode full but sampling disabled', () => {
618+
before(() => {
619+
tracer._tracer.configure({ env: 'tester', sampler: { sampleRate: 0 } })
620+
621+
return agent.load('mongodb-core', {
622+
dbmPropagationMode: 'full'
623+
})
624+
})
625+
626+
after(() => {
627+
tracer._tracer.configure({ env: 'tester', sampler: { sampleRate: 1 } })
628+
629+
return agent.close({ ritmReset: false })
630+
})
631+
632+
beforeEach(async () => {
633+
client = await createClient()
634+
db = client.db('test')
635+
collection = db.collection(collectionName)
636+
637+
startSpy = sinon.spy(MongodbCorePlugin.prototype, 'start')
638+
})
639+
640+
afterEach(() => {
641+
startSpy?.restore()
642+
})
643+
644+
it(
645+
'DBM propagation should inject full mode with traceparent as comment and the rejected sampling decision',
646+
done => {
647+
agent
648+
.assertSomeTraces(traces => {
649+
const span = traces[0][0]
650+
const traceId = span.meta['_dd.p.tid'] + span.trace_id.toString(16).padStart(16, '0')
651+
const spanId = span.span_id.toString(16).padStart(16, '0')
652+
653+
expect(startSpy.called).to.be.true
654+
const { comment } = startSpy.getCall(0).args[0].ops
655+
expect(comment).to.include(`traceparent='00-${traceId}-${spanId}-00'`)
656+
})
657+
.then(done)
658+
.catch(done)
659+
660+
collection.find({
661+
_id: Buffer.from('1234')
662+
}).toArray()
663+
})
664+
})
665+
617666
describe('with heartbeatEnabled configuration', () => {
618667
describe('when heartbeat tracing is disabled via config', () => {
619668
before(() => {

‎packages/datadog-plugin-mysql/test/index.spec.js‎

Lines changed: 32 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -468,6 +468,8 @@ describe('Plugin', () => {
468468
connection.end(() => {
469469
agent.close({ ritmReset: false }).then(done)
470470
})
471+
472+
tracer._tracer.configure({ env: 'tester', sampler: { sampleRate: 1 } })
471473
})
472474

473475
beforeEach(async () => {
@@ -491,13 +493,26 @@ describe('Plugin', () => {
491493

492494
expect(queryText).to.equal(
493495
`/*dddb='db',dddbs='post',dde='tester',ddh='127.0.0.1',ddps='test',ddpv='${ddpv}',` +
494-
`traceparent='00-${traceId}-${spanId}-00'*/ SELECT 1 + 1 AS solution`)
496+
`traceparent='00-${traceId}-${spanId}-01'*/ SELECT 1 + 1 AS solution`)
495497
}).then(done, done)
496498
connection.query('SELECT 1 + 1 AS solution', () => {
497499
queryText = connection._protocol._queue[0].sql
498500
})
499501
})
500502

503+
it('query text should contain rejected sampling decision in the traceparent', done => {
504+
tracer._tracer.configure({ env: 'tester', sampler: { sampleRate: 0 } })
505+
let queryText = ''
506+
507+
agent.assertSomeTraces(traces => {
508+
expect(queryText).to.include('-00\'*/ SELECT 1 + 1 AS solution')
509+
}).then(done, done)
510+
511+
connection.query('SELECT 1 + 1 AS solution', () => {
512+
queryText = connection._protocol._queue[0].sql
513+
})
514+
})
515+
501516
it('query should inject _dd.dbm_trace_injected into span', done => {
502517
agent.assertSomeTraces(traces => {
503518
expect(traces[0][0].meta).to.have.property('_dd.dbm_trace_injected', 'true')
@@ -548,6 +563,8 @@ describe('Plugin', () => {
548563
pool.end(() => {
549564
agent.close({ ritmReset: false }).then(done)
550565
})
566+
567+
tracer._tracer.configure({ env: 'tester', sampler: { sampleRate: 1 } })
551568
})
552569

553570
beforeEach(async () => {
@@ -571,8 +588,21 @@ describe('Plugin', () => {
571588

572589
expect(queryText).to.equal(
573590
`/*dddb='db',dddbs='post',dde='tester',ddh='127.0.0.1',ddps='test',ddpv='${ddpv}',` +
574-
`traceparent='00-${traceId}-${spanId}-00'*/ SELECT 1 + 1 AS solution`)
591+
`traceparent='00-${traceId}-${spanId}-01'*/ SELECT 1 + 1 AS solution`)
592+
}).then(done, done)
593+
pool.query('SELECT 1 + 1 AS solution', () => {
594+
queryText = pool._allConnections[0]._protocol._queue[0].sql
595+
})
596+
})
597+
598+
it('query text should contain rejected sampling decision in the traceparent', done => {
599+
tracer._tracer.configure({ env: 'tester', sampler: { sampleRate: 0 } })
600+
let queryText = ''
601+
602+
agent.assertSomeTraces(() => {
603+
expect(queryText).to.include('-00\'*/ SELECT 1 + 1 AS solution')
575604
}).then(done, done)
605+
576606
pool.query('SELECT 1 + 1 AS solution', () => {
577607
queryText = pool._allConnections[0]._protocol._queue[0].sql
578608
})

‎packages/datadog-plugin-mysql2/test/index.spec.js‎

Lines changed: 40 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -459,6 +459,8 @@ describe('Plugin', () => {
459459
connection.end(() => {
460460
agent.close({ ritmReset: false }).then(done)
461461
})
462+
463+
tracer._tracer.configure({ env: 'tester', sampler: { sampleRate: 1 } })
462464
})
463465

464466
beforeEach(async () => {
@@ -482,13 +484,30 @@ describe('Plugin', () => {
482484

483485
expect(queryText).to.equal(
484486
`/*dddb='db',dddbs='post',dde='tester',ddh='127.0.0.1',ddps='test',ddpv='${ddpv}',` +
485-
`traceparent='00-${traceId}-${spanId}-00'*/ SELECT 1 + 1 AS solution`)
487+
`traceparent='00-${traceId}-${spanId}-01'*/ SELECT 1 + 1 AS solution`)
486488
}).then(done, done)
487489
const connect = connection.query('SELECT 1 + 1 AS solution', () => {
488490
queryText = connect.sql
489491
})
490492
})
491493

494+
it('query text should contain rejected sampling decision in the traceparent', done => {
495+
tracer._tracer.configure({ env: 'tester', sampler: { sampleRate: 0 } })
496+
let queryText = ''
497+
498+
agent.assertSomeTraces(traces => {
499+
const expectedTimePrefix = traces[0][0].meta['_dd.p.tid'].toString(16).padStart(16, '0')
500+
const traceId = expectedTimePrefix + traces[0][0].trace_id.toString(16).padStart(16, '0')
501+
const spanId = traces[0][0].span_id.toString(16).padStart(16, '0')
502+
503+
expect(queryText).to.include(`traceparent='00-${traceId}-${spanId}-00'*/ SELECT 1 + 1 AS solution`)
504+
}).then(done, done)
505+
506+
const connect = connection.query('SELECT 1 + 1 AS solution', () => {
507+
queryText = connect.sql
508+
})
509+
})
510+
492511
it('query should inject _dd.dbm_trace_injected into span', done => {
493512
agent.assertSomeTraces(traces => {
494513
expect(traces[0][0].meta).to.have.property('_dd.dbm_trace_injected', 'true')
@@ -539,6 +558,8 @@ describe('Plugin', () => {
539558
pool.end(() => {
540559
agent.close({ ritmReset: false }).then(done)
541560
})
561+
562+
tracer._tracer.configure({ env: 'tester', sampler: { sampleRate: 1 } })
542563
})
543564

544565
beforeEach(async () => {
@@ -562,13 +583,30 @@ describe('Plugin', () => {
562583

563584
expect(queryText).to.equal(
564585
`/*dddb='db',dddbs='post',dde='tester',ddh='127.0.0.1',ddps='test',ddpv='${ddpv}',` +
565-
`traceparent='00-${traceId}-${spanId}-00'*/ SELECT 1 + 1 AS solution`)
586+
`traceparent='00-${traceId}-${spanId}-01'*/ SELECT 1 + 1 AS solution`)
566587
}).then(done, done)
567588
const queryPool = pool.query('SELECT 1 + 1 AS solution', () => {
568589
queryText = queryPool.sql
569590
})
570591
})
571592

593+
it('query text should contain rejected sampling decision in the traceparent', done => {
594+
tracer._tracer.configure({ env: 'tester', sampler: { sampleRate: 0 } })
595+
let queryText = ''
596+
597+
agent.assertSomeTraces(traces => {
598+
const expectedTimePrefix = traces[0][0].meta['_dd.p.tid'].toString(16).padStart(16, '0')
599+
const traceId = expectedTimePrefix + traces[0][0].trace_id.toString(16).padStart(16, '0')
600+
const spanId = traces[0][0].span_id.toString(16).padStart(16, '0')
601+
602+
expect(queryText).to.include(`traceparent='00-${traceId}-${spanId}-00'*/ SELECT 1 + 1 AS solution`)
603+
}).then(done, done)
604+
605+
const queryPool = pool.query('SELECT 1 + 1 AS solution', () => {
606+
queryText = queryPool.sql
607+
})
608+
})
609+
572610
it('query should inject _dd.dbm_trace_injected into span', done => {
573611
agent.assertSomeTraces(traces => {
574612
expect(traces[0][0].meta).to.have.property('_dd.dbm_trace_injected', 'true')

‎packages/datadog-plugin-pg/test/index.spec.js‎

Lines changed: 19 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -666,6 +666,8 @@ describe('Plugin', () => {
666666

667667
afterEach((done) => {
668668
client.end(done)
669+
670+
tracer._tracer.configure({ env: 'tester', sampler: { sampleRate: 1 } })
669671
})
670672

671673
it('query config objects should be handled', async () => {
@@ -685,7 +687,23 @@ describe('Plugin', () => {
685687

686688
expect(queryText).to.equal(
687689
`/*dddb='postgres',dddbs='post',dde='tester',ddh='127.0.0.1',ddps='test',ddpv='${ddpv}',` +
688-
`traceparent='00-${traceId}-${spanId}-00'*/ SELECT $1::text as message`)
690+
`traceparent='00-${traceId}-${spanId}-01'*/ SELECT $1::text as message`)
691+
})
692+
})
693+
694+
it('query text should contain rejected sampling decision in the traceparent', async () => {
695+
tracer._tracer.configure({ env: 'tester', sampler: { sampleRate: 0 } })
696+
const query = {
697+
text: 'SELECT $1::text as message'
698+
}
699+
700+
const queryPromise = client.query(query, ['Hello world!'])
701+
const queryText = client.queryQueue[0].text
702+
703+
await queryPromise
704+
705+
await agent.assertSomeTraces(() => {
706+
expect(queryText).to.include('-00\'*/ SELECT $1::text as message')
689707
})
690708
})
691709

‎packages/dd-trace/src/plugins/database.js‎

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -77,6 +77,7 @@ class DatabasePlugin extends StoragePlugin {
7777
return servicePropagation
7878
} else if (mode === 'full') {
7979
span.setTag('_dd.dbm_trace_injected', 'true')
80+
span._processor.sample(span)
8081
const traceparent = span._spanContext.toTraceparent()
8182
return `${servicePropagation},traceparent='${traceparent}'`
8283
}

‎packages/dd-trace/src/span_processor.js‎

Lines changed: 7 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -26,6 +26,12 @@ class SpanProcessor {
2626
this._gitMetadataTagger = new GitMetadataTagger(config)
2727
}
2828

29+
sample (span) {
30+
const spanContext = span.context()
31+
this._prioritySampler.sample(spanContext)
32+
this._spanSampler.sample(spanContext)
33+
}
34+
2935
process (span) {
3036
const spanContext = span.context()
3137
const active = []
@@ -40,8 +46,7 @@ class SpanProcessor {
4046
return
4147
}
4248
if (started.length === finished.length || finished.length >= flushMinSpans) {
43-
this._prioritySampler.sample(spanContext)
44-
this._spanSampler.sample(spanContext)
49+
this.sample(span)
4550
this._gitMetadataTagger.tagGitMetadata(spanContext)
4651

4752
let isChunkRoot = true

‎packages/dd-trace/test/span_processor.spec.js‎

Lines changed: 6 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -72,6 +72,12 @@ describe('SpanProcessor', () => {
7272
expect(prioritySampler.sample).to.have.been.calledWith(finishedSpan.context())
7373
})
7474

75+
it('should generate sampling priority when sampling manually', () => {
76+
processor.sample(finishedSpan)
77+
78+
expect(prioritySampler.sample).to.have.been.calledWith(finishedSpan.context())
79+
})
80+
7581
it('should erase the trace once finished', () => {
7682
trace.started = [finishedSpan]
7783
trace.finished = [finishedSpan]

0 commit comments

Comments
 (0)