fix: ensure correct run context for 'elasticsearch' (legacy) instrume…

…ntation (#2551)

Refs: #2430

CHANGELOG.asciidoc

@@ -61,6 +61,9 @@ Notes:
 * Fixes for run context handling for 'mongodb-core' instrumentation.
   ({issues}2430[#2430])
 
+* Fixes for run context handling for 'elasticsearch' instrumentation.
+  ({issues}2430[#2430])
+
 
 [[release-notes-3.27.0]]
 ==== 3.27.0 2022/01/17

lib/instrumentation/modules/elasticsearch.js

@@ -56,14 +56,16 @@ function getHostAndPortFromTransportConfig (config) {
 module.exports = function (elasticsearch, agent, { enabled }) {
   if (!enabled) return elasticsearch
 
+  const ins = agent._instrumentation
+
   agent.logger.debug('shimming elasticsearch.Transport.prototype.request')
   shimmer.wrap(elasticsearch.Transport && elasticsearch.Transport.prototype, 'request', wrapRequest)
 
   return elasticsearch
 
   function wrapRequest (original) {
     return function wrappedRequest (params, cb) {
-      var span = agent.startSpan(null, 'db', 'elasticsearch', 'request')
+      var span = ins.createSpan(null, 'db', 'elasticsearch', 'request')
       var id = span && span.transaction.id
       var method = params && params.method
       var path = params && params.path
@@ -85,15 +87,17 @@ module.exports = function (elasticsearch, agent, { enabled }) {
         }
         span.setDestinationContext(getDBDestination(span, host, port))
 
+        const parentRunContext = ins.currRunContext()
+        const spanRunContext = parentRunContext.enterSpan(span)
         if (typeof cb === 'function') {
           var args = Array.prototype.slice.call(arguments)
           args[1] = function () {
             span.end()
-            return cb.apply(this, arguments)
+            ins.withRunContext(parentRunContext, cb, this, ...arguments)
           }
-          return original.apply(this, args)
+          return ins.withRunContext(spanRunContext, original, this, ...args)
         } else {
-          const originalPromise = original.apply(this, arguments)
+          const originalPromise = ins.withRunContext(spanRunContext, original, this, ...arguments)
 
           const descriptors = Object.getOwnPropertyDescriptors(originalPromise)
           delete descriptors.domain

test/instrumentation/modules/elasticsearch.test.js

@@ -35,6 +35,7 @@ test('client.ping with callback', function userLandCode (t) {
     agent.endTransaction()
     agent.flush()
   })
+  t.ok(agent.currentSpan === null, 'no currentSpan in sync code after elasticsearch client command')
 })
 
 test('client.ping with promise', function userLandCode (t) {
@@ -50,6 +51,7 @@ test('client.ping with promise', function userLandCode (t) {
   }, function (err) {
     t.error(err)
   })
+  t.ok(agent.currentSpan === null, 'no currentSpan in sync code after elasticsearch client command')
 })
 
 test('client.search with callback', function userLandCode (t) {
@@ -65,6 +67,7 @@ test('client.search with callback', function userLandCode (t) {
     agent.endTransaction()
     agent.flush()
   })
+  t.ok(agent.currentSpan === null, 'no currentSpan in sync code after elasticsearch client command')
 })
 
 test('client.search with abort', function userLandCode (t) {
@@ -76,6 +79,7 @@ test('client.search with abort', function userLandCode (t) {
   var query = { q: 'pants' }
 
   var req = client.search(query)
+  t.ok(agent.currentSpan === null, 'no currentSpan in sync code after elasticsearch client command')
 
   setImmediate(() => {
     req.abort()
@@ -110,6 +114,7 @@ if (semver.satisfies(pkg.version, '>= 10')) {
       agent.endTransaction()
       agent.flush()
     })
+    t.ok(agent.currentSpan === null, 'no currentSpan in sync code after elasticsearch client command')
   })
 }
 
@@ -139,6 +144,7 @@ if (semver.satisfies(pkg.version, '>= 13')) {
       agent.endTransaction()
       agent.flush()
     })
+    t.ok(agent.currentSpan === null, 'no currentSpan in sync code after elasticsearch client command')
   })
 
   test('client.msearchTempate with callback', function userLandCode (t) {
@@ -171,6 +177,7 @@ if (semver.satisfies(pkg.version, '>= 13')) {
       agent.endTransaction()
       agent.flush()
     })
+    t.ok(agent.currentSpan === null, 'no currentSpan in sync code after elasticsearch client command')
   })
 }
 
@@ -185,6 +192,7 @@ test('client.count with callback', function userLandCode (t) {
     agent.endTransaction()
     agent.flush()
   })
+  t.ok(agent.currentSpan === null, 'no currentSpan in sync code after elasticsearch client command')
 })
 
 test('client with host=<array of host:port>', function userLandCode (t) {
@@ -196,6 +204,7 @@ test('client with host=<array of host:port>', function userLandCode (t) {
     agent.endTransaction()
     agent.flush()
   })
+  t.ok(agent.currentSpan === null, 'no currentSpan in sync code after elasticsearch client command')
 })
 
 test('client with hosts=<array of host:port>', function userLandCode (t) {
@@ -207,6 +216,7 @@ test('client with hosts=<array of host:port>', function userLandCode (t) {
     agent.endTransaction()
     agent.flush()
   })
+  t.ok(agent.currentSpan === null, 'no currentSpan in sync code after elasticsearch client command')
 })
 
 test('client with hosts="http://host:port"', function userLandCode (t) {
@@ -222,6 +232,7 @@ test('client with hosts="http://host:port"', function userLandCode (t) {
     agent.endTransaction()
     agent.flush()
   })
+  t.ok(agent.currentSpan === null, 'no currentSpan in sync code after elasticsearch client command')
 })
 
 function done (t, method, path, query, abort = false) {