refactor: Separated context classes for agent in standard and opentelemetry bridge mode (#2967)

Author: bizob2828

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

@@ -7,6 +7,7 @@
 
 const { AsyncLocalStorage } = require('async_hooks')
 const Context = require('./context')
+const OtelContext = require('../otel/context')
 
 /**
  * Class for managing state in the agent.
@@ -18,7 +19,8 @@ const Context = require('./context')
  * @class
  */
 class AsyncLocalContextManager {
-  constructor() {
+  constructor(isOtelBridgeMode) {
+    this.isOtelBridgeMode = isOtelBridgeMode
     this._asyncLocalStorage = new AsyncLocalStorage()
   }
 
@@ -28,7 +30,7 @@ class AsyncLocalContextManager {
    * @returns {object} The current active context.
    */
   getContext() {
-    return this._asyncLocalStorage.getStore() || new Context()
+    return this._asyncLocalStorage.getStore() || (this.isOtelBridgeMode ? new OtelContext() : new Context())
   }
 
   /**

lib/context-manager/context.js

@@ -4,14 +4,11 @@
  */
 
 'use strict'
-const { otelSynthesis } = require('../symbols')
-const FakeSpan = require('../otel/fake-span')
 
 module.exports = class Context {
-  constructor(transaction, segment, parentContext) {
+  constructor(transaction, segment) {
     this._transaction = transaction
     this._segment = segment
-    this._otelCtx = parentContext ? new Map(parentContext) : new Map()
   }
 
   get segment() {
@@ -26,83 +23,23 @@ module.exports = class Context {
    * Constructs a new context from segment about to be bound to context manager
    * along with the current transaction.
    *
-   * If agent is in otel bridge mode it will also bind a FakeSpan to the otel ctx.
-   *
    * @param {object} params to function
    * @param {TraceSegment} params.segment segment to bind to context
    * @param {Transaction} params.transaction active transaction
    * @returns {Context} a newly constructed context
    */
   enterSegment({ segment, transaction = this._transaction }) {
-    if (transaction?.agent?.otelSpanKey) {
-      this._otelCtx.set(transaction.agent.otelSpanKey, new FakeSpan(segment, transaction))
-    }
     return new this.constructor(transaction, segment, this._otelCtx)
   }
 
   /**
    * Constructs a new context from transaction about to be bound to context manager.
    * It uses the trace root segment as the segment in context.
    *
-   * If agent is in otel bridge mode it will also bind a FakeSpan to the otel ctx.
-   *
-   * @param {object} params to function
-   * @param {TraceSegment} params.segment transaction trace root segment
-   * @param {Transaction} params.transaction transaction to bind to context
-   * @param transaction
+   * @param {Transaction} transaction transaction to bind to context
    * @returns {Context} a newly constructed context
    */
   enterTransaction(transaction) {
-    if (transaction?.agent?.otelSpanKey) {
-      this._otelCtx.set(transaction.agent.otelSpanKey, new FakeSpan(transaction.trace.root, transaction))
-    }
-    return new this.constructor(transaction, transaction.trace.root, this._otelCtx)
-  }
-
-  /**
-   * Required for bridging OTEL data into the agent.
-   *
-   * @param {string} key Stored entity name to retrieve.
-   *
-   * @returns {*} The stored value.
-   */
-  getValue(key) {
-    return this._otelCtx.get(key)
-  }
-
-  /**
-   * Required for bridging OTEL data into the agent.
-   *
-   * @param {string} key Name for stored value.
-   * @param {*} value Value to store.
-   *
-   * @returns {object} The context manager object.
-   */
-  setValue(key, value) {
-    let ctx
-
-    if (value[otelSynthesis] && value[otelSynthesis].segment && value[otelSynthesis].transaction) {
-      const { segment, transaction } = value[otelSynthesis]
-      segment.start()
-      ctx = new this.constructor(transaction, segment, this._otelCtx)
-    } else {
-      ctx = new this.constructor(this._transaction, this._segment, this._otelCtx)
-    }
-
-    ctx._otelCtx.set(key, value)
-    return ctx
-  }
-
-  /**
-   * Required for bridging OTEL data into the agent.
-   *
-   * @param {string} key Named value to remove from the store.
-   *
-   * @returns {object} The context manager object.
-   */
-  deleteValue(key) {
-    const ctx = new this.constructor(this._transaction, this._segment, this._otelCtx)
-    ctx._otelCtx.delete(key)
-    return ctx
+    return new this.constructor(transaction, transaction.trace.root)
   }
 }

lib/otel/context.js

@@ -0,0 +1,89 @@
+/*
+ * Copyright 2025 New Relic Corporation. All rights reserved.
+ * SPDX-License-Identifier: Apache-2.0
+ */
+
+'use strict'
+const { otelSynthesis } = require('../symbols')
+const FakeSpan = require('./fake-span')
+const Context = require('../context-manager/context')
+
+module.exports = class OtelContext extends Context {
+  constructor(transaction, segment, parentContext) {
+    super(transaction, segment)
+    this._otelCtx = parentContext ? new Map(parentContext) : new Map()
+  }
+
+  /**
+   * Constructs a new context from segment about to be bound to context manager
+   * along with the current transaction. It will also bind a FakeSpan to the `_otelCtx`
+   *
+   * @param {object} params to function
+   * @param {TraceSegment} params.segment segment to bind to context
+   * @param {Transaction} params.transaction active transaction
+   * @returns {OtelContext} a newly constructed context
+   */
+  enterSegment({ segment, transaction = this._transaction }) {
+    this._otelCtx.set(transaction.agent.otelSpanKey, new FakeSpan(segment, transaction))
+    return new this.constructor(transaction, segment, this._otelCtx)
+  }
+
+  /**
+   * Constructs a new context from transaction about to be bound to context manager.
+   * It uses the trace root segment as the segment in context. It will also bind a FakeSpan to the `_otelCtx`.
+   *
+   * @param {Transaction} transaction transaction to bind to context
+   * @returns {OtelContext} a newly constructed context
+   */
+  enterTransaction(transaction) {
+    this._otelCtx.set(transaction.agent.otelSpanKey, new FakeSpan(transaction.trace.root, transaction))
+    return new this.constructor(transaction, transaction.trace.root, this._otelCtx)
+  }
+
+  /**
+   * Used to retrieve data from `_otelCtx`
+   *
+   * @param {string} key Stored entity name to retrieve.
+   *
+   * @returns {*} The stored value.
+   */
+  getValue(key) {
+    return this._otelCtx.get(key)
+  }
+
+  /**
+   * Used to set data on `_otelCtx`
+   *
+   * @param {string} key Name for stored value.
+   * @param {*} value Value to store.
+   *
+   * @returns {OtelContext} The context object.
+   */
+  setValue(key, value) {
+    let ctx
+
+    if (value[otelSynthesis] && value[otelSynthesis].segment && value[otelSynthesis].transaction) {
+      const { segment, transaction } = value[otelSynthesis]
+      segment.start()
+      ctx = new this.constructor(transaction, segment, this._otelCtx)
+    } else {
+      ctx = new this.constructor(this._transaction, this._segment, this._otelCtx)
+    }
+
+    ctx._otelCtx.set(key, value)
+    return ctx
+  }
+
+  /**
+   * Used to remove data from `_otelCtx`
+   *
+   * @param {string} key Named value to remove from the store.
+   *
+   * @returns {OtelContext} The context object.
+   */
+  deleteValue(key) {
+    const ctx = new this.constructor(this._transaction, this._segment, this._otelCtx)
+    ctx._otelCtx.delete(key)
+    return ctx
+  }
+}

lib/transaction/tracer/index.js

@@ -23,7 +23,7 @@ function Tracer(agent) {
   }
 
   this.agent = agent
-  this._contextManager = new AsyncLocalContextManager()
+  this._contextManager = new AsyncLocalContextManager(agent.config.feature_flag.opentelemetry_bridge)
 }
 
 Tracer.prototype.getContext = getContext

test/unit/context-manager/async-local-context-manager.test.js

@@ -88,7 +88,7 @@ test('runInContext()', async (t) => {
   })
 
   await t.test('should restore previous context on exception', () => {
-    const contextManager = new AsyncLocalContextManager({})
+    const contextManager = new AsyncLocalContextManager()
 
     const context = contextManager.getContext()
     const previousContext = context.enterSegment({ segment: 'previous', transaction: 'tx' })
@@ -109,7 +109,7 @@ test('runInContext()', async (t) => {
   })
 
   await t.test('should apply `cbThis` arg to execution', (t, end) => {
-    const contextManager = new AsyncLocalContextManager({})
+    const contextManager = new AsyncLocalContextManager()
 
     const context = contextManager.getContext()
     const expectedThis = () => {}
@@ -123,7 +123,7 @@ test('runInContext()', async (t) => {
   })
 
   await t.test('should apply args array to execution', (t, end) => {
-    const contextManager = new AsyncLocalContextManager({})
+    const contextManager = new AsyncLocalContextManager()
 
     const context = contextManager.getContext()
     const expectedArg1 = 'first arg'
@@ -140,7 +140,7 @@ test('runInContext()', async (t) => {
   })
 
   await t.test('should apply arguments construct to execution', (t, end) => {
-    const contextManager = new AsyncLocalContextManager({})
+    const contextManager = new AsyncLocalContextManager()
     const context = contextManager.getContext()
     const expectedArg1 = 'first arg'
     const expectedArg2 = 'second arg'

test/unit/lib/otel/context.test.js

@@ -72,3 +72,28 @@ test('should remove otelSynthesis symbol when it exists on value', () => {
   const otelData = newCtx.getValue(key)
   assert.deepEqual(otelData, { test: 'value' })
 })
+
+test('should add transaction and trace root to otel ctx', (t) => {
+  const { agent } = t.nr
+  const ctx = otel.context.active()
+  const transaction = { agent, traceId: 'traceId', trace: { root: { id: 'segmentId' } } }
+  const newContext = ctx.enterTransaction(transaction)
+  const fakeSpan = newContext.getValue(agent.otelSpanKey)
+  assert.deepEqual(fakeSpan, {
+    segment: transaction.trace.root,
+    transaction
+  })
+})
+
+test('should add segment to otel ctx', (t) => {
+  const { agent } = t.nr
+  const ctx = otel.context.active()
+  ctx._transaction = { agent, traceId: 'traceId' }
+  const segment = { id: 'segmentId' }
+  const newContext = ctx.enterSegment({ segment })
+  const fakeSpan = newContext.getValue(agent.otelSpanKey)
+  assert.deepEqual(fakeSpan, {
+    segment,
+    transaction: newContext.transaction
+  })
+})