feat: Removed transaction from segment. Introduced a new enterSegment and enterTransaction to make context propagation more clear

api.js

@@ -961,19 +961,21 @@
 
   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 @@
     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 @@
   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 @@
     )
 
     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)) {
@@ -1524,12 +1540,13 @@
   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) {
@@ -1547,7 +1564,7 @@
  * @param {string} params.traceId Identifier for the feedback event.
  * Obtained from {@link getTraceMetadata}.
  * @param {string} params.category A tag for the event.
  * @param {string} params.rating A indicator of how useful the message was.
  * @param {string} [params.message] The message that triggered the event.
  * @param {object} [params.metadata] Additional key-value pairs to associate
  * with the recorded event.
@@ -1901,7 +1918,7 @@
   transaction.ignoreApdex = true
 }
 
 /**
  * Run a function with the passed in LLM context as the active context and return its return value.
  *
  * @example

lib/agent.js

@@ -29,7 +29,6 @@
 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,
@@ -207,7 +206,7 @@
       periodMs: config.event_harvest_config.report_period_ms,
       limit: config.event_harvest_config.harvest_limits.analytic_event_data,
       config,
       enabled: (config) => config.transaction_events.enabled
     },
     this
   )
@@ -218,7 +217,7 @@
       limit: config.event_harvest_config.harvest_limits.custom_event_data,
       metricNames: NAMES.CUSTOM_EVENTS,
       config,
       enabled: (config) => config.custom_insights_events.enabled
     },
     this
   )
@@ -228,7 +227,7 @@
       periodMs: DEFAULT_HARVEST_INTERVAL_MS,
       limit: MAX_ERROR_TRACES_DEFAULT,
       config,
       enabled: (config) => config.error_collector.enabled && config.collect_errors
     },
     this.collector,
     this.harvester
@@ -239,7 +238,7 @@
       periodMs: config.event_harvest_config.report_period_ms,
       limit: config.event_harvest_config.harvest_limits.error_event_data,
       config,
       enabled: (config) => config.error_collector.enabled && config.error_collector.capture_events
     },
     this
   )
@@ -257,9 +256,8 @@
 
   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,
@@ -293,13 +291,7 @@
   // 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.
@@ -598,11 +590,33 @@
     )
   }
 
-  // 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
@@ -753,9 +767,7 @@
     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) {
@@ -834,6 +846,7 @@
  */
 Agent.prototype.getLinkingMetadata = function getLinkingMetadata() {
   const segment = this.tracer.getSegment()
+  const transaction = this.tracer.getTransaction()
   const config = this.config
 
   const linkingMetadata = {
@@ -842,8 +855,8 @@
     'hostname': config.getHostnameSafe()
   }
 
-  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/attributes.js

@@ -22,7 +22,7 @@ class Attributes {
    * @param {string} scope
    *  The scope of the attributes this will collect. Must be `transaction` or
    *  `segment`.
-   * @param {number} [limit=Infinity]
+   * @param {number} [limit]
    *  The maximum number of attributes to retrieve for each destination.
    */
   constructor(scope, limit = Infinity) {
@@ -104,7 +104,7 @@ class Attributes {
    * @param {DESTINATIONS}  destinations  - The default destinations for this key.
    * @param {string}        key           - The attribute name.
    * @param {string}        value         - The attribute value.
-   * @param {boolean} [truncateExempt=false] - Flag marking value exempt from truncation
+   * @param {boolean} [truncateExempt] - Flag marking value exempt from truncation
    */
   addAttribute(destinations, key, value, truncateExempt = false) {
     if (this.attributeCount + 1 > this.limit) {

lib/collector/remote-method.js

@@ -92,7 +92,7 @@ RemoteMethod.prototype._reportDataUsage = function reportDataUsage(sent, receive
  * you're doing it wrong.
  *
  * @param {object}   payload           Serializable payload.
- * @param {object}   [nrHeaders=null]  NR request headers from connect response.
+ * @param {object}   [nrHeaders]  NR request headers from connect response.
  * @param {Function} callback          What to do next. Gets passed any error.
  */
 RemoteMethod.prototype.invoke = function invoke(payload, nrHeaders, callback) {

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/parsed-statement.js

@@ -21,10 +21,9 @@ function ParsedStatement(type, operation, collection, raw) {
   }
 }
 
-ParsedStatement.prototype.recordMetrics = function recordMetrics(segment, scope) {
+ParsedStatement.prototype.recordMetrics = function recordMetrics(segment, scope, transaction) {
   const duration = segment.getDurationInMillis()
   const exclusive = segment.getExclusiveDurationInMillis()
-  const transaction = segment.transaction
   const type = transaction.isWeb() ? DB.WEB : DB.OTHER
   const thisTypeSlash = this.type + '/'
   const operation = DB.OPERATION + '/' + thisTypeSlash + this.operation
@@ -68,7 +67,13 @@ ParsedStatement.prototype.recordMetrics = function recordMetrics(segment, scope)
   }
 
   if (this.raw) {
-    transaction.agent.queries.add(segment, this.type.toLowerCase(), this.raw, this.trace)
+    transaction.agent.queries.add({
+      segment,
+      transaction,
+      type: this.type.toLowerCase(),
+      query: this.raw,
+      trace: this.trace
+    })
   }
 }
 

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