fix: Updated undici instrumentation to fix crash with trying to calculate exclusive duration on a segment that no longer exists (#2884)

Author: bizob2828

documentation/feature-flags.md

@@ -20,12 +20,6 @@ Any prerelease flags can be enabled or disabled in your agent config by adding a
 * Environment Variable: `NEW_RELIC_FEATURE_FLAG_REVERSE_NAMING_RULES`
 * Description: Naming rules are in forward order by default.  
 
-#### undici_async_tracking
-* Enabled by default: `true`
-* Configuration: `{ feature_flag: { undici_async_tracking: true|false }}`
-* Environment Variable: `NEW_RELIC_FEATURE_FLAG_UNDICI_ASYNC_TRACKING`
-* Description: If you have multiple undici requests being made in parallel, you may find some state issues if requests to an app are made with keep-alive. If so, *disabling* this flag will avoid these state issues, though at the cost of some broken segment nesting.
-
 #### unresolved_promise_cleanup
 * Enabled by default: `true`
 * Configuration: `{ feature_flag: { unresolved_promise_cleanup: true|false }}`

lib/feature_flags.js

@@ -13,7 +13,6 @@ exports.prerelease = {
 
   promise_segments: false,
   reverse_naming_rules: false,
-  undici_async_tracking: true,
   unresolved_promise_cleanup: true,
   kafkajs_instrumentation: false
 }
@@ -36,6 +35,7 @@ exports.released = [
   'await_support',
   'certificate_bundle',
   'async_local_context',
+  'undici_async_tracking',
   'undici_instrumentation',
   'aws_bedrock_instrumentation',
   'langchain_instrumentation'

lib/instrumentation/undici.js

@@ -11,8 +11,6 @@ const logger = require('../logger').child({ component: 'undici' })
 const NAMES = require('../metrics/names')
 const symbols = require('../symbols')
 // eslint-disable-next-line n/no-unsupported-features/node-builtins
-const { executionAsyncResource } = require('async_hooks')
-// eslint-disable-next-line n/no-unsupported-features/node-builtins
 const diagnosticsChannel = require('diagnostics_channel')
 const synthetics = require('../synthetics')
 const urltils = require('../util/urltils')
@@ -49,42 +47,6 @@ module.exports.unsubscribe = function unsubscribe() {
   })
 }
 
