fix: ensure correct run-context for 'memcached' instrumentation

This ensures that a created 'memcached' span is not the currentSpan
in user code after the call.

Refs: #2430

CHANGELOG.asciidoc

@@ -39,6 +39,10 @@ Notes:
 [float]
 ===== Bug fixes
 
+* Fix run-context handling for 'memcached' instrumentation so that the
+  automatically created Memcached span is never the `currentSpan` in user
+  code.
+
 * Fixes for 'ioredis' instrumentation ({pull}2460[#2460]):
 +
 **  Fix run-context so that a span created in the same tick as an ioredis

examples/trace-memcached.js

@@ -1,43 +1,46 @@
 #!/usr/bin/env node
 
-// A small example showing Elastic APM tracing of a script using `memcached`.
+// A small example showing Elastic APM tracing the 'memcached' package.
 //
-// By default this will use a Memcached on localhost. You can use:
-//    npm run docker:start
-// to start a Memcached container (and other containers used for testing of
-// this project).
+// This assumes a Memcached server running on localhost. You can use:
+//    npm run docker:start memcached
+// to start a Memcached container. Then `npm run docker:stop` to stop it.
 
 const apm = require('../').start({ // elastic-apm-node
   serviceName: 'example-trace-memcached'
 })
 
 const Memcached = require('memcached')
-
-const HOST = process.env.MEMCACHED_HOST || '127.0.0.1'
-const PORT = 11211
-const client = new Memcached(`${HOST}:${PORT}`, { timeout: 500 })
+const memcached = new Memcached('localhost:11211', { timeout: 500 })
 
 // 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 a transaction. More details at:
 // https://www.elastic.co/guide/en/apm/agent/nodejs/current/custom-transactions.html
-apm.startTransaction('t0')
+const t1 = apm.startTransaction('t1')
+
+memcached.touch('foo', 10, function (err, res) {
+  console.log('touch foo: err=%s res=%j', err && err.message, res)
+})
+memcached.set('foo', 'bar', 10, function (err, res) {
+  console.log('set foo: err=%s res=%j', err && err.message, res)
 
-client.version(function (err, data) {
-  console.log('Version: data=%j (err=%s)', data, err)
+  memcached.get('foo', function (err, res) {
+    console.log('get foo: err=%s res=%j', err && err.message, res)
+  })
 
-  client.set('foo', 'bar', 10, function (err) {
-    console.log('Set: foo (err=%s)', err)
+  memcached.gets('foo', function (err, res) {
+    console.log('gets foo: err=%s res=%j', err && err.message, res)
 
-    client.get('foo', function (err, data) {
-      console.log('Get foo: %s (err=%s)', data, err)
+    memcached.cas('foo', 'baz', res.cas, 10, function (casErr, casRes) {
+      console.log('cas foo: err=%s res=%j', casErr && casErr.message, casRes)
 
-      client.get('foo', function (err, data) {
-        console.log('Get foo (again): %s (err=%s)', data, err)
+      memcached.get('foo', function (err, res) {
+        console.log('get foo: err=%s res=%j', err && err.message, res)
 
-        apm.endTransaction()
-        client.end()
+        t1.end()
+        memcached.end()
       })
     })
   })

lib/instrumentation/modules/memcached.js

@@ -13,14 +13,17 @@ module.exports = function (memcached, agent, { version, enabled }) {
     agent.logger.debug('Memcached version %s not supported - aborting...', version)
     return memcached
   }
+
+  const ins = agent._instrumentation
+
   agent.logger.debug('shimming memcached.prototype.command')
   shimmer.wrap(memcached.prototype, 'command', wrapCommand)
   shimmer.wrap(memcached.prototype, 'connect', wrapConnect)
   return memcached
 
   function wrapConnect (original) {
     return function wrappedConnect () {
-      const currentSpan = agent._instrumentation.currSpan()
+      const currentSpan = ins.currSpan()
       const server = arguments[0]
       agent.logger.debug('intercepted call to memcached.prototype.connect %o', { server })
 
@@ -34,28 +37,37 @@ module.exports = function (memcached, agent, { version, enabled }) {
 
   // Wrap the generic command that is used to build touch, get, gets etc
   function wrapCommand (original) {
-    return function wrappedCommand () {
-      if (typeof arguments[0] === 'function') {
-        var query = arguments[0]()
-        // If the callback is not a function the user doesn't care about result
-        if (query && typeof query.callback === 'function') {
-          var span = agent.startSpan(`memcached.${query.type}`, 'db', 'memcached', query.type)
-          agent.logger.debug('intercepted call to memcached.prototype.command %o', { id: span && span.id, type: query.type })
-          if (span) {
-            span.setDbContext({ statement: `${query.type} ${query.key}`, type: 'memcached' })
-            const origCallback = query.callback
-            query.callback = agent._instrumentation.bindFunction(function tracedCallback () {
-              span.end()
-              return origCallback.apply(this, arguments)
-            })
-            // Rewrite the query compiler with the wrapped callback
-            arguments[0] = function queryCompiler () {
-              return query
-            }
-          }
-        }
+    return function wrappedCommand (queryCompiler, _server) {
+      if (typeof queryCompiler !== 'function') {
+        return original.apply(this, arguments)
       }
-      return original.apply(this, arguments)
+
+      var query = queryCompiler()
+      // Replace the queryCompiler function so it isn't called a second time.
+      arguments[0] = function prerunQueryCompiler () {
+        return query
+      }
+
+      // If the callback is not a function the user doesn't care about result.
+      if (!query && typeof query.callback !== 'function') {
+        return original.apply(this, arguments)
+      }
+
+      const span = ins.createSpan(`memcached.${query.type}`, 'db', 'memcached', query.type)
+      if (!span) {
+        return original.apply(this, arguments)
+      }
+
+      agent.logger.debug('intercepted call to memcached.prototype.command %o', { id: span.id, type: query.type })
+      span.setDbContext({ statement: `${query.type} ${query.key}`, type: 'memcached' })
+
+      const spanRunContext = ins.currRunContext().enterSpan(span)
+      const origCallback = query.callback
+      query.callback = ins.bindFunctionToRunContext(spanRunContext, function tracedCallback () {
+        span.end()
+        return origCallback.apply(this, arguments)
+      })
+      return ins.withRunContext(spanRunContext, original, this, ...arguments)
     }
   }
 }

test/instrumentation/modules/memcached.test.js

@@ -1,7 +1,10 @@
 'use strict'
 
 // Memcached isn't supported on Windows
-if (process.platform === 'win32') process.exit()
+if (process.platform === 'win32') {
+  console.log('# SKIP memcached does not support Windows')
+  process.exit()
+}
 
 var agent = require('../../..').start({
   serviceName: 'test-memcached',
@@ -78,6 +81,7 @@ test('memcached', function (t) {
   var cache = new Memcached(`${host}:11211`, { timeout: 500 })
   agent.startTransaction('myTrans')
   cache.set('foo', 'bar', 300, (err) => {
+    t.ok(agent.currentSpan === null, 'memcached span should not be currentSpan in callback')
     t.error(err)
     cache.get('foo', (err, data) => {
       t.error(err)
@@ -96,17 +100,15 @@ test('memcached', function (t) {
                 t.strictEqual(data, undefined)
                 cache.end()
                 agent.endTransaction()
-                setTimeout(function () {
-                  // Wait for spans to encode and be sent, before flush.
-                  agent.flush()
-                }, 200)
+                agent.flush()
               })
             })
           })
         })
       })
     })
   })
+  t.ok(agent.currentSpan === null, 'memcached span should not be currentSpan in same tick after client method call')
 })
 
 function resetAgent (cb) {