fixup! chore(vitest): emit benchmark markers

Author: not-matthias

packages/benchmark.js-plugin/src/index.ts

@@ -9,6 +9,8 @@ import {
   SetupInstrumentsResponse,
   teardownCore,
   tryIntrospect,
+  wrapWithRootFrame,
+  wrapWithRootFrameSync,
 } from "@codspeed/core";
 import Benchmark from "benchmark";
 import buildSuiteAdd from "./buildSuiteAdd";
@@ -195,7 +197,7 @@ async function runBenchmarks({
       await optimizeFunction(benchPayload);
       await mongoMeasurement.start(uri);
       global.gc?.();
-      await (async function __codspeed_root_frame__() {
+      await wrapWithRootFrame(async () => {
         InstrumentHooks.startBenchmark();
         await benchPayload();
         InstrumentHooks.stopBenchmark();
@@ -205,7 +207,7 @@ async function runBenchmarks({
     } else {
       optimizeFunctionSync(benchPayload);
       await mongoMeasurement.start(uri);
-      (function __codspeed_root_frame__() {
+      wrapWithRootFrameSync(() => {
         InstrumentHooks.startBenchmark();
         benchPayload();
         InstrumentHooks.stopBenchmark();

packages/core/src/index.ts

@@ -77,6 +77,7 @@ export type {
 } from "./generated/openapi";
 export { getV8Flags, tryIntrospect } from "./introspection";
 export { optimizeFunction, optimizeFunctionSync } from "./optimization";
+export { wrapWithRootFrame, wrapWithRootFrameSync } from "./rootFrame";
 export * from "./utils";
 export * from "./walltime";
 export type { InstrumentMode };

packages/core/src/rootFrame.ts

@@ -0,0 +1,19 @@
+/**
+ * Wrap a benchmark function so it executes under a frame named
+ * `__codspeed_root_frame__`. CodSpeed uses this frame to locate the
+ * benchmark root in collected call stacks; samples without it cannot be
+ * attributed to a benchmark.
+ */
+export function wrapWithRootFrame<T>(
+  fn: () => T | Promise<T>,
+): () => Promise<T> {
+  return async function __codspeed_root_frame__() {
+    return await fn();
+  };
+}
+
+export function wrapWithRootFrameSync<T>(fn: () => T): () => T {
+  return function __codspeed_root_frame__() {
+    return fn();
+  };
+}

packages/tinybench-plugin/src/analysis.ts

@@ -3,6 +3,8 @@ import {
   InstrumentHooks,
   mongoMeasurement,
   optimizeFunction,
+  wrapWithRootFrame,
+  wrapWithRootFrameSync,
 } from "@codspeed/core";
 import { Bench, Fn, FnOptions, Task } from "tinybench";
 import { BaseBenchRunner } from "./shared";
@@ -25,18 +27,6 @@ class AnalysisBenchRunner extends BaseBenchRunner {
     return InstrumentHooks.isInstrumented() ? "Measured" : "Checked";
   }
 
-  private wrapFunctionWithFrame(fn: Fn, isAsync: boolean): Fn {
-    if (isAsync) {
-      return async function __codspeed_root_frame__() {
-        await fn();
-      };
-    } else {
-      return function __codspeed_root_frame__() {
-        fn();
-      };
-    }
-  }
-
   protected async runTaskAsync(task: Task, uri: string): Promise<void> {
     const { fnOpts, fn } = task as unknown as { fnOpts?: FnOptions; fn: Fn };
 
@@ -50,10 +40,7 @@ class AnalysisBenchRunner extends BaseBenchRunner {
     await mongoMeasurement.start(uri);
 
     global.gc?.();
-    await this.wrapWithInstrumentHooksAsync(
-      this.wrapFunctionWithFrame(fn, true),
-      uri,
-    );
+    await this.wrapWithInstrumentHooksAsync(wrapWithRootFrame(fn), uri);
 
     await mongoMeasurement.stop(uri);
     await fnOpts?.afterEach?.call(task, "run");
@@ -68,7 +55,7 @@ class AnalysisBenchRunner extends BaseBenchRunner {
     fnOpts?.beforeAll?.call(task, "run");
     fnOpts?.beforeEach?.call(task, "run");
 
-    this.wrapWithInstrumentHooks(this.wrapFunctionWithFrame(fn, false), uri);
+    this.wrapWithInstrumentHooks(wrapWithRootFrameSync(fn), uri);
 
     fnOpts?.afterEach?.call(task, "run");
     fnOpts?.afterAll?.call(task, "run");

packages/tinybench-plugin/src/walltime.ts

@@ -3,6 +3,8 @@ import {
   mongoMeasurement,
   msToNs,
   msToS,
+  wrapWithRootFrame,
+  wrapWithRootFrameSync,
   writeWalltimeResults,
   type BenchmarkStats,
   type Benchmark as CodspeedBenchmark,
@@ -64,21 +66,10 @@ class WalltimeBenchRunner extends BaseBenchRunner {
 
   private wrapTaskFunction(task: Task, isAsync: boolean): void {
     const { fn } = task as unknown as { fn: Fn };
-    if (isAsync) {
-      // eslint-disable-next-line no-inner-declarations
-      async function __codspeed_root_frame__() {
-        await fn();
-      }
-      // eslint-disable-next-line @typescript-eslint/no-explicit-any
-      (task as any).fn = __codspeed_root_frame__;
-    } else {
-      // eslint-disable-next-line no-inner-declarations
-      function __codspeed_root_frame__() {
-        fn();
-      }
-      // eslint-disable-next-line @typescript-eslint/no-explicit-any
-      (task as any).fn = __codspeed_root_frame__;
-    }
+    // eslint-disable-next-line @typescript-eslint/no-explicit-any
+    (task as any).fn = isAsync
+      ? wrapWithRootFrame(fn)
+      : wrapWithRootFrameSync(fn);
   }
 
   private registerCodspeedBenchmarkFromTask(task: Task): void {

packages/vitest-plugin/src/analysis.ts

@@ -5,6 +5,7 @@ import {
   optimizeFunction,
   setupCore,
   teardownCore,
+  wrapWithRootFrame,
 } from "@codspeed/core";
 import { Benchmark, type RunnerTestSuite } from "vitest";
 import { NodeBenchmarkRunner } from "vitest/runners";
@@ -47,7 +48,7 @@ async function runAnalysisBench(
   await callSuiteHook(suite, benchmark, "beforeEach");
   await mongoMeasurement.start(uri);
   global.gc?.();
-  await (async function __codspeed_root_frame__() {
+  await wrapWithRootFrame(async () => {
     InstrumentHooks.startBenchmark();
     // @ts-expect-error we do not need to bind the function to an instance of tinybench's Bench
     await fn();

packages/vitest-plugin/src/walltime/index.ts

@@ -3,6 +3,7 @@ import {
   MARKER_TYPE_BENCHMARK_END,
   MARKER_TYPE_BENCHMARK_START,
   setupCore,
+  wrapWithRootFrame,
   writeWalltimeResults,
 } from "@codspeed/core";
 import {
@@ -79,9 +80,18 @@ export class WalltimeRunner extends NodeBenchmarkRunner {
     tinybench.Task.prototype.run = async function () {
       const suiteUri = getSuiteUri();
 
+      // eslint-disable-next-line @typescript-eslint/no-explicit-any
+      const task = this as any;
+      const originalFn = task.fn;
+      task.fn = wrapWithRootFrame(() => originalFn.call(task));
+
       InstrumentHooks.startBenchmark();
       const runStart = InstrumentHooks.currentTimestamp();
-      await originalRun.call(this);
+      try {
+        await originalRun.call(this);
+      } finally {
+        task.fn = originalFn;
+      }
       const runEnd = InstrumentHooks.currentTimestamp();
       InstrumentHooks.stopBenchmark();