diff --git a/README.md b/README.md index d71a5a0..7c76790 100644 --- a/README.md +++ b/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 diff --git a/src/environment.ts b/src/environment.ts index 0576bd7..24ce3c1 100644 --- a/src/environment.ts +++ b/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; } diff --git a/src/index.ts b/src/index.ts index 1bf754e..a1bfb85 100644 --- a/src/index.ts +++ b/src/index.ts @@ -42,12 +42,22 @@ function setLabels(labels: Record): void { getProfiler().wallProfiler.profiler.setLabels(labels); } -export function wrapWithLabels( +export function wrapWithLabels( lbls: Record, - fn: () => void, + fn: (...args: TArgs) => R, + ...args: TArgs +): R; +export function wrapWithLabels( + lbls: Record, + fn: () => R, ...args: unknown[] -): void { - getProfiler().wallProfiler.profiler.wrapWithLabels(lbls, fn, ...args); +): R; +export function wrapWithLabels( + lbls: Record, + fn: (...args: TArgs) => R, + ...args: TArgs +): R { + return getProfiler().wallProfiler.profiler.wrapWithLabels(lbls, fn, ...args); } function startWallProfiling(): void { diff --git a/src/profilers/heap-profiler.ts b/src/profilers/heap-profiler.ts index 03ef1fa..ec361c6 100644 --- a/src/profilers/heap-profiler.ts +++ b/src/profilers/heap-profiler.ts @@ -25,7 +25,7 @@ export class HeapProfiler implements Profiler { throw new Error("heap profiler doesn't support labels"); } - public wrapWithLabels(): void { + public wrapWithLabels(): never { throw new Error("heap profiler doesn't support labels"); } diff --git a/src/profilers/profiler.ts b/src/profilers/profiler.ts index 4c57e2f..58d65d4 100644 --- a/src/profilers/profiler.ts +++ b/src/profilers/profiler.ts @@ -5,11 +5,11 @@ export interface Profiler { setLabels(labels: Record): void; - wrapWithLabels( + wrapWithLabels( labels: Record, - fn: () => void, - ...args: unknown[] - ): void; + fn: (...args: TArgs) => R, + ...args: TArgs + ): R; start(args: TStartArgs): void; diff --git a/src/profilers/pyroscope-profiler.ts b/src/profilers/pyroscope-profiler.ts index 65ec089..4b75813 100644 --- a/src/profilers/pyroscope-profiler.ts +++ b/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: diff --git a/src/profilers/wall-profiler.ts b/src/profilers/wall-profiler.ts index 643592e..a7bae92 100644 --- a/src/profilers/wall-profiler.ts +++ b/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 { private lastProfiledAt: Date; private lastContext: ProfilerContext; private lastSamplingIntervalMicros!: number; + private readonly contextBoundary = new AsyncLocalStorage(); - 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 { return this.innerProfile(true); } - public wrapWithLabels( + public wrapWithLabels( 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 { } 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(); diff --git a/src/pyroscope-config.ts b/src/pyroscope-config.ts index 024f395..c6923d9 100644 --- a/src/pyroscope-config.ts +++ b/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; diff --git a/src/utils/check-pyroscope-config.ts b/src/utils/check-pyroscope-config.ts index db1ec15..526a147 100644 --- a/src/utils/check-pyroscope-config.ts +++ b/src/utils/check-pyroscope-config.ts @@ -253,7 +253,9 @@ function hasValidWall(config: Record): boolean { return ( hasValidSamplingDurationMs(wallConfig) && - hasValidSamplingIntervalMicros(wallConfig) + hasValidSamplingIntervalMicros(wallConfig) && + (wallConfig.asyncContext === undefined || + typeof wallConfig.asyncContext === 'boolean') ); } diff --git a/src/utils/get-env.ts b/src/utils/get-env.ts index 11eacb2..f138bcf 100644 --- a/src/utils/get-env.ts +++ b/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'] ), diff --git a/src/utils/process-config.ts b/src/utils/process-config.ts index c3a00a8..2636734 100644 --- a/src/utils/process-config.ts +++ b/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: diff --git a/test/fixtures/async-context.test.ts b/test/fixtures/async-context.test.ts new file mode 100644 index 0000000..e1005c6 --- /dev/null +++ b/test/fixtures/async-context.test.ts @@ -0,0 +1,380 @@ +import { strict as assert } from 'node:assert'; +import { AsyncLocalStorage, AsyncResource } from 'node:async_hooks'; +import { createServer } from 'node:http'; +import { it } from 'node:test'; +import { setImmediate, setTimeout } from 'node:timers/promises'; +import { gunzipSync } from 'node:zlib'; +import { time } from '@datadog/pprof'; +import { Profile } from 'pprof-format'; +import Pyroscope, { wrapWithLabels } from '../../src/index.js'; +import { getProfiler } from '../../src/utils/pyroscope-profiler.js'; + +const profiles: Profile[] = []; +let uploaded!: () => void; +const firstUpload = new Promise((resolve) => { + uploaded = resolve; +}); +let setterUploaded!: () => void; +const setterUpload = new Promise((resolve) => { + setterUploaded = resolve; +}); +const server = createServer(async (req, res) => { + assert.ok(req.url?.includes('format=pprof')); + const chunks: Buffer[] = []; + for await (const chunk of req) chunks.push(chunk); + const profile = Profile.decode(gunzipSync(Buffer.concat(chunks))); + profiles.push(profile); + res.end('ok'); + if ( + profile.stringTable.strings.some((name) => + name.includes(':itemBeforeAwait:') + ) + ) { + uploaded(); + } + if ( + profile.stringTable.strings.some((name) => name.includes(':setterSibling:')) + ) { + setterUploaded(); + } +}); +it('native async-context lifecycle and sample attribution', async (t) => { + await new Promise((resolve, reject) => { + server.once('error', reject); + server.listen(0, '127.0.0.1', () => { + server.off('error', reject); + resolve(); + }); + }); + t.after(async () => { + try { + await Pyroscope.stop(); + } finally { + await new Promise((resolve, reject) => + server.close((error) => (error ? reject(error) : resolve())) + ); + } + }); + const address = server.address(); + assert.ok(address && typeof address !== 'string'); + const config = { + appName: 'async-label-test', + serverAddress: `http://127.0.0.1:${address.port}`, + flushIntervalMs: 40, + wall: { + asyncContext: true, + samplingIntervalMicros: 1000, + collectCpuTime: true, + }, + }; + + Pyroscope.init(config); + if (process.argv[2] === 'unsupported') { + assert.throws( + () => Pyroscope.start(), + /wall.asyncContext requires native AsyncContextFrame/ + ); + assert.equal(time.isStarted(), false); + Pyroscope.init({ + ...config, + wall: { ...config.wall, asyncContext: false }, + }); + Pyroscope.startWallProfiling(); + assert.equal( + wrapWithLabels({ page: 'item' }, () => 42), + 42 + ); + } else { + assert.throws(() => Pyroscope.setLabels({ phase: 'early' }), /not started/); + Pyroscope.startWallProfiling(); + await t.test( + 'preserves callback results, nested scopes and tracing', + verifyContract + ); + await t.test( + 'isolates persistent setters from existing async work', + verifySetLabels + ); + Pyroscope.setLabels({}); + await t.test( + 'exports historical and concurrent labels across flushes', + verifySamples + ); + await t.test( + 'clears labels on stop and supports setters after restart', + async () => { + assert.deepEqual(Pyroscope.getLabels(), {}); + assert.throws( + () => Pyroscope.setLabels({ phase: 'late' }), + /not started/ + ); + Pyroscope.startWallProfiling(); + assert.deepEqual(Pyroscope.getLabels(), {}); + Pyroscope.setLabels({ phase: 'restarted' }); + await setImmediate(); + assert.deepEqual(Pyroscope.getLabels(), { phase: 'restarted' }); + await wrapWithLabels({ page: 'new' }, async () => { + await setImmediate(); + assert.deepEqual(Pyroscope.getLabels(), { + phase: 'restarted', + page: 'new', + }); + }); + assert.deepEqual(Pyroscope.getLabels(), { phase: 'restarted' }); + } + ); + } +}); + +async function verifyContract() { + assert.deepEqual(Pyroscope.getLabels(), {}); + const result: number = wrapWithLabels( + { page: 'item' }, + (a: number, b: number) => a + b, + 2, + 3 + ); + assert.equal(result, 5); + const original = Promise.resolve(8); + const returned: Promise = wrapWithLabels({}, () => original, 'extra'); + assert.equal(returned, original); + assert.equal(await returned, 8); + + const error = new Error('callback failure'); + assert.throws( + () => + wrapWithLabels({ page: 'item' }, () => { + throw error; + }), + (caught) => caught === error + ); + assert.deepEqual(Pyroscope.getLabels(), {}); + const tracing = new AsyncLocalStorage(); + await tracing.run('trace-id', () => + wrapWithLabels({ page: 'item', shard: 3 }, async () => { + await setTimeout(1); + assert.equal(tracing.getStore(), 'trace-id'); + await assert.rejects( + wrapWithLabels({ page: 'search' }, async () => { + await Promise.resolve(); + assert.deepEqual(Pyroscope.getLabels(), { page: 'search', shard: 3 }); + throw error; + }), + (caught) => caught === error + ); + assert.deepEqual(Pyroscope.getWallLabels(), { page: 'item', shard: 3 }); + assert.throws( + () => + wrapWithLabels({ page: 'search' }, () => { + throw error; + }), + (caught) => caught === error + ); + assert.deepEqual(Pyroscope.getLabels(), { page: 'item', shard: 3 }); + + getProfiler().wallProfiler.profiler.profile(); + assert.deepEqual(Pyroscope.getLabels(), { page: 'item', shard: 3 }); + await setImmediate(); + assert.deepEqual(Pyroscope.getLabels(), { page: 'item', shard: 3 }); + }) + ); + assert.equal(tracing.getStore(), undefined); + assert.deepEqual(Pyroscope.getLabels(), {}); + + const resource = wrapWithLabels( + { page: 'search' }, + () => new AsyncResource('scoped-work') + ); + try { + await wrapWithLabels({ page: 'item' }, async () => { + await setImmediate(); + resource.runInAsyncScope(() => + assert.deepEqual(Pyroscope.getLabels(), { page: 'search' }) + ); + assert.deepEqual(Pyroscope.getLabels(), { page: 'item' }); + }); + } finally { + resource.emitDestroy(); + } +} + +async function verifySetLabels() { + Pyroscope.setLabels({ phase: 'root' }); + assert.deepEqual(Pyroscope.getLabels(), { phase: 'root' }); + await setImmediate(); + assert.deepEqual(Pyroscope.getLabels(), { phase: 'root' }); + const tracing = new AsyncLocalStorage(); + await tracing.run('setter-trace', () => + wrapWithLabels({ page: 'item' }, async () => { + const parent = { phase: 'root', page: 'item' }; + beforeSetter(); + let release!: () => void; + const gate = new Promise((resolve) => { + release = resolve; + }); + const writer = (async () => { + await gate; + Pyroscope.setLabels({ phase: 'first' }); + firstSetter(); + Pyroscope.setWallLabels({ phase: 'second' }); + secondSetter(); + await setImmediate(); + assert.deepEqual(Pyroscope.getLabels(), { phase: 'second' }); + assert.equal(tracing.getStore(), 'setter-trace'); + await assert.rejects( + wrapWithLabels({ page: 'nested' }, async () => { + assert.deepEqual(Pyroscope.getLabels(), { + phase: 'second', + page: 'nested', + }); + Pyroscope.setLabels({ phase: 'nested-replaced' }); + await setImmediate(); + assert.deepEqual(Pyroscope.getLabels(), { + phase: 'nested-replaced', + }); + throw new Error('nested setter'); + }), + /nested setter/ + ); + assert.deepEqual(Pyroscope.getLabels(), { phase: 'second' }); + })(); + const sibling = (async () => { + await gate; + await setImmediate(); + assert.deepEqual(Pyroscope.getLabels(), parent); + setterSibling(); + })(); + release(); + await Promise.all([writer, sibling]); + assert.deepEqual(Pyroscope.getLabels(), parent); + }) + ); + assert.deepEqual(Pyroscope.getLabels(), { phase: 'root' }); + await setterUpload; + assert.deepEqual(Pyroscope.getLabels(), { phase: 'root' }); +} + +function beforeSetter() { + burn(); +} +function firstSetter() { + burn(); +} +function secondSetter() { + burn(); +} +function setterSibling() { + burn(); +} + +function burn() { + const end = performance.now() + 120; + while (performance.now() < end) Math.sqrt(Math.random()); +} +function itemBeforeAwait() { + burn(); +} +function itemAfterAwait() { + burn(); +} +function itemAfterFlush() { + burn(); +} +function searchAfterAwait() { + burn(); +} +function nestedSearch() { + burn(); +} +function itemAfterNested() { + burn(); +} +function unrelatedWork() { + burn(); +} + +async function verifySamples() { + await Promise.all([ + wrapWithLabels({ page: 'item', shard: 3 }, async () => { + itemBeforeAwait(); + await Promise.resolve(); + assert.deepEqual(Pyroscope.getLabels(), { page: 'item', shard: 3 }); + itemAfterAwait(); + await firstUpload; + assert.deepEqual(Pyroscope.getLabels(), { page: 'item', shard: 3 }); + itemAfterFlush(); + await wrapWithLabels({ page: 'search' }, async () => { + await setImmediate(); + nestedSearch(); + }); + itemAfterNested(); + }), + wrapWithLabels({ page: 'search' }, async () => { + await setTimeout(1); + assert.deepEqual(Pyroscope.getLabels(), { page: 'search' }); + searchAfterAwait(); + }), + (async () => { + await setImmediate(); + assert.deepEqual(Pyroscope.getLabels(), {}); + unrelatedWork(); + })(), + ]); + await Pyroscope.stopWallProfiling(); + assert.ok( + profiles.length >= 2, + 'must include a periodic export and final flush' + ); + + const expected: Record> = { + itemBeforeAwait: { page: 'item', shard: 3 }, + itemAfterAwait: { page: 'item', shard: 3 }, + itemAfterFlush: { page: 'item', shard: 3 }, + searchAfterAwait: { page: 'search' }, + nestedSearch: { page: 'search', shard: 3 }, + itemAfterNested: { page: 'item', shard: 3 }, + unrelatedWork: {}, + beforeSetter: { phase: 'root', page: 'item' }, + firstSetter: { phase: 'first' }, + secondSetter: { phase: 'second' }, + setterSibling: { phase: 'root', page: 'item' }, + }; + const seen = new Set(); + for (const profile of profiles) { + const string = (id: number | bigint) => + profile.stringTable.strings[Number(id)]; + assert.ok(profile.sampleType.some((type) => string(type.type) === 'cpu')); + const functions = new Map( + profile.function.map((fn) => [Number(fn.id), string(fn.name)]) + ); + const locations = new Map( + profile.location.map((loc) => [ + Number(loc.id), + loc.line.map((line) => functions.get(Number(line.functionId)) ?? ''), + ]) + ); + for (const sample of profile.sample) { + const names = sample.locationId.flatMap( + (id) => locations.get(Number(id)) ?? [] + ); + const work = Object.keys(expected).find((name) => + names.some((frame) => frame.includes(`:${name}:`)) + ); + if (!work) continue; + const labels = Object.fromEntries( + sample.label.map((label) => [ + string(label.key), + Number(label.str) === 0 ? Number(label.num) : string(label.str), + ]) + ); + assert.deepEqual(labels, expected[work], `${work} sample attribution`); + assert.ok(Number(sample.value[0]) > 0); + seen.add(work); + } + } + assert.deepEqual( + [...seen].sort(), + Object.keys(expected).sort(), + 'every named workload must have real samples' + ); +} diff --git a/test/labels.test.ts b/test/labels.test.ts new file mode 100644 index 0000000..16880b3 --- /dev/null +++ b/test/labels.test.ts @@ -0,0 +1,134 @@ +import { describe, it } from 'node:test'; +import { strict as assert } from 'node:assert'; +import { time } from '@datadog/pprof'; +import Pyroscope, { wrapWithLabels } from '../src/index.js'; +import { getProfiler } from '../src/utils/pyroscope-profiler.js'; +import { checkPyroscopeConfig } from '../src/utils/check-pyroscope-config.js'; +import { processConfig } from '../src/utils/process-config.js'; +import { getEnv } from '../src/utils/get-env.js'; + +function start() { + Pyroscope.init(); + const { profiler, startArgs } = getProfiler().wallProfiler; + profiler.start(startArgs); +} + +describe('synchronous labels (default mode)', () => { + it('returns the callback result and forwards typed arguments', (t) => { + start(); + t.after(() => time.stop()); + const result: number = wrapWithLabels( + { page: 'item' }, + (a: number, b: number) => { + assert.deepEqual(Pyroscope.getLabels(), { page: 'item' }); + return a + b; + }, + 2, + 3 + ); + assert.equal(result, 5); + assert.deepEqual(Pyroscope.getLabels(), {}); + }); + + it('forwards extra arguments to callbacks without declared parameters', (t) => { + start(); + t.after(() => time.stop()); + const callback = t.mock.fn(() => 42); + const result: number = wrapWithLabels({}, callback, 'extra'); + assert.equal(result, 42); + assert.deepEqual(callback.mock.calls[0].arguments, ['extra']); + }); + + it('restores parent labels on nested synchronous exceptions', (t) => { + start(); + t.after(() => time.stop()); + Pyroscope.setLabels({ region: 'eu' }); + const error = new Error('callback failure'); + Pyroscope.wrapWithLabels({ page: 'item' }, () => { + assert.throws( + () => + Pyroscope.wrapWithLabels({ page: 'search' }, () => { + throw error; + }), + (caught) => caught === error + ); + assert.deepEqual(Pyroscope.getLabels(), { region: 'eu', page: 'item' }); + }); + assert.deepEqual(Pyroscope.getLabels(), { region: 'eu' }); + }); + + it('keeps setLabels process-wide for existing async work', async (t) => { + start(); + t.after(() => time.stop()); + Pyroscope.setLabels({ phase: 'before' }); + let release!: () => void; + const gate = new Promise((resolve) => { + release = resolve; + }); + const reader = (async () => { + await gate; + assert.deepEqual(Pyroscope.getLabels(), { phase: 'after' }); + })(); + Pyroscope.setLabels({ phase: 'after' }); + release(); + await reader; + }); + + it('preserves promise identity without extending the synchronous scope', async (t) => { + start(); + t.after(() => time.stop()); + let original!: Promise; + const returned: Promise = Pyroscope.wrapWithLabels( + { page: 'item' }, + () => { + original = (async () => { + assert.deepEqual(Pyroscope.getLabels(), { page: 'item' }); + await Promise.resolve(); + assert.deepEqual(Pyroscope.getLabels(), {}); + return 7; + })(); + return original; + } + ); + assert.equal(returned, original); + assert.equal(await returned, 7); + }); +}); + +describe('async-context configuration', () => { + for (const value of ['true', 1, null, {}, []]) { + it(`rejects non-boolean asyncContext ${JSON.stringify(value)}`, () => { + assert.throws( + () => checkPyroscopeConfig({ wall: { asyncContext: value } }), + /valid wall options/ + ); + }); + } + + it('allows explicit false to override environment opt-in', () => { + const previous = process.env.PYROSCOPE_WALL_ASYNC_CONTEXT; + try { + for (const value of ['true', '1']) { + process.env.PYROSCOPE_WALL_ASYNC_CONTEXT = value; + assert.equal(processConfig({}, getEnv()).wall?.asyncContext, true); + assert.equal( + processConfig({ wall: { asyncContext: false } }, getEnv()).wall + ?.asyncContext, + false + ); + } + process.env.PYROSCOPE_WALL_ASYNC_CONTEXT = 'false'; + assert.equal( + processConfig({ wall: { asyncContext: true } }, getEnv()).wall + ?.asyncContext, + true + ); + delete process.env.PYROSCOPE_WALL_ASYNC_CONTEXT; + assert.equal(processConfig({}, getEnv()).wall?.asyncContext, undefined); + } finally { + if (previous === undefined) + delete process.env.PYROSCOPE_WALL_ASYNC_CONTEXT; + else process.env.PYROSCOPE_WALL_ASYNC_CONTEXT = previous; + } + }); +}); diff --git a/test/profiler.test.ts b/test/profiler.test.ts index 77c3a56..effa42a 100644 --- a/test/profiler.test.ts +++ b/test/profiler.test.ts @@ -1,6 +1,8 @@ import { describe, it, type TestContext } from 'node:test'; import { strict as assert } from 'node:assert'; import process from 'node:process'; +import { execFileSync } from 'node:child_process'; +import { fileURLToPath } from 'node:url'; import Pyroscope from '../src/index.js'; import { VERSION } from '../src/version.js'; @@ -430,3 +432,41 @@ describe('common behaviour of profilers', () => { assert.strictEqual(req.headers['x-scope-orgid'], 'my-tenant-id'); }); }); + +describe('async-context labels', () => { + const fixture = fileURLToPath( + new URL('./fixtures/async-context.test.js', import.meta.url) + ); + const major = Number(process.versions.node.split('.')[0]); + + function run(flags: string[], mode: string) { + execFileSync(process.execPath, [...flags, fixture, mode], { + timeout: 30_000, + env: { + ...process.env, + NODE_OPTIONS: '', + NODE_TEST_CONTEXT: undefined, + PYROSCOPE_WALL_ASYNC_CONTEXT: '', + }, + stdio: 'pipe', + }); + } + + if (major >= 24) { + it('propagates async labels and exports correctly attributed native samples', () => + run([], 'supported')); + it('rejects disabled AsyncContextFrame and permits a subsequent default-mode start', () => + run(['--no-async-context-frame'], 'unsupported')); + } else { + it('rejects unavailable AsyncContextFrame and permits a subsequent default-mode start', () => + run([], 'unsupported')); + if ( + process.allowedNodeEnvironmentFlags.has( + '--experimental-async-context-frame' + ) + ) { + it('propagates async labels and exports correctly attributed native samples with the runtime flag', () => + run(['--experimental-async-context-frame'], 'supported')); + } + } +});