feat: async context labels

README.md

@@ -16,6 +16,13 @@ Visit the [npm package][releases] to find the latest version of this package.
 
 Visit [docs](https://grafana.com/docs/pyroscope/latest/configure-client/language-sdks/nodejs/) page for usage and configuration documentation.
 
+Async labels are opt-in: configure `wall: { asyncContext: true }`, then use
+`wrapWithLabels` to scope work across `await`. Requires Node.js 24+ with
+AsyncContextFrame enabled, or Node.js 22 with
+`--experimental-async-context-frame`. In this mode, `setLabels` replaces labels
+in the current async context without changing sibling contexts. Heap profiles
+are unchanged.
+
 [Grafana Pyroscope]: https://grafana.com/oss/pyroscope/
 [@datadog/pprof]: https://github.com/DataDog/pprof-nodejs
 [v8-prof]: https://v8.dev/docs/profile

src/environment.ts

@@ -11,6 +11,7 @@ export interface Environment {
   wallSamplingDurationMs: number | undefined;
   wallSamplingIntervalMicros: number | undefined;
   wallCollectCpuTime: boolean | undefined;
+  wallAsyncContext: boolean | undefined;
   stripFilenames: StripFilenamesMode | undefined;
   shortenPaths: boolean | undefined;
 }

src/index.ts

@@ -42,12 +42,22 @@ function setLabels(labels: Record<string, number | string>): void {
   getProfiler().wallProfiler.profiler.setLabels(labels);
 }
 
-export function wrapWithLabels(
+export function wrapWithLabels<R, TArgs extends unknown[]>(
   lbls: Record<string, string | number>,
-  fn: () => void,
+  fn: (...args: TArgs) => R,
+  ...args: TArgs
+): R;
+export function wrapWithLabels<R>(
+  lbls: Record<string, string | number>,
+  fn: () => R,
   ...args: unknown[]
-): void {
-  getProfiler().wallProfiler.profiler.wrapWithLabels(lbls, fn, ...args);
+): R;
+export function wrapWithLabels<R, TArgs extends unknown[]>(
+  lbls: Record<string, string | number>,
+  fn: (...args: TArgs) => R,
+  ...args: TArgs
+): R {
+  return getProfiler().wallProfiler.profiler.wrapWithLabels(lbls, fn, ...args);
 }
 
 function startWallProfiling(): void {

src/profilers/heap-profiler.ts

@@ -25,7 +25,7 @@ export class HeapProfiler implements Profiler<HeapProfilerStartArgs> {
     throw new Error("heap profiler doesn't support labels");
   }
 
-  public wrapWithLabels(): void {
+  public wrapWithLabels(): never {
     throw new Error("heap profiler doesn't support labels");
   }
 

src/profilers/profiler.ts

@@ -5,11 +5,11 @@ export interface Profiler<TStartArgs> {
 
   setLabels(labels: Record<string, number | string>): void;
 
-  wrapWithLabels(
+  wrapWithLabels<R, TArgs extends unknown[]>(
     labels: Record<string, number | string>,
-    fn: () => void,
-    ...args: unknown[]
-  ): void;
+    fn: (...args: TArgs) => R,
+    ...args: TArgs
+  ): R;
 
   start(args: TStartArgs): void;
 

src/profilers/pyroscope-profiler.ts

@@ -107,7 +107,7 @@ export class PyroscopeProfiler {
     return new ContinuousProfiler({
       exporter,
       flushIntervalMs: flushIntervalMs,
-      profiler: new WallProfiler(),
+      profiler: new WallProfiler(config.wall?.asyncContext ?? false),
       startArgs: {
         sourceMapper: this.toDDSourceMapper(config.sourceMapper),
         samplingDurationMs:

src/profilers/wall-profiler.ts

@@ -1,5 +1,6 @@
 import { time, SourceMapper, LabelSet, TimeProfileNode } from '@datadog/pprof';
 import { Profile } from 'pprof-format';
+import { AsyncLocalStorage } from 'node:async_hooks';
 
 import { ProfileExport } from '../profile-exporter.js';
 import { Profiler } from './profiler.js';
@@ -32,20 +33,24 @@ export interface ProfilerContext {
   labels?: LabelSet;
 }
 
-// eslint-disable-next-line @typescript-eslint/no-explicit-any
-type AnyWrappedCallback = any;
-
 export class WallProfiler implements Profiler<WallProfilerStartArgs> {
   private lastProfiledAt: Date;
   private lastContext: ProfilerContext;
   private lastSamplingIntervalMicros!: number;
+  private readonly contextBoundary = new AsyncLocalStorage<ProfilerContext>();
 
-  constructor() {
+  constructor(private readonly asyncContext = false) {
     this.lastContext = {};
     this.lastProfiledAt = new Date();
   }
 
   public getLabels(): LabelSet {
+    if (this.asyncContext) {
+      const context: ProfilerContext | undefined = time.isStarted()
+        ? time.getContext()
+        : undefined;
+      return context?.labels ?? {};
+    }
     return this.lastContext.labels ?? {};
   }
 
@@ -54,43 +59,91 @@ export class WallProfiler implements Profiler<WallProfilerStartArgs> {
     return this.innerProfile(true);
   }
 
-  public wrapWithLabels(
+  public wrapWithLabels<R, TArgs extends unknown[]>(
     lbls: LabelSet,
-    fn: () => void,
-    ...args: unknown[]
-  ): void {
+    fn: (...args: TArgs) => R,
+    ...args: TArgs
+  ): R {
     const oldLabels = this.getLabels();
-    this.setLabels({
-      ...oldLabels,
-      ...lbls,
-    });
-    (fn as AnyWrappedCallback)(...args);
-    this.setLabels({
-      ...oldLabels,
-    });
+    const labels = { ...oldLabels, ...lbls };
+    if (this.asyncContext) {
+      return time.runWithContext({ labels }, fn, ...args);
+    }
+    this.setLabels(labels);
+    try {
+      return fn(...args);
+    } finally {
+      this.setLabels({ ...oldLabels });
+    }
   }
 
   public setLabels(labels: LabelSet): void {
-    this.newContext({
-      labels: labels,
-    });
+    const context = { labels };
+    if (this.asyncContext) {
+      if (!time.isStarted()) {
+        throw new Error('Wall profiler is not started');
+      }
+      // The native setter mutates a shared frame in place; copy it first so
+      // already-created async work retains its previous labels.
+      this.contextBoundary.enterWith(context);
+      time.setContext(context);
+    } else {
+      this.newContext(context);
+    }
   }
 
   public start(args: WallProfilerStartArgs): void {
-    if (!time.isStarted()) {
+    if (time.isStarted()) {
+      if (this.asyncContext) {
+        this.checkAsyncContext();
+      }
+    } else {
       log('start');
 
       this.lastProfiledAt = new Date();
       this.lastSamplingIntervalMicros = args.samplingDurationMs;
-      time.start({
-        sourceMapper: args.sourceMapper,
-        durationMillis: args.samplingDurationMs,
-        intervalMicros: args.samplingIntervalMicros,
-        withContexts: true,
-        workaroundV8Bug: true,
-        collectCpuTime: args.collectCpuTime,
-      });
-      this.newContext({});
+      try {
+        time.start({
+          sourceMapper: args.sourceMapper,
+          durationMillis: args.samplingDurationMs,
+          intervalMicros: args.samplingIntervalMicros,
+          withContexts: true,
+          useCPED: this.asyncContext,
+          workaroundV8Bug: true,
+          collectCpuTime: args.collectCpuTime,
+        });
+        if (this.asyncContext) {
+          this.checkAsyncContext();
+        } else {
+          this.newContext({});
+        }
+      } catch (error) {
+        if (time.isStarted()) {
+          time.stop();
+        }
+        if (this.asyncContext) {
+          throw new Error(
+            'wall.asyncContext requires native AsyncContextFrame support. ' +
+              'On Node.js 22, enable --experimental-async-context-frame; ' +
+              'on Node.js 24+, do not disable AsyncContextFrame. ' +
+              'Otherwise leave wall.asyncContext disabled.',
+            { cause: error }
+          );
+        }
+        throw error;
+      }
+    }
+  }
+
+  private checkAsyncContext(): void {
+    // useCPED can start successfully even when V8 cannot read its contexts.
+    const context: ProfilerContext = {};
+    const supported = time.runWithContext(
+      context,
+      () => time.getContext() === context
+    );
+    if (!supported) {
+      throw new Error('The native profiler could not read the scoped context');
     }
   }
 
@@ -110,8 +163,13 @@ export class WallProfiler implements Profiler<WallProfilerStartArgs> {
   }
 
   private innerProfile(restart: boolean): ProfileExport {
-    this.newContext({});
+    if (!this.asyncContext) {
+      this.newContext({});
+    }
     const profile: Profile = time.stop(restart, this.generateLabels);
+    if (!restart) {
+      this.contextBoundary.disable();
+    }
 
     const lastProfileStartedAt: Date = this.lastProfiledAt;
     this.lastProfiledAt = new Date();

src/pyroscope-config.ts

@@ -27,6 +27,7 @@ export interface PyroscopeConfig {
 export type StripFilenamesMode = 'all' | 'dependencies';
 
 export interface PyroscopeWallConfig {
+  asyncContext?: boolean | undefined;
   samplingDurationMs?: number | undefined;
   samplingIntervalMicros?: number | undefined;
   collectCpuTime?: boolean | undefined;

src/utils/check-pyroscope-config.ts

@@ -253,7 +253,9 @@ function hasValidWall(config: Record<string | symbol, unknown>): boolean {
 
   return (
     hasValidSamplingDurationMs(wallConfig) &&
-    hasValidSamplingIntervalMicros(wallConfig)
+    hasValidSamplingIntervalMicros(wallConfig) &&
+    (wallConfig.asyncContext === undefined ||
+      typeof wallConfig.asyncContext === 'boolean')
   );
 }
 

src/utils/get-env.ts

@@ -20,6 +20,9 @@ export function getEnv(): Environment {
     wallSamplingIntervalMicros: parseNumericEnv(
       process.env['PYROSCOPE_WALL_SAMPLING_INTERVAL_MICROS']
     ),
+    wallAsyncContext: parseBooleanEnv(
+      process.env['PYROSCOPE_WALL_ASYNC_CONTEXT']
+    ),
     wallCollectCpuTime: parseBooleanEnv(
       process.env['PYROSCOPE_WALL_COLLECT_CPU_TIME']
     ),

src/utils/process-config.ts

@@ -18,6 +18,7 @@ export function processConfig(
       env.adhocServerAddress ?? config.serverAddress ?? env.serverAddress,
     tags: config.tags,
     wall: {
+      asyncContext: config.wall?.asyncContext ?? env.wallAsyncContext,
       samplingDurationMs:
         config.wall?.samplingDurationMs ?? env.wallSamplingDurationMs,
       samplingIntervalMicros: