fix(vitest-plugin): bracket walltime window around the measured loop only

The runner wrapped the entire tinybench Task.run() in the instrumentation
window, folding the run-mode setup/teardown hooks and processRunResult (which
sorts the latency samples twice and computes the full statistics) into the
recorded sample. That framework overhead showed up in the walltime flamegraph.

Drive the window from each bench's run-mode setup/teardown hooks instead, so it
brackets only the measured loop. Warmup already runs as a separate Vitest call
and is now excluded, and the post-loop statistics computation falls outside the
window. User-provided hooks are preserved.

The tinybench module namespace is frozen, so the instrumented Bench is handed
back through a fresh module-shaped object that Vitest destructures from rather
than reassigning the export.

Closes COD-2925
Co-Authored-By: Claude <noreply@anthropic.com>

Author: GuillaumeLagrange

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

@@ -6,6 +6,7 @@ import {
   wrapWithRootFrame,
   writeWalltimeResults,
 } from "@codspeed/core";
+import type * as tinybench from "tinybench";
 import {
   RunnerTaskEventPack,
   RunnerTaskResultPack,
@@ -15,6 +16,27 @@ import { NodeBenchmarkRunner } from "vitest/runners";
 import { patchRootSuiteWithFullFilePath } from "../common";
 import { extractBenchmarkResults } from "./utils";
 
+type Tinybench = typeof tinybench;
+
+/** A tinybench task, exposing the `fn` the runner wraps with the root frame. */
+interface TinybenchTask {
+  name: string;
+  // eslint-disable-next-line @typescript-eslint/no-explicit-any
+  fn: (...args: any[]) => any;
+}
+
+/** tinybench's per-task setup/teardown hook signature. */
+type TinybenchHook = (
+  task: TinybenchTask,
+  mode: "run" | "warmup",
+) => Promise<void> | void;
+
+/** The mutable subset of a tinybench Bench the runner reaches into. */
+interface TinybenchBench {
+  setup: TinybenchHook;
+  teardown: TinybenchHook;
+}
+
 /**
  * WalltimeRunner uses Vitest's default benchmark execution
  * and extracts results from the suite after completion
@@ -24,6 +46,9 @@ export class WalltimeRunner extends NodeBenchmarkRunner {
   private suiteUris = new Map<string, string>();
   /// Suite ID of the currently running suite, to allow constructing the URI in the context of tinybench tasks
   private currentSuiteId: string | null = null;
+  // Carries the window start timestamp from the setup hook to the teardown
+  // hook. Tasks run strictly sequentially, so a single field is enough.
+  private runStart: bigint | null = null;
 
   async runSuite(suite: RunnerTestSuite): Promise<void> {
     patchRootSuiteWithFullFilePath(suite);
@@ -59,59 +84,118 @@ export class WalltimeRunner extends NodeBenchmarkRunner {
     }
   }
 
-  async importTinybench() {
+  private getBenchmarkUri(taskName: string): string {
+    if (this.currentSuiteId === null) {
+      throw new Error("currentSuiteId is null - something went wrong");
+    }
+    const suiteUri = this.suiteUris.get(this.currentSuiteId) || "";
+    return `${suiteUri}::${taskName}`;
+  }
+
+  async importTinybench(): Promise<Tinybench> {
     const tinybench = await super.importTinybench();
 
-    if (this.isTinybenchHookedWithCodspeed) {
-      return tinybench;
+    // `tinybench` is a frozen ES module namespace, so the `Bench` export cannot
+    // be reassigned. Mutating the shared `Task.prototype` in place is allowed
+    // and only needs to happen once; the instrumented `Bench` is handed back
+    // through a fresh module-shaped object that Vitest destructures from.
+    if (!this.isTinybenchHookedWithCodspeed) {
+      this.isTinybenchHookedWithCodspeed = true;
+      this.patchTaskWithRootFrame(tinybench);
     }
-    this.isTinybenchHookedWithCodspeed = true;
 
-    const originalRun = tinybench.Task.prototype.run;
-    const pid = process.pid;
-
-    const getSuiteUri = (): string => {
-      if (this.currentSuiteId === null) {
-        throw new Error("currentSuiteId is null - something went wrong");
-      }
-      return this.suiteUris.get(this.currentSuiteId) || "";
+    return {
+      ...tinybench,
+      Bench: this.createInstrumentedBench(tinybench),
     };
+  }
 
-    tinybench.Task.prototype.run = async function () {
-      const suiteUri = getSuiteUri();
+  /**
+   * Wrap each task's function with the root frame so collected stacks can be
+   * attributed to a benchmark. The window itself is driven by the bench's
+   * setup/teardown hooks (see createInstrumentedBench).
+   */
+  private patchTaskWithRootFrame(tinybench: Tinybench): void {
+    const originalRun = tinybench.Task.prototype.run;
 
-      // eslint-disable-next-line @typescript-eslint/no-explicit-any
-      const task = this as any;
+    tinybench.Task.prototype.run = async function () {
+      const task = this as unknown as TinybenchTask;
       const originalFn = task.fn;
       task.fn = wrapWithRootFrame(() => originalFn.call(task));
 
-      InstrumentHooks.startBenchmark();
-      const runStart = InstrumentHooks.currentTimestamp();
       try {
         await originalRun.call(this);
       } finally {
-        const runEnd = InstrumentHooks.currentTimestamp();
         task.fn = originalFn;
+      }
 
-        // Benchmark markers must land inside the sample window opened by
-        // startBenchmark(), so they have to be emitted before stopBenchmark()
-        // closes it. The runner consumes the FIFO stream in order, so a marker
-        // sent after StopBenchmark falls outside the sample and breaks the
-        // expected SampleStart > BenchmarkStart > BenchmarkEnd > SampleEnd nesting.
-        InstrumentHooks.addMarker(pid, MARKER_TYPE_BENCHMARK_START, runStart);
-        InstrumentHooks.addMarker(pid, MARKER_TYPE_BENCHMARK_END, runEnd);
+      return this;
+    };
+  }
 
-        InstrumentHooks.stopBenchmark();
+  /**
+   * Drive the instrumentation window from each bench's run-mode setup/teardown
+   * hooks so it brackets only tinybench's measured loop, excluding the warmup
+   * that Vitest runs beforehand and the statistics computation tinybench
+   * performs after the loop. Wrapping the whole `Task.run()` would otherwise
+   * fold all of that framework overhead into the recorded sample.
+   *
+   * User-provided hooks are preserved and keep their order relative to the work
+   * under test.
+   */
+  private createInstrumentedBench(
+    tinybench: Tinybench,
+  ): typeof tinybench.Bench {
+    // eslint-disable-next-line @typescript-eslint/no-this-alias
+    const runner = this;
+    const OriginalBench = tinybench.Bench;
+
+    class InstrumentedBench extends OriginalBench {
+      // eslint-disable-next-line @typescript-eslint/no-explicit-any
+      constructor(...benchArgs: any[]) {
+        super(...benchArgs);
+        runner.installInstrumentHooks(this as unknown as TinybenchBench);
+      }
+    }
 
-        // Look up the URI by task name
-        const uri = `${suiteUri}::${this.name}`;
-        InstrumentHooks.setExecutedBenchmark(pid, uri);
+    return InstrumentedBench;
+  }
+
+  private installInstrumentHooks(bench: TinybenchBench): void {
+    const userSetup = bench.setup;
+    const userTeardown = bench.teardown;
+
+    bench.setup = async (task, mode) => {
+      await userSetup(task, mode);
+      if (mode === "run") {
+        InstrumentHooks.startBenchmark();
+        this.runStart = InstrumentHooks.currentTimestamp();
       }
+    };
 
-      return this;
+    bench.teardown = async (task, mode) => {
+      if (mode === "run") {
+        this.closeInstrumentWindow(this.getBenchmarkUri(task.name));
+      }
+      await userTeardown(task, mode);
     };
+  }
+
+  private closeInstrumentWindow(uri: string): void {
+    const runEnd = InstrumentHooks.currentTimestamp();
+    const pid = process.pid;
 
-    return tinybench;
+    // Benchmark markers must land inside the sample window opened by
+    // startBenchmark(), so they have to be emitted before stopBenchmark()
+    // closes it. The runner consumes the FIFO stream in order, so a marker
+    // sent after StopBenchmark falls outside the sample and breaks the
+    // expected SampleStart > BenchmarkStart > BenchmarkEnd > SampleEnd nesting.
+    InstrumentHooks.addMarker(pid, MARKER_TYPE_BENCHMARK_START, this.runStart!);
+    InstrumentHooks.addMarker(pid, MARKER_TYPE_BENCHMARK_END, runEnd);
+
+    InstrumentHooks.stopBenchmark();
+    InstrumentHooks.setExecutedBenchmark(pid, uri);
+    this.runStart = null;
   }
 
   // Allow tinybench to retrieve the path to the currently running suite