feat(profiler): add custom profiling labels API (#7879)

* feat(profiler): add custom profiling labels API

Add tracer.profiling.runWithLabels() and tracer.profiling.setCustomLabelKeys()
APIs that allow users to attach custom labels to wall profiler samples, enabling
flame graph filtering by business logic dimensions (e.g. customer, region) in
the Datadog UI. This brings Node.js profiler to parity with Java and Go.

The implementation stores custom labels alongside span profiling context in
pprof's CPED AsyncLocalStorage as a 2-element array [profilingContext,
customLabels]. A monotonic flag avoids getContext() overhead when custom labels
have never been used. Labels propagate correctly through async continuations.
Internal labels (span id, thread name, etc.) always take precedence over custom
labels with the same key.

Requires AsyncContextFrame (ACF) to be enabled; gracefully falls back to a
passthrough (just calls fn()) when ACF is off or profiling is disabled.

Author: szegedi

docs/test.ts

@@ -522,6 +522,14 @@ const res = {} as OutgoingMessage
 resBlockRequest = tracer.appsec.blockRequest(req, res)
 tracer.appsec.setUser(user)
 
+// Profiling custom labels
+tracer.profiling.setCustomLabelKeys(['customer', 'region'])
+tracer.profiling.setCustomLabelKeys(new Set(['customer', 'region']))
+const labelResult: number = tracer.profiling.runWithLabels({ customer: 'acme', region: 'us-east' }, () => 42)
+tracer.profiling.runWithLabels({ tier: 'premium' }, () => {
+  tracer.profiling.runWithLabels({ region: 'eu-west' }, () => {})
+})
+
 // OTel TracerProvider registers and provides a tracer
 const provider: opentelemetry.TracerProvider = new tracer.TracerProvider();
 provider.register();

index.d.ts

