fix(core): retire finished spans without losing context

Long-lived async resources can retain finished middleware spans and their parent chains after export, making retained heap grow with request count.

The original span remains active until serialization and response listeners complete. Retained resources then keep a compact propagation facade, so late work preserves trace and baggage context without pinning the span graph.

Fixes: #9227

Author: BridgeAR

.github/CODEOWNERS

@@ -247,6 +247,8 @@
 /packages/dd-trace/test/opentracing/ @DataDog/apm-sdk-capabilities-js
 /packages/dd-trace/src/remote_config/ @DataDog/apm-sdk-capabilities-js
 /packages/dd-trace/test/remote_config/ @DataDog/apm-sdk-capabilities-js
+/packages/dd-trace/src/active-span.js @DataDog/apm-sdk-capabilities-js
+/packages/dd-trace/test/active-span.spec.js @DataDog/apm-sdk-capabilities-js
 /packages/dd-trace/src/baggage.js @DataDog/apm-sdk-capabilities-js
 /packages/dd-trace/test/baggage.spec.js @DataDog/apm-sdk-capabilities-js
 /packages/dd-trace/src/sampler.js @DataDog/apm-sdk-capabilities-js

packages/datadog-core/src/storage.js

@@ -2,6 +2,8 @@
 
 const { AsyncLocalStorage } = require('async_hooks')
 
