Update PR #2275

Author: lforst

js/src/auto-instrumentations/bundler/plugin.ts

@@ -25,15 +25,13 @@ export interface BundlerPluginOptions {
   browser?: boolean;
 
   /**
-   * Previously replaced Node's `diagnostics_channel` implementation with a
-   * browser-compatible shim in transformed code.
+   * Marks transformed source as targeting a browser or edge-like environment.
    *
-   * Global instrumentation hooks do not use `diagnostics_channel`, so this
-   * option no longer does anything. It is retained for backwards
-   * compatibility.
+   * This retains the previous browser-target behavior of the option. Global
+   * instrumentation hooks are runtime-independent, so no diagnostics-channel
+   * compatibility shim is injected.
    *
-   * @deprecated This option no longer does anything. Use `browser` to mark
-   * browser or edge-like bundles.
+   * @deprecated Use `browser` instead.
    */
   useDiagnosticChannelCompatShim?: boolean;
 }
@@ -56,7 +54,8 @@ function getModuleVersion(basedir: string): string | undefined {
 }
 
 export const unplugin = createUnplugin<BundlerPluginOptions>((options = {}) => {
-  const browser = options.browser ?? false;
+  const browser =
+    options.browser ?? options.useDiagnosticChannelCompatShim ?? false;
   const allInstrumentations = getDefaultInstrumentationConfigs({
     additionalInstrumentations: options.instrumentations,
   });

js/src/auto-instrumentations/bundler/webpack-loader.ts

@@ -128,7 +128,7 @@ function codeTransformerLoader(
     modulePath: normalizedModulePath,
     source: code,
     format: isModule ? "esm" : "cjs",
-    browser: options.browser ?? false,
+    browser: options.browser ?? options.useDiagnosticChannelCompatShim ?? false,
   });
   if (patched !== null) {
     return callback(null, patched);

js/src/auto-instrumentations/orchestrion-js/transformer.ts

@@ -88,7 +88,6 @@ export class Transformer {
       const query = astQuery || this.fromFunctionQuery(resolvedFunctionQuery);
       const state: TransformState = {
         ...config,
-        moduleType,
         moduleVersion: this.version,
         functionQuery: resolvedFunctionQuery,
         operator: this.getOperator(resolvedFunctionQuery.kind),

js/src/auto-instrumentations/orchestrion-js/transforms.ts

@@ -11,7 +11,7 @@ import {
   GLOBAL_INSTRUMENTATION_HOOKS_PROTOCOL_VERSION,
   GLOBAL_INSTRUMENTATION_HOOKS_REGISTRY_BRAND,
 } from "../../global-instrumentation-hooks";
-import type { FunctionQuery, InstrumentationConfig, ModuleType } from "./types";
+import type { FunctionQuery, InstrumentationConfig } from "./types";
 
 type AnyNode = any;
 type TransformFn = (
@@ -23,14 +23,14 @@ type TransformFn = (
 type TraceOperator = "traceCallback" | "tracePromise" | "traceSync";
 
 export interface TransformState extends InstrumentationConfig {
-  moduleType: ModuleType;
   moduleVersion: string;
   functionQuery: FunctionQuery;
   operator: TraceOperator;
   functionIndex?: number;
 }
 
 const CHANNEL_REGEX = /[^\w]/g;
+const SHARED_HOOK_LOOKUP = "tr_ch_apm$get_hook";
 
 function formatChannelVariable(channelName: string): string {
   return `tr_ch_apm$${channelName.replace(CHANNEL_REGEX, "_")}`;
@@ -49,6 +49,10 @@ export const transforms: Record<string, TransformFn> = {
     } = state;
     const channelVariable = formatChannelVariable(channelName);
     const channelGetter = formatChannelGetter(channelName);
+    const hasSharedHookLookup = node.body.some(
+      (child: AnyNode) =>
+        child.declarations?.[0]?.id?.name === SHARED_HOOK_LOOKUP,
+    );
 
     if (
       node.body.some(
@@ -58,12 +62,10 @@ export const transforms: Record<string, TransformFn> = {
       return;
     }
 
-    const index = node.body.findIndex(
-      (child: AnyNode) => child.directive === "use strict",
-    );
-    const code = `
-      let ${channelVariable};
-      const ${channelGetter} = () => {
+    const sharedHookLookup = hasSharedHookLookup
+      ? ""
+      : `
+      const ${SHARED_HOOK_LOOKUP} = (__apm$hookName, __apm$operator) => {
         try {
           const __apm$hooks = globalThis[${JSON.stringify(
             GLOBAL_INSTRUMENTATION_HOOKS_KEY,
@@ -76,7 +78,7 @@ export const transforms: Record<string, TransformFn> = {
           ) return undefined;
           const __apm$hook = Map.prototype.get.call(
             __apm$hooks,
-            "orchestrion:${name}:${channelName}"
+            __apm$hookName
           );
           if (
             (__apm$hook === null ||
@@ -86,15 +88,34 @@ export const transforms: Record<string, TransformFn> = {
               GLOBAL_INSTRUMENTATION_HOOK_BRAND,
             )})] !== ${GLOBAL_INSTRUMENTATION_HOOKS_PROTOCOL_VERSION} ||
             typeof __apm$hook.hasSubscribers !== "boolean" ||
-            typeof __apm$hook[${JSON.stringify(operator)}] !== "function"
+            typeof __apm$hook[__apm$operator] !== "function"
           ) return undefined;
-          return ${channelVariable} ??= __apm$hook;
+          return __apm$hook;
         } catch {
           return undefined;
         }
       };
     `;
+    const code = `
+      ${sharedHookLookup}
+      let ${channelVariable};
+      const ${channelGetter} = () =>
+        ${channelVariable} ??= ${SHARED_HOOK_LOOKUP}(
+          ${JSON.stringify(`orchestrion:${name}:${channelName}`)},
+          ${JSON.stringify(operator)}
+        );
+    `;
 
+    const sharedHookLookupIndex = node.body.findIndex(
+      (child: AnyNode) =>
+        child.declarations?.[0]?.id?.name === SHARED_HOOK_LOOKUP,
+    );
+    const index =
+      sharedHookLookupIndex === -1
+        ? node.body.findIndex(
+            (child: AnyNode) => child.directive === "use strict",
+          )
+        : sharedHookLookupIndex;
     node.body.splice(index + 1, 0, ...parse(code).body);
   },
 
@@ -139,18 +160,14 @@ function traceFunction(
     (!methodName && !privateMethodName && !functionName);
   const type = isConstructor ? "ArrowFunctionExpression" : "FunctionExpression";
 
-  node.body = wrap(
-    state,
-    {
-      type,
-      params: node.params,
-      body: node.body,
-      async: node.async,
-      expression: false,
-      generator: node.generator,
-    },
-    program,
-  );
+  node.body = wrap(state, {
+    type,
+    params: node.params,
+    body: node.body,
+    async: node.async,
+    expression: false,
+    generator: node.generator,
+  });
 
   node.generator = false;
   node.async = false;
@@ -202,22 +219,17 @@ function traceInstanceMethod(
   const fn = ctorBody[1].expression.right;
 
   fn.async = operator === "tracePromise";
-  fn.body = wrap(
-    state,
-    { type: "Identifier", name: `__apm$${methodName}` },
-    program,
-  );
+  fn.body = wrap(state, {
+    type: "Identifier",
+    name: `__apm$${methodName}`,
+  });
 
   wrapSuper(fn);
 
   ctor.value.body.body.push(...ctorBody);
 }
 
-function wrap(
-  state: TransformState,
-  node: AnyNode,
-  program?: AnyNode,
-): AnyNode {
+function wrap(state: TransformState, node: AnyNode): AnyNode {
   const { operator, moduleVersion } = state;
 
   const wrapper =

js/src/global-instrumentation-hooks.ts

@@ -137,20 +137,16 @@ function setContextValue(
   }
 }
 
-function defaultTransform<M>(message: M): M {
-  return message;
-}
-
 function wrapStoreRun<M>(
   store: GlobalHookAsyncLocalStorage<unknown>,
   message: M,
   next: () => unknown,
-  transform: GlobalHookTransformFunction<M, unknown> = defaultTransform,
+  transform?: GlobalHookTransformFunction<M, unknown>,
 ): () => unknown {
   return () => {
     let context: unknown;
     try {
-      context = transform(message);
+      context = transform ? transform(message) : message;
     } catch (error) {
       reportError(error);
       return next();

js/tests/auto-instrumentations/transformation.test.ts

@@ -369,6 +369,48 @@ describe("Orchestrion Transformation Tests", () => {
       );
     });
 
+    it("injects one shared global hook lookup for multiple channels", () => {
+      const matcher = create([
+        {
+          ...testConfig({ functionName: "first", kind: "Sync" }),
+          channelName: "first",
+        },
+        {
+          ...testConfig({ functionName: "second", kind: "Async" }),
+          channelName: "second",
+        },
+      ]);
+      const transformer = matcher.getTransformer(
+        "test-sdk",
+        "1.0.0",
+        "index.mjs",
+      );
+
+      expect(transformer).toBeDefined();
+      const result = transformer!.transform(
+        `
+          export function first() {
+            return "first";
+          }
+          export async function second() {
+            return "second";
+          }
+        `,
+        "esm",
+      );
+
+      expect(result.code.match(/const tr_ch_apm\$get_hook =/g)).toHaveLength(1);
+      expect(
+        result.code.match(
+          /braintrust\.global-instrumentation-hooks\.registry/g,
+        ),
+      ).toHaveLength(1);
+      expect(result.code).toContain("orchestrion:test-sdk:first");
+      expect(result.code).toContain("orchestrion:test-sdk:second");
+      expect(result.code).toContain("__apm$hook.traceSync");
+      expect(result.code).toContain("__apm$hook.tracePromise");
+    });
+
     it("generates source maps", () => {
       const result = transformTestCode(
         { functionName: "query", kind: "Sync" },
@@ -449,6 +491,7 @@ describe("Orchestrion Transformation Tests", () => {
 
     it.each([
       ["browser", "browser", { browser: true }],
+      ["legacy-browser", "browser", { useDiagnosticChannelCompatShim: true }],
       ["edge", "neutral", { browser: true }],
     ] as const)(
       "should keep Mastra %s bundles free of Node-only patches",
@@ -715,8 +758,15 @@ describe("Orchestrion Transformation Tests", () => {
     });
 
     it.each([
-      ["apply", "node", "node", { useDiagnosticChannelCompatShim: true }, true],
+      [
+        "skip",
+        "legacy browser",
+        "web",
+        { useDiagnosticChannelCompatShim: true },
+        false,
+      ],
       ["skip", "browser", "web", { browser: true }, false],
+      ["apply", "node", "node", { browser: false }, true],
     ] as const)(
       "should %s special-case patches for %s turbopack loader targets",
       async (_action, runtime, target, options, shouldPatch) => {