wip: updated shim/tracer to get transaction from context, this is sti…

…ll a mess but all versioned tests are passing

lib/context-manager/async-local-context-manager.js

@@ -56,6 +56,14 @@ class AsyncLocalContextManager {
    * @returns {*} Returns the value returned by the callback function.
    */
   runInContext(context, callback, cbThis, args) {
+    if (!(context.hasOwnProperty('segment') && context.hasOwnProperty('transaction'))) {
+      const ctx = this.getContext()
+      context = {
+        segment: context,
+        transaction: ctx?.transaction
+      }
+    }
+
     const toInvoke = cbThis ? callback.bind(cbThis) : callback
 
     if (args) {

lib/instrumentation/core/net.js

@@ -44,6 +44,7 @@ module.exports = function initialize(agent, net, moduleName, shim) {
   function wrapListen2(shim, fn) {
     return function wrappedListen2() {
       const segment = shim.getActiveSegment()
+      const transaction = shim.tracer.getTransaction()
       const emit = this.emit
 
       if (!segment || !emit) {
@@ -62,10 +63,10 @@ module.exports = function initialize(agent, net, moduleName, shim) {
         const child = shim.createSegment('net.Server.onconnection', segment)
 
         if (socket._handle.onread) {
-          shim.bindSegment(socket._handle, 'onread', child)
+          shim.bindSegment(socket._handle, 'onread', { segment: child, transaction })
         }
 
-        return shim.applySegment(emit, child, true, this, arguments)
+        return shim.applySegment(emit, { segment: child, transaction }, true, this, arguments)
       }
     }
   }

lib/instrumentation/when/contextualizer.js

@@ -6,8 +6,9 @@
 'use strict'
 const symbols = require('../../symbols')
 
-function Context(segment) {
+function Context(segment, transaction) {
   this.segments = [segment]
+  this.transaction = transaction
 }
 
 Context.prototype = Object.create(null)
@@ -87,9 +88,10 @@ function bindChild(ctxlzr, next) {
  * @param {Function} prev previous function in chain
  * @param {Function} next next function in chain
  * @param {object} segment proper segment to bind
+ * @param transaction
  * @returns {void}
  */
-Contextualizer.link = function link(prev, next, segment) {
+Contextualizer.link = function link(prev, next, segment, transaction) {
   let ctxlzr = prev && prev[symbols.context]
   if (ctxlzr && !ctxlzr.isActive()) {
     ctxlzr = prev[symbols.context] = null
@@ -100,7 +102,7 @@ Contextualizer.link = function link(prev, next, segment) {
   } else if (segment) {
     // This next promise is the root of a chain. Either there was no previous
     // promise or the promise was created out of context.
-    next[symbols.context] = new Contextualizer(0, new Context(segment))
+    next[symbols.context] = new Contextualizer(0, new Context(segment, transaction))
   }
 }
 
@@ -114,7 +116,7 @@ Contextualizer.prototype = Object.create(null)
 Contextualizer.prototype.isActive = function isActive() {
   const segments = this.context.segments
   const segment = segments[this.idx] || segments[this.parentIdx] || segments[0]
-  return segment && segment.transaction.isActive()
+  return segment && this.context.transaction.isActive()
 }
 
 /**
@@ -133,6 +135,10 @@ Contextualizer.prototype.getSegment = function getSegment() {
   return segment
 }
 
+Contextualizer.prototype.getTransaction = function getTransaction() {
+  return this.context.transaction
+}
+
 /**
  * Sets the set to the appropriate index
  *

lib/instrumentation/when/index.js

@@ -82,10 +82,11 @@ module.exports = function initialize(shim, when) {
     }
 
     const parent = agent.tracer.getSegment()
+    const transaction = agent.tracer.getTransaction()
     let promise = null
     if (
       !parent ||
-      !parent.transaction.isActive() ||
+      !transaction.isActive() ||
       typeof executor !== 'function' ||
       arguments.length !== 1
     ) {
@@ -105,11 +106,13 @@ module.exports = function initialize(shim, when) {
       promise = new Promise(wrapExecutorContext(context))
       context.promise = promise
       const segment = _createSegment(segmentName)
-      Contextualizer.link(null, promise, segment)
+      Contextualizer.link(null, promise, segment, transaction)
 
       try {
         // Must run after promise is defined so that `__NR_wrapper` can be set.
-        agent.tracer.bindFunction(executor, segment, true).apply(context.self, context.args)
+        agent.tracer
+          .bindFunction(executor, { transaction, segment }, true)
+          .apply(context.self, context.args)
       } catch (e) {
         context.args[1](e)
       }
@@ -234,6 +237,7 @@ module.exports = function initialize(shim, when) {
 
       const segmentNamePrefix = 'Promise#' + name + ' '
       const thenSegment = agent.tracer.getSegment()
+      const transaction = agent.tracer.getTransaction()
       const promise = this
       const ctx = { next: undefined, useAllParams, isWrapped: false, segmentNamePrefix }
 
@@ -243,7 +247,7 @@ module.exports = function initialize(shim, when) {
 
       // If we got a promise (which we should have), link the parent's context.
       if (!ctx.isWrapped && ctx.next instanceof Promise && ctx.next !== promise) {
-        Contextualizer.link(promise, ctx.next, thenSegment)
+        Contextualizer.link(promise, ctx.next, thenSegment, transaction)
       }
       return ctx.next
     }
@@ -277,14 +281,17 @@ module.exports = function initialize(shim, when) {
       let promSegment = ctx.next[symbols.context].getSegment()
       const segmentName = ctx.segmentNamePrefix + (fn.name || ANONYMOUS)
       const segment = _createSegment(segmentName, promSegment)
+      const transaction = ctx.next[symbols.context].getTransaction()
       if (segment && segment !== promSegment) {
         ctx.next[symbols.context].setSegment(segment)
         promSegment = segment
       }
 
       let ret = null
       try {
-        ret = agent.tracer.bindFunction(fn, promSegment, true).apply(this, arguments)
+        ret = agent.tracer
+          .bindFunction(fn, { transaction, segment: promSegment }, true)
+          .apply(this, arguments)
       } finally {
         if (ret && typeof ret.then === 'function') {
           ret = ctx.next[symbols.context].continue(ret)
@@ -310,9 +317,10 @@ module.exports = function initialize(shim, when) {
     // eslint-disable-next-line camelcase
     return function __NR_wrappedCast() {
       const segment = _createSegment(CAST_SEGMENT_NAME)
+      const transaction = agent.tracer.getTransaction()
       const prom = cast.apply(this, arguments)
       if (segment) {
-        Contextualizer.link(null, prom, segment)
+        Contextualizer.link(null, prom, segment, transaction)
       }
       return prom
     }

lib/shim/promise-shim.js

@@ -116,6 +116,7 @@ class PromiseShim extends Shim {
           _wrapExecutorContext(shim, args)
         },
         post: function postPromise(shim, Promise, name, args) {
+          const transaction = shim.tracer.getTransaction()
           // This extra property is added by `_wrapExecutorContext` in the pre step.
           const executor = args[0]
           const context = executor && executor[symbols.executorContext]
@@ -124,7 +125,7 @@ class PromiseShim extends Shim {
           }
 
           context.promise = this
-          Contextualizer.link(null, this, shim.getSegment())
+          Contextualizer.link(null, this, shim.getSegment(), transaction)
           try {
             // Must run after promise is defined so that `__NR_wrapper` can be set.
             context.executor.apply(context.self, context.args)
@@ -163,12 +164,13 @@ class PromiseShim extends Shim {
 
       return function wrappedExecutorCaller(executor) {
         const parent = shim.getActiveSegment()
+        const transaction = shim.tracer.getTransaction()
         if (!this || !parent) {
           return caller.apply(this, arguments)
         }
 
         if (!this[symbols.context]) {
-          Contextualizer.link(null, this, parent)
+          Contextualizer.link(null, this, parent, transaction)
         }
 
         const args = shim.argsToArray.apply(shim, arguments)
@@ -220,9 +222,10 @@ class PromiseShim extends Shim {
 
       return function __NR_wrappedCast() {
         const segment = shim.getSegment()
+        const transaction = shim.tracer.getTransaction()
         const prom = cast.apply(this, arguments)
         if (segment) {
-          Contextualizer.link(null, prom, segment)
+          Contextualizer.link(null, prom, segment, transaction)
         }
         return prom
       }
@@ -320,12 +323,19 @@ class PromiseShim extends Shim {
          */
         function __NR_wrappedPromisified() {
           const segment = shim.getActiveSegment()
+          const transaction = shim.tracer.getTransaction()
           if (!segment) {
             return promisified.apply(this, arguments)
           }
 
-          const prom = shim.applySegment(promisified, segment, true, this, arguments)
-          Contextualizer.link(null, prom, segment)
+          const prom = shim.applySegment(
+            promisified,
+            { transaction, segment },
+            true,
+            this,
+            arguments
+          )
+          Contextualizer.link(null, prom, segment, transaction)
           return prom
         }
       }
@@ -413,10 +423,11 @@ function wrapHandler({ handler, index, argsLength, useAllParams, ctx, shim }) {
       ctx.handler[symbols.context].setSegment(segment)
       promSegment = segment
     }
+    const transaction = ctx.handler[symbols.context].getTransaction()
 
     let ret = null
     try {
-      ret = shim.applySegment(handler, promSegment, true, this, arguments)
+      ret = shim.applySegment(handler, { transaction, segment: promSegment }, true, this, arguments)
     } finally {
       if (ret && typeof ret.then === 'function') {
         ret = ctx.handler[symbols.context].continueContext(ret)
@@ -445,6 +456,8 @@ function _wrapThen(shim, fn, _name, useAllParams) {
       return fn.apply(this, arguments)
     }
 
+    const transaction = shim.tracer.getTransaction()
+
     const thenSegment = shim.getSegment()
     const promise = this
 
@@ -466,7 +479,7 @@ function _wrapThen(shim, fn, _name, useAllParams) {
 
     // If we got a promise (which we should have), link the parent's context.
     if (!ctx.isWrapped && ctx.handler instanceof shim._class && ctx.handler !== promise) {
-      Contextualizer.link(promise, ctx.handler, thenSegment)
+      Contextualizer.link(promise, ctx.handler, thenSegment, transaction)
     }
     return ctx.handler
   }
@@ -476,8 +489,9 @@ function _wrapThen(shim, fn, _name, useAllParams) {
  * @private
  */
 class Context {
-  constructor(segment) {
+  constructor(segment, transaction) {
     this.segments = [segment]
+    this.transaction = transaction
   }
 
   branch() {
@@ -535,7 +549,7 @@ class Contextualizer {
     ctxlzr.child = false
   }
 
-  static link(prev, next, segment) {
+  static link(prev, next, segment, transaction) {
     let ctxlzr = prev && prev[symbols.context]
     if (ctxlzr && !ctxlzr.isActive()) {
       ctxlzr = prev[symbols.context] = null
@@ -562,14 +576,18 @@ class Contextualizer {
     } else if (segment) {
       // This next promise is the root of a chain. Either there was no previous
       // promise or the promise was created out of context.
-      next[symbols.context] = new Contextualizer(0, new Context(segment))
+      next[symbols.context] = new Contextualizer(0, new Context(segment, transaction))
     }
   }
 
   isActive() {
     const segments = this.context.segments
     const segment = segments[this.idx] || segments[this.parentIdx] || segments[0]
-    return segment && segment.transaction.isActive()
+    return segment && this.context.transaction.isActive()
+  }
+
+  getTransaction() {
+    return this.context.transaction
   }
 
   getSegment() {