feat: Updated shim/when instrumentation to use tracer to run in conte…

…xt (#2642)
newrelic · Oct 8, 2024 · 1a80ad5 · 1a80ad5
1 parent bb0a1ef
commit 1a80ad5
Show file tree

Hide file tree

Showing 2 changed files with 10 additions and 28 deletions.
diff --git a/lib/instrumentation/when/index.js b/lib/instrumentation/when/index.js
@@ -21,7 +21,6 @@ const Contextualizer = require('./contextualizer')
  */
 module.exports = function initialize(shim, when) {
   const agent = shim.agent
-  const contextManager = agent._contextManager
   const spec = WHEN_SPEC
 
   // Wrap library-level methods.
@@ -82,7 +81,7 @@ module.exports = function initialize(shim, when) {
       return Promise(executor) // eslint-disable-line new-cap
     }
 
-    const parent = contextManager.getContext()
+    const parent = agent.tracer.getSegment()
     let promise = null
     if (
       !parent ||
@@ -108,14 +107,11 @@ module.exports = function initialize(shim, when) {
       const segment = _createSegment(segmentName)
       Contextualizer.link(null, promise, segment)
 
-      segment.start()
       try {
         // Must run after promise is defined so that `__NR_wrapper` can be set.
-        contextManager.runInContext(segment, executor, context.self, context.args)
+        agent.tracer.bindFunction(executor, segment, true).apply(context.self, context.args)
       } catch (e) {
         context.args[1](e)
-      } finally {
-        segment.touch()
       }
     }
 

diff --git a/lib/shim/shim.js b/lib/shim/shim.js
@@ -1083,7 +1083,7 @@ function getSegment(obj) {
     return obj[symbols.segment]
   }
 
-  return this._contextManager.getContext()
+  return this.tracer.getSegment()
 }
 
 /**
@@ -1172,32 +1172,18 @@ function applySegment(func, segment, full, context, args, inContextCB) {
 
   this.logger.trace('Applying segment %s', segment.name)
 
-  const contextManager = this._contextManager
-  const prevSegment = contextManager.getContext()
-
-  return contextManager.runInContext(segment, function runInContextCb() {
-    if (full) {
-      segment.start()
-    }
-
+  /**
+   *
+   */
+  function runInContextCb() {
     if (typeof inContextCB === 'function') {
       inContextCB(segment)
     }
 
-    try {
-      return fnApply.call(func, context, args)
-    } catch (error) {
-      if (prevSegment === null && process.domain != null) {
-        process.domain[symbols.segment] = contextManager.getContext()
-      }
+    return fnApply.call(func, this, arguments)
+  }
 
-      throw error // Re-throwing application error, this is not an agent error.
-    } finally {
-      if (full) {
-        segment.touch()
-      }
-    }
-  })
+  return this.tracer.bindFunction(runInContextCb, segment, full).apply(context, args)
 }
 /* eslint-enable max-params */