From 710dd2e46b3a6f7aecaf60fd28502dca94182f90 Mon Sep 17 00:00:00 2001 From: isaacs Date: Tue, 23 Jun 2026 22:40:59 -0700 Subject: [PATCH 1/2] feat: allow message.result to be changed Allow an `end` handler to modify the return value of a synchronous function. This is needed for the NestJS instrumentation, because its request handler mechanism is a factory that returns an inner arrow function, capturing values that parity with the OTel-based instrumentation requires us to report on. (These values are not present in the inner function's arguments; they are captured by a closure only.) This *could* be done by wrapping both the inner and outer methods, and correlating them with a WeakMap, but this would be more brittle and complicated, for no benefit. --- README.md | 48 ++++++++++++++++++++++++++-- index.d.ts | 2 +- lib/transforms.js | 45 ++++++++++++++++++++++++-- tests/mutable_result_cjs/mod.js | 18 +++++++++++ tests/mutable_result_cjs/test.js | 54 ++++++++++++++++++++++++++++++++ tests/tests.test.mjs | 37 ++++++++++++++++++++++ 6 files changed, 198 insertions(+), 6 deletions(-) create mode 100644 tests/mutable_result_cjs/mod.js create mode 100644 tests/mutable_result_cjs/test.js diff --git a/README.md b/README.md index db94428..53f5403 100644 --- a/README.md +++ b/README.md @@ -87,6 +87,47 @@ const instrumentation = { This also works for class exports (e.g., `export { MyClass as PublicClass }`). +### Mutable Result + +By default a subscriber can observe a function's return value +(via `message.result`) but cannot change it. Setting +`mutableResult: true` on a **synchronous** `FunctionQuery` makes +the wrapper return whatever `message.result` holds after the +`end` event has published, so an `end` handler can reassign +`message.result` to mutate the return value. This is useful when +a function returns another function (or object) that you need to +wrap, like a factory that returns a per-request handler: + +```js +const instrumentation = { + channelName: "create-handler", + module: { name: "my-framework", versionRange: ">=1.0.0", filePath: "lib/router.js" }, + functionQuery: { methodName: "create", kind: "Sync", mutableResult: true }, +}; +``` + +```js +const { tracingChannel } = require("node:diagnostics_channel"); + +tracingChannel("orchestrion:my-framework:create-handler").subscribe({ + end(message) { + const original = message.result; + // Replace the returned handler with a wrapped version. + message.result = function wrapped(...args) { + // ...start a span, etc. + return original.apply(this, args); + }; + }, +}); +``` + +`mutableResult` is only valid with `kind: "Sync"`. Combining it +with any other kind throws at transform time. On the throw path +the original error still propagates; the substituted return +value only applies when the function returns normally. If no +subscriber reassigns `message.result`, the original return value +is preserved unchanged. + ### API Reference ```ts @@ -108,13 +149,14 @@ type FunctionQuery = index?: number | null; callbackIndex?: number; isExportAlias?: boolean; + mutableResult?: boolean; } | // Match method on objects - { methodName: string; kind: FunctionKind; index?: number | null; callbackIndex?: number } + { methodName: string; kind: FunctionKind; index?: number | null; callbackIndex?: number; mutableResult?: boolean } | // Match standalone function - { functionName: string; kind: FunctionKind; index?: number | null; callbackIndex?: number; isExportAlias?: boolean } + { functionName: string; kind: FunctionKind; index?: number | null; callbackIndex?: number; isExportAlias?: boolean; mutableResult?: boolean } | // Match arrow function or function expression - { expressionName: string; kind: FunctionKind; index?: number | null; callbackIndex?: number; isExportAlias?: boolean }; + { expressionName: string; kind: FunctionKind; index?: number | null; callbackIndex?: number; isExportAlias?: boolean; mutableResult?: boolean }; | // Match private class methods { className: string; privateMethodName: string; kind: FunctionKind; index?: number | null; callbackIndex?: number }; ``` diff --git a/index.d.ts b/index.d.ts index 380e404..9f32f2c 100644 --- a/index.d.ts +++ b/index.d.ts @@ -27,7 +27,7 @@ export type FunctionKind = "Sync" | "Async" | "Callback" | "Auto"; /** * Describes which function to instrument */ -export type FunctionQuery = { className: string; methodName: string; kind: FunctionKind; index?: number | null; isExportAlias?: boolean } | { className: string; privateMethodName: string; kind: FunctionKind; index?: number | null } | { className: string; index?: number | null; isExportAlias?: boolean } | { methodName: string; kind: FunctionKind; index?: number | null } | { functionName: string; kind: FunctionKind; index?: number | null; isExportAlias?: boolean } | { expressionName: string; kind: FunctionKind; index?: number | null; isExportAlias?: boolean }; +export type FunctionQuery = { className: string; methodName: string; kind: FunctionKind; index?: number | null; isExportAlias?: boolean; mutableResult?: boolean } | { className: string; privateMethodName: string; kind: FunctionKind; index?: number | null } | { className: string; index?: number | null; isExportAlias?: boolean } | { methodName: string; kind: FunctionKind; index?: number | null; mutableResult?: boolean } | { functionName: string; kind: FunctionKind; index?: number | null; isExportAlias?: boolean; mutableResult?: boolean } | { expressionName: string; kind: FunctionKind; index?: number | null; isExportAlias?: boolean; mutableResult?: boolean }; /** * A custom transform function registered via `addTransform`. diff --git a/lib/transforms.js b/lib/transforms.js index 9f08134..be94f6b 100644 --- a/lib/transforms.js +++ b/lib/transforms.js @@ -210,7 +210,7 @@ function traceInstanceMethod (state, node, program) { const __apm$${methodName} = this["${methodName}"] this["${methodName}"] = function () {} if (typeof __apm$${methodName} === 'function') { - Object.defineProperty(this["${methodName}"], 'length', { + Object.defineProperty(this["${methodName}"], 'length', { value: __apm$${methodName}.length, configurable: true }) @@ -242,7 +242,18 @@ function traceInstanceMethod (state, node, program) { */ function wrap (state, node, program) { const { operator, moduleVersion } = state - const { returnKind } = state.functionQuery + const { returnKind, mutableResult } = state.functionQuery + + // `mutableResult` lets a subscriber replace the function's synchronous + // return value by reassigning `message.result` in its `end` handler. + // This is only meaningful for the synchronous wrapper. The + // callback/promise wrappers return a callback result or a (possibly + // chained) promise, neither of which is the value a subscriber would + // substitute here. Fail loudly rather than silently ignoring the option + // on the wrong `kind`. + if (mutableResult && operator !== 'traceSync') { + throw new Error(`functionQuery.mutableResult is only supported with kind: 'Sync' (got operator '${operator}')`) + } const iterPatch = returnKind ? generateIterPatch(state, returnKind, program) : '' @@ -544,8 +555,38 @@ function wrapPromise (state, node, iterPatch = '') { */ function wrapSync (state, node, iterPatch = '') { const { channelName } = state + const { mutableResult } = state.functionQuery const channelVariable = formatChannelVariable(channelName) + // When `mutableResult` is set, the return value is whatever + // `__apm$ctx.result` holds *after* `end` has published. + // So a subscriber's `end` handler can eassign `message.result` + // to substitute it (e.g. to wrap a returned handler). The `return` + // therefore moves below the `try`/`finally`: on the throw path the + // `catch` re-throws, so the trailing `return` is unreachable. + if (mutableResult) { + return parse(` + function wrapper () { + if (!tr_ch_apm_hasSubscribers(${channelVariable})) return __apm$traced(); + + return ${channelVariable}.start.runStores(__apm$ctx, () => { + try { + __apm$ctx.result = __apm$traced(); + ${iterPatch} + } catch (err) { + __apm$ctx.error = err; + ${channelVariable}.error.publish(__apm$ctx); + throw err; + } finally { + __apm$ctx.self ??= this; + ${channelVariable}.end.publish(__apm$ctx); + } + return __apm$ctx.result; + }); + } + `) + } + return parse(` function wrapper () { if (!tr_ch_apm_hasSubscribers(${channelVariable})) return __apm$traced(); diff --git a/tests/mutable_result_cjs/mod.js b/tests/mutable_result_cjs/mod.js new file mode 100644 index 0000000..2f7458a --- /dev/null +++ b/tests/mutable_result_cjs/mod.js @@ -0,0 +1,18 @@ +function create (instance, callback) { + // Returns a per-request handler that closes over `instance`/`callback`, + // mirroring NestJS `RouterExecutionContext.create`. A subscriber needs to + // replace this returned function with a span-wrapping version. + return function handler (req) { + return callback.call(instance, req) + } +} + +function compute (x) { + return x + 1 +} + +function boom () { + throw new Error('boom') +} + +module.exports = { create, compute, boom } diff --git a/tests/mutable_result_cjs/test.js b/tests/mutable_result_cjs/test.js new file mode 100644 index 0000000..f851a91 --- /dev/null +++ b/tests/mutable_result_cjs/test.js @@ -0,0 +1,54 @@ +const { create, compute, boom } = require('./instrumented.js') +const assert = require('node:assert') +const { tracingChannel } = require('node:diagnostics_channel') + +// Scenario 1: a subscriber substitutes the returned handler by reassigning +// `message.result` in `end`. The wrapper must return the substituted function. +const events = [] +tracingChannel('orchestrion:undici:create_mutable').subscribe({ + start () { + events.push('start') + }, + end (message) { + events.push('end') + const original = message.result + assert.strictEqual(typeof original, 'function') + message.result = function wrapped (...args) { + events.push('wrapped') + return `[${original.apply(this, args)}]` + } + } +}) + +const instance = { name: 'Ctrl' } +const callback = function (req) { + return `${this.name}:${req}` +} +const handler = create(instance, callback) +assert.strictEqual(typeof handler, 'function') +// The handler returned to the caller is the substituted wrapper. +assert.strictEqual(handler('GET'), '[Ctrl:GET]') +assert.deepStrictEqual(events, ['start', 'end', 'wrapped']) + +// Scenario 2: a subscriber observes `message.result` but does NOT reassign it. +// The original return value must be preserved unchanged. +let observed +tracingChannel('orchestrion:undici:compute_mutable').subscribe({ + end (message) { + observed = message.result + } +}) +assert.strictEqual(compute(41), 42) +assert.strictEqual(observed, 42) + +// Scenario 3 (throw path): the restructured try/finally must still propagate +// the error and publish it; the trailing `return message.result` is unreachable. +let errored +tracingChannel('orchestrion:undici:boom_mutable').subscribe({ + error (message) { + errored = message.error + } +}) +assert.throws(() => boom(), /boom/) +assert.ok(errored instanceof Error) +assert.strictEqual(errored.message, 'boom') diff --git a/tests/tests.test.mjs b/tests/tests.test.mjs index e2eccdc..2b16053 100644 --- a/tests/tests.test.mjs +++ b/tests/tests.test.mjs @@ -394,6 +394,43 @@ describe('ast_query_cjs', () => { }) }) +describe('mutable_result_cjs', () => { + test('lets a subscriber substitute the synchronous return value via message.result', () => { + runTest('mutable_result_cjs', [ + { + channelName: 'create_mutable', + module: { name: TEST_MODULE_NAME, versionRange: '>=0.0.1', filePath: TEST_MODULE_PATH }, + functionQuery: { functionName: 'create', kind: 'Sync', mutableResult: true }, + }, + { + channelName: 'compute_mutable', + module: { name: TEST_MODULE_NAME, versionRange: '>=0.0.1', filePath: TEST_MODULE_PATH }, + functionQuery: { functionName: 'compute', kind: 'Sync', mutableResult: true }, + }, + { + channelName: 'boom_mutable', + module: { name: TEST_MODULE_NAME, versionRange: '>=0.0.1', filePath: TEST_MODULE_PATH }, + functionQuery: { functionName: 'boom', kind: 'Sync', mutableResult: true }, + }, + ]) + }) + + test('throws when mutableResult is combined with a non-Sync kind', () => { + const instrumentor = create([ + { + channelName: 'fetch_mutable_async', + module: { name: TEST_MODULE_NAME, versionRange: '>=0.0.1', filePath: TEST_MODULE_PATH }, + functionQuery: { functionName: 'fetch', kind: 'Async', mutableResult: true }, + }, + ]) + const transformer = instrumentor.getTransformer(TEST_MODULE_NAME, TEST_MODULE_VERSION, TEST_MODULE_PATH) + assert.throws( + () => transformer.transform('async function fetch () { return 42 }', 'cjs'), + /mutableResult is only supported with kind: 'Sync'/ + ) + }) +}) + describe('polyfill_cjs', () => { test('instruments with a custom dc module (cjs)', () => { runTest('polyfill_cjs', [ From c7417fc98bd33a05c558db0003c920c3d287a3c5 Mon Sep 17 00:00:00 2001 From: isaacs Date: Wed, 24 Jun 2026 08:03:31 -0700 Subject: [PATCH 2/2] feat: make astQuery a first-class feature This is the cleanest way to do certain things that cannot be reached by name alone. However, it's brittle to rely on a feature, even a stable one, which is not documented or tested. So, treat it as a first-class feature. --- README.md | 72 ++++++++++++++++++++-- index.d.ts | 68 ++++++++++++++------ lib/transformer.js | 8 ++- lib/transforms.js | 2 +- tests/ast_query_returned_arrow_cjs/mod.js | 12 ++++ tests/ast_query_returned_arrow_cjs/test.js | 26 ++++++++ tests/tests.test.mjs | 29 +++++++++ 7 files changed, 192 insertions(+), 25 deletions(-) create mode 100644 tests/ast_query_returned_arrow_cjs/mod.js create mode 100644 tests/ast_query_returned_arrow_cjs/test.js diff --git a/README.md b/README.md index 53f5403..3fe4e5f 100644 --- a/README.md +++ b/README.md @@ -128,6 +128,52 @@ value only applies when the function returns normally. If no subscriber reassigns `message.result`, the original return value is preserved unchanged. +### AST Query + +The name-based `FunctionQuery` variants cover the common cases +(named functions, class/object methods, expressions). When you +need to target a node they can't express, such as an **anonymous +function returned by a factory** (a decorator factory, a +per-request handler), you can set `astQuery` to a raw +[esquery](https://github.com/estools/esquery) selector instead. + +When present, `astQuery` chooses the nodes to instrument and +takes precedence over `functionQuery`'s matching fields; +`functionQuery` then only supplies behaviour (`kind`, `index`, +`callbackIndex`, `mutableResult`) and may be omitted (it defaults +to `kind: "Sync"`). + +For example, to instrument the decorator returned by a factory: + +```js +function Injectable(options) { + return (target) => { /* applied to the decorated class */ }; +} +``` + +```js +const instrumentation = { + channelName: "injectable-apply", + module: { name: "@nestjs/common", versionRange: ">=8.0.0", filePath: "decorators/core/injectable.decorator.js" }, + // Match the arrow returned from `Injectable`. There is no name to target! + astQuery: 'FunctionDeclaration[id.name="Injectable"] ReturnStatement > ArrowFunctionExpression', + functionQuery: { kind: "Sync" }, +}; +``` + +The channel then fires each time the decorator is applied, with +the decorated target available as `message.arguments[0]`, which a +subscriber can mutate, for example to wrap prototype methods. + +An `astQuery` is used verbatim, so it can match any node, +including ones the name-based variants don't expose, such as +anonymous or deeply nested functions. (Both name-based and +`astQuery` matching work on synchronous and async functions +alike.) + +If an `astQuery` matches no nodes, the "failed to find injection +points" error includes the selector so it can be debugged. + ### API Reference ```ts @@ -174,11 +220,29 @@ type ModuleMatcher = { #### **`InstrumentationConfig`** ```ts -type InstrumentationConfig = { - channelName: string; // Name of the diagnostics channel - module: ModuleMatcher; - functionQuery: FunctionQuery; +// Behaviour-only fields, used when `astQuery` does the matching. +type FunctionBehavior = { + kind?: FunctionKind; + index?: number | null; + callbackIndex?: number; + mutableResult?: boolean; }; + +type InstrumentationConfig = + | { + channelName: string; // Name of the diagnostics channel + module: ModuleMatcher; + functionQuery: FunctionQuery; // Name-based matching + astQuery?: string; // Raw esquery selector; takes precedence over functionQuery matching + transform?: string; // Name of a custom transform registered via addTransform + } + | { + channelName: string; + module: ModuleMatcher; + astQuery: string; // Raw esquery selector chooses the node(s) + functionQuery?: FunctionBehavior; // Behaviour only; matching fields ignored + transform?: string; + }; ``` ### Functions diff --git a/index.d.ts b/index.d.ts index 9f32f2c..ab59a2a 100644 --- a/index.d.ts +++ b/index.d.ts @@ -36,28 +36,58 @@ export type FunctionQuery = { className: string; methodName: string; kind: Funct export type CustomTransform = (state: unknown, node: Node, parent: Node, ancestry: Node[]) => void; /** - * Configuration for injecting instrumentation code + * The behaviour-only fields of a `FunctionQuery`. Used together with `astQuery`, + * where the raw selector chooses the node and these fields drive how it is + * wrapped (the name-based matching fields are ignored). */ -export interface InstrumentationConfig { - /** - * The name of the diagnostics channel to publish to - */ - channelName: string; - /** - * The module matcher to identify the module and file to instrument - */ - module: ModuleMatcher; - /** - * The function query to identify the function to instrument - */ - functionQuery: FunctionQuery; - /** - * The name of a custom transform registered via `addTransform`. - * When set, takes precedence over `functionQuery.kind`. - */ - transform?: string; +export interface FunctionBehavior { + kind?: FunctionKind; + index?: number | null; + callbackIndex?: number; + mutableResult?: boolean; } +/** + * Configuration for injecting instrumentation code. + * + * Either `functionQuery` (name-based matching) or `astQuery` (a raw esquery + * selector) must identify the node(s) to instrument. When `astQuery` is set it + * takes precedence over `functionQuery`'s matching fields, and `functionQuery` + * becomes an optional bag of behaviour options ({@link FunctionBehavior}). + */ +export type InstrumentationConfig = + | { + /** The name of the diagnostics channel to publish to */ + channelName: string; + /** The module matcher to identify the module and file to instrument */ + module: ModuleMatcher; + /** The function query to identify the function to instrument */ + functionQuery: FunctionQuery; + /** + * A raw esquery selector that chooses the node(s) to instrument. When + * set, it takes precedence over `functionQuery`'s matching fields. + */ + astQuery?: string; + /** + * The name of a custom transform registered via `addTransform`. + * When set, takes precedence over `functionQuery.kind`. + */ + transform?: string; + } + | { + channelName: string; + module: ModuleMatcher; + /** + * A raw esquery selector that chooses the node(s) to instrument. This is + * the escape hatch for shapes the name-based `functionQuery` can't + * express, e.g. an anonymous arrow returned by a factory function. + */ + astQuery: string; + /** Behaviour options for the matched node(s); matching fields are ignored. */ + functionQuery?: FunctionBehavior; + transform?: string; + }; + /** * Describes the module and file path you would like to match */ diff --git a/lib/transformer.js b/lib/transformer.js index 76d6c4d..3f7a5c7 100644 --- a/lib/transformer.js +++ b/lib/transformer.js @@ -101,6 +101,11 @@ class Transformer { } const resolvedFunctionQuery = this.#resolveExportAlias(functionQuery, aliases) + // A raw `astQuery` selector takes precedence over the name-based + // `functionQuery`: it chooses which node(s) to instrument, while + // `functionQuery` still supplies behaviour (kind/index/callbackIndex/ + // mutableResult). This is the escape hatch for shapes the name-based + // queries can't express (e.g. an anonymous arrow returned by a factory). const query = astQuery || this.#fromFunctionQuery(resolvedFunctionQuery) const state = { ...config, @@ -119,7 +124,8 @@ class Transformer { } if (injectionCount === 0 && this.#configs.length > 0) { - const names = this.#configs.map(({ functionQuery = {} }) => { + const names = this.#configs.map(({ astQuery, functionQuery = {} }) => { + if (astQuery) return astQuery const resolvedQuery = this.#resolveExportAlias(functionQuery, aliases) const queryName = (q) => q.methodName || q.privateMethodName || q.functionName || q.expressionName || 'constructor' const originalName = queryName(functionQuery) diff --git a/lib/transforms.js b/lib/transforms.js index be94f6b..c7cc7bc 100644 --- a/lib/transforms.js +++ b/lib/transforms.js @@ -560,7 +560,7 @@ function wrapSync (state, node, iterPatch = '') { // When `mutableResult` is set, the return value is whatever // `__apm$ctx.result` holds *after* `end` has published. - // So a subscriber's `end` handler can eassign `message.result` + // So a subscriber's `end` handler can reassign `message.result` // to substitute it (e.g. to wrap a returned handler). The `return` // therefore moves below the `try`/`finally`: on the throw path the // `catch` re-throws, so the trailing `return` is unreachable. diff --git a/tests/ast_query_returned_arrow_cjs/mod.js b/tests/ast_query_returned_arrow_cjs/mod.js new file mode 100644 index 0000000..4834abb --- /dev/null +++ b/tests/ast_query_returned_arrow_cjs/mod.js @@ -0,0 +1,12 @@ +// Mirrors a NestJS-style decorator factory: `Decorator(options)` returns the +// actual decorator `(target) => {...}` that is applied to a class. The returned +// arrow is anonymous and synchronous, so the name-based FunctionQuery variants +// cannot target it. Only a raw `astQuery` can do that. +function Decorator (options) { + return (target) => { + target.decorated = options + return target + } +} + +module.exports = { Decorator } diff --git a/tests/ast_query_returned_arrow_cjs/test.js b/tests/ast_query_returned_arrow_cjs/test.js new file mode 100644 index 0000000..2992473 --- /dev/null +++ b/tests/ast_query_returned_arrow_cjs/test.js @@ -0,0 +1,26 @@ +const { Decorator } = require('./instrumented.js') +const assert = require('node:assert') +const { tracingChannel } = require('node:diagnostics_channel') + +// The channel fires when the (anonymous, sync) returned decorator is applied. +// A subscriber can mutate the decorated target (arg 0): building block for +// instrumenting decorated classes. +let started = false +tracingChannel('orchestrion:undici:decorator_apply').subscribe({ + start (message) { + started = true + const target = message.arguments[0] + target.instrumented = true + } +}) + +class MyService {} +const decorate = Decorator({ scope: 'request' }) +const result = decorate(MyService) + +assert.strictEqual(started, true) +// Original decorator still ran. +assert.strictEqual(result, MyService) +assert.deepStrictEqual(MyService.decorated, { scope: 'request' }) +// Subscriber's mutation of the target took effect. +assert.strictEqual(MyService.instrumented, true) diff --git a/tests/tests.test.mjs b/tests/tests.test.mjs index 2b16053..d89839b 100644 --- a/tests/tests.test.mjs +++ b/tests/tests.test.mjs @@ -431,6 +431,35 @@ describe('mutable_result_cjs', () => { }) }) +describe('ast_query_returned_arrow_cjs', () => { + test('instruments an anonymous arrow returned by a factory via astQuery', () => { + runTest('ast_query_returned_arrow_cjs', [ + { + channelName: 'decorator_apply', + module: { name: TEST_MODULE_NAME, versionRange: '>=0.0.1', filePath: TEST_MODULE_PATH }, + astQuery: 'FunctionDeclaration[id.name="Decorator"] ReturnStatement > ArrowFunctionExpression', + functionQuery: { kind: 'Sync' }, + }, + ]) + }) + + test('reports the astQuery selector when no injection points are found', () => { + const instrumentor = create([ + { + channelName: 'no_match', + module: { name: TEST_MODULE_NAME, versionRange: '>=0.0.1', filePath: TEST_MODULE_PATH }, + astQuery: 'FunctionDeclaration[id.name="doesNotExist"]', + functionQuery: { kind: 'Sync' }, + }, + ]) + const transformer = instrumentor.getTransformer(TEST_MODULE_NAME, TEST_MODULE_VERSION, TEST_MODULE_PATH) + assert.throws( + () => transformer.transform('function fetch () { return 42 }', 'cjs'), + /doesNotExist/ + ) + }) +}) + describe('polyfill_cjs', () => { test('instruments with a custom dc module (cjs)', () => { runTest('polyfill_cjs', [