-/**
- * Retrieves the current segment(parent in our context) and transaction from executionAsyncResource
- * or from context manager then adds to the executionAsyncResource for future
- * undici requests within same async context.
- *
- * It was found that when running concurrent undici requests
- * within a transaction that the parent segment would get out of sync
- * depending on the async context of the transaction.  By using
- * `async_hooks.executionResource` it is more reliable.
- *
- * Note: However, if you have concurrent undici requests in a transaction
- * and the request to the transaction is using a keep alive there is a chance the
- * executionAsyncResource may be incorrect because of shared connections.  To revert to a more
- * naive tracking of parent set `config.feature_flag.undici_async_tracking: false` and
- * it will just get the segment and transaction from the context manager.
- *
- * @param {Shim} shim instance of shim
- * @returns {TraceSegment} parent segment
- */
-function getParentContext(shim) {
-  const { config } = shim.agent
-  if (config.feature_flag.undici_async_tracking) {
-    const resource = executionAsyncResource()
-
-    if (!resource[symbols.parentSegment]) {
-      const parent = shim.getSegment()
-      resource[symbols.parentSegment] = parent
-      const tx = shim.tracer.getTransaction()
-      resource[symbols.transaction] = tx
-    }
-    return { segment: resource[symbols.parentSegment], transaction: resource[symbols.transaction] }
-  }
-
-  return shim.tracer.getContext()
-}
-
 /**
  * Injects relevant DT headers for the external request
  *
@@ -160,7 +122,7 @@ function createExternalSegment({ shim, request, segment }) {
  */
 function requestCreateHook(shim, { request }) {
   const { config } = shim.agent
-  const { transaction, segment } = getParentContext(shim)
+  const { transaction, segment } = shim.tracer.getContext()
   request[symbols.parentSegment] = segment
   request[symbols.transaction] = transaction
   if (!(segment || transaction) || segment?.opaque) {
@@ -221,9 +183,8 @@ function requestHeadersHook(shim, { request, response }) {
 }
 
 /**
- * Gets the active and parent from given ctx(request, client connector)
- * and ends active and restores parent to active.  If an error exists
- * it will add the error to the transaction
+ * Gets the active segment, parent segment and transaction from given ctx(request, client connector)
+ * and ends segment and sets the previous parent segment as the active segment.  If an error exists it will add the error to the transaction
  *
  * @param {Shim} shim instance of shim
  * @param {object} params object from undici hook
@@ -236,14 +197,14 @@ function endAndRestoreSegment(shim, { request, error }) {
   const tx = request[symbols.transaction]
   if (activeSegment) {
     activeSegment.end()
+  }
 
-    if (error) {
-      handleError(shim, tx, error)
-    }
+  if (error && tx) {
+    handleError(shim, tx, error)
+  }
 
-    if (parentSegment) {
-      shim.setActiveSegment(parentSegment)
-    }
+  if (parentSegment) {
+    shim.setActiveSegment(parentSegment)
   }
 }
 

test/integration/instrumentation/fetch.test.js

@@ -144,13 +144,57 @@ test('fetch', async function (t) {
       const [{ status }, { status: status2 }] = await Promise.all([req1, req2])
       assert.equal(status, 200)
       assert.equal(status2, 200)
-      assertSegments(tx.trace, tx.trace.root, [`External/${HOST}/post`, `External/${HOST}/put`], {
+      const postName = `External/${HOST}/post`
+      const putName = `External/${HOST}/put`
+      const postSegment = metrics.findSegment(tx.trace, tx.trace.root, postName)
+      assert.equal(postSegment.parentId, tx.trace.root.id)
+      const putSegment = metrics.findSegment(tx.trace, tx.trace.root, putName)
+      // parent of put is the post segment because it is still the active one
+      // not ideal, but our instrumentation does not play nice with diagnostic_channel
+      // we're setting the active segment in the `undici:request:create` and restoring
+      // the parent segment in the request end
+      assert.equal(putSegment.parentId, postSegment.id)
+      assertSegments(tx.trace, tx.trace.root, [postSegment, putSegment], {
         exact: false
       })
       tx.end()
     })
   })
 
+  await t.test('concurrent requests in diff transaction', async () => {
+    const tx1 = helper.runInTransaction(agent, async (tx) => {
+      const { status } = await fetch(`${REQUEST_URL}/post`, {
+        method: 'POST',
+        headers: {
+          'Content-Type': 'application.json'
+        },
+        body: Buffer.from('{"key":"value"}')
+      })
+      assert.equal(status, 200)
+      const postName = `External/${HOST}/post`
+      const postSegment = metrics.findSegment(tx.trace, tx.trace.root, postName)
+      assert.equal(postSegment.parentId, tx.trace.root.id)
+      tx.end()
+    })
+
+    const tx2 = helper.runInTransaction(agent, async(tx) => {
+      const { status } = await fetch(`${REQUEST_URL}/put`, {
+        method: 'PUT',
+        headers: {
+          'Content-Type': 'application.json'
+        },
+        body: Buffer.from('{"key":"value"}')
+      })
+      assert.equal(status, 200)
+      const putName = `External/${HOST}/put`
+      const putSegment = metrics.findSegment(tx.trace, tx.trace.root, putName)
+      assert.equal(putSegment.parentId, tx.trace.root.id)
+      tx.end()
+    })
+
+    await Promise.all([tx1, tx2])
+  })
+
   await t.test('invalid host', async () => {
     await helper.runInTransaction(agent, async (tx) => {
       try {

test/unit/instrumentation/undici.test.js

@@ -27,9 +27,6 @@ test('undici instrumentation', async function (t) {
   const agent = helper.loadMockedAgent()
   agent.config.distributed_tracing.enabled = false
   agent.config.cross_application_tracer.enabled = false
-  agent.config.feature_flag = {
-    undici_async_tracking: true
-  }
   const shim = new TransactionShim(agent, 'undici')
   const loggerMock = require('../mocks/logger')(sandbox)
   const undiciInstrumentation = proxyquire('../../../lib/instrumentation/undici', {
@@ -43,7 +40,6 @@ test('undici instrumentation', async function (t) {
     sandbox.resetHistory()
     agent.config.distributed_tracing.enabled = false
     agent.config.cross_application_tracer.enabled = false
-    agent.config.feature_flag.undici_async_tracking = true
     helper.unloadAgent(agent)
   })
 
@@ -130,84 +126,6 @@ test('undici instrumentation', async function (t) {
       }
     )
 
-    await t.test(
-      'should get the parent segment executionAsyncResource when it already exists',
-      function (t, end) {
-        helper.runInTransaction(agent, function (tx) {
-          const addHeader = sandbox.stub()
-          const request = { origin: HOST, path: '/foo-2', addHeader }
-          channels.create.publish({ request })
-          const segment = tx.trace.add('another segment')
-          segment.start()
-          shim.setActiveSegment(segment)
-          const request2 = { path: '/path', addHeader, origin: HOST }
-          channels.create.publish({ request: request2 })
-          assert.equal(
-            request[symbols.parentSegment].id,
-            request2[symbols.parentSegment].id,
-            'parent segment should be same'
-          )
-          assert.equal(
-            request[symbols.transaction].id,
-            request2[symbols.transaction].id,
-            'tx should be same'
-          )
-          tx.end()
-          end()
-        })
-      }
-    )
-
-    await t.test(
-      'should get diff parent segment across diff async execution contexts',
-      function (t, end) {
-        helper.runInTransaction(agent, function (tx) {
-          const request = { origin: HOST, path: '/request1', addHeader: sandbox.stub() }
-          channels.create.publish({ request })
-          Promise.resolve('test').then(() => {
-            const segment = tx.trace.add('another segment')
-            segment.start()
-            shim.setActiveSegment(segment)
-            const request2 = { path: '/request2', addHeader: sandbox.stub(), origin: HOST }
-            channels.create.publish({ request: request2 })
-            assert.notEqual(request[symbols.parentSegment], request2[symbols.parentSegment])
-            assert.equal(request[symbols.transaction], request2[symbols.transaction])
-            tx.end()
-            end()
-          })
-        })
-      }
-    )
-
-    await t.test(
-      'should get the parent segment shim when `undici_async_tracking` is false',
-      function (t, end) {
-        agent.config.feature_flag.undici_async_tracking = false
-        helper.runInTransaction(agent, function (tx) {
-          const addHeader = sandbox.stub()
-          const request = { path: '/foo-2', addHeader, origin: HOST }
-          channels.create.publish({ request })
-          const segment = tx.trace.add('another segment')
-          segment.start()
-          shim.setActiveSegment(segment)
-          const request2 = { path: '/path', addHeader, origin: HOST }
-          channels.create.publish({ request: request2 })
-          assert.notEqual(
-            request[symbols.parentSegment].name,
-            request2[symbols.parentSegment].name,
-            'parent segment should not be same'
-          )
-          assert.equal(
-            request[symbols.transaction].id,
-            request2[symbols.transaction].id,
-            'tx should be the same'
-          )
-          tx.end()
-          end()
-        })
-      }
-    )
-
     await t.test(
       'should name segment with appropriate attrs based on request.path',
       function (t, end) {

test/versioned/undici/requests.test.js

@@ -194,13 +194,59 @@ test('Undici request tests', async (t) => {
       const [{ statusCode }, { statusCode: statusCode2 }] = await Promise.all([req1, req2])
       assert.equal(statusCode, 200)
       assert.equal(statusCode2, 200)
-      assertSegments(tx.trace, tx.trace.root, [`External/${HOST}/post`, `External/${HOST}/put`], {
+      const postName = `External/${HOST}/post`
+      const putName = `External/${HOST}/put`
+      const postSegment = metrics.findSegment(tx.trace, tx.trace.root, postName)
+      assert.equal(postSegment.parentId, tx.trace.root.id)
+      const putSegment = metrics.findSegment(tx.trace, tx.trace.root, putName)
+      // parent of put is the post segment because it is still the active one
+      // not ideal, but our instrumentation does not play nice with diagnostic_channel
+      // we're setting the active segment in the `undici:request:create` and restoring
+      // the parent segment in the request end
+      assert.equal(putSegment.parentId, postSegment.id)
+      assertSegments(tx.trace, tx.trace.root, [postSegment, putSegment], {
         exact: false
       })
       tx.end()
     })
   })
 
+  await t.test('concurrent requests in diff transaction', async () => {
+    const tx1 = helper.runInTransaction(agent, async (tx) => {
+      const { statusCode } = await undici.request(REQUEST_URL, {
+        path: '/post',
+        method: 'POST',
+        headers: {
+          'Content-Type': 'application.json'
+        },
+        body: Buffer.from('{"key":"value"}')
+      })
+      assert.equal(statusCode, 200)
+      const postName = `External/${HOST}/post`
+      const postSegment = metrics.findSegment(tx.trace, tx.trace.root, postName)
+      assert.equal(postSegment.parentId, tx.trace.root.id)
+      tx.end()
+    })
+
+    const tx2 = helper.runInTransaction(agent, async(tx) => {
+      const { statusCode } = await undici.request(REQUEST_URL, {
+        path: '/put',
+        method: 'PUT',
+        headers: {
+          'Content-Type': 'application.json'
+        },
+        body: Buffer.from('{"key":"value"}')
+      })
+      assert.equal(statusCode, 200)
+      const putName = `External/${HOST}/put`
+      const putSegment = metrics.findSegment(tx.trace, tx.trace.root, putName)
+      assert.equal(putSegment.parentId, tx.trace.root.id)
+      tx.end()
+    })
+
+    await Promise.all([tx1, tx2])
+  })
+
   await t.test('invalid host', async () => {
     await helper.runInTransaction(agent, async (tx) => {
       try {