chore(vitest,tinybench): emit benchmark markers

Emit benchmark start/end markers for the tinybench plugin and the
vitest walltime runner, wrapping the measured function in a root frame.

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/shared.ts

@@ -1,4 +1,11 @@
-import { InstrumentHooks, setupCore, teardownCore } from "@codspeed/core";
+import {
+  getInstrumentMode,
+  InstrumentHooks,
+  MARKER_TYPE_BENCHMARK_END,
+  MARKER_TYPE_BENCHMARK_START,
+  setupCore,
+  teardownCore,
+} from "@codspeed/core";
 import { Bench, Fn, Task } from "tinybench";
 import { getTaskUri } from "./uri";
 
@@ -63,25 +70,43 @@ export abstract class BaseBenchRunner {
   protected abstract finalizeAsyncRun(): Task[];
   protected abstract finalizeSyncRun(): Task[];
 
+  private sendRunMarkers(runStart: bigint, runEnd: bigint): void {
+    if (getInstrumentMode() !== "walltime") {
+      return;
+    }
+    InstrumentHooks.addMarker(
+      process.pid,
+      MARKER_TYPE_BENCHMARK_START,
+      runStart,
+    );
+    InstrumentHooks.addMarker(process.pid, MARKER_TYPE_BENCHMARK_END, runEnd);
+  }
+
   public setupBenchMethods(): void {
     this.bench.run = async () => {
       this.setupBenchRun();
 
+      const runStart = InstrumentHooks.currentTimestamp();
       for (const task of this.bench.tasks) {
         const uri = this.getTaskUri(task);
         await this.runTaskAsync(task, uri);
       }
+      const runEnd = InstrumentHooks.currentTimestamp();
+      this.sendRunMarkers(runStart, runEnd);
 
       return this.finalizeAsyncRun();
     };
 
     this.bench.runSync = () => {
       this.setupBenchRun();
 
+      const runStart = InstrumentHooks.currentTimestamp();
       for (const task of this.bench.tasks) {
         const uri = this.getTaskUri(task);
         this.runTaskSync(task, uri);
       }
+      const runEnd = InstrumentHooks.currentTimestamp();
+      this.sendRunMarkers(runStart, runEnd);
 
       return this.finalizeSyncRun();
     };

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/tinybench-plugin/tests/index.integ.test.ts

@@ -16,6 +16,8 @@ const mockCore = vi.hoisted(() => {
       startBenchmark: vi.fn(),
       stopBenchmark: vi.fn(),
       setExecutedBenchmark: vi.fn(),
+      currentTimestamp: vi.fn().mockReturnValue(0n),
+      addMarker: vi.fn(),
     },
     optimizeFunction: vi
       .fn()

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,9 +3,9 @@ import {
   MARKER_TYPE_BENCHMARK_END,
   MARKER_TYPE_BENCHMARK_START,
   setupCore,
+  wrapWithRootFrame,
   writeWalltimeResults,
 } from "@codspeed/core";
-import { Fn } from "tinybench";
 import {
   RunnerTaskEventPack,
   RunnerTaskResultPack,
@@ -78,37 +78,27 @@ export class WalltimeRunner extends NodeBenchmarkRunner {
     };
 
     tinybench.Task.prototype.run = async function () {
-      const { fn } = this as { fn: Fn };
       const suiteUri = getSuiteUri();
 
-      const finishRound = (roundStart: bigint): void => {
-        const roundEnd = InstrumentHooks.currentTimestamp();
-        InstrumentHooks.addMarker(pid, MARKER_TYPE_BENCHMARK_START, roundStart);
-        InstrumentHooks.addMarker(pid, MARKER_TYPE_BENCHMARK_END, roundEnd);
-      };
-
-      function __codspeed_root_frame__(): unknown {
-        const roundStart = InstrumentHooks.currentTimestamp();
-        const result = fn();
-        if (
-          result !== null &&
-          typeof result === "object" &&
-          typeof (result as PromiseLike<unknown>).then === "function"
-        ) {
-          return (result as PromiseLike<unknown>).then((value) => {
-            finishRound(roundStart);
-            return value;
-          });
-        }
-        finishRound(roundStart);
-        return result;
-      }
-      (this as { fn: Fn }).fn = __codspeed_root_frame__ as Fn;
+      // 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();
-      await originalRun.call(this);
+      const runStart = InstrumentHooks.currentTimestamp();
+      try {
+        await originalRun.call(this);
+      } finally {
+        task.fn = originalFn;
+      }
+      const runEnd = InstrumentHooks.currentTimestamp();
       InstrumentHooks.stopBenchmark();
 
+      // Emit a single marker pair covering the whole measurement run
+      InstrumentHooks.addMarker(pid, MARKER_TYPE_BENCHMARK_START, runStart);
+      InstrumentHooks.addMarker(pid, MARKER_TYPE_BENCHMARK_END, runEnd);
+
       // Look up the URI by task name
       const uri = `${suiteUri}::${this.name}`;
       InstrumentHooks.setExecutedBenchmark(pid, uri);