fix: ensure correct run context for mysql2 instrumentation

CHANGELOG.asciidoc

@@ -82,6 +82,9 @@ Notes:
     internally-queued before the RedisClient is "ready". ({issues}2459[#2459])
 **  Add destination context so Redis shows up on the Service Map.
 
+* Fix run-context handling for 'mysql2' instrumentation to avoid accidental
+  creation of *child* spans of the automatic mysql spans.
+  ({issues}2430[#2430]})
 
 [[release-notes-3.25.0]]
 ==== 3.25.0 2021/11/24

examples/trace-mysql2.js

@@ -0,0 +1,88 @@
+#!/usr/bin/env node --unhandled-rejections=strict
+
+// A small example showing Elastic APM tracing of a script the `mysql2` package.
+//
+// By default this will use a MySQL on localhost with user 'root'. You can use:
+//    npm run docker:start mysql
+// to start a MySQL container. Then `npm run docker:stop` to stop it.
+
+const apm = require('../').start({ // elastic-apm-node
+  serviceName: 'example-trace-mysql2'
+})
+
+const mysql = require('mysql2')
+const mysqlPromise = require('mysql2/promise')
+
+const conn = mysql.createConnection({ user: 'root' })
+
+// 1. Simple queries, callback-style.
+// Note: For tracing spans to be created, there must be an active transaction.
+// Typically, a transaction is automatically started for incoming HTTP
+// requests to a Node.js server. However, because this script is not running
+// an HTTP server, we manually start transactions. More details at:
+// https://www.elastic.co/guide/en/apm/agent/nodejs/current/custom-transactions.html
+const t1 = apm.startTransaction('t1-callback-style')
+conn.query('SELECT 1 + 1 AS solution', (err, results, _fields) => {
+  console.log('SELECT 1+1: err=%s results=%o', err, results)
+})
+conn.query('SELECT ? + ? AS solution', [2, 2], (err, results, _fields) => {
+  console.log('SELECT 2+2: err=%s results=%o', err, results)
+  t1.end()
+})
+
+// 2. Tracing of prepared statements can show that subsequent executions of the
+// same query can be much faster. This example is derived from:
+// https://github.com/sidorares/node-mysql2/blob/master/examples/execute.js
+const t2 = apm.startTransaction('t2-prepared-statements')
+conn.execute(
+  'select ?+1 as qqq, ? as rrr, ? as yyy',
+  [1, null, 3],
+  (err, rows, fields) => {
+    console.log('execute 1: err=%s results=%o', err, rows)
+    conn.execute(
+      'select ?+1 as qqq, ? as rrr, ? as yyy',
+      [3, null, 3],
+      (err, rows, fields) => {
+        console.log('execute 2: err=%s results=%o', err, rows)
+        conn.unprepare('select ?+1 as qqq, ? as rrr, ? as yyy')
+        conn.execute(
+          'select ?+1 as qqq, ? as rrr, ? as yyy',
+          [3, null, 3],
+          (err, rows, fields) => {
+            console.log('execute 3: err=%s results=%o', err, rows)
+            t2.end()
+          }
+        )
+      }
+    )
+  }
+)
+
+// 3. Promise style
+async function promiseStyle () {
+  const conn2 = await mysqlPromise.createConnection({ user: 'root' })
+  const t3 = apm.startTransaction('t3-promise-style')
+
+  const [rows] = await conn2.query('select 3 + 3 as solution')
+  console.log('select 3+3: rows=%o', rows)
+
+  // "upgrade" from non-promise connection
+  conn.promise().query('select 4 + 4 as solution')
+    .then(([rows, _fields]) => {
+      console.warn('XXX in then: %s', apm._instrumentation._runCtxMgr)
+      console.log('select 4+4: rows=%o', rows)
+    })
+    .catch(err => {
+      console.log('select 4+4: err=%s', err)
+    })
+
+  t3.end()
+  conn2.close()
+}
+promiseStyle()
+
+// Lazily shutdown client after everything above is finished.
+setTimeout(() => {
+  console.log('Done')
+  conn.end()
+}, 1000)

lib/instrumentation/modules/mysql2.js

@@ -8,42 +8,52 @@ var symbols = require('../../symbols')
 var { getDBDestination } = require('../context')
 
 module.exports = function (mysql2, agent, { version, enabled }) {
-  var ins = agent._instrumentation
+  if (!enabled) {
+    return mysql2
+  }
   if (!semver.satisfies(version, '>=1 <3')) {
     agent.logger.debug('mysql2 version %s not supported - aborting...', version)
     return mysql2
   }
 
+  var ins = agent._instrumentation
+
   shimmer.wrap(mysql2.Connection.prototype, 'query', wrapQuery)
   shimmer.wrap(mysql2.Connection.prototype, 'execute', wrapQuery)
 
   return mysql2
 
   function wrapQuery (original) {
     return function wrappedQuery (sql, values, cb) {
-      var span = enabled && agent.startSpan(null, 'db', 'mysql', 'query')
-      var id = span && span.transaction.id
-      var hasCallback = false
-      var sqlStr
+      agent.logger.debug('intercepted call to mysql2.%s', original.name)
 
-      if (span) {
-        if (this[symbols.knexStackObj]) {
-          span.customStackTrace(this[symbols.knexStackObj])
-          this[symbols.knexStackObj] = null
-        }
-        // get connection parameters from mysql config
-        let host, port
-        if (typeof this.config === 'object') {
-          ({ host, port } = this.config)
-        }
-        span.setDestinationContext(getDBDestination(span, host, port))
+      var span = ins.createSpan(null, 'db', 'mysql', 'query')
+      if (!span) {
+        return original.apply(this, arguments)
       }
 
+      if (this[symbols.knexStackObj]) {
+        span.customStackTrace(this[symbols.knexStackObj])
+        this[symbols.knexStackObj] = null
+      }
+
+      let hasCallback = false
+      const wrapCallback = function (origCallback) {
+        hasCallback = true
+        return ins.bindFunction(function wrappedCallback (_err) {
+          span.end()
+          return origCallback.apply(this, arguments)
+        })
+      }
+
+      let sqlStr
       switch (typeof sql) {
         case 'string':
           sqlStr = sql
           break
         case 'object':
+          // `mysql2.{query,execute}` support the only arg being an instance
+          // of the internal mysql2 `Command` object.
           if (typeof sql.onResult === 'function') {
             sql.onResult = wrapCallback(sql.onResult)
           }
@@ -53,47 +63,43 @@ module.exports = function (mysql2, agent, { version, enabled }) {
           arguments[0] = wrapCallback(sql)
           break
       }
-
-      if (span && sqlStr) {
-        agent.logger.debug('extracted sql from mysql2 query %o', { id: id, sql: sqlStr })
+      if (sqlStr) {
         span.setDbContext({ statement: sqlStr, type: 'sql' })
         span.name = sqlSummary(sqlStr)
+      } else {
+        span.setDbContext({ type: 'sql' })
       }
 
+      let host, port
+      if (typeof this.config === 'object') {
+        ({ host, port } = this.config)
+      }
+      span.setDestinationContext(getDBDestination(span, host, port))
+
       if (typeof values === 'function') {
         arguments[1] = wrapCallback(values)
       } else if (typeof cb === 'function') {
         arguments[2] = wrapCallback(cb)
       }
+      const spanRunContext = ins.currRunContext().enterSpan(span)
+      const result = ins.withRunContext(spanRunContext, original, this, ...arguments)
 
-      var result = original.apply(this, arguments)
       if (result && !hasCallback) {
         ins.bindEmitter(result)
-        if (span) {
-          shimmer.wrap(result, 'emit', function (original) {
-            return function (event) {
-              switch (event) {
-                case 'error':
-                case 'close':
-                case 'end':
-                  span.end()
-              }
-              return original.apply(this, arguments)
+        shimmer.wrap(result, 'emit', function (origEmit) {
+          return function (event) {
+            switch (event) {
+              case 'error':
+              case 'close':
+              case 'end':
+                span.end()
             }
-          })
-        }
+            return origEmit.apply(this, arguments)
+          }
+        })
       }
 
       return result
-
-      function wrapCallback (cb) {
-        hasCallback = true
-        return agent._instrumentation.bindFunction(span ? wrappedCallback : cb)
-        function wrappedCallback () {
-          span.end()
-          return cb.apply(this, arguments)
-        }
-      }
     }
   }
 }

test/instrumentation/modules/mysql2/mysql.test.js

@@ -98,6 +98,7 @@ factories.forEach(function (f) {
               factory(function () {
                 agent.startTransaction('foo')
                 queryable[executor].apply(queryable, args)
+                t.ok(agent.currentSpan === null, 'mysql2 span should not spill into calling code')
               })
             })
           })
@@ -121,6 +122,7 @@ factories.forEach(function (f) {
                 factory(function () {
                   agent.startTransaction('foo')
                   var promise = queryablePromise[executor].apply(queryablePromise, args)
+                  t.ok(agent.currentSpan === null, 'mysql2 span should not spill into calling code')
                   basicQueryPromise(t, promise)
                 })
               })
@@ -146,6 +148,7 @@ factories.forEach(function (f) {
                 factory(function () {
                   agent.startTransaction('foo')
                   var stream = queryable[executor].apply(queryable, args)
+                  t.ok(agent.currentSpan === null, 'mysql2 span should not spill into calling code')
                   basicQueryStream(stream, t)
                 })
               })
@@ -167,6 +170,7 @@ factories.forEach(function (f) {
 
           data.spans.forEach(function (span) {
             assertSpan(t, span, sql)
+            t.equal(span.parent_id, trans.id, 'each mysql2 span is a child of the transaction')
           })
 
           t.end()
@@ -211,6 +215,7 @@ factories.forEach(function (f) {
 
           data.spans.forEach(function (span) {
             assertSpan(t, span, sql)
+            t.equal(span.parent_id, trans.id, 'each mysql2 span is a child of the transaction')
           })
 
           t.end()
@@ -324,6 +329,7 @@ factories.forEach(function (f) {
             queryable.getConnection(function (err, conn) {
               t.error(err)
               conn.query(sql, basicQueryCallback(t))
+              t.ok(agent.currentSpan === null, 'mysql2 span should not spill into calling code')
             })
           })
         })
@@ -349,6 +355,7 @@ function basicQueryPromise (t, p) {
 
 function basicQueryCallback (t) {
   return function (err, rows, fields) {
+    t.ok(agent.currentSpan === null, 'mysql2 span should not spill into calling code')
     t.error(err)
     t.strictEqual(rows[0].solution, 2)
     agent.endTransaction()
@@ -358,13 +365,16 @@ function basicQueryCallback (t) {
 function basicQueryStream (stream, t) {
   var results = 0
   stream.on('error', function (err) {
+    t.ok(agent.currentSpan === null, 'mysql2 span should not be active in user code')
     t.error(err)
   })
   stream.on('result', function (row) {
+    t.ok(agent.currentSpan === null, 'mysql2 span should not be active in user code')
     results++
     t.strictEqual(row.solution, 2)
   })
   stream.on('end', function () {
+    t.ok(agent.currentSpan === null, 'mysql2 span should not be active in user code')
     t.strictEqual(results, 1)
     agent.endTransaction()
   })