fix: Updated cassandra-driver instrumentation to properly trace promise based executions (#2351)

Author: svetlanabrennan

lib/instrumentation/cassandra-driver.js

@@ -29,7 +29,7 @@ module.exports = function initialize(_agent, cassandra, _moduleName, shim) {
     ClientProto,
     ['connect', 'shutdown'],
     function operationSpec(shim, _fn, name) {
-      return new OperationSpec({ callback: shim.LAST, name })
+      return new OperationSpec({ callback: shim.LAST, name, promise: true })
     }
   )
 
@@ -39,7 +39,8 @@ module.exports = function initialize(_agent, cassandra, _moduleName, shim) {
       '_execute',
       new QuerySpec({
         query: shim.FIRST,
-        callback: shim.LAST
+        callback: shim.LAST,
+        promise: true
       })
     )
 
@@ -64,7 +65,8 @@ module.exports = function initialize(_agent, cassandra, _moduleName, shim) {
       '_innerExecute',
       new QuerySpec({
         query: shim.FIRST,
-        callback: shim.LAST
+        callback: shim.LAST,
+        promise: true
       })
     )
 
@@ -109,7 +111,8 @@ module.exports = function initialize(_agent, cassandra, _moduleName, shim) {
     'batch',
     new QuerySpec({
       query: findBatchQueryArg,
-      callback: shim.LAST
+      callback: shim.LAST,
+      promise: true
     })
   )
 }

test/versioned/cassandra-driver/query.tap.js

@@ -11,6 +11,7 @@ const helper = require('../../lib/agent_helper')
 
 const agent = helper.instrumentMockedAgent()
 const cassandra = require('cassandra-driver')
+const { findSegment } = require('../../lib/metrics_helper')
 
 // constants for keyspace and table creation
 const KS = 'test'
@@ -25,6 +26,26 @@ const client = new cassandra.Client({
   localDataCenter: 'datacenter1'
 })
 
+const colValArr = ['Jim', 'Bob', 'Joe']
+const pkValArr = [111, 222, 333]
+let insQuery = 'INSERT INTO ' + KS + '.' + FAM + ' (' + PK + ',' + COL
+insQuery += ') VALUES(?, ?);'
+
+const insArr = [
+  { query: insQuery, params: [pkValArr[0], colValArr[0]] },
+  { query: insQuery, params: [pkValArr[1], colValArr[1]] },
+  { query: insQuery, params: [pkValArr[2], colValArr[2]] }
+]
+
+const hints = [
+  ['int', 'varchar'],
+  ['int', 'varchar'],
+  ['int', 'varchar']
+]
+
+let selQuery = 'SELECT * FROM ' + KS + '.' + FAM + ' WHERE '
+selQuery += PK + ' = 111;'
+
 /**
  * Deletion of testing keyspace if already exists,
  * then recreation of a testable keyspace and table
@@ -33,7 +54,7 @@ const client = new cassandra.Client({
  * @param Callback function to set off running the tests
  */
 
