From a5c42f81b72f074f77229c56f05f88a15a0bd131 Mon Sep 17 00:00:00 2001
From: Bob Evans <robert.evans25@gmail.com>
Date: Wed, 9 Oct 2024 17:14:47 -0400
Subject: [PATCH] wip: updated shim/tracer to get transaction from context,
 this is still a mess but all versioned tests are passing

---
 .../async-local-context-manager.js            |  8 ++
 lib/instrumentation/core/net.js               |  5 +-
 lib/instrumentation/when/contextualizer.js    | 14 ++-
 lib/instrumentation/when/index.js             | 20 +++--
 lib/shim/promise-shim.js                      | 40 ++++++---
 lib/shim/shim.js                              | 87 ++++++++++++++-----
 lib/shim/transaction-shim.js                  |  6 +-
 lib/shim/webframework-shim/middleware.js      | 10 ++-
 lib/transaction/tracer/index.js               | 45 ++++++----
 test/benchmark/shim/record.bench.js           |  4 +-
 test/benchmark/tracer/bindFunction.bench.js   |  2 +-
 test/benchmark/tracer/segments.bench.js       |  2 +-
 test/benchmark/tracer/transaction.bench.js    |  2 +-
 test/benchmark/tracer/utilities.bench.js      |  2 +-
 .../recordMiddleware.bench.js                 |  4 +-
 .../native-promises/native-promises.tap.js    |  8 +-
 test/integration/core/net.tap.js              |  2 +-
 test/integration/transaction/tracer.tap.js    |  8 +-
 .../async-local-context-manager.test.js       | 28 +++---
 .../instrumentation/http/outbound.test.js     |  2 +-
 test/unit/shim/shim.test.js                   | 79 ++++++-----------
 test/unit/transaction.test.js                 | 18 ++--
 22 files changed, 239 insertions(+), 157 deletions(-)

diff --git a/lib/context-manager/async-local-context-manager.js b/lib/context-manager/async-local-context-manager.js
index 9ad596b550..277b3705f7 100644
--- a/lib/context-manager/async-local-context-manager.js
+++ b/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) {
diff --git a/lib/instrumentation/core/net.js b/lib/instrumentation/core/net.js
index c7df7090ce..6e9a94ce84 100644
--- a/lib/instrumentation/core/net.js
+++ b/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)
       }
     }
   }
diff --git a/lib/instrumentation/when/contextualizer.js b/lib/instrumentation/when/contextualizer.js
index ed95351e1c..7ba49fdc81 100644
--- a/lib/instrumentation/when/contextualizer.js
+++ b/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
  *
diff --git a/lib/instrumentation/when/index.js b/lib/instrumentation/when/index.js
index ddace24ee9..22aa025fb2 100644
--- a/lib/instrumentation/when/index.js
+++ b/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,6 +281,7 @@ 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
@@ -284,7 +289,9 @@ module.exports = function initialize(shim, when) {
 
       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
     }