+const kStoreRetirement = Symbol('dd-trace.store-retirement')
+
 /**
  * `AsyncLocalStorage` with a `getHandle()` escape hatch: a span stashes the
  * active handle at creation (see opentracing/span.js) so a later context can
@@ -19,6 +21,15 @@ const { AsyncLocalStorage } = require('async_hooks')
  * @typedef {Record<string, T>} Store
  */
 class DatadogStorage extends AsyncLocalStorage {
+  /**
+   * @param {Store<unknown>} [store]
+   * @override
+   */
+  enterWith (store) {
+    store?.[kStoreRetirement]?.add(store)
+    return super.enterWith(store)
+  }
+
   /**
    * Passthrough to the real `getStore()`. A span stashes this handle and feeds
    * it back to `getStore(handle)` later. Identical in both modes: under ACF the
@@ -53,6 +64,7 @@ const isACFActive = (() => {
   return active
 })()
 
+/* istanbul ignore if: legacy async_hooks backend is covered on supported Node releases without ACF */
 if (!isACFActive) {
   const superGetStore = AsyncLocalStorage.prototype.getStore
   const superEnterWith = AsyncLocalStorage.prototype.enterWith
@@ -66,6 +78,7 @@ if (!isACFActive) {
    * @param {Store<unknown>} [store]
    */
   DatadogStorage.prototype.enterWith = function enterWith (store) {
+    store?.[kStoreRetirement]?.add(store)
     const handle = { noop: store?.noop }
     stores.set(handle, store)
     superEnterWith.call(this, handle)
@@ -128,4 +141,4 @@ function storage (namespace) {
   return storages[namespace]
 }
 
-module.exports = { storage, isACFActive }
+module.exports = { storage, isACFActive, kStoreRetirement }

packages/datadog-core/test/storage.spec.js

@@ -6,7 +6,7 @@ const { executionAsyncResource } = require('async_hooks')
 const { describe, it, beforeEach, afterEach } = require('mocha')
 
 require('../../dd-trace/test/setup/core')
-const { storage } = require('../src/storage')
+const { kStoreRetirement, storage } = require('../src/storage')
 
 describe('storage', () => {
   let testStorage
@@ -51,6 +51,21 @@ describe('storage', () => {
     assert.strictEqual(storage('test'), testStorage)
   })
 
+  it('should register entered stores with their retirement group', () => {
+    const entered = []
+    const store = {
+      [kStoreRetirement]: {
+        add (value) {
+          entered.push(value)
+        },
+      },
+    }
+
+    testStorage.enterWith(store)
+
+    assert.deepStrictEqual(entered, [store])
+  })
+
   it('should not have its store referenced by the underlying async resource', () => {
     const resource = executionAsyncResource()
 

packages/datadog-instrumentations/src/http/server.js

@@ -10,6 +10,7 @@ const startServerCh = channel('apm:http:server:request:start')
 const exitServerCh = channel('apm:http:server:request:exit')
 const errorServerCh = channel('apm:http:server:request:error')
 const finishServerCh = channel('apm:http:server:request:finish')
+const postFinishServerCh = channel('apm:http:server:request:postfinish')
 const startWriteHeadCh = channel('apm:http:server:response:writeHead:start')
 const finishSetHeaderCh = channel('datadog:http:server:response:set-header:finish')
 const startSetHeaderCh = channel('datadog:http:server:response:set-header:start')
@@ -52,14 +53,20 @@ function wrapResponseEmit (originalEmit) {
   // skips its name rewrite; rest params keep the per-event forwarding
   // allocation-free.
   return function emit (...args) {
-    if (!finishServerCh.hasSubscribers) {
+    if (!finishServerCh.hasSubscribers && !postFinishServerCh.hasSubscribers) {
       return Reflect.apply(originalEmit, this, args)
     }
 
     const eventName = args[0]
     if ((eventName === 'finish' || eventName === 'close') && !requestFinishedSet.has(this)) {
-      finishServerCh.publish({ req: this.req })
+      const ctx = { req: this.req }
+      finishServerCh.publish(ctx)
       requestFinishedSet.add(this)
+      try {
+        return Reflect.apply(originalEmit, this, args)
+      } finally {
+        postFinishServerCh.publish(ctx)
+      }
     }
 
     return Reflect.apply(originalEmit, this, args)

packages/datadog-instrumentations/src/light-my-request.js

@@ -22,6 +22,7 @@ const startServerCh = channel('apm:http:server:request:start')
 const exitServerCh = channel('apm:http:server:request:exit')
 const errorServerCh = channel('apm:http:server:request:error')
 const finishServerCh = channel('apm:http:server:request:finish')
+const postFinishServerCh = channel('apm:http:server:request:postfinish')
 
 addHook({ name: 'light-my-request', versions: ['>=3'] }, (lightMyRequest) => {
   // Wrap the inject function
@@ -51,29 +52,22 @@ function wrapDispatchFunc (dispatchFunc) {
     // Publish start event (same as HTTP server)
     startServerCh.publish({ req, res, abortController })
 
-    // Track when response finishes via 'finish' event (like HTTP instrumentation)
     let finishCalled = false
-    const onFinish = () => {
-      if (!finishCalled) {
-        finishCalled = true
-        finishServerCh.publish({ req })
-      }
-    }
-
-    // light-my-request Response emits 'finish' when done
-    if (res.on && typeof res.on === 'function') {
-      res.once('finish', onFinish)
-      res.once('close', onFinish)
-    }
-
-    // Also wrap end() as fallback
-    const originalEnd = res.end
-    if (originalEnd) {
-      res.end = function wrappedEnd (...args) {
-        const result = originalEnd.apply(this, args)
-        // Trigger finish if events don't fire
-        setImmediate(onFinish)
-        return result
+    const originalEmit = res.emit
+    if (typeof originalEmit === 'function') {
+      res.emit = function emit (...args) {
+        const eventName = args[0]
+        if ((eventName === 'finish' || eventName === 'close') && !finishCalled) {
+          finishCalled = true
+          const ctx = { req }
+          finishServerCh.publish(ctx)
+          try {
+            return Reflect.apply(originalEmit, this, args)
+          } finally {
+            postFinishServerCh.publish(ctx)
+          }
+        }
+        return Reflect.apply(originalEmit, this, args)
       }
     }
 

packages/datadog-instrumentations/test/http.spec.js

@@ -10,17 +10,19 @@ const sinon = require('sinon')
 const agent = require('../../dd-trace/test/plugins/agent')
 
 describe('client', () => {
-  let url, http, startChannelCb, endChannelCb, asyncStartChannelCb, errorChannelCb
+  let tracer, url, http, startChannelCb, endChannelCb, asyncStartChannelCb, errorChannelCb
 
   const startChannel = dc.channel('apm:http:client:request:start')
   const endChannel = dc.channel('apm:http:client:request:end')
   const asyncStartChannel = dc.channel('apm:http:client:request:asyncStart')
   const errorChannel = dc.channel('apm:http:client:request:error')
   const responseFinishChannel = dc.channel('apm:http:client:response:finish')
   const responseStartChannel = dc.channel('apm:http:client:response:start')
+  const serverFinishChannel = dc.channel('apm:http:server:request:finish')
+  const serverPostFinishChannel = dc.channel('apm:http:server:request:postfinish')
 
   before(async () => {
-    await agent.load('http')
+    tracer = await agent.load('http')
   })
 
   after(() => {
@@ -291,6 +293,60 @@ describe('client', () => {
         const describeIfHttp = httpSchema === 'http' ? describe : describe.skip
 
         describeIfHttp('with local http server', () => {
+          it('keeps response events unchanged without lifecycle subscribers', async () => {
+            tracer.use('http', false)
+            let finished = false
+            const server = http.createServer((req, res) => {
+              res.on('finish', () => {
+                finished = true
+              })
+              res.end()
+            })
+
+            try {
+              await new Promise(resolve => server.listen(0, resolve))
+              await new Promise((resolve, reject) => {
+                http.get(`http://127.0.0.1:${server.address().port}/`, res => {
+                  res.resume()
+                  res.on('end', resolve)
+                }).on('error', reject)
+              })
+            } finally {
+              await new Promise(resolve => server.close(resolve))
+              tracer.use('http', {})
+            }
+
+            assert.strictEqual(finished, true)
+          })
+
+          it('publishes server postfinish after response listeners', async () => {
+            const calls = []
+            const onFinish = () => calls.push('finish')
+            const onPostFinish = () => calls.push('postfinish')
+            serverFinishChannel.subscribe(onFinish)
+            serverPostFinishChannel.subscribe(onPostFinish)
+            const server = http.createServer((req, res) => {
+              res.on('finish', () => calls.push('listener'))
+              res.end()
+            })
+
+            try {
+              await new Promise(resolve => server.listen(0, resolve))
+              await new Promise((resolve, reject) => {
+                http.get(`http://127.0.0.1:${server.address().port}/`, res => {
+                  res.resume()
+                  res.on('end', resolve)
+                }).on('error', reject)
+              })
+            } finally {
+              serverFinishChannel.unsubscribe(onFinish)
+              serverPostFinishChannel.unsubscribe(onPostFinish)
+              await new Promise(resolve => server.close(resolve))
+            }
+
+            assert.deepStrictEqual(calls, ['finish', 'listener', 'postfinish'])
+          })
+
           function requestWithLocalServer ({ responseHeaders, responseBody, onResponse }, done) {
             const server = http.createServer((req, res) => {
               res.writeHead(200, responseHeaders)

packages/datadog-instrumentations/test/light-my-request.spec.js

@@ -12,9 +12,10 @@ withVersions('light-my-request', 'light-my-request', version => describe('light-
   const startServerCh = dc.channel('apm:http:server:request:start')
   const exitServerCh = dc.channel('apm:http:server:request:exit')
   const finishServerCh = dc.channel('apm:http:server:request:finish')
+  const postFinishServerCh = dc.channel('apm:http:server:request:postfinish')
   const errorServerCh = dc.channel('apm:http:server:request:error')
 
-  let startStub, exitStub, finishStub, errorStub
+  let startStub, exitStub, finishStub, postFinishStub, errorStub
   let inject, Fastify
 
   before(async () => {
@@ -31,18 +32,21 @@ withVersions('light-my-request', 'light-my-request', version => describe('light-
     startStub = sinon.stub()
     exitStub = sinon.stub()
     finishStub = sinon.stub()
+    postFinishStub = sinon.stub()
     errorStub = sinon.stub()
 
     startServerCh.subscribe(startStub)
     exitServerCh.subscribe(exitStub)
     finishServerCh.subscribe(finishStub)
+    postFinishServerCh.subscribe(postFinishStub)
     errorServerCh.subscribe(errorStub)
   })
 
   afterEach(() => {
     startServerCh.unsubscribe(startStub)
     exitServerCh.unsubscribe(exitStub)
     finishServerCh.unsubscribe(finishStub)
+    postFinishServerCh.unsubscribe(postFinishStub)
     errorServerCh.unsubscribe(errorStub)
   })
 
@@ -102,7 +106,7 @@ withVersions('light-my-request', 'light-my-request', version => describe('light-
       assert(testCall, 'exit channel should be called for /test request')
     })
 
-    it('should publish to finish channel when response completes', async () => {
+    it('should publish finish and postfinish when the response completes', async () => {
       await app.inject({
         method: 'GET',
         url: '/test',
@@ -119,6 +123,13 @@ withVersions('light-my-request', 'light-my-request', version => describe('light-
       })
 
       assert(testCall, 'finish channel should be called for /test request')
+      const postFinishCall = postFinishStub.getCalls().find(call => {
+        const { req } = call.args[0]
+        return req.url === '/test'
+      })
+
+      assert(postFinishCall, 'postfinish channel should be called for /test request')
+      assert(testCall.calledBefore(postFinishCall))
     })
 
     it('should link res.req for context tracking', async () => {

packages/datadog-plugin-google-cloud-pubsub/src/pubsub-push-subscription.js

@@ -1,6 +1,7 @@
 'use strict'
 
 const TracingPlugin = require('../../dd-trace/src/plugins/tracing')
+const { createStoreRetirement, enterSpanForRetirement } = require('../../dd-trace/src/active-span')
 const SpanContext = require('../../dd-trace/src/opentracing/span_context')
 const id = require('../../dd-trace/src/id')
 const log = require('../../dd-trace/src/log')
@@ -32,12 +33,28 @@ class GoogleCloudPubsubPushSubscriptionPlugin extends TracingPlugin {
     this.addSub('apm:http:server:request:finish', ({ req }) => {
       this.#finishPushReceiveSpan(req)
     })
+    this.addSub('apm:http:server:request:postfinish', ({ req }) => {
+      this.#retirePushReceiveSpan(req)
+    })
   }
 
+  /**
+   * @param {import('node:http').IncomingMessage} req
+   */
   #finishPushReceiveSpan (req) {
-    const pushReceiveSpan = pushReceiveSpans.get(req)
-    if (pushReceiveSpan && !pushReceiveSpan._duration) {
-      pushReceiveSpan.finish()
+    const entry = pushReceiveSpans.get(req)
+    if (entry && !entry.span._duration) {
+      entry.span.finish()
+    }
+  }
+
+  /**
+   * @param {import('node:http').IncomingMessage} req
+   */
+  #retirePushReceiveSpan (req) {
+    const entry = pushReceiveSpans.get(req)
+    if (entry) {
+      entry.retirement.retire()
       pushReceiveSpans.delete(req)
     }
   }
@@ -61,6 +78,9 @@ class GoogleCloudPubsubPushSubscriptionPlugin extends TracingPlugin {
     return false
   }
 
+  /**
+   * @param {{ req: import('node:http').IncomingMessage, res: import('node:http').ServerResponse }} ctx
+   */
   #createPushReceiveSpanAndActivate (ctx) {
     const { req, res } = ctx
     const messageData = this.#parseMessage(req)
@@ -90,8 +110,9 @@ class GoogleCloudPubsubPushSubscriptionPlugin extends TracingPlugin {
       return
     }
 
-    this.enter(pushReceiveSpan, { req, res })
-    pushReceiveSpans.set(req, pushReceiveSpan)
+    const retirement = createStoreRetirement(pushReceiveSpan.context())
+    enterSpanForRetirement(pushReceiveSpan, { req, res }, retirement)
+    pushReceiveSpans.set(req, { span: pushReceiveSpan, retirement })
   }
 
   #parseMessage (req) {