Add perf profiling to vitest

packages/vitest-plugin/package.json

@@ -31,12 +31,14 @@
     "@codspeed/core": "workspace:^5.0.0"
   },
   "peerDependencies": {
+    "tinybench": "^2.9.0",
     "vite": "^4.2.0 || ^5.0.0 || ^6.0.0 || ^7.0.0",
     "vitest": ">=3.2"
   },
   "devDependencies": {
     "@total-typescript/shoehorn": "^0.1.1",
     "execa": "^8.0.1",
+    "tinybench": "^2.9.0",
     "vite": "^7.0.0",
     "vitest": "^3.2.4"
   }

packages/vitest-plugin/rollup.config.ts

@@ -27,7 +27,7 @@ export default defineConfig([
     external: ["@codspeed/core", /^vitest/],
   },
   {
-    input: "src/walltime.ts",
+    input: "src/walltime/index.ts",
     output: { file: "dist/walltime.mjs", format: "es" },
     plugins: jsPlugins(pkg.version),
     external: ["@codspeed/core", /^vitest/],

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

@@ -0,0 +1,115 @@
+import {
+  InstrumentHooks,
+  setupCore,
+  writeWalltimeResults,
+} from "@codspeed/core";
+import { Fn } from "tinybench";
+import {
+  RunnerTaskEventPack,
+  RunnerTaskResultPack,
+  type RunnerTestSuite,
+} from "vitest";
+import { NodeBenchmarkRunner } from "vitest/runners";
+import { patchRootSuiteWithFullFilePath } from "../common";
+import { extractBenchmarkResults } from "./utils";
+
+/**
+ * WalltimeRunner uses Vitest's default benchmark execution
+ * and extracts results from the suite after completion
+ */
+export class WalltimeRunner extends NodeBenchmarkRunner {
+  private isTinybenchHookedWithCodspeed = false;
+  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;
+
+  async runSuite(suite: RunnerTestSuite): Promise<void> {
+    patchRootSuiteWithFullFilePath(suite);
+    this.populateBenchmarkUris(suite);
+
+    setupCore();
+
+    await super.runSuite(suite);
+
+    const benchmarks = await extractBenchmarkResults(suite);
+
+    if (benchmarks.length > 0) {
+      writeWalltimeResults(benchmarks);
+      console.log(
+        `[CodSpeed] Done collecting walltime data for ${benchmarks.length} benches.`
+      );
+    } else {
+      console.warn(
+        `[CodSpeed] No benchmark results found after suite execution`
+      );
+    }
+  }
+
+  private populateBenchmarkUris(suite: RunnerTestSuite, parentPath = ""): void {
+    const currentPath =
+      parentPath !== "" ? `${parentPath}::${suite.name}` : suite.name;
+
+    for (const task of suite.tasks) {
+      if (task.type === "suite") {
+        this.suiteUris.set(task.id, `${currentPath}::${task.name}`);
+        this.populateBenchmarkUris(task, currentPath);
+      }
+    }
+  }
+
+  async importTinybench(): Promise<typeof import("tinybench")> {
+    const tinybench = await super.importTinybench();
+
+    if (this.isTinybenchHookedWithCodspeed) {
+      return tinybench;
+    }
+    this.isTinybenchHookedWithCodspeed = true;
+
+    const originalRun = tinybench.Task.prototype.run;
+
+    const getSuiteUri = (): string => {
+      if (this.currentSuiteId === null) {
+        throw new Error("currentSuiteId is null - something went wrong");
+      }
+      return this.suiteUris.get(this.currentSuiteId) || "";
+    };
+
+    tinybench.Task.prototype.run = async function () {
+      const { fn } = this as { fn: Fn };
+      const suiteUri = getSuiteUri();
+
+      function __codspeed_root_frame__() {
+        return fn();
+      }
+      (this as { fn: Fn }).fn = __codspeed_root_frame__;
+
+      InstrumentHooks.startBenchmark();
+      await originalRun.call(this);
+      InstrumentHooks.stopBenchmark();
+
+      // Look up the URI by task name
+      const uri = `${suiteUri}::${this.name}`;
+      InstrumentHooks.setExecutedBenchmark(process.pid, uri);
+
+      return this;
+    };
+
+    return tinybench;
+  }
+
+  // Allow tinybench to retrieve the path to the currently running suite
+  async onTaskUpdate(
+    _: RunnerTaskResultPack[],
+    events: RunnerTaskEventPack[]
+  ): Promise<void> {
+    events.map((event) => {
+      const [id, eventName] = event;
+
+      if (eventName === "suite-prepare") {
+        this.currentSuiteId = id;
+      }
+    });
+  }
+}
+
+export default WalltimeRunner;