From 1a80ad56a6d502182a0db368f40443467f7943df Mon Sep 17 00:00:00 2001
From: Bob Evans <revans@newrelic.com>
Date: Tue, 8 Oct 2024 15:46:45 -0400
Subject: [PATCH] feat: Updated shim/when instrumentation to use tracer to run
 in context (#2642)

---
 lib/instrumentation/when/index.js |  8 ++------
 lib/shim/shim.js                  | 30 ++++++++----------------------
 2 files changed, 10 insertions(+), 28 deletions(-)

diff --git a/lib/instrumentation/when/index.js b/lib/instrumentation/when/index.js
index 9b77068984..ddace24ee9 100644
--- 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
index 9505a28872..0a393500f1 100644
--- 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 */