feat: make astQuery a first-class feature

README.md

@@ -87,6 +87,93 @@ 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.
+
+### 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
@@ -108,13 +195,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 };
 ```
@@ -132,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

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`.
@@ -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
  */

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)

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 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.
+  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();

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 }

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)

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 }