@@ -130,6 +130,11 @@ interface Tracer extends opentracing.Tracer {
 
   appsec: tracer.Appsec;
 
+  /**
+   * Profiling API for attaching custom labels to profiler samples.
+   */
+  profiling: tracer.Profiling;
+
   TracerProvider: tracer.opentelemetry.TracerProvider;
 
   dogstatsd: tracer.DogStatsD;
@@ -1570,6 +1575,35 @@ declare namespace tracer {
     trackUserLoginFailure(login: string, metadata?: any): void;
   }
 
+  export interface Profiling {
+    /**
+     * Declares the set of custom label keys that will be used with
+     * {@link runWithLabels}. This is used for profile upload metadata
+     * (so the Datadog UI knows which keys to index for filtering) and
+     * for pprof serialization optimization.
+     *
+     * @param keys Custom label key names.
+     */
+    setCustomLabelKeys(keys: Iterable<string>): void;
+
+    /**
+     * Runs a function with custom profiling labels attached to all wall profiler
+     * samples taken during its execution. Labels are key-value pairs that appear
+     * in the pprof output and can be used to filter flame graphs in the Datadog UI.
+     *
+     * Requires AsyncContextFrame (ACF) to be enabled. Supports nesting: inner
+     * calls merge labels with outer calls, with inner values taking precedence.
+     *
+     * When profiling is not enabled or ACF is not active, the function is still
+     * called but labels are silently dropped.
+     *
+     * @param labels Custom labels to attach to profiler samples.
+     * @param fn Function to execute with the labels.
+     * @returns The return value of fn.
+     */
+    runWithLabels<T>(labels: Record<string, string | number>, fn: () => T): T;
+  }
+
   export interface Appsec {
     /**
      * Links a successful login event to the current trace. Will link the passed user to the current trace with Appsec.setUser() internally.

packages/dd-trace/src/noop/proxy.js

@@ -13,6 +13,10 @@ const noopDogStatsDClient = new NoopDogStatsDClient()
 const noopLLMObs = new NoopLLMObsSDK(noop)
 const noopOpenFeatureProvider = new NoopFlaggingProvider()
 const noopAIGuard = new NoopAIGuardSDK()
+const noopProfiling = {
+  setCustomLabelKeys () {},
+  runWithLabels (labels, fn) { return fn() },
+}
 
 /** @type {import('../../src/index')} Proxy */
 class NoopProxy {
@@ -98,6 +102,10 @@ class NoopProxy {
     return this
   }
 
+  get profiling () {
+    return noopProfiling
+  }
+
   get TracerProvider () {
     return require('../opentelemetry/tracer_provider')
   }

packages/dd-trace/src/profiler.js

@@ -14,4 +14,26 @@ module.exports = {
   stop: () => {
     profiler.stop()
   },
+
+  /**
+   * Declares the set of custom label keys that will be used with
+   * {@link runWithLabels}.
+   *
+   * @param {Iterable<string>} keys - Custom label key names
+   */
+  setCustomLabelKeys: (keys) => {
+    profiler.setCustomLabelKeys(keys)
+  },
+
+  /**
+   * Runs a function with custom profiling labels attached to wall profiler samples.
+   *
+   * @param {Record<string, string | number>} labels - Custom labels to attach
+   * @param {function(): T} fn - Function to execute with the labels
+   * @returns {T} The return value of fn
+   * @template T
+   */
+  runWithLabels: (labels, fn) => {
+    return profiler.runWithLabels(labels, fn)
+  },
 }

packages/dd-trace/src/profiling/exporters/event_serializer.js

@@ -22,7 +22,7 @@ class EventSerializer {
     return `${type}.pprof`
   }
 
-  getEventJSON ({ profiles, infos, start, end, tags = {}, endpointCounts }) {
+  getEventJSON ({ profiles, infos, start, end, tags = {}, endpointCounts, customAttributes }) {
     const event = {
       attachments: Object.keys(profiles).map(t => this.typeToFile(t)),
       start: start.toISOString(),
@@ -80,6 +80,10 @@ class EventSerializer {
       },
     }
 
+    if (customAttributes) {
+      event.custom_attributes = customAttributes
+    }
+
     if (processTags.serialized) {
       event[processTags.PROFILING_FIELD_NAME] = processTags.serialized
     }

packages/dd-trace/src/profiling/profiler.js

@@ -51,6 +51,7 @@ class Profiler extends EventEmitter {
   #compressionFnInitialized = false
   #compressionOptions
   #config
+  #customLabelKeys = new Set()
   #enabled = false
   #endpointCounts = new Map()
   #lastStart
@@ -135,6 +136,45 @@ class Profiler extends EventEmitter {
     return this.#enabled
   }
 
+  /**
+   * Declares the set of custom label keys that will be used with
+   * {@link runWithLabels}. This is used for profile upload metadata and
+   * for pprof serialization optimization (low-cardinality deduplication).
+   *
+   * @param {Iterable<string>} keys - Custom label key names
+   */
+  setCustomLabelKeys (keys) {
+    this.#customLabelKeys.clear()
+    for (const key of keys) {
+      this.#customLabelKeys.add(key)
+    }
+    if (this.#config) {
+      for (const profiler of this.#config.profilers) {
+        profiler.setCustomLabelKeys?.(this.#customLabelKeys)
+      }
+    }
+  }
+
+  /**
+   * Runs a function with custom profiling labels attached to wall profiler samples.
+   *
+   * @param {Record<string, string | number>} labels - Custom labels to attach
+   * @param {function(): T} fn - Function to execute with the labels
+   * @returns {T} The return value of fn
+   * @template T
+   */
+  runWithLabels (labels, fn) {
+    if (!this.#enabled || !this.#config) {
+      return fn()
+    }
+    for (const profiler of this.#config.profilers) {
+      if (profiler.runWithLabels) {
+        return profiler.runWithLabels(labels, fn)
+      }
+    }
+    return fn()
+  }
+
   #logError (err) {
     logError(this.#logger, err)
   }
@@ -410,7 +450,10 @@ class Profiler extends EventEmitter {
 
     tags.snapshot = snapshotKind
     tags.profile_seq = this.#profileSeq++
-    const exportSpec = { profiles, infos, start, end, tags, endpointCounts }
+    const customAttributes = this.#customLabelKeys.size > 0
+      ? [...this.#customLabelKeys]
+      : undefined
+    const exportSpec = { profiles, infos, start, end, tags, endpointCounts, customAttributes }
     const tasks = this.#config.exporters.map(exporter =>
       exporter.export(exportSpec).catch(err => {
         if (this.#logger) {

packages/dd-trace/src/profiling/profilers/wall.js

@@ -108,6 +108,8 @@ class NativeWallProfiler {
   #captureSpanData = false
   #codeHotspotsEnabled = false
   #cpuProfilingEnabled = false
+  #customLabelsActive = false
+  #customLabelKeys
   #endpointCollectionEnabled = false
   #flushIntervalMillis = 0
   #logger
@@ -245,7 +247,24 @@ class NativeWallProfiler {
     // context -- we simply can't tell which one it might've been across all
     // possible async context frames.
     if (this.#asyncContextFrameEnabled) {
-      this.#pprof.time.setContext(sampleContext)
+      if (this.#customLabelsActive) {
+        // Custom labels may be active in this async context. The current CPED
+        // context could be a 2-element array [profilingContext, customLabels].
+        // Replace the profiling context while preserving the custom labels.
+        // This flag is monotonic (once set, stays true) because async
+        // continuations from runWithLabels can fire at any time after the
+        // synchronous runWithLabels call has returned.
+        const current = this.#pprof.time.getContext()
+        if (Array.isArray(current)) {
+          if (current[0] !== sampleContext) {
+            this.#pprof.time.setContext([sampleContext, current[1]])
+          }
+        } else if (current !== sampleContext) {
+          this.#pprof.time.setContext(sampleContext)
+        }
+      } else {
+        this.#pprof.time.setContext(sampleContext)
+      }
     } else {
       const sampleCount = this._profilerState[kSampleCount]
       if (sampleCount !== this._lastSampleCount) {
@@ -344,6 +363,13 @@ class NativeWallProfiler {
     const lowCardinalityLabels = Object.keys(getThreadLabels())
     lowCardinalityLabels.push(TRACE_ENDPOINT_LABEL)
 
+    // Custom labels are expected to be low-cardinality (e.g. customer tier, region)
+    if (this.#customLabelKeys) {
+      for (const key of this.#customLabelKeys) {
+        lowCardinalityLabels.push(key)
+      }
+    }
+
     const profile = this.#pprof.time.stop(restart, this.#boundGenerateLabels, lowCardinalityLabels)
 
     if (restart) {
@@ -383,7 +409,29 @@ class NativeWallProfiler {
       return getThreadLabels()
     }
 
-    const labels = { ...getThreadLabels() }
+    // 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. Note that with ACF, we don't
+    // use the "ref" indirection.
+    let ref
+    let customLabels
+    const cctx = context.context
+    if (this.#asyncContextFrameEnabled) {
+      // When custom labels are active with ACF, context.context is a 2-element array:
+      // [profilingContext, customLabels]. Otherwise it's a plain object.
+      if (Array.isArray(cctx)) {
+        [ref, customLabels] = cctx
+      } else {
+        ref = cctx
+      }
+    } else {
+      ref = cctx?.ref
+    }
+
+    // Custom labels are spread first so that internal labels always take
+    // precedence and overwrite them.
+    const labels = customLabels === undefined
+      ? { ...getThreadLabels() }
+      : { ...customLabels, ...getThreadLabels() }
 
     if (this.#timelineEnabled) {
       // Incoming timestamps are in microseconds, we emit nanos.
@@ -395,10 +443,6 @@ class NativeWallProfiler {
       labels['async id'] = asyncId
     }
 
-    // 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. Note that with async context
-    // frame, we don't use the "ref" indirection.
-    const ref = this.#asyncContextFrameEnabled ? context.context : context.context?.ref
     if (typeof ref !== 'object') {
       return labels
     }
@@ -421,6 +465,45 @@ class NativeWallProfiler {
     return labels
   }
 
+  /**
+   * Sets the custom label keys used for pprof low-cardinality deduplication.
+   * Called once by the top-level Profiler when keys are declared.
+   *
+   * @param {Iterable<string>} keys
+   */
+  setCustomLabelKeys (keys) {
+    this.#customLabelKeys = keys
+  }
+
+  /**
+   * Runs a function with custom profiling labels attached to all wall profiler
+   * samples taken during its execution. Labels are key-value pairs that appear
+   * in the pprof output and can be used to filter flame graphs in the Datadog UI.
+   *
+   * Requires AsyncContextFrame (ACF) to be enabled. Supports nesting: inner
+   * calls merge labels with outer calls, with inner values taking precedence.
+   *
+   * @param {Record<string, string | number>} labels - Custom labels to attach
+   * @param {function(): T} fn - Function to execute with the labels
+   * @returns {T} The return value of fn
+   * @template T
+   */
+  runWithLabels (labels, fn) {
+    if (!this.#asyncContextFrameEnabled || !this.#withContexts) {
+      return fn()
+    }
+
+    // Read current context; merge custom labels if already in a runWithLabels scope
+    const current = this.#pprof.time.getContext()
+    const isCurrentArray = Array.isArray(current)
+    const customLabels = isCurrentArray ? { ...current[1], ...labels } : labels
+
+    const profilingContext = (isCurrentArray ? current[0] : current) ?? {}
+
+    this.#customLabelsActive = true
+    return this.#pprof.time.runWithContext([profilingContext, customLabels], fn)
+  }
+
   profile (restart) {
     return this.#stop(restart)
   }

packages/dd-trace/src/proxy.js

@@ -330,6 +330,25 @@ class Tracer extends NoopProxy {
     }
   }
 
+  /**
+   * @override
+   */
+  get profiling () {
+    // Lazily require the profiler module and cache the result. If profiling
+    // is not enabled, runWithLabels still works as a passthrough (just calls fn()).
+    const profilerModule = require('./profiler')
+    const profiling = {
+      setCustomLabelKeys (keys) {
+        profilerModule.setCustomLabelKeys(keys)
+      },
+      runWithLabels (labels, fn) {
+        return profilerModule.runWithLabels(labels, fn)
+      },
+    }
+    Reflect.defineProperty(this, 'profiling', { value: profiling, configurable: true, enumerable: true })
+    return profiling
+  }
+
   /**
    * @override
    */