Inject trace info as comment to MongoDB operation when dbm propagation is enabled.

docker-compose.yml

@@ -37,7 +37,7 @@ services:
     ports:
       - "127.0.0.1:6379:6379"
   mongo:
-    image: circleci/mongo:3.6
+    image: circleci/mongo:4.4
     platform: linux/amd64
     ports:
       - "127.0.0.1:27017:27017"

packages/datadog-plugin-mongodb-core/src/index.js

@@ -11,8 +11,9 @@ class MongodbCorePlugin extends DatabasePlugin {
   start ({ ns, ops, options = {}, name }) {
     const query = getQuery(ops)
     const resource = truncate(getResource(this, ns, query, name))
-    this.startSpan(this.operationName(), {
-      service: this.serviceName({ pluginConfig: this.config }),
+    const service = this.serviceName({ pluginConfig: this.config })
+    const span = this.startSpan(this.operationName(), {
+      service,
       resource,
       type: 'mongodb',
       kind: 'client',
@@ -24,6 +25,7 @@ class MongodbCorePlugin extends DatabasePlugin {
         'out.port': options.port
       }
     })
+    ops = this.injectDbmCommand(span, ops, service)
   }
 
   getPeerService (tags) {
@@ -34,6 +36,30 @@ class MongodbCorePlugin extends DatabasePlugin {
     }
     return super.getPeerService(tags)
   }
+
+  injectDbmCommand (span, command, serviceName) {
+    const dbmTraceComment = this.createDbmComment(span, serviceName)
+
+    if (!dbmTraceComment) {
+      return command
+    }
+
+    // create a copy of the command to avoid mutating the original
+    const dbmTracedCommand = { ...command }
+
+    if (dbmTracedCommand.comment) {
+      // if the command already has a comment, append the dbm trace comment
+      if (typeof dbmTracedCommand.comment === 'string') {
+        dbmTracedCommand.comment += `,${dbmTraceComment}`
+      } else if (Array.isArray(command.comment)) {
+        dbmTracedCommand.comment.push(dbmTraceComment)
+      } // do nothing if the comment is not a string or an array
+    } else {
+      dbmTracedCommand.comment = dbmTraceComment
+    }
+
+    return dbmTracedCommand
+  }
 }
 
 function sanitizeBigInt (data) {

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

@@ -1,10 +1,14 @@
 'use strict'
 
+const sinon = require('sinon')
 const semver = require('semver')
 const agent = require('../../dd-trace/test/plugins/agent')
 const { ERROR_MESSAGE, ERROR_TYPE, ERROR_STACK } = require('../../dd-trace/src/constants')
 const { expectedSchema, rawExpectedSchema } = require('./naming')
 
+const MongodbCorePlugin = require('../../datadog-plugin-mongodb-core/src/index')
+const ddpv = require('mocha/package.json').version
+
 const withTopologies = fn => {
   withVersions('mongodb-core', ['mongodb-core', 'mongodb'], '<4', (version, moduleName) => {
     describe('using the server topology', () => {
@@ -29,6 +33,7 @@
   let id
   let tracer
   let collection
+  let injectDbmCommandSpy
 
   describe('mongodb-core (core)', () => {
     withTopologies(getServer => {
@@ -397,6 +402,181 @@
           }
         )
       })
+
+      describe('with dbmPropagationMode service', () => {
+        before(() => {
+          return agent.load('mongodb-core', { dbmPropagationMode: 'service' })
+        })
+
+        after(() => {
+          return agent.close({ ritmReset: false })
+        })
+
+        beforeEach(done => {
+          const Server = getServer()
+
+          server = new Server({
+            host: '127.0.0.1',
+            port: 27017,
+            reconnect: false
+          })
+
+          server.on('connect', () => done())
+          server.on('error', done)
+
+          server.connect()
+
+          injectDbmCommandSpy = sinon.spy(MongodbCorePlugin.prototype, 'injectDbmCommand')
+        })
+
+        afterEach(() => {
+          injectDbmCommandSpy?.restore()
+        })
+
+        it('DBM propagation should inject service mode as comment', done => {
+          agent
+            .use(traces => {
+              const span = traces[0][0]
+
+              expect(injectDbmCommandSpy.called).to.be.true
+              const instrumentedCommand = injectDbmCommandSpy.getCall(0).returnValue
+              expect(instrumentedCommand).to.have.property('comment')
+              expect(instrumentedCommand.comment).to.equal(
+                `dddb='${encodeURIComponent(span.meta['db.name'])}',` +
+                'dddbs=\'test-mongodb\',' +
+                'dde=\'tester\',' +
+                `ddh='${encodeURIComponent(span.meta['out.host'])}',` +
+                `ddps='${encodeURIComponent(span.meta.service)}',` +
+                `ddpv='${ddpv}',` +
+                `ddprs='${encodeURIComponent(span.meta['peer.service'])}'`
+              )
+            })
+            .then(done)
+            .catch(done)
+
+          server.insert(`test.${collection}`, [{ a: 1 }], () => {})
+        })
+
+        it('DBM propagation should inject service mode after eixsting str comment', done => {
+          agent
+            .use(traces => {
+              const span = traces[0][0]
+
+              expect(injectDbmCommandSpy.called).to.be.true
+              const instrumentedCommand = injectDbmCommandSpy.getCall(0).returnValue
+              expect(instrumentedCommand).to.have.property('comment')
+              expect(instrumentedCommand.comment).to.equal(
+                `test comment,` +
+                `dddb='${encodeURIComponent(span.meta['db.name'])}',` +
+                'dddbs=\'test-mongodb\',' +
+                'dde=\'tester\',' +
+                `ddh='${encodeURIComponent(span.meta['out.host'])}',` +
+                `ddps='${encodeURIComponent(span.meta.service)}',` +
+                `ddpv='${ddpv}',` +
+                `ddprs='${encodeURIComponent(span.meta['peer.service'])}'`
+              )
+            })
+            .then(done)
+            .catch(done)
+
+            server.command(`test.${collection}`, {
+              find: `test.${collection}`,
+              query: {
+                _id: Buffer.from('1234')
+              },
+              comment: "test comment"
+            }, () => {})
+        })
+
+        it('DBM propagation should inject service mode after eixsting array comment', done => {
+          agent
+            .use(traces => {
+              const span = traces[0][0]
+
+              expect(injectDbmCommandSpy.called).to.be.true
+              const instrumentedCommand = injectDbmCommandSpy.getCall(0).returnValue
+              expect(instrumentedCommand).to.have.property('comment')
+              expect(instrumentedCommand.comment).to.deep.equal([
+                `test comment`,
+                `dddb='${encodeURIComponent(span.meta['db.name'])}',` +
+                'dddbs=\'test-mongodb\',' +
+                'dde=\'tester\',' +
+                `ddh='${encodeURIComponent(span.meta['out.host'])}',` +
+                `ddps='${encodeURIComponent(span.meta.service)}',` +
+                `ddpv='${ddpv}',` +
+                `ddprs='${encodeURIComponent(span.meta['peer.service'])}'`
+              ])
+            })
+            .then(done)
+            .catch(done)
+
+            server.command(`test.${collection}`, {
+              find: `test.${collection}`,
+              query: {
+                _id: Buffer.from('1234')
+              },
+              comment: ["test comment"]
+            }, () => {})
+        })
+      })
+
+      describe('with dbmPropagationMode full', () => {
+        before(() => {
+          return agent.load('mongodb-core', { dbmPropagationMode: 'full' })
+        })
+
+        after(() => {
+          return agent.close({ ritmReset: false })
+        })
+
+        beforeEach(done => {
+          const Server = getServer()
+
+          server = new Server({
+            host: '127.0.0.1',
+            port: 27017,
+            reconnect: false
+          })
+
+          server.on('connect', () => done())
+          server.on('error', done)
+
+          server.connect()
+
+          injectDbmCommandSpy = sinon.spy(MongodbCorePlugin.prototype, 'injectDbmCommand')
+        })
+
+        afterEach(() => {
+          injectDbmCommandSpy?.restore()
+        })
+
+        it('DBM propagation should inject full mode with traceparent as comment', done => {
+          agent
+            .use(traces => {
+              const span = traces[0][0]
+              const traceId = span.meta['_dd.p.tid'] + span.trace_id.toString(16).padStart(16, '0')
+              const spanId = span.span_id.toString(16).padStart(16, '0')
+
+              expect(injectDbmCommandSpy.called).to.be.true
+              const instrumentedCommand = injectDbmCommandSpy.getCall(0).returnValue
+              expect(instrumentedCommand).to.have.property('comment')
+              expect(instrumentedCommand.comment).to.equal(
+                `dddb='${encodeURIComponent(span.meta['db.name'])}',` +
+                'dddbs=\'test-mongodb\',' +
+                'dde=\'tester\',' +
+                `ddh='${encodeURIComponent(span.meta['out.host'])}',` +
+                `ddps='${encodeURIComponent(span.meta.service)}',` +
+                `ddpv='${ddpv}',` +
+                `ddprs='${encodeURIComponent(span.meta['peer.service'])}',` +
+                `traceparent='00-${traceId}-${spanId}-00'`
+              )
+            })
+            .then(done)
+            .catch(done)
+
+          server.insert(`test.${collection}`, [{ a: 1 }], () => {})
+        })
+      })
     })
   })
 })

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

@@ -1,9 +1,13 @@
 'use strict'
 
+const sinon = require('sinon')
 const semver = require('semver')
 const agent = require('../../dd-trace/test/plugins/agent')
 const { expectedSchema, rawExpectedSchema } = require('./naming')
 
+const MongodbCorePlugin = require('../../datadog-plugin-mongodb-core/src/index')
+const ddpv = require('mocha/package.json').version
+
 const withTopologies = fn => {
   const isOldNode = semver.satisfies(process.version, '<=14')
   const range = isOldNode ? '>=2 <6' : '>=2' // TODO: remove when 3.x support is removed.
@@ -44,6 +48,7 @@ describe('Plugin', () => {
   let collection
   let db
   let BSON
+  let injectDbmCommandSpy
 
   describe('mongodb-core', () => {
     withTopologies(createClient => {
@@ -334,6 +339,109 @@ describe('Plugin', () => {
           }
         )
       })
+
+      describe('with dbmPropagationMode service', () => {
+        before(() => {
+          return agent.load('mongodb-core', {
+            dbmPropagationMode: 'service'
+          })
+        })
+
+        after(() => {
+          return agent.close({ ritmReset: false })
+        })
+
+        beforeEach(async () => {
+          client = await createClient()
+          db = client.db('test')
+          collection = db.collection(collectionName)
+
+          injectDbmCommandSpy = sinon.spy(MongodbCorePlugin.prototype, 'injectDbmCommand')
+        })
+
+        afterEach(() => {
+          injectDbmCommandSpy?.restore()
+        })
+
+        it('DBM propagation should inject service mode as comment', done => {
+          agent
+            .use(traces => {
+              const span = traces[0][0]
+
+              expect(injectDbmCommandSpy.called).to.be.true
+              const instrumentedCommand = injectDbmCommandSpy.getCall(0).returnValue
+              expect(instrumentedCommand).to.have.property('comment')
+              expect(instrumentedCommand.comment).to.equal(
+                `dddb='${encodeURIComponent(span.meta['db.name'])}',` +
+                'dddbs=\'test-mongodb\',' +
+                'dde=\'tester\',' +
+                `ddh='${encodeURIComponent(span.meta['out.host'])}',` +
+                `ddps='${encodeURIComponent(span.meta.service)}',` +
+                `ddpv='${ddpv}',` +
+                `ddprs='${encodeURIComponent(span.meta['peer.service'])}'`
+              )
+            })
+            .then(done)
+            .catch(done)
+
+          collection.find({
+            _id: Buffer.from('1234')
+          }).toArray()
+        })
+      })
+
+      describe('with dbmPropagationMode full', () => {
+        before(() => {
+          return agent.load('mongodb-core', {
+            dbmPropagationMode: 'full'
+          })
+        })
+
+        after(() => {
+          return agent.close({ ritmReset: false })
+        })
+
+        beforeEach(async () => {
+          client = await createClient()
+          db = client.db('test')
+          collection = db.collection(collectionName)
+
+          injectDbmCommandSpy = sinon.spy(MongodbCorePlugin.prototype, 'injectDbmCommand')
+        })
+
+        afterEach(() => {
+          injectDbmCommandSpy?.restore()
+        })
+
+        it('DBM propagation should inject full mode with traceparent as comment', done => {
+          agent
+            .use(traces => {
+              const span = traces[0][0]
+              const traceId = span.meta['_dd.p.tid'] + span.trace_id.toString(16).padStart(16, '0')
+              const spanId = span.span_id.toString(16).padStart(16, '0')
+
+              expect(injectDbmCommandSpy.called).to.be.true
+              const instrumentedCommand = injectDbmCommandSpy.getCall(0).returnValue
+              expect(instrumentedCommand).to.have.property('comment')
+              expect(instrumentedCommand.comment).to.equal(
+                `dddb='${encodeURIComponent(span.meta['db.name'])}',` +
+                'dddbs=\'test-mongodb\',' +
+                'dde=\'tester\',' +
+                `ddh='${encodeURIComponent(span.meta['out.host'])}',` +
+                `ddps='${encodeURIComponent(span.meta.service)}',` +
+                `ddpv='${ddpv}',` +
+                `ddprs='${encodeURIComponent(span.meta['peer.service'])}',` +
+                `traceparent='00-${traceId}-${spanId}-00'`
+              )
+            })
+            .then(done)
+            .catch(done)
+
+          collection.find({
+            _id: Buffer.from('1234')
+          }).toArray()
+        })
+      })
     })
   })
 })