feat: Removed transaction from segment. Introduced a new enterSegment…

… and enterTransaction to make context propagation more clear (#2646)

api.js

@@ -961,19 +961,21 @@ API.prototype.startWebTransaction = function startWebTransaction(url, handle) {
 
   const shim = this.shim
   const tracer = this.agent.tracer
-  const parent = tracer.getTransaction()
+  const parentTx = tracer.getTransaction()
 
   assignCLMSymbol(shim, handle)
   return tracer.transactionNestProxy('web', function startWebSegment() {
-    const tx = tracer.getTransaction()
+    const context = tracer.getContext()
+    const tx = context?.transaction
+    const parent = context?.segment
 
     if (!tx) {
       return handle.apply(this, arguments)
     }
 
-    if (tx === parent) {
+    if (tx === parentTx) {
       logger.debug('not creating nested transaction %s using transaction %s', url, tx.id)
-      return tracer.addSegment(url, null, null, true, handle)
+      return tracer.addSegment(url, null, parent, true, handle)
     }
 
     logger.debug(
@@ -985,10 +987,16 @@ API.prototype.startWebTransaction = function startWebTransaction(url, handle) {
     tx.nameState.setName(NAMES.CUSTOM, null, NAMES.ACTION_DELIMITER, url)
     tx.url = url
     tx.applyUserNamingRules(tx.url)
-    tx.baseSegment = tracer.createSegment(url, recordWeb)
+    tx.baseSegment = tracer.createSegment({
+      name: url,
+      recorder: recordWeb,
+      transaction: tx,
+      parent
+    })
+    const newContext = context.enterSegment({ transaction: tx, segment: tx.baseSegment })
     tx.baseSegment.start()
 
-    const boundHandle = tracer.bindFunction(handle, tx.baseSegment)
+    const boundHandle = tracer.bindFunction(handle, newContext)
     maybeAddCLMAttributes(handle, tx.baseSegment)
     let returnResult = boundHandle.call(this)
     if (returnResult && shim.isPromise(returnResult)) {
@@ -1061,19 +1069,21 @@ function startBackgroundTransaction(name, group, handle) {
   const tracer = this.agent.tracer
   const shim = this.shim
   const txName = group + '/' + name
-  const parent = tracer.getTransaction()
+  const parentTx = tracer.getTransaction()
 
   assignCLMSymbol(shim, handle)
   return tracer.transactionNestProxy('bg', function startBackgroundSegment() {
-    const tx = tracer.getTransaction()
+    const context = tracer.getContext()
+    const tx = context?.transaction
+    const parent = context?.segment
 
     if (!tx) {
       return handle.apply(this, arguments)
     }
 
-    if (tx === parent) {
+    if (tx === parentTx) {
       logger.debug('not creating nested transaction %s using transaction %s', txName, tx.id)
-      return tracer.addSegment(txName, null, null, true, handle)
+      return tracer.addSegment(txName, null, parent, true, handle)
     }
 
     logger.debug(
@@ -1085,11 +1095,17 @@ function startBackgroundTransaction(name, group, handle) {
     )
 
     tx._partialName = txName
-    tx.baseSegment = tracer.createSegment(name, recordBackground)
+    tx.baseSegment = tracer.createSegment({
+      name,
+      recorder: recordBackground,
+      transaction: tx,
+      parent
+    })
+    const newContext = context.enterSegment({ transaction: tx, segment: tx.baseSegment })
     tx.baseSegment.partialName = group
     tx.baseSegment.start()
 
-    const boundHandle = tracer.bindFunction(handle, tx.baseSegment)
+    const boundHandle = tracer.bindFunction(handle, newContext)
     maybeAddCLMAttributes(handle, tx.baseSegment)
     let returnResult = boundHandle.call(this)
     if (returnResult && shim.isPromise(returnResult)) {
@@ -1525,12 +1541,13 @@ API.prototype.getTraceMetadata = function getTraceMetadata() {
   const metadata = {}
 
   const segment = this.agent.tracer.getSegment()
-  if (!segment) {
+  const transaction = this.agent.tracer.getTransaction()
+  if (!(segment || transaction)) {
     logger.debug('No transaction found when calling API#getTraceMetadata')
   } else if (!this.agent.config.distributed_tracing.enabled) {
     logger.debug('Distributed tracing disabled when calling API#getTraceMetadata')
   } else {
-    metadata.traceId = segment.transaction.traceId
+    metadata.traceId = transaction.traceId
 
     const spanId = segment.getSpanId()
     if (spanId) {

lib/agent.js

@@ -29,7 +29,6 @@ const TxSegmentNormalizer = require('./metrics/normalizer/tx_segment')
 const uninstrumented = require('./uninstrumented')
 const util = require('util')
 const createSpanEventAggregator = require('./spans/create-span-event-aggregator')
-const createContextManager = require('./context-manager/create-context-manager')
 const {
   maybeAddDatabaseAttributes,
   maybeAddExternalAttributes,
@@ -258,9 +257,8 @@ function Agent(config) {
 
   this.errors = new ErrorCollector(config, errorTraceAggregator, errorEventAggregator, this.metrics)
 
-  this._contextManager = createContextManager(this.config)
   // Transaction tracing.
-  this.tracer = new Tracer(this, this._contextManager)
+  this.tracer = new Tracer(this)
   this.traces = new TransactionTraceAggregator(
     {
       periodMs: DEFAULT_HARVEST_INTERVAL_MS,
@@ -294,13 +292,7 @@ function Agent(config) {
   // Set up all the configuration events the agent needs to listen for.
   this._listenForConfigChanges()
 
-  // Entity tracking metrics.
-  this.totalActiveSegments = 0
-  this.segmentsCreatedInHarvest = 0
-  this.segmentsClearedInHarvest = 0
-  // Used by shutdown code as well as entity tracking stats
-  this.activeTransactions = 0
-
+  this.initCounters()
   this.llm = {}
 
   // Finally, add listeners for the agent's own events.
@@ -599,11 +591,33 @@ Agent.prototype._generateEntityStatsAndClear = function _generateHarvestMetrics(
     )
   }
 
-  // Reset the counters.
+  this.resetCounters()
+}
+
+Agent.prototype.initCounters = function initCounters() {
+  // Entity tracking metrics.
+  this.totalActiveSegments = 0
   this.segmentsCreatedInHarvest = 0
   this.segmentsClearedInHarvest = 0
+  // Used by shutdown code as well as entity tracking stats
+  this.activeTransactions = 0
 }
 
+Agent.prototype.incrementCounters = function incrementCounters() {
+  ++this.totalActiveSegments
+  ++this.segmentsCreatedInHarvest
+}
+
+Agent.prototype.decrementCounters = function decrementCounters(transaction) {
+  --this.activeTransactions
+  this.totalActiveSegments -= transaction.numSegments
+  this.segmentsClearedInHarvest += transaction.numSegments
+}
+
+Agent.prototype.resetCounters = function resetCounters() {
+  this.segmentsCreatedInHarvest = 0
+  this.segmentsClearedInHarvest = 0
+}
 /**
  * Public interface for passing configuration data from the collector
  * on to the configuration, in an effort to keep them at least somewhat
@@ -755,9 +769,7 @@ Agent.prototype._transactionFinished = function _transactionFinished(transaction
     logger.debug('Ignoring %s (%s).', transaction.name, transaction.id)
   }
 
-  --this.activeTransactions
-  this.totalActiveSegments -= transaction.numSegments
-  this.segmentsClearedInHarvest += transaction.numSegments
+  this.decrementCounters(transaction)
 }
 
 Agent.prototype.setLambdaArn = function setLambdaArn(arn) {
@@ -839,12 +851,13 @@ Agent.prototype._listenForConfigChanges = function _listenForConfigChanges() {
  */
 Agent.prototype.getLinkingMetadata = function getLinkingMetadata(excludeServiceLinks = false) {
   const segment = this.tracer.getSegment()
+  const transaction = this.tracer.getTransaction()
   const config = this.config
 
   const linkingMetadata = {}
 
-  if (config.distributed_tracing.enabled && segment) {
-    linkingMetadata['trace.id'] = segment.transaction.traceId
+  if (config.distributed_tracing.enabled && segment && transaction) {
+    linkingMetadata['trace.id'] = transaction.traceId
     const spanId = segment.getSpanId()
     if (spanId) {
       linkingMetadata['span.id'] = spanId

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

@@ -6,6 +6,7 @@
 'use strict'
 
 const { AsyncLocalStorage } = require('async_hooks')
+const Context = require('./context')
 
 /**
  * Class for managing state in the agent.
@@ -17,12 +18,7 @@ const { AsyncLocalStorage } = require('async_hooks')
  * @class
  */
 class AsyncLocalContextManager {
-  /**
-   * @param {object} config New Relic config instance
-   */
-  constructor(config) {
-    this._config = config
-
+  constructor() {
     this._asyncLocalStorage = new AsyncLocalStorage()
   }
 
@@ -32,7 +28,7 @@ class AsyncLocalContextManager {
    * @returns {object} The current active context.
    */
   getContext() {
-    return this._asyncLocalStorage.getStore() || null
+    return this._asyncLocalStorage.getStore() || new Context()
   }
 
   /**

lib/context-manager/context.js

@@ -0,0 +1,29 @@
+/*
+ * Copyright 2024 New Relic Corporation. All rights reserved.
+ * SPDX-License-Identifier: Apache-2.0
+ */
+
+'use strict'
+
+module.exports = class Context {
+  constructor(transaction, segment) {
+    this._transaction = transaction
+    this._segment = segment
+  }
+
+  get segment() {
+    return this._segment
+  }
+
+  get transaction() {
+    return this._transaction
+  }
+
+  enterSegment({ segment, transaction = this.transaction }) {
+    return new this.constructor(transaction, segment)
+  }
+
+  enterTransaction(transaction) {
+    return new this.constructor(transaction, transaction.trace.root)
+  }
+}

lib/db/query-sample.js

@@ -34,7 +34,7 @@ QuerySample.prototype.merge = function merge(sample) {
 }
 
 QuerySample.prototype.prepareJSON = function prepareJSON(done) {
-  const transaction = this.trace.segment.transaction
+  const transaction = this.trace.transaction
   const sample = this
   const trace = sample.trace
 
@@ -56,7 +56,7 @@ QuerySample.prototype.prepareJSON = function prepareJSON(done) {
 }
 
 QuerySample.prototype.prepareJSONSync = function prepareJSONSync() {
-  const transaction = this.trace.segment.transaction
+  const transaction = this.trace.transaction
   const sample = this
   const trace = sample.trace
 
@@ -99,7 +99,7 @@ QuerySample.prototype.getParams = function getParams() {
   }
 
   if (this.tracer.config.distributed_tracing.enabled) {
-    this.trace.segment.transaction.addDistributedTraceIntrinsics(params)
+    this.trace.transaction.addDistributedTraceIntrinsics(params)
   }
 
   return params

lib/db/query-trace-aggregator.js

@@ -48,7 +48,7 @@ class QueryTraceAggregator extends Aggregator {
     }
   }
 
-  add(segment, type, query, trace) {
+  add({ segment, transaction, type, query, trace }) {
     const ttConfig = this.config.transaction_tracer
 
     // If DT is enabled and the segment is part of a sampled transaction
@@ -57,12 +57,12 @@ class QueryTraceAggregator extends Aggregator {
     let slowQuery
     switch (ttConfig.record_sql) {
       case 'raw':
-        slowQuery = new SlowQuery(segment, type, query, trace)
+        slowQuery = new SlowQuery({ segment, transaction, type, query, trace })
         logger.trace('recording raw sql')
         segment.addAttribute('sql', slowQuery.query, true)
         break
       case 'obfuscated':
-        slowQuery = new SlowQuery(segment, type, query, trace)
+        slowQuery = new SlowQuery({ transaction, segment, type, query, trace })
         logger.trace('recording obfuscated sql')
         segment.addAttribute('sql_obfuscated', slowQuery.obfuscated, true)
         break
@@ -78,7 +78,7 @@ class QueryTraceAggregator extends Aggregator {
       return
     }
 
-    slowQuery = slowQuery || new SlowQuery(segment, type, query, trace)
+    slowQuery = slowQuery || new SlowQuery({ segment, transaction, type, query, trace })
 
     segment.addAttribute('backtrace', slowQuery.trace)
 

lib/db/slow-query.js

@@ -10,14 +10,15 @@ const crypto = require('crypto')
 const path = require('path')
 const NR_ROOT = path.resolve(__dirname, '..')
 
-function SlowQuery(segment, type, query, trace) {
+function SlowQuery({ segment, transaction, type, query, trace }) {
   this.obfuscated = obfuscate(query, type)
   this.normalized = this.obfuscated.replace(/\?\s*,\s*|\s*/g, '')
   this.id = normalizedHash(this.normalized)
   this.segment = segment
   this.query = query
   this.metric = segment.name
   this.trace = formatTrace(trace)
+  this.transaction = transaction
   this.duration = segment.getDurationInMillis()
 }
 

lib/instrumentation/@nestjs/core.js

@@ -20,8 +20,7 @@ module.exports = function initialize(agent, core, moduleName, shim) {
 
   shim.wrap(core.BaseExceptionFilter.prototype, 'handleUnknownError', (shim, original) => {
     return function wrappedHandleUnknownError(exception) {
-      const segment = shim.getActiveSegment()
-      const transaction = segment?.transaction
+      const transaction = shim.tracer.getTransaction()
       if (transaction) {
         shim.agent.errors.add(transaction, exception)
         logger.trace(exception, 'Captured error handled by Nest.js exception filter.')