refactor: Updated subscribers to have a common createSegment that creates, assigns attributes, starts segment (newrelic#3317)

Author: bizob2828

lib/subscribers/base.js

@@ -7,6 +7,12 @@
 // eslint-disable-next-line n/no-unsupported-features/node-builtins
 const { tracingChannel } = require('node:diagnostics_channel')
 
+/**
+ * Base class for defining a subscriber.
+ * events property is an array with the following possible event names:
+ *  `start`, `end`, `asyncStart`, `asyncEnd`, `error`
+ *  @link https://nodejs.org/api/diagnostics_channel.html#class-tracingchannel
+ */
 class Subscriber {
   constructor({ agent, logger, packageName, channelName }) {
     this.agent = agent
@@ -23,31 +29,88 @@ class Subscriber {
     this.store = agent.tracer._contextManager._asyncLocalStorage
   }
 
+  /**
+   * Creates a segment with a name, parent, transaction and optional recorder.
+   * If the segment is successfully created, it will be started and added to the context.
+   * @param {Object} params - Parameters for creating the segment
+   * @param {string} params.name - The name of the segment
+   * @param {Object} params.recorder - Optional recorder for the segment
+   * @param {Context} params.ctx - The context containing the parent segment and transaction
+   * @returns {Context} - The updated context with the new segment or existing context if segment creation fails
+   */
+  createSegment({ name, recorder, ctx }) {
+    const segment = this.agent.tracer.createSegment({
+      name,
+      parent: ctx?.segment,
+      recorder,
+      transaction: ctx?.transaction,
+    })
+
+    if (segment) {
+      segment.opaque = this.opaque
+      segment.start()
+      this.logger.trace('Created segment %s', name)
+      this.addAttributes(segment)
+      const newCtx = ctx.enterSegment({ segment })
+      return newCtx
+    } else {
+      this.logger.trace('Failed to create segment for %s', name)
+      return ctx
+    }
+  }
+
+  /**
+   * By default this is a no-op, but can be overridden by subclasses
+   * @param {Segment} segment - The segment to which attributes will be added
+   * @returns {void}
+   */
+  addAttributes(segment) {
+
+  }
+
+  /**
+   * Checks if the subscriber is enabled based on the agent's configuration.
+   */
   get enabled() {
     return this.config.instrumentation[this.packageName].enabled === true
   }
 
+  /**
+   * Enables the subscriber by binding the store to the channel and setting up the handler.
+   * If the subscriber requires an active transaction, it will check the context before passing the event to the handler.
+   * @returns {Context} - The context after processing the event
+   */
   enable() {
     this.channel.start.bindStore(this.store, (data) => {
       const ctx = this.agent.tracer.getContext()
       if (this.requireActiveTx && !ctx?.transaction?.isActive()) {
         this.logger.debug('Not recording event for %s, transaction is not active', this.package)
-        return
+        return ctx
       }
 
       return this.handler(data, ctx)
     })
   }
 
+  /**
+   * Disables the subscriber by unbinding the store from the channel.
+   */
   disable() {
     this.channel.start.unbindStore(this.store)
   }
 
+  /**
+   * Common handler for when async events end.
+   * It gets the context and touches the segment if it exists.
+   */
   asyncEnd() {
     const ctx = this.agent.tracer.getContext()
     ctx?.segment?.touch()
   }
 
+  /*
+   * Subscribes to the events defined in the `events` array.
+   */
   subscribe() {
     this.subscriptions = this.events.reduce((events, curr) => {
       events[curr] = this[curr].bind(this)
@@ -57,8 +120,12 @@ class Subscriber {
     this.channel.subscribe(this.subscriptions)
   }
 
+  /**
+   * Unsubscribes from the events defined in the `events` array..
+   */
   unsubscribe() {
     this.channel.unsubscribe(this.subscriptions)
+    this.subscriptions = null
   }
 }
 

lib/subscribers/db-operation.js

@@ -10,16 +10,11 @@ const { DB } = require('../metrics/names')
 class DbOperationSubscriber extends DbSubscriber {
   handler(data, ctx) {
     const name = `${DB.OPERATION}/${this.system}/${this.operation}`
-    const segment = this.agent.tracer.createSegment({
+    return this.createSegment({
       name,
-      parent: ctx.segment,
+      ctx,
       recorder: recordOperationMetrics.bind(this),
-      transaction: ctx.transaction
     })
-
-    this.addAttributes(segment)
-    const newCtx = ctx.enterSegment({ segment })
-    return newCtx
   }
 }
 

lib/subscribers/db-query.js

@@ -14,18 +14,11 @@ class DbQuerySubscriber extends DbSubscriber {
     const queryString = this.queryString
     const parsed = this.parseQueryString(queryString)
     const name = `${DB.STATEMENT}/${this.system}/${parsed.collection}/${parsed.operation}`
-    const segment = this.agent.tracer.createSegment({
+    return this.createSegment({
       name,
-      parent: ctx.segment,
+      ctx,
       recorder: recordQueryMetrics.bind(parsed),
-      transaction: ctx.transaction
     })
-
-    segment.opaque = this.opaque
-
-    this.addAttributes(segment)
-    const newCtx = ctx.enterSegment({ segment })
-    return newCtx
   }
 
   parseQueryString(queryString) {

lib/subscribers/mcp-sdk/client.js

@@ -11,24 +11,18 @@ class McpClientSubscriber extends Subscriber {
   constructor({ agent, logger, channelName }) {
     super({ agent, logger, packageName: '@modelcontextprotocol/sdk', channelName })
     this.events = ['asyncEnd']
-    this.requireActiveTx = true
     this.segmentName = 'Unknown'
   }
 
   get enabled() {
-    return this.config.instrumentation[this.packageName].enabled === true &&
-    this.config.ai_monitoring.enabled === true
+    return super.enabled && this.config.ai_monitoring.enabled === true
   }
 
   handler(ctx) {
-    const segment = this.agent.tracer.createSegment({
+    return this.createSegment({
       name: this.segmentName,
-      parent: ctx.segment,
-      transaction: ctx.transaction
+      ctx
     })
-
-    const newCtx = ctx.enterSegment({ segment })
-    return newCtx
   }
 }
 

test/unit/lib/subscribers/base.test.js

@@ -0,0 +1,130 @@
+/*
+ * Copyright 2025 New Relic Corporation. All rights reserved.
+ * SPDX-License-Identifier: Apache-2.0
+ */
+
+'use strict'
+const assert = require('node:assert')
+const test = require('node:test')
+const Subscriber = require('#agentlib/subscribers/base.js')
+const helper = require('#testlib/agent_helper.js')
+const loggerMock = require('../../mocks/logger')
+const { tspl } = require('@matteo.collina/tspl')
+
+test.beforeEach((ctx) => {
+  const agent = helper.loadMockedAgent()
+  const logger = loggerMock()
+  const subscriber = new Subscriber({ agent, logger, packageName: 'test-package', channelName: 'test-channel' })
+  ctx.nr = { agent, subscriber }
+})
+
+test.afterEach((ctx) => {
+  helper.unloadAgent(ctx.nr.agent)
+})
+
+test('should define default properties to subscriber', (t) => {
+  const { subscriber } = t.nr
+  assert.ok(subscriber.agent)
+  assert.ok(subscriber.logger)
+  assert.ok(subscriber.config)
+  assert.equal(subscriber.packageName, 'test-package')
+  assert.equal(subscriber.channelName, 'test-channel')
+  assert.deepEqual(subscriber.events, [])
+  assert.equal(subscriber.opaque, false)
+  assert.equal(subscriber.prefix, 'orchestrion:')
+  assert.equal(subscriber.requireActiveTx, true)
+  assert.equal(subscriber.id, 'orchestrion:test-package:test-channel')
+  assert.ok(subscriber.channel)
+  assert.ok(subscriber.store)
+})
+
+test('addAttributes should not crash', (t) => {
+  const { subscriber } = t.nr
+  assert.doesNotThrow(() => {
+    subscriber.addAttributes({})
+  })
+})
+
+test('enabled should return true if package is enabled', (t) => {
+  const { subscriber } = t.nr
+  subscriber.config.instrumentation['test-package'] = { enabled: true }
+  assert.equal(subscriber.enabled, true)
+})
+
+test('enabled should return false if package is not enabled', (t) => {
+  const { subscriber } = t.nr
+  subscriber.config.instrumentation['test-package'] = { enabled: false }
+  assert.equal(subscriber.enabled, false)
+})
+
+test('should create segment if active tx with proper parent', async (t) => {
+  const { agent, subscriber } = t.nr
+  const plan = tspl(t, { plan: 7 })
+  subscriber.addAttributes = (segment) => {
+    plan.equal(segment.name, 'test-segment')
+  }
+  helper.runInTransaction(agent, async () => {
+    const ctx = agent.tracer.getContext()
+    const newCtx = subscriber.createSegment({
+      name: 'test-segment',
+      ctx,
+    })
+
+    plan.ok(newCtx)
+    plan.equal(newCtx.transaction.id, ctx.transaction.id)
+    const segment = newCtx.segment
+    plan.equal(segment.name, 'test-segment')
+    plan.equal(segment.parentId, ctx.segment.id)
+    plan.equal(segment.opaque, false)
+    // indicates that the segment timer is running
+    plan.equal(segment.timer.state, 2)
+  })
+
+  await plan.completed
+})
+
+test('should not create segment if no active tx', (t) => {
+  const { agent, subscriber } = t.nr
+  const ctx = agent.tracer.getContext()
+  const newCtx = subscriber.createSegment({
+    name: 'test-segment',
+    ctx,
+  })
+
+  assert.deepEqual(newCtx, ctx)
+  assert.ok(!newCtx.segment)
+})
+
+test('should touch segment when asyncEnd is called', (t, end) => {
+  const { agent, subscriber } = t.nr
+  helper.runInTransaction(agent, () => {
+    const ctx = agent.tracer.getContext()
+    const segment = ctx.segment
+    assert.equal(segment.timer.state, 2)
+    assert.equal(segment.timer.touched, false)
+    subscriber.asyncEnd()
+    assert.equal(segment.timer.state, 2)
+    assert.equal(segment.timer.touched, true)
+    end()
+  })
+})
+
+test('should subscribe/unsubscribe to specific events on channel', (t) => {
+  const { subscriber } = t.nr
+  subscriber.bogus = () => {}
+  subscriber.start = () => {}
+  subscriber.end = () => {}
+  subscriber.events = ['bogus', 'start', 'end']
+  subscriber.subscribe()
+  assert.equal(subscriber.channel.start.hasSubscribers, true)
+  assert.equal(subscriber.channel.end.hasSubscribers, true)
+  assert.ok(!subscriber.channel.bogus)
+  assert.ok(subscriber.subscriptions.bogus)
+  assert.ok(subscriber.subscriptions.start)
+  assert.ok(subscriber.subscriptions.end)
+  subscriber.unsubscribe()
+  assert.equal(subscriber.channel.start.hasSubscribers, false)
+  assert.equal(subscriber.channel.end.hasSubscribers, false)
+  assert.ok(!subscriber.channel.bogus)
+  assert.equal(subscriber.subscriptions, null)
+})