-async function cassSetup(runTest) {
+async function cassSetup() {
   const setupClient = new cassandra.Client({
     contactPoints: [params.cassandra_host],
     protocolOptions: params.cassandra_port,
@@ -66,146 +87,151 @@ async function cassSetup(runTest) {
   await runCommand(famCreate)
 
   setupClient.shutdown()
-  runTest()
 }
 
 test('Cassandra instrumentation', { timeout: 5000 }, async function testInstrumentation(t) {
-  t.plan(1)
-  await cassSetup(runTest)
-
-  function runTest() {
-    t.test('executeBatch', function (t) {
-      t.notOk(agent.getTransaction(), 'no transaction should be in play')
-      helper.runInTransaction(agent, function transactionInScope(tx) {
-        const transaction = agent.getTransaction()
-        t.ok(transaction, 'transaction should be visible')
-        t.equal(tx, transaction, 'We got the same transaction')
-        const colValArr = ['Jim', 'Bob', 'Joe']
-        const pkValArr = [111, 222, 333]
-        let insQuery = 'INSERT INTO ' + KS + '.' + FAM + ' (' + PK + ',' + COL
-        insQuery += ') VALUES(?, ?);'
-
-        const insArr = [
-          { query: insQuery, params: [pkValArr[0], colValArr[0]] },
-          { query: insQuery, params: [pkValArr[1], colValArr[1]] },
-          { query: insQuery, params: [pkValArr[2], colValArr[2]] }
-        ]
-
-        const hints = [
-          ['int', 'varchar'],
-          ['int', 'varchar'],
-          ['int', 'varchar']
-        ]
-
-        client.batch(insArr, { hints: hints }, function done(error, ok) {
+  t.before(async function () {
+    await cassSetup()
+  })
+
+  t.teardown(function tearDown() {
+    helper.unloadAgent(agent)
+    client.shutdown()
+  })
+
+  t.afterEach(() => {
+    agent.queries.clear()
+    agent.metrics.clear()
+  })
+
+  t.test('executeBatch - callback style', function (t) {
+    t.notOk(agent.getTransaction(), 'no transaction should be in play')
+    helper.runInTransaction(agent, function transactionInScope(tx) {
+      const transaction = agent.getTransaction()
+      t.ok(transaction, 'transaction should be visible')
+      t.equal(tx, transaction, 'We got the same transaction')
+
+      client.batch(insArr, { hints: hints }, function done(error, ok) {
+        if (error) {
+          t.error(error)
+          return t.end()
+        }
+
+        t.ok(agent.getTransaction(), 'transaction should still be visible')
+        t.ok(ok, 'everything should be peachy after setting')
+
+        client.execute(selQuery, function (error, value) {
           if (error) {
-            t.fail(error)
-            return t.end()
+            return t.error(error)
           }
 
-          t.ok(agent.getTransaction(), 'transaction should still be visible')
-          t.ok(ok, 'everything should be peachy after setting')
-
-          let selQuery = 'SELECT * FROM ' + KS + '.' + FAM + ' WHERE '
-          selQuery += PK + ' = 111;'
-          client.execute(selQuery, function (error, value) {
-            if (error) {
-              return t.fail(error)
-            }
+          t.ok(agent.getTransaction(), 'transaction should still still be visible')
+          t.equal(value.rows[0][COL], colValArr[0], 'Cassandra client should still work')
+
+          t.equal(
+            transaction.trace.root.children.length,
+            1,
+            'there should be only one child of the root'
+          )
+          verifyTrace(t, transaction.trace, KS + '.' + FAM)
+          transaction.end()
+          checkMetric(t)
+          t.end()
+        })
+      })
+    })
+  })
 
+  t.test('executeBatch - promise style', function (t) {
+    t.notOk(agent.getTransaction(), 'no transaction should be in play')
+    helper.runInTransaction(agent, function transactionInScope(tx) {
+      const transaction = agent.getTransaction()
+      t.ok(transaction, 'transaction should be visible')
+      t.equal(tx, transaction, 'We got the same transaction')
+
+      client.batch(insArr, { hints: hints }).then(function () {
+        client
+          .execute(selQuery)
+          .then((result) => {
             t.ok(agent.getTransaction(), 'transaction should still still be visible')
-            t.equal(value.rows[0][COL], colValArr[0], 'Cassandra client should still work')
-
-            const trace = transaction.trace
-            t.ok(trace, 'trace should exist')
-            t.ok(trace.root, 'root element should exist')
-
-            t.equal(trace.root.children.length, 1, 'there should be only one child of the root')
-
-            const setSegment = trace.root.children[0]
-            t.ok(setSegment, 'trace segment for insert should exist')
-            if (setSegment) {
-              t.equal(
-                setSegment.name,
-                'Datastore/statement/Cassandra/test.testFamily/insert/batch',
-                'should register the executeBatch'
-              )
-              t.ok(
-                setSegment.children.length >= 2,
-                'set should have atleast a dns lookup and callback child'
-              )
-
-              const setSegmentAttributes = setSegment.getAttributes()
-              t.equal(setSegmentAttributes.product, 'Cassandra', 'should set product attribute')
-              t.equal(setSegmentAttributes.port_path_or_id, '9042', 'should set port attribute')
-              t.equal(
-                setSegmentAttributes.database_name,
-                'test',
-                'should set database_name attribute'
-              )
-              t.equal(
-                setSegmentAttributes.host,
-                agent.config.getHostnameSafe(),
-                'should set host attribute'
-              )
-
-              const childIndex = setSegment.children.length - 1
-              const getSegment = setSegment.children[childIndex].children[0]
-              t.ok(getSegment, 'trace segment for select should exist')
-              if (getSegment) {
-                t.equal(
-                  getSegment.name,
-                  'Datastore/statement/Cassandra/test.testFamily/select',
-                  'should register the execute'
-                )
-
-                t.ok(getSegment.children.length >= 1, 'get should have a callback segment')
-
-                const getSegmentAttributes = getSegment.getAttributes()
-                t.equal(
-                  getSegmentAttributes.product,
-                  'Cassandra',
-                  'get should set product attribute'
-                )
-                t.equal(
-                  getSegmentAttributes.port_path_or_id,
-                  '9042',
-                  'get should set port attribute'
-                )
-                t.equal(
-                  getSegmentAttributes.database_name,
-                  'test',
-                  'get should set database_name attribute'
-                )
-                t.equal(
-                  getSegmentAttributes.host,
-                  agent.config.getHostnameSafe(),
-                  'get should set host attribute'
-                )
-
-                t.ok(getSegment.timer.hrDuration, 'trace segment should have ended')
-              }
-            }
-
+            t.equal(result.rows[0][COL], colValArr[0], 'Cassandra client should still work')
+
+            t.equal(
+              transaction.trace.root.children.length,
+              2,
+              'there should be two children of the root'
+            )
+            verifyTrace(t, transaction.trace, KS + '.' + FAM)
             transaction.end()
-            checkMetric('Datastore/operation/Cassandra/insert', 1)
-            checkMetric('Datastore/allWeb', 2)
-            checkMetric('Datastore/Cassandra/allWeb', 2)
-            checkMetric('Datastore/Cassandra/all', 2)
-            checkMetric('Datastore/all', 2)
-            checkMetric('Datastore/statement/Cassandra/test.testFamily/insert', 1)
-            checkMetric('Datastore/operation/Cassandra/select', 1)
-            checkMetric('Datastore/statement/Cassandra/test.testFamily/select', 1)
-
+            checkMetric(t)
+          })
+          .catch((error) => {
+            t.error(error)
+          })
+          .finally(() => {
             t.end()
           })
+      })
+    })
+  })
+
+  t.test('executeBatch - slow query', function (t) {
+    t.notOk(agent.getTransaction(), 'no transaction should be in play')
+    helper.runInTransaction(agent, function transactionInScope(tx) {
+      // enable slow queries
+      agent.config.transaction_tracer.explain_threshold = 1
+      agent.config.transaction_tracer.record_sql = 'raw'
+      agent.config.slow_sql.enabled = true
+
+      const transaction = agent.getTransaction()
+      t.ok(transaction, 'transaction should be visible')
+      t.equal(tx, transaction, 'We got the same transaction')
+
+      client.batch(insArr, { hints: hints }, function done(error, ok) {
+        if (error) {
+          t.error(error)
+          return t.end()
+        }
+
+        const slowQuery = 'SELECT * FROM ' + KS + '.' + FAM
+        t.ok(agent.getTransaction(), 'transaction should still be visible')
+        t.ok(ok, 'everything should be peachy after setting')
+
+        client.execute(slowQuery, function (error) {
+          if (error) {
+            return t.error(error)
+          }
+
+          verifyTrace(t, transaction.trace, KS + '.' + FAM)
+          transaction.end()
+          t.ok(agent.queries.samples.size > 0, 'there should be a slow query')
+          checkMetric(t)
+          t.end()
         })
       })
+    })
+  })
 
-      function checkMetric(name, count, scoped) {
-        const agentMetrics = agent.metrics._metrics
-        const metric = agentMetrics[scoped ? 'scoped' : 'unscoped'][name]
-        t.ok(metric, 'metric "' + name + '" should exist')
+  function checkMetric(t, scoped) {
+    const agentMetrics = agent.metrics._metrics
+
+    const expected = {
+      'Datastore/operation/Cassandra/insert': 1,
+      'Datastore/allWeb': 2,
+      'Datastore/Cassandra/allWeb': 2,
+      'Datastore/Cassandra/all': 2,
+      'Datastore/all': 2,
+      'Datastore/statement/Cassandra/test.testFamily/insert': 1,
+      'Datastore/operation/Cassandra/select': 1,
+      'Datastore/statement/Cassandra/test.testFamily/select': 1
+    }
+
+    for (const expectedMetric in expected) {
+      if (expected.hasOwnProperty(expectedMetric)) {
+        const count = expected[expectedMetric]
+
+        const metric = agentMetrics[scoped ? 'scoped' : 'unscoped'][expectedMetric]
+        t.ok(metric, 'metric "' + expectedMetric + '" should exist')
         if (!metric) {
           return
         }
@@ -217,11 +243,54 @@ test('Cassandra instrumentation', { timeout: 5000 }, async function testInstrume
         t.ok(metric.max, 'should have set max')
         t.ok(metric.sumOfSquares, 'should have set sumOfSquares')
       }
-    })
+    }
+  }
 
-    t.teardown(function tearDown() {
-      helper.unloadAgent(agent)
-      client.shutdown()
-    })
+  function verifyTrace(t, trace, table) {
+    t.ok(trace, 'trace should exist')
+    t.ok(trace.root, 'root element should exist')
+
+    const setSegment = findSegment(
+      trace.root,
+      'Datastore/statement/Cassandra/' + table + '/insert/batch'
+    )
+
+    t.ok(setSegment, 'trace segment for insert should exist')
+
+    if (setSegment) {
+      verifyTraceSegment(t, setSegment, 'insert/batch')
+
+      t.ok(
+        setSegment.children.length >= 2,
+        'set should have at least a dns lookup and callback/promise child'
+      )
+
+      const getSegment = findSegment(
+        trace.root,
+        'Datastore/statement/Cassandra/' + table + '/select'
+      )
+      t.ok(getSegment, 'trace segment for select should exist')
+
+      if (getSegment) {
+        verifyTraceSegment(t, getSegment, 'select')
+
+        t.ok(getSegment.children.length >= 1, 'get should have a callback/promise segment')
+        t.ok(getSegment.timer.hrDuration, 'trace segment should have ended')
+      }
+    }
+  }
+
+  function verifyTraceSegment(t, segment, queryType) {
+    t.equal(
+      segment.name,
+      'Datastore/statement/Cassandra/' + KS + '.' + FAM + '/' + queryType,
+      'should register the execute'
+    )
+
+    const segmentAttributes = segment.getAttributes()
+    t.equal(segmentAttributes.product, 'Cassandra', 'should set product attribute')
+    t.equal(segmentAttributes.port_path_or_id, '9042', 'should set port attribute')
+    t.equal(segmentAttributes.database_name, 'test', 'should set database_name attribute')
+    t.equal(segmentAttributes.host, agent.config.getHostnameSafe(), 'should set host attribute')
   }
 })