Allow the CPU profiler to store sample context in the async context frame

szegedi · szegedi · commit 13c325ff5352 · 2025-05-16T13:22:12.000+02:00
diff --git a/packages/dd-trace/src/profiling/config.js b/packages/dd-trace/src/profiling/config.js
@@ -15,6 +15,7 @@ const { GIT_REPOSITORY_URL, GIT_COMMIT_SHA } = require('../plugins/util/tags')
 const { tagger } = require('./tagger')
 const { isFalse, isTrue } = require('../util')
 const { getAzureTagsFromMetadata, getAzureAppMetadata } = require('../azure_metadata')
+const satisfies = require('semifies')
 
 class Config {
   constructor (options = {}) {
@@ -41,6 +42,7 @@ class Config {
       DD_PROFILING_TIMELINE_ENABLED,
       DD_PROFILING_UPLOAD_PERIOD,
       DD_PROFILING_UPLOAD_TIMEOUT,
+      DD_PROFILING_USE_ASYNC_CONTEXT_FRAME,
       DD_PROFILING_V8_PROFILER_BUG_WORKAROUND,
       DD_PROFILING_WALLTIME_ENABLED,
       DD_SERVICE,
@@ -191,6 +193,27 @@ class Config {
     logExperimentalVarDeprecation('CPU_ENABLED')
     checkOptionWithSamplingContextAllowed(this.cpuProfilingEnabled, 'CPU profiling')
 
+    function turnOffAsyncContextFrame (that, msg) {
+      logger.warn(
+        `DD_PROFILING_USE_ASYNC_CONTEXT_FRAME was set ${msg}, it will have no effect.`)
+      that.useAsyncContextFrame = false
+    }
+
+    this.useAsyncContextFrame = isTrue(coalesce(options.useAsyncContextFrame,
+      DD_PROFILING_USE_ASYNC_CONTEXT_FRAME, false))
+    if (this.useAsyncContextFrame) {
+      if (satisfies(process.versions.node, '>=24.0.0')) {
+        if (process.execArgv.includes('--no-async-context-frame')) {
+          turnOffAsyncContextFrame(this, 'with --no-async-context-frame')
+        }
+      } else if (satisfies(process.versions.node, '>=23.0.0')) {
+        if (!process.execArgv.includes('--experimental-async-context-frame')) {
+          turnOffAsyncContextFrame(this, 'without --experimental-async-context-frame')
+        }
+      } else {
+        turnOffAsyncContextFrame(this, 'but it requires at least Node 23')
+      }
+    }
     this.profilers = ensureProfilers(profilers, this)
   }
 }
diff --git a/packages/dd-trace/src/profiling/profilers/wall.js b/packages/dd-trace/src/profiling/profilers/wall.js
@@ -15,7 +15,7 @@ const {
 
 const { isWebServerSpan, endpointNameFromTags, getStartedSpans } = require('../webspan-utils')
 
-const beforeCh = dc.channel('dd-trace:storage:before')
+let beforeCh
 const enterCh = dc.channel('dd-trace:storage:enter')
 const spanFinishCh = dc.channel('dd-trace:span:finish')
 const profilerTelemetryMetrics = telemetryMetrics.manager.namespace('profilers')
@@ -29,40 +29,94 @@ function getActiveSpan () {
   return store && store.span
 }
 
+function updateContext (context) {
+  // Converting spanIDs to strings is not necessary as generateLabels will do it
+  // too. When we don't use async context frame, we can convert them when the
+  // sample is taken though so we amortize the latency of operations. It is an
+  // optimization.
+  if (typeof context.spanId === 'object') {
+    context.spanId = context.spanId.toString(10)
+  }
+  if (typeof context.rootSpanId === 'object') {
+    context.rootSpanId = context.rootSpanId.toString(10)
+  }
+  if (context.webTags !== undefined && context.endpoint === undefined) {
+    // endpoint may not be determined yet, but keep it as fallback
+    // if tags are not available anymore during serialization
+    context.endpoint = endpointNameFromTags(context.webTags)
+  }
+}
+
 let channelsActivated = false
-function ensureChannelsActivated () {
+function ensureChannelsActivated (useAsyncContextFrame) {
   if (channelsActivated) return
 
-  const { AsyncLocalStorage, createHook } = require('async_hooks')
   const shimmer = require('../../../../datadog-shimmer')
 
-  createHook({ before: () => beforeCh.publish() }).enable()
+  let interceptingAsyncContextFrameSet = false
+  // When using the async context frame to store sample context (available with
+  // Node 24), we do not need to use the async hooks anymore.
+  if (!useAsyncContextFrame) {
+    const { createHook } = require('async_hooks')
+    beforeCh = dc.channel('dd-trace:storage:before')
+    createHook({ before: () => beforeCh.publish() }).enable()
+  } else if (process.execArgv.includes('--expose-internals')) {
+    const AsyncContextFrame = require('internal/async_context_frame')
+    shimmer.wrap(AsyncContextFrame, 'set', function (original) {
+      return function (frame) {
+        original.call(this, frame)
+        enterCh.publish()
+      }
+    })
+    interceptingAsyncContextFrameSet = true
+  }
 
-  let inRun = false
-  shimmer.wrap(AsyncLocalStorage.prototype, 'enterWith', function (original) {
-    return function (...args) {
-      const retVal = original.apply(this, args)
-      if (!inRun) enterCh.publish()
-      return retVal
-    }
-  })
+  if (!interceptingAsyncContextFrameSet) {
+    const { AsyncLocalStorage, AsyncResource } = require('async_hooks')
 
-  shimmer.wrap(AsyncLocalStorage.prototype, 'run', function (original) {
-    return function (store, callback, ...args) {
-      const wrappedCb = shimmer.wrapFunction(callback, cb => function (...args) {
-        inRun = false
+    let inRun = false
+    shimmer.wrap(AsyncLocalStorage.prototype, 'enterWith', function (original) {
+      return function (...args) {
+        const retVal = original.apply(this, args)
+        if (!inRun) enterCh.publish()
+        return retVal
+      }
+    })
+
+    shimmer.wrap(AsyncLocalStorage.prototype, 'run', function (original) {
+      return function (store, callback, ...args) {
+        const wrappedCb = shimmer.wrapFunction(callback, cb => function (...args) {
+          inRun = false
+          enterCh.publish()
+          const retVal = cb.apply(this, args)
+          inRun = true
+          return retVal
+        })
+        inRun = true
+        const retVal = original.call(this, store, wrappedCb, ...args)
         enterCh.publish()
-        const retVal = cb.apply(this, args)
+        inRun = false
+        return retVal
+      }
+    })
+
+    shimmer.wrap(AsyncResource.prototype, 'runInAsyncScope', function (original) {
+      return function (callback, thisArg, ...args) {
+        const wrappedCb = shimmer.wrapFunction(callback, cb => function (...args) {
+          inRun = false
+          enterCh.publish()
+          const retVal = cb.apply(this, args)
+          inRun = true
+          return retVal
+        })
         inRun = true
+        const retVal = original.call(this, wrappedCb, thisArg, ...args)
+        enterCh.publish()
+        inRun = false
         return retVal
-      })
-      inRun = true
-      const retVal = original.call(this, store, wrappedCb, ...args)
-      enterCh.publish()
-      inRun = false
-      return retVal
-    }
-  })
+      }
+    })
+  }
 
   channelsActivated = true
 }
@@ -76,6 +130,8 @@ class NativeWallProfiler {
     this._endpointCollectionEnabled = !!options.endpointCollectionEnabled
     this._timelineEnabled = !!options.timelineEnabled
     this._cpuProfilingEnabled = !!options.cpuProfilingEnabled
+    this._useAsyncContextFrame = !!options.useAsyncContextFrame
+
     // We need to capture span data into the sample context for either code hotspots
     // or endpoint collection.
     this._captureSpanData = this._codeHotspotsEnabled || this._endpointCollectionEnabled
@@ -131,19 +187,24 @@ class NativeWallProfiler {
       withContexts: this._withContexts,
       lineNumbers: false,
       workaroundV8Bug: this._v8ProfilerBugWorkaroundEnabled,
-      collectCpuTime: this._cpuProfilingEnabled
+      collectCpuTime: this._cpuProfilingEnabled,
+      useCPED: this._useAsyncContextFrame
     })
 
     if (this._withContexts) {
-      this._setNewContext()
+      if (!this._useAsyncContextFrame) {
+        this._setNewContext()
+      }
 
       if (this._captureSpanData) {
         this._profilerState = this._pprof.time.getState()
         this._lastSampleCount = 0
 
-        ensureChannelsActivated()
+        ensureChannelsActivated(this._useAsyncContextFrame)
 
-        beforeCh.subscribe(this._enter)
+        if (!this._useAsyncContextFrame) {
+          beforeCh.subscribe(this._enter)
+        }
         enterCh.subscribe(this._enter)
         spanFinishCh.subscribe(this._spanFinished)
       }
@@ -155,17 +216,37 @@ class NativeWallProfiler {
   _enter () {
     if (!this._started) return
 
-    const sampleCount = this._profilerState[kSampleCount]
-    if (sampleCount !== this._lastSampleCount) {
-      this._lastSampleCount = sampleCount
-      const context = this._currentContext.ref
-      this._setNewContext()
+    const span = getActiveSpan()
+    const sampleContext = span ? this._getProfilingContext(span) : {}
+
+    // Note that we store the sample context differently with and without the
+    // async context frame. With the async context frame, we tell the profiler
+    // to store the sample context directly in the frame on each enterWith.
+    // Without the async context frame, we store one holder object as the
+    // profiler's single sample context, and reassign its "ref" property on
+    // every async context change. Then when we detect that the profiler took a
+    // sample (and thus bound the holder as that sample's context), we create a
+    // new holder object so that we no longer mutate the old one. This is really
+    // an optimization to avoid going to profiler's native SetContext every
+    // time. With async context frame however, we can't have that optimization,
+    // as we can't tell from which async context frame was the sampling context
+    // taken. For the same reason we can't call updateContext() on the old
+    // context -- we simply can't tell which one it might've been across all
+    // possible async context frames.
+    if (this._useAsyncContextFrame) {
+      this._pprof.time.setContext(sampleContext)
+    } else {
+      const sampleCount = this._profilerState[kSampleCount]
+      if (sampleCount !== this._lastSampleCount) {
+        this._lastSampleCount = sampleCount
+        const context = this._currentContext.ref
+        this._setNewContext()
 
-      this._updateContext(context)
-    }
+        updateContext(context)
+      }
 
-    const span = getActiveSpan()
-    this._currentContext.ref = span ? this._getProfilingContext(span) : {}
+      this._currentContext.ref = sampleContext
+    }
   }
 
   _getProfilingContext (span) {
@@ -246,7 +327,7 @@ class NativeWallProfiler {
   _stop (restart) {
     if (!this._started) return
 
-    if (this._captureSpanData) {
+    if (this._captureSpanData && !this._useAsyncContextFrame) {
       // update last sample context if needed
       this._enter()
       this._lastSampleCount = 0
@@ -260,7 +341,9 @@ class NativeWallProfiler {
       }
     } else {
       if (this._captureSpanData) {
-        beforeCh.unsubscribe(this._enter)
+        if (!this._useAsyncContextFrame) {
+          beforeCh.unsubscribe(this._enter)
+        }
         enterCh.unsubscribe(this._enter)
         spanFinishCh.unsubscribe(this._spanFinished)
         this._profilerState = undefined
@@ -297,19 +380,20 @@ class NativeWallProfiler {
     }
 
     // Native profiler doesn't set context.context for some samples, such as idle samples or when
-    // the context was otherwise unavailable when the sample was taken.
-    const ref = context.context?.ref
+    // the context was otherwise unavailable when the sample was taken. Note that with async context
+    // frame, we don't use the "ref" indirection.
+    const ref = this._useAsyncContextFrame ? context.context : context.context?.ref
     if (typeof ref !== 'object') {
       return labels
     }
 
     const { spanId, rootSpanId, webTags, endpoint } = ref
 
     if (spanId !== undefined) {
-      labels[SPAN_ID_LABEL] = spanId
+      labels[SPAN_ID_LABEL] = typeof spanId === 'object' ? spanId.toString(10) : spanId
     }
     if (rootSpanId !== undefined) {
-      labels[LOCAL_ROOT_SPAN_ID_LABEL] = rootSpanId
+      labels[LOCAL_ROOT_SPAN_ID_LABEL] = typeof rootSpanId === 'object' ? rootSpanId.toString(10) : rootSpanId
     }
     if (webTags !== undefined && Object.keys(webTags).length !== 0) {
       labels['trace endpoint'] = endpointNameFromTags(webTags)