diff --git a/lib/shim/promise-shim.js b/lib/shim/promise-shim.js
index ef3c864927..3f2e2cf4d9 100644
--- a/lib/shim/promise-shim.js
+++ b/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() {
diff --git a/lib/shim/shim.js b/lib/shim/shim.js
index 0a393500f1..a940ffc280 100644
--- a/lib/shim/shim.js
+++ b/lib/shim/shim.js
@@ -15,6 +15,7 @@ const symbols = require('../symbols')
 const { addCLMAttributes: maybeAddCLMAttributes } = require('../util/code-level-metrics')
 const { makeId } = require('../util/hashes')
 const { isBuiltin } = require('module')
+const TraceSegment = require('../transaction/trace/segment')
 
 // Some modules do terrible things, like change the prototype of functions. To
 // avoid crashing things we'll use a cached copy of apply everywhere.
@@ -39,7 +40,6 @@ function Shim(agent, moduleName, resolvedName, shimName, pkgVersion) {
 
   this._logger = logger.child({ module: moduleName })
   this._agent = agent
-  this._contextManager = agent._contextManager
   this._toExport = null
   this._debug = false
   this.defineProperty(this, 'moduleName', moduleName)
@@ -674,17 +674,19 @@ function recordWrapper({ shim, fn, name, recordNamer }) {
       return fnApply.call(fn, this, args)
     }
 
+    let transaction
     // See if we're in an active transaction.
-    let parent
     if (segDesc.parent) {
+      segDesc.parent = segDesc.parent?.transaction?.isActive() && segDesc.parent?.segment ? segDesc.parent.segment : shim.getActiveSegment() 
+      transaction = segDesc?.parent?.transaction
       // We only want to continue recording in a transaction if the
       // transaction is active.
-      parent = segDesc.parent.transaction.isActive() ? segDesc.parent : null
     } else {
-      parent = shim.getActiveSegment()
+      segDesc.parent = shim.getActiveSegment()
+      transaction = shim.tracer.getTransaction()
     }
 
-    if (!parent) {
+    if (!segDesc.parent) {
       shim.logger.debug('Not recording function %s, not in a transaction.', name)
       return fnApply.call(fn, this, arguments)
     }
@@ -697,14 +699,23 @@ function recordWrapper({ shim, fn, name, recordNamer }) {
     //  - We are _not_ making an internal segment.
     //  - OR the parent segment is either not internal or not from this shim.
     const shouldCreateSegment = !(
-      parent.opaque ||
-      (segDesc.internal && parent.internal && shim === parent.shim)
+      segDesc.parent.opaque ||
+      (segDesc.internal && segDesc.parent.internal && shim === segDesc.parent.shim)
     )
 
-    const segment = shouldCreateSegment ? _rawCreateSegment(shim, segDesc) : parent
+    const segment = shouldCreateSegment ? _rawCreateSegment(shim, segDesc) : segDesc.parent
     maybeAddCLMAttributes(fn, segment)
 
-    return _doRecord.call(this, { segment, args, segDesc, shouldCreateSegment, shim, fn, name })
+    return _doRecord.call(this, {
+      transaction,
+      segment,
+      args,
+      segDesc,
+      shouldCreateSegment,
+      shim,
+      fn,
+      name
+    })
   }
 }
 
@@ -742,15 +753,16 @@ function _hasValidCallbackArg(shim, args, specCallback) {
  * @param {boolean} params.shouldCreateSegment Whether the recorder should create a segment
  * @param {Shim} params.shim instance of shim
  * @param {Function} params.fn function being wrapped
+ * @param params.transaction
  * @param {string} params.name name of function being wrapped
  * @returns {shim|promise} Returns a shim or promise with recorder segment and
  * bound callbacks, if applicable
  */
-function _doRecord({ segment, args, segDesc, shouldCreateSegment, shim, fn, name }) {
+function _doRecord({ segment, transaction, args, segDesc, shouldCreateSegment, shim, fn, name }) {
   // Now bind any callbacks specified in the segment descriptor.
   _bindAllCallbacks.call(this, shim, fn, name, args, {
     spec: segDesc,
-    segment: segment,
+    segment,
     shouldCreateSegment: shouldCreateSegment
   })
 
@@ -758,7 +770,16 @@ function _doRecord({ segment, args, segDesc, shouldCreateSegment, shim, fn, name
   // The reason there is no check for `segment` is because it should
   // be guaranteed by the parent and active transaction check
   // at the beginning of this function.
-  let ret = _applyRecorderSegment({ segment, ctx: this, args, segDesc, shim, fn, name })
+  let ret = _applyRecorderSegment({
+    segment,
+    transaction,
+    ctx: this,
+    args,
+    segDesc,
+    shim,
+    fn,
+    name
+  })
   if (ret) {
     if (segDesc.stream) {
       shim.logger.trace('Binding return value as stream.')
@@ -785,15 +806,16 @@ function _doRecord({ segment, args, segDesc, shouldCreateSegment, shim, fn, name
  * @param {Spec} params.segDesc Segment descriptor spec
  * @param {Shim} params.shim instance of shim
  * @param {Function} params.fn function being wrapped
+ * @param params.transaction
  * @param {string} params.name name of function being wrapped
  * @returns {*} return value of wrapped function
  */
-function _applyRecorderSegment({ segment, ctx, args, segDesc, shim, fn, name }) {
+function _applyRecorderSegment({ transaction, segment, ctx, args, segDesc, shim, fn, name }) {
   let error = null
   let promised = false
   let ret
   try {
-    ret = shim.applySegment(fn, segment, true, ctx, args, segDesc.inContext)
+    ret = shim.applySegment(fn, { transaction, segment }, true, ctx, args, segDesc.inContext)
     if (segDesc.after && segDesc.promise && shim.isPromise(ret)) {
       promised = true
       return ret.then(
@@ -948,12 +970,18 @@ function bindSegment(nodule, property, segment, full) {
     segment = property
     property = null
   }
+  
+  let transaction
+  if (segment?.segment instanceof TraceSegment) {
+    segment = segment?.segment
+    transaction = segment?.transaction
+  } 
 
   // This protects against the `bindSegment(func, null, true)` case, where the
   // segment is `null`, and thus `true` (the full param) is detected as the
   // segment.
   if (segment != null && !this.isObject(segment)) {
-    this.logger.debug({ segment: segment }, 'Segment is not a segment, not binding.')
+    this.logger.debug({ segment }, 'Segment is not a segment, not binding.')
     return nodule
   }
 
@@ -964,11 +992,18 @@ function bindSegment(nodule, property, segment, full) {
 
     // Wrap up the function with this segment.
     segment = segment || shim.getSegment()
+    transaction = transaction || shim.tracer.getTransaction()
     if (!segment) {
       return func
     }
 
-    const binder = _makeBindWrapper(shim, func, segment, full || false)
+    let ctx
+    if (transaction) {
+      ctx = { transaction, segment }
+    } else {
+      ctx = segment
+    }
+    const binder = _makeBindWrapper(shim, func, ctx, full || false)
     shim.storeSegment(binder, segment)
     return binder
   })
@@ -1037,6 +1072,9 @@ function bindCallbackSegment(spec, args, cbIdx, parentSegment) {
 function wrapCallback({ shim, args, cbIdx, parentSegment, spec }) {
   const cb = args[cbIdx]
   const realParent = parentSegment || shim.getSegment()
+  const ctx = shim.agent._contextManager.getContext()
+  const transaction = ctx?.transaction
+
   args[cbIdx] = shim.wrap(cb, null, function callbackWrapper(shim, fn, name) {
     return function wrappedCallback() {
       if (realParent) {
@@ -1060,7 +1098,7 @@ function wrapCallback({ shim, args, cbIdx, parentSegment, spec }) {
 
       // CB may end the transaction so update the parent's time preemptively.
       realParent && realParent.touch()
-      return shim.applySegment(cb, segment, true, this, arguments)
+      return shim.applySegment(cb, { transaction, segment }, true, this, arguments)
     }
   })
   shim.storeSegment(args[cbIdx], realParent)
@@ -1121,7 +1159,8 @@ function getActiveSegment(obj) {
  * @returns {TraceSegment} - The segment set as active on the context.
  */
 function setActiveSegment(segment) {
-  this._contextManager.setContext(segment)
+  const transaction = this.tracer.getTransaction()
+  this.tracer.setSegment({ segment, transaction })
   return segment
 }
 
@@ -1165,19 +1204,23 @@ function applySegment(func, segment, full, context, args, inContextCB) {
     return
   }
 
-  if (!segment) {
+  if (!(segment instanceof TraceSegment || segment?.segment instanceof TraceSegment)) {
     this.logger.trace('No segment to apply to function.')
     return fnApply.call(func, context, args)
   }
 
   this.logger.trace('Applying segment %s', segment.name)
+  let active = segment
+  if (segment?.segment) {
+    active = segment.segment
+  }
 
   /**
    *
    */
   function runInContextCb() {
     if (typeof inContextCB === 'function') {
-      inContextCB(segment)
+      inContextCB(active)
     }
 
     return fnApply.call(func, this, arguments)
@@ -1916,7 +1959,7 @@ function wrapStreamEmit({ stream, shim, segment, specEvent, shouldCreateSegment,
   // Wrap emit such that each event handler is executed within context of this
   // segment or the event-specific segment.
   shim.wrap(stream, 'emit', function wrapEmit(shim, emit) {
-    const tx = segment.transaction
+    const tx = shim.tracer.getTransaction()
     const streamBoundEmit = shim.bindSegment(emit, segment, true)
     let eventSegment = null
     let eventBoundEmit = null
@@ -1931,7 +1974,7 @@ function wrapStreamEmit({ stream, shim, segment, specEvent, shouldCreateSegment,
       if (evnt === specEvent && tx.isActive()) {
         if (!eventBoundEmit) {
           eventSegment = shim.createSegment(segmentName, segment)
-          eventBoundEmit = shim.bindSegment(emit, eventSegment, true)
+          eventBoundEmit = shim.bindSegment(emit, { transaction: tx, segment: eventSegment }, true)
         }
         eventSegment.addAttribute('count', ++emitCount)
         emitToCall = eventBoundEmit
diff --git a/lib/shim/transaction-shim.js b/lib/shim/transaction-shim.js
index d54a8b2341..157881e4cd 100644
--- a/lib/shim/transaction-shim.js
+++ b/lib/shim/transaction-shim.js
@@ -376,7 +376,8 @@ function _makeNestedTransWrapper(shim, fn, name, spec) {
       shim.logger.trace('Creating new nested %s transaction for %s', spec.type, name)
       transaction = new Transaction(shim.agent)
       transaction.type = spec.type
-      segment = transaction.trace.root
+      // segment = transaction.trace.root
+      segment = { transaction, segment: transaction.trace.root }
     }
 
     return shim.applySegment(fn, segment, false, this, arguments)
@@ -411,6 +412,7 @@ function _makeTransWrapper(shim, fn, name, spec) {
     shim.logger.trace('Creating new %s transaction for %s', spec.type, name)
     const transaction = new Transaction(shim.agent)
     transaction.type = spec.type
-    return shim.applySegment(fn, transaction.trace.root, false, this, arguments)
+    const ctx = { transaction, segment: transaction.trace.root }
+    return shim.applySegment(fn, ctx, false, this, arguments)
   }
 }
diff --git a/lib/shim/webframework-shim/middleware.js b/lib/shim/webframework-shim/middleware.js
index d76f4cab86..2d412cbab5 100644
--- a/lib/shim/webframework-shim/middleware.js
+++ b/lib/shim/webframework-shim/middleware.js
@@ -173,7 +173,10 @@ function middlewareWithCallbackRecorder({ spec, typeDetails, metricName, isError
     return new RecorderSpec({
       name: segmentName,
       callback: nextWrapper,
-      parent: txInfo.segmentStack[txInfo.segmentStack.length - 1],
+      parent: {
+        transaction: txInfo.transaction,
+        segment: txInfo.segmentStack[txInfo.segmentStack.length - 1]
+      },
       recorder,
       parameters: params,
       after: function afterExec({ shim, error }) {
@@ -238,7 +241,10 @@ function middlewareWithPromiseRecorder({ spec, typeDetails, metricName, isErrorW
     // Finally, return the segment descriptor.
     return new RecorderSpec({
       name: segmentName,
-      parent: txInfo.segmentStack[txInfo.segmentStack.length - 1],
+      parent: {
+        transaction: txInfo.transaction,
+        segment: txInfo.segmentStack[txInfo.segmentStack.length - 1]
+      },
       promise: spec.promise,
       callback: nextWrapper,
       recorder,
diff --git a/lib/transaction/tracer/index.js b/lib/transaction/tracer/index.js
index 2797eb79e3..a0ba1e96c0 100644
--- a/lib/transaction/tracer/index.js
+++ b/lib/transaction/tracer/index.js
@@ -45,9 +45,9 @@ Tracer.prototype.wrapSyncFunction = wrapSyncFunction
 Tracer.prototype.wrapCallback = wrapCallback
 
 function getTransaction() {
-  const currentSegment = this._contextManager.getContext()
-  if (currentSegment && currentSegment.transaction && currentSegment.transaction.isActive()) {
-    return currentSegment.transaction
+  const ctx = this._contextManager.getContext()
+  if (ctx?.transaction && ctx?.transaction?.isActive()) {
+    return ctx.transaction
   }
 
   return null
@@ -55,11 +55,13 @@ function getTransaction() {
 
 // TODO: Remove/replace external uses to tracer.getSegment()
 function getSegment() {
-  return this._contextManager.getContext()
+  const ctx = this._contextManager.getContext()
+  return ctx?.segment || null
 }
 
-function setSegment(segment) {
-  this._contextManager.setContext(segment)
+// TODO: update to setNewContext or something like that
+function setSegment({ segment, transaction }) {
+  this._contextManager.setContext({ transaction, segment })
 }
 
 // TODO: Remove/replace external uses to tracer.getSpanContext()
@@ -128,7 +130,8 @@ function transactionProxy(handler) {
       return handler.apply(this, arguments)
     }
     const transaction = new Transaction(tracer.agent)
-    return tracer.bindFunction(handler, transaction.trace.root, true).apply(this, arguments)
+    const ctx = { transaction, segment: transaction.trace.root }
+    return tracer.bindFunction(handler, ctx, true).apply(this, arguments)
   }
 
   wrapped[symbols.original] = handler
@@ -176,7 +179,8 @@ function transactionNestProxy(type, handler) {
     if (createNew) {
       transaction = new Transaction(tracer.agent)
       transaction.type = type
-      segment = transaction.trace.root
+      // segment = transaction.trace.root
+      segment = { transaction, segment: transaction.trace.root }
     }
 
     return tracer.bindFunction(handler, segment).apply(this, arguments)
@@ -192,23 +196,32 @@ function bindFunction(handler, segment, full) {
     return handler
   }
 
-  return _makeWrapped(this, handler, segment || this.getSegment(), !!full)
+  let transaction
+  if (segment?.segment) {
+    segment = segment?.segment
+    transaction = segment?.transaction
+  } else {
+    segment = segment || this.getSegment()
+    transaction = this.getTransaction()
+  }
+
+  return _makeWrapped({ tracer: this, handler, segment, transaction, full: !!full })
 }
-function _makeWrapped(tracer, handler, active, full) {
+function _makeWrapped({ tracer, handler, segment, transaction, full }) {
   wrapped[symbols.original] = getOriginal(handler)
-  wrapped[symbols.segment] = active
+  wrapped[symbols.segment] = segment
 
   return wrapped
 
   function wrapped() {
     const prev = tracer.getSegment()
 
-    if (active && full) {
-      active.start()
+    if (segment && full) {
+      segment.start()
     }
 
     try {
-      return tracer._contextManager.runInContext(active, handler, this, arguments)
+      return tracer._contextManager.runInContext({ transaction, segment }, handler, this, arguments)
     } catch (err) {
       logger.trace(err, 'Error from wrapped function:')
 
@@ -218,8 +231,8 @@ function _makeWrapped(tracer, handler, active, full) {
 
       throw err // Re-throwing application error, this is not an agent error.
     } finally {
-      if (active && full) {
-        active.touch()
+      if (segment && full) {
+        segment.touch()
       }
     }
   }
diff --git a/test/benchmark/shim/record.bench.js b/test/benchmark/shim/record.bench.js
index db8354c4fc..e7003e3640 100644
--- a/test/benchmark/shim/record.bench.js
+++ b/test/benchmark/shim/record.bench.js
@@ -41,7 +41,7 @@ const wrapped = shim.record(getTest(), 'func', function () {
 suite.add({
   name: 'wrapper - no transaction',
   fn: function () {
-    tracer.setSegment(null)
+    tracer.setSegment({ segment: null, transaction: null })
     wrapped.func(noop)
   }
 })
@@ -49,7 +49,7 @@ suite.add({
 suite.add({
   name: 'wrapper - in transaction',
   fn: function () {
-    tracer.setSegment(transaction.trace.root)
+    tracer.setSegment({ transaction, segment: transaction.trace.root })
     wrapped.func(noop)
   }
 })
diff --git a/test/benchmark/tracer/bindFunction.bench.js b/test/benchmark/tracer/bindFunction.bench.js
index 3feb0b1433..cdb059b94f 100644
--- a/test/benchmark/tracer/bindFunction.bench.js
+++ b/test/benchmark/tracer/bindFunction.bench.js
@@ -14,7 +14,7 @@ const tracer = helper.getTracer()
 const tx = helper.runInTransaction(s.agent, function (_tx) {
   return _tx
 })
-tracer.setSegment(tx.root)
+tracer.setSegment({ transaction: tx, segment: tx.root })
 
 preOptBind()
 const bound = tracer.bindFunction(shared.getTest().func, tx.root, true)
diff --git a/test/benchmark/tracer/segments.bench.js b/test/benchmark/tracer/segments.bench.js
index 2d636dfde6..b3834f3e26 100644
--- a/test/benchmark/tracer/segments.bench.js
+++ b/test/benchmark/tracer/segments.bench.js
@@ -16,7 +16,7 @@ const tx = helper.runInTransaction(s.agent, function (_tx) {
   return _tx
 })
 
-tracer.setSegment(tx.root)
+tracer.setSegment({ transaction: tx.root, segment: tx.root })
 
 suite.add({
   name: 'tracer.getSegment',
diff --git a/test/benchmark/tracer/transaction.bench.js b/test/benchmark/tracer/transaction.bench.js
index c2e668db7c..7ce74e470f 100644
--- a/test/benchmark/tracer/transaction.bench.js
+++ b/test/benchmark/tracer/transaction.bench.js
@@ -15,7 +15,7 @@ const tx = helper.runInTransaction(s.agent, function (_tx) {
   return _tx
 })
 
-tracer.setSegment(tx.root)
+tracer.setSegment({ transaction: tx, segment: tx.root })
 
 suite.add({
   name: 'tracer.getTransaction',
diff --git a/test/benchmark/tracer/utilities.bench.js b/test/benchmark/tracer/utilities.bench.js
index 4866e07e95..46e85bcd24 100644
--- a/test/benchmark/tracer/utilities.bench.js
+++ b/test/benchmark/tracer/utilities.bench.js
@@ -16,7 +16,7 @@ const tx = helper.runInTransaction(s.agent, function (_tx) {
   return _tx
 })
 
-tracer.setSegment(tx.root)
+tracer.setSegment({ transaction: tx, segment: tx.root })
 
 suite.add({
   name: 'tracer.slice',
diff --git a/test/benchmark/webframework-shim/recordMiddleware.bench.js b/test/benchmark/webframework-shim/recordMiddleware.bench.js
index 3c72ddf7bf..d08d0eac6e 100644
--- a/test/benchmark/webframework-shim/recordMiddleware.bench.js
+++ b/test/benchmark/webframework-shim/recordMiddleware.bench.js
@@ -59,7 +59,7 @@ function addTests(name, speccer) {
   suite.add({
     name: name + ' - wrapper (no tx)    ',
     fn: function () {
-      tracer.setSegment(null)
+      tracer.setSegment({ segment: null, transaction: null })
       middleware(getReqd(), {}, noop)
     }
   })
@@ -67,7 +67,7 @@ function addTests(name, speccer) {
   suite.add({
     name: name + ' - wrapper (tx)       ',
     fn: function () {
-      tracer.setSegment(transaction.trace.root)
+      tracer.setSegment({ transaction, segment: transaction.trace.root })
       middleware(getReqd(), {}, noop)
     }
   })
diff --git a/test/integration/core/native-promises/native-promises.tap.js b/test/integration/core/native-promises/native-promises.tap.js
index c7cac141dc..4d3af5877b 100644
--- a/test/integration/core/native-promises/native-promises.tap.js
+++ b/test/integration/core/native-promises/native-promises.tap.js
@@ -287,19 +287,19 @@ function createPromiseTests(t, config) {
   t.test('handles multi-entry callbacks correctly', function (t) {
     const { agent, tracer } = setupAgent(t, config)
 
-    helper.runInTransaction(agent, function () {
+    helper.runInTransaction(agent, function (tx) {
       const root = tracer.getSegment()
 
       const aSeg = agent.tracer.createSegment('A')
-      tracer.setSegment(aSeg)
+      tracer.setSegment({ segment: aSeg, transaction: tx })
 
       const resA = new TestResource(1)
 
       const bSeg = agent.tracer.createSegment('B')
-      tracer.setSegment(bSeg)
+      tracer.setSegment({ segment: bSeg, transaction: tx })
       const resB = new TestResource(2)
 
-      tracer.setSegment(root)
+      tracer.setSegment({ segment: root, transaction: tx })
 
       resA.doStuff(() => {
         t.equal(
diff --git a/test/integration/core/net.tap.js b/test/integration/core/net.tap.js
index 78237cdf01..56b9b9653b 100644
--- a/test/integration/core/net.tap.js
+++ b/test/integration/core/net.tap.js
@@ -21,7 +21,7 @@ test('createServer', function createServerTest(t) {
 
     server.listen(4123, function listening() {
       // leave transaction
-      tracer.setSegment(null)
+      tracer.setSegment({ segment: null, transaction: null })
       const socket = net.connect({ port: 4123 })
       socket.write('test123')
       socket.end()
diff --git a/test/integration/transaction/tracer.tap.js b/test/integration/transaction/tracer.tap.js
index 0a97294c21..7c19df2e02 100644
--- a/test/integration/transaction/tracer.tap.js
+++ b/test/integration/transaction/tracer.tap.js
@@ -26,7 +26,7 @@ test('bind in transaction', function testBind(t) {
     t.equal(tracer.getSegment(), root, 'should start at root segment')
     let bound = tracer.bindFunction(compare)
 
-    tracer.setSegment(null)
+    tracer.setSegment({ transaction: null, segment: null })
     bound.call(context, root)
     t.equal(tracer.getSegment(), null, 'should reset segment after being called')
 
@@ -35,7 +35,7 @@ test('bind in transaction', function testBind(t) {
     bound.call(context, other)
 
     t.comment('null segment bind')
-    tracer.setSegment(root)
+    tracer.setSegment({ transaction: transaction, segment: root })
     bound = tracer.bindFunction(compare, null)
 
     t.equal(tracer.getSegment(), root, 'should be back to root segment')
@@ -89,7 +89,7 @@ test('bind + throw', function testThrows(t) {
     compare(dangerous(null, root), root)
 
     t.comment('null is active')
-    tracer.setSegment(null)
+    tracer.setSegment({ transaction: null, segment: null })
     compare(dangerous(root, root), null)
     compare(dangerous(null, null), null)
 
@@ -580,7 +580,7 @@ test('wrapFunctionNoSegment', function testwrapFunctionNoSegment(t) {
     t.equal(this, outer)
     t.equal(tracer.getSegment(), seg)
     process.nextTick(function next() {
-      tracer.setSegment(null)
+      tracer.setSegment({ transaction: null, segment: null })
       callback.apply(inner, args)
     })
   }
diff --git a/test/unit/context-manager/async-local-context-manager.test.js b/test/unit/context-manager/async-local-context-manager.test.js
index a935bc30fb..dc0b151bfe 100644
--- a/test/unit/context-manager/async-local-context-manager.test.js
+++ b/test/unit/context-manager/async-local-context-manager.test.js
@@ -43,10 +43,10 @@ test('runInContext()', async (t) => {
   await t.test('should set context to active for life of callback', (t, end) => {
     const contextManager = new AsyncLocalContextManager({})
 
-    const previousContext = { name: 'previous' }
+    const previousContext = { segment: 'previous', transaction: 'tx' }
     contextManager.setContext(previousContext)
 
-    const newContext = { name: 'new' }
+    const newContext = { segment: 'new', transaction: 'tx1' }
 
     contextManager.runInContext(newContext, () => {
       const context = contextManager.getContext()
@@ -59,24 +59,24 @@ test('runInContext()', async (t) => {
   await t.test('should restore previous context when callback completes', () => {
     const contextManager = new AsyncLocalContextManager({})
 
-    const previousContext = { name: 'previous' }
+    const previousContext = { segment: 'previous', transaction: 'tx' }
     contextManager.setContext(previousContext)
 
-    const newContext = { name: 'new' }
+    const newContext = { segment: 'new', transaction: 'tx1' }
     contextManager.runInContext(newContext, () => {})
 
     const context = contextManager.getContext()
 
-    assert.equal(context, previousContext)
+    assert.deepEqual(context, previousContext)
   })
 
   await t.test('should restore previous context on exception', () => {
     const contextManager = new AsyncLocalContextManager({})
 
-    const previousContext = { name: 'previous' }
+    const previousContext = { segment: 'previous', transaction: 'tx' }
     contextManager.setContext(previousContext)
 
-    const newContext = { name: 'new' }
+    const newContext = { segment: 'new', transaction: 'tx1' }
 
     try {
       contextManager.runInContext(newContext, () => {
@@ -89,16 +89,16 @@ test('runInContext()', async (t) => {
 
     const context = contextManager.getContext()
 
-    assert.equal(context, previousContext)
+    assert.deepEqual(context, previousContext)
   })
 
   await t.test('should apply `cbThis` arg to execution', (t, end) => {
     const contextManager = new AsyncLocalContextManager({})
 
-    const previousContext = { name: 'previous' }
+    const previousContext = { segment: 'previous', transaction: 'tx' }
     contextManager.setContext(previousContext)
 
-    const newContext = { name: 'new' }
+    const newContext = { segment: 'new', transaction: 'tx1' }
     const expectedThis = () => {}
 
     contextManager.runInContext(newContext, functionRunInContext, expectedThis)
@@ -112,10 +112,10 @@ test('runInContext()', async (t) => {
   await t.test('should apply args array to execution', (t, end) => {
     const contextManager = new AsyncLocalContextManager({})
 
-    const previousContext = { name: 'previous' }
+    const previousContext = { segment: 'previous', transaction: 'tx' }
     contextManager.setContext(previousContext)
 
-    const newContext = { name: 'new' }
+    const newContext = { segment: 'new', transaction: 'tx1' }
     const expectedArg1 = 'first arg'
     const expectedArg2 = 'second arg'
     const args = [expectedArg1, expectedArg2]
@@ -132,10 +132,10 @@ test('runInContext()', async (t) => {
   await t.test('should apply arguments construct to execution', (t, end) => {
     const contextManager = new AsyncLocalContextManager({})
 
-    const previousContext = { name: 'previous' }
+    const previousContext = { segment: 'previous', transaction: 'tx' }
     contextManager.setContext(previousContext)
 
-    const newContext = { name: 'new' }
+    const newContext = { segment: 'new', transaction: 'tx1' }
     const expectedArg1 = 'first arg'
     const expectedArg2 = 'second arg'
 
diff --git a/test/unit/instrumentation/http/outbound.test.js b/test/unit/instrumentation/http/outbound.test.js
index d8c9785577..4463e83218 100644
--- a/test/unit/instrumentation/http/outbound.test.js
+++ b/test/unit/instrumentation/http/outbound.test.js
@@ -516,7 +516,7 @@ test('when working with http.request', async (t) => {
       const parentSegment = agent.tracer.createSegment('ParentSegment')
       parentSegment.opaque = true
 
-      tracer.setSegment(parentSegment) // make the current active segment
+      tracer.setSegment({ transaction, segment: parentSegment }) // make the current active segment
 
       http.get(`${host}${path}`, (res) => {
         const segment = tracer.getSegment()
diff --git a/test/unit/shim/shim.test.js b/test/unit/shim/shim.test.js
index 71238d557d..57b0150126 100644
--- a/test/unit/shim/shim.test.js
+++ b/test/unit/shim/shim.test.js
@@ -20,6 +20,8 @@ const {
 const promiseResolvers = require('../../lib/promise-resolvers')
 const { tspl } = require('@matteo.collina/tspl')
 const tempOverrideUncaught = require('../../lib/temp-override-uncaught')
+const Transaction = require('../../../lib/transaction')
+const TraceSegment = require('../../../lib/transaction/trace/segment')
 
 test('Shim', async function (t) {
   function beforeEach(ctx) {
@@ -362,21 +364,8 @@ test('Shim', async function (t) {
     t.beforeEach(function (ctx) {
       beforeEach(ctx)
 
-      ctx.nr.segment = {
-        started: false,
-        touched: false,
-        probed: false,
-        start: function () {
-          this.started = true
-        },
-        touch: function () {
-          this.touched = true
-        },
-        probe: function () {
-          this.probed = true
-        }
-      }
-
+      const transaction = new Transaction(ctx.nr.agent)
+      ctx.nr.segment = new TraceSegment(transaction, 'test')
       ctx.nr.startingSegment = ctx.nr.tracer.getSegment()
     })
 
@@ -446,9 +435,9 @@ test('Shim', async function (t) {
       // no segment is passed in.  To get around this we set the
       // active segment to an object known not to be null then do the
       // wrapping.
-      tracer.setSegment(segment)
+      tracer.setSegment({ segment })
       const wrapped = shim.bindSegment(wrappable.getActiveSegment)
-      tracer.setSegment(startingSegment)
+      tracer.setSegment({ segment: startingSegment })
 
       assert.equal(wrapped(), segment)
       assert.equal(tracer.getSegment(), startingSegment)
@@ -459,8 +448,8 @@ test('Shim', async function (t) {
       shim.bindSegment(wrappable, 'getActiveSegment', segment)
       wrappable.getActiveSegment()
 
-      assert.equal(segment.started, false)
-      assert.equal(segment.touched, false)
+      assert.equal(segment.timer.state, 1)
+      assert.equal(segment.timer.touched, false)
     })
 
     await t.test('should start and touch the segment if `full` is `true`', function (t) {
@@ -468,13 +457,13 @@ test('Shim', async function (t) {
       shim.bindSegment(wrappable, 'getActiveSegment', segment, true)
       wrappable.getActiveSegment()
 
-      assert.equal(segment.started, true)
-      assert.equal(segment.touched, true)
+      assert.equal(segment.timer.state, 2)
+      assert.equal(segment.timer.touched, true)
     })
 
     await t.test('should default to the current segment', function (t) {
       const { tracer, segment, shim, wrappable } = t.nr
-      tracer.setSegment(segment)
+      tracer.setSegment({ segment })
       shim.bindSegment(wrappable, 'getActiveSegment')
       const activeSegment = wrappable.getActiveSegment()
       assert.equal(activeSegment, segment)
@@ -1912,7 +1901,7 @@ test('Shim', async function (t) {
 
     await t.test('should return the current segment if the function is not bound', function (t) {
       const { tracer, segment, shim } = t.nr
-      tracer.setSegment(segment)
+      tracer.setSegment({ segment })
       assert.equal(
         shim.getSegment(function () {}),
         segment
@@ -1921,7 +1910,7 @@ test('Shim', async function (t) {
 
     await t.test('should return the current segment if no object is provided', function (t) {
       const { tracer, segment, shim } = t.nr
-      tracer.setSegment(segment)
+      tracer.setSegment({ segment })
       assert.equal(shim.getSegment(), segment)
     })
   })
@@ -1954,7 +1943,7 @@ test('Shim', async function (t) {
       'should return the current segment if the function is not bound when transaction is active',
       function (t) {
         const { tracer, segment, shim } = t.nr
-        tracer.setSegment(segment)
+        tracer.setSegment({ segment })
         assert.equal(
           shim.getActiveSegment(function () {}),
           segment
@@ -1966,7 +1955,7 @@ test('Shim', async function (t) {
       'should return the current segment if no object is provided when transaction is active',
       function (t) {
         const { tracer, segment, shim } = t.nr
-        tracer.setSegment(segment)
+        tracer.setSegment({ segment })
         assert.equal(shim.getActiveSegment(), segment)
       }
     )
@@ -1986,7 +1975,7 @@ test('Shim', async function (t) {
       function (t) {
         const { tracer, segment, shim } = t.nr
         segment.transaction.active = false
-        tracer.setSegment(segment)
+        tracer.setSegment({ segment })
         assert.equal(
           shim.getActiveSegment(function () {}),
           null
@@ -1999,7 +1988,7 @@ test('Shim', async function (t) {
       function (t) {
         const { tracer, segment, shim } = t.nr
         segment.transaction.active = false
-        tracer.setSegment(segment)
+        tracer.setSegment({ segment })
         assert.equal(shim.getActiveSegment(), null)
       }
     )
@@ -2019,7 +2008,7 @@ test('Shim', async function (t) {
     await t.test('should default to the current segment', function (t) {
       const { tracer, shim, wrappable } = t.nr
       const segment = { probe: function () {} }
-      tracer.setSegment(segment)
+      tracer.setSegment({ segment })
       shim.storeSegment(wrappable)
       assert.equal(shim.getSegment(wrappable), segment)
     })
@@ -2177,20 +2166,8 @@ test('Shim', async function (t) {
   await t.test('#applySegment', async function (t) {
     t.beforeEach(function (ctx) {
       beforeEach(ctx)
-      ctx.nr.segment = {
-        name: 'segment',
-        started: false,
-        touched: false,
-        start: function () {
-          this.started = true
-        },
-        touch: function () {
-          this.touched = true
-        },
-        probe: function () {
-          this.probed = true
-        }
-      }
+      const transaction = new Transaction(ctx.nr.agent)
+      ctx.nr.segment = new TraceSegment(transaction, 'test')
     })
     t.afterEach(afterEach)
 
@@ -2238,25 +2215,25 @@ test('Shim', async function (t) {
     await t.test('should make the segment active for the duration of execution', function (t) {
       const { tracer, segment, shim, wrappable } = t.nr
       const prevSegment = { name: 'prevSegment', probe: function () {} }
-      tracer.setSegment(prevSegment)
+      tracer.setSegment({ segment: prevSegment })
 
       const activeSegment = shim.applySegment(wrappable.getActiveSegment, segment)
       assert.equal(tracer.getSegment(), prevSegment)
       assert.equal(activeSegment, segment)
-      assert.equal(segment.touched, false)
-      assert.equal(segment.started, false)
+      assert.equal(segment.timer.touched, false)
+      assert.equal(segment.timer.state, 1)
     })
 
     await t.test('should start and touch the segment if `full` is `true`', function (t) {
       const { segment, shim, wrappable } = t.nr
       shim.applySegment(wrappable.getActiveSegment, segment, true)
-      assert.equal(segment.touched, true)
-      assert.equal(segment.started, true)
+      assert.equal(segment.timer.touched, true)
+      assert.equal(segment.timer.state,  2)
     })
 
     await t.test('should not change the active segment if `segment` is `null`', function (t) {
       const { tracer, segment, shim, wrappable } = t.nr
-      tracer.setSegment(segment)
+      tracer.setSegment({ segment })
       let activeSegment = null
       assert.doesNotThrow(function () {
         activeSegment = shim.applySegment(wrappable.getActiveSegment, null)
@@ -2300,7 +2277,7 @@ test('Shim', async function (t) {
           throw new Error('test error')
         }
         const prevSegment = { name: 'prevSegment', probe: function () {} }
-        tracer.setSegment(prevSegment)
+        tracer.setSegment({ segment: prevSegment })
 
         assert.throws(function () {
           shim.applySegment(func, segment)
@@ -2320,7 +2297,7 @@ test('Shim', async function (t) {
           shim.applySegment(func, segment, true)
         }, 'Error: test error')
 
-        assert.equal(segment.touched, true)
+        assert.equal(segment.timer.touched, true)
       }
     )
   })
diff --git a/test/unit/transaction.test.js b/test/unit/transaction.test.js
index 8a643a45d0..2a78fb1e4c 100644
--- a/test/unit/transaction.test.js
+++ b/test/unit/transaction.test.js
@@ -1142,11 +1142,11 @@ test('_createDistributedTracePayload', async (t) => {
   await t.test('adds the current span id as the parent span id', (t) => {
     const { agent, txn, tracer } = t.nr
     agent.config.span_events.enabled = true
-    tracer.setSegment(txn.trace.root)
+    tracer.setSegment({ segment: txn.trace.root, transaction: txn })
     txn.sampled = true
     const payload = JSON.parse(txn._createDistributedTracePayload().text())
     assert.equal(payload.d.id, txn.trace.root.id)
-    tracer.setSegment(null)
+    tracer.setSegment({ segment: null, transaction: null })
     agent.config.span_events.enabled = false
   })
 
@@ -1155,10 +1155,10 @@ test('_createDistributedTracePayload', async (t) => {
     agent.config.span_events.enabled = true
     txn._calculatePriority()
     txn.sampled = false
-    tracer.setSegment(txn.trace.root)
+    tracer.setSegment({ segment: txn.trace.root, transaction: txn })
     const payload = JSON.parse(txn._createDistributedTracePayload().text())
     assert.equal(payload.d.id, undefined)
-    tracer.setSegment(null)
+    tracer.setSegment({ segment: null, transaction: null })
     agent.config.span_events.enabled = false
   })
 
@@ -1458,7 +1458,7 @@ test('insertDistributedTraceHeaders', async (t) => {
 
     const txn = new Transaction(agent)
 
-    tracer.setSegment(txn.trace.root)
+    tracer.setSegment({ transaction: txn, segment: txn.trace.root })
 
     const outboundHeaders = createHeadersAndInsertTrace(txn)
     const traceparent = outboundHeaders.traceparent
@@ -1485,7 +1485,7 @@ test('insertDistributedTraceHeaders', async (t) => {
     const txn = new Transaction(agent)
     const lowercaseHexRegex = /^[a-f0-9]+/
 
-    tracer.setSegment(txn.trace.root)
+    tracer.setSegment({ transaction: txn, segment:  txn.trace.root })
 
     const outboundHeaders = createHeadersAndInsertTrace(txn)
     const traceparent = outboundHeaders.traceparent
@@ -1504,7 +1504,7 @@ test('insertDistributedTraceHeaders', async (t) => {
 
     const txn = new Transaction(agent)
 
-    tracer.setSegment(txn.trace.root)
+    tracer.setSegment({ transaction: txn, segment:  txn.trace.root })
     txn.sampled = true
 
     const outboundHeaders = createHeadersAndInsertTrace(txn)
@@ -1526,7 +1526,7 @@ test('insertDistributedTraceHeaders', async (t) => {
 
     txn.acceptTraceContextPayload(traceparent, tracestate)
 
-    tracer.setSegment(txn.trace.root)
+    tracer.setSegment({ transaction: txn, segment:  txn.trace.root })
 
     const outboundHeaders = createHeadersAndInsertTrace(txn)
     const traceparentParts = outboundHeaders.traceparent.split('-')
@@ -2054,7 +2054,7 @@ function createHeadersAndInsertTrace(transaction) {
 
 function addSegmentInContext(tracer, transaction, name) {
   const segment = new Segment(transaction, name)
-  tracer.setSegment(segment)
+  tracer.setSegment({ transaction, segment })
 
   return segment
 }