add feature to disable mongo heartbeats spans (#5526)

* add feature to disable mongodb heartbeat tracing

Author: wconti27

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

@@ -1,14 +1,29 @@
 'use strict'
 
 const DatabasePlugin = require('../../dd-trace/src/plugins/database')
+const coalesce = require('koalas')
 
 class MongodbCorePlugin extends DatabasePlugin {
   static get id () { return 'mongodb-core' }
   static get component () { return 'mongodb' }
   // avoid using db.name for peer.service since it includes the collection name
   // should be removed if one day this will be fixed
   static get peerServicePrecursors () { return [] }
+
+  configure (config) {
+    super.configure(config)
+    this.config.heartbeatEnabled = coalesce(
+      config.heartbeatEnabled,
+      process.env.DD_TRACE_MONGODB_HEARTBEAT_ENABLED,
+      true
+    )
+  }
+
   start ({ ns, ops, options = {}, name }) {
+    // heartbeat commands can be disabled if this.config.heartbeatEnabled is false
+    if (!this.config.heartbeatEnabled && isHeartbeat(ops, this.config)) {
+      return
+    }
     const query = getQuery(ops)
     const resource = truncate(getResource(this, ns, query, name))
     const service = this.serviceName({ pluginConfig: this.config })
@@ -153,4 +168,9 @@ function isBinary (val) {
   return val && val._bsontype === 'Binary'
 }
 
+function isHeartbeat (ops, config) {
+  // Check if it's a heartbeat command hello: 1 or helloOk: 1
+  return ops && typeof ops === 'object' && (ops.hello === 1 || ops.helloOk === true)
+}
+
 module.exports = MongodbCorePlugin

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

@@ -459,6 +459,95 @@ describe('Plugin', () => {
           }).toArray()
         })
       })
+
+      describe('with heartbeatEnabled configuration', () => {
+        describe('when heartbeat tracing is disabled', () => {
+          before(() => {
+            return agent.load('mongodb-core', {
+              heartbeatEnabled: false
+            })
+          })
+
+          after(() => {
+            return agent.close({ ritmReset: false })
+          })
+
+          beforeEach(async () => {
+            client = await createClient()
+            db = client.db('test')
+          })
+
+          it('should NOT create a span for heartbeat commands', (done) => {
+            const parentSpan = tracer.startSpan('test.parent')
+
+            agent
+              .use(traces => {
+                // Should only receive the trace for the parent span
+                expect(traces[0]).to.have.length(1)
+                const span = traces[0][0]
+                expect(span.name).to.equal('test.parent')
+              })
+              .then(done)
+
+            // Activate parent span scope and trigger heartbeat command
+            tracer.scope().activate(parentSpan, async () => {
+              // Admin connect should be all that is needed to trigger heartbeat command for newer versions of mongo
+              client = await createClient()
+              db = client.db('test')
+
+              // but we should send a test heartbeat command since older versions of mongo don't auto-send heartbeats
+              db.command({ hello: 1 })
+              setTimeout(() => parentSpan.finish(), 50)
+            })
+          })
+        })
+
+        describe('when heartbeat tracing is enabled (default)', () => {
+          before(() => {
+            return agent.load('mongodb-core', {
+              heartbeatEnabled: true
+            })
+          })
+
+          after(() => {
+            return agent.close({ ritmReset: false })
+          })
+
+          beforeEach(async () => {
+          })
+
+          it('should create a child span for heartbeat commands', (done) => {
+            const parentSpan = tracer.startSpan('test.parent')
+
+            agent
+              .use(traces => {
+                expect(traces[0]).to.have.length.at.least(2)
+                const rootSpan = traces[0][0]
+
+                expect(rootSpan.name).to.equal('test.parent')
+
+                // assert that some child spans were created, these are the heartbeat spans
+                // don't assert on exact number of spans because it's dynamic
+                for (const childSpan of traces[0].slice(1)) {
+                  expect(childSpan.name).to.equal(expectedSchema.outbound.opName)
+                  expect(childSpan.parent_id.toString()).to.equal(rootSpan.span_id.toString()) // Verify parent-child
+                }
+              })
+              .then(done)
+
+            // Activate parent span scope and trigger heartbeat command
+            tracer.scope().activate(parentSpan, async () => {
+              // Admin connect should be all that is needed to trigger heartbeat command for newer versions of mongo
+              client = await createClient()
+              db = client.db('test')
+
+              // but we should send a test heartbeat command since older versions of mongo don't auto-send heartbeats
+              db.command({ hello: 1 })
+              setTimeout(() => parentSpan.finish(), 200)
+            })
+          })
+        })
+      })
     })
   })
 })