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

…ntation Refs: #2430
elastic · Jan 28, 2022 · a29b223 · a29b223
1 parent bb32bc2
commit a29b223
Show file tree

Hide file tree

Showing 2 changed files with 19 additions and 4 deletions.
diff --git a/lib/instrumentation/modules/elasticsearch.js b/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

diff --git a/test/instrumentation/modules/elasticsearch.test.js b/test/instrumentation/modules/elasticsearch.test.js
@@ -34,6 +34,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) {
@@ -49,6 +50,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) {
@@ -64,6 +66,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) {
@@ -75,6 +78,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()
@@ -109,6 +113,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')
   })
 }
 
@@ -138,6 +143,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) {
@@ -170,6 +176,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')
   })
 }
 
@@ -184,6 +191,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) {
@@ -195,6 +203,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) {
@@ -206,6 +215,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) {
@@ -221,6 +231,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) {