feat(vitest-plugin): allow perf profiling for vitest plugin

Author: GuillaumeLagrange

packages/vitest-plugin/package.json

@@ -28,7 +28,8 @@
     "bench": "vitest bench"
   },
   "dependencies": {
-    "@codspeed/core": "workspace:^5.0.0"
+    "@codspeed/core": "workspace:^5.0.0",
+    "tinybench": "^2.9.0"
   },
   "peerDependencies": {
     "vite": "^4.2.0 || ^5.0.0 || ^6.0.0 || ^7.0.0",

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

@@ -1,7 +1,15 @@
-import { setupCore, writeWalltimeResults } from "@codspeed/core";
+import {
+  InstrumentHooks,
+  setupCore,
+  writeWalltimeResults,
+} from "@codspeed/core";
+import { Fn } from "tinybench";
 import { type RunnerTestSuite } from "vitest";
 import { NodeBenchmarkRunner } from "vitest/runners";
-import { patchRootSuiteWithFullFilePath } from "../common";
+import {
+  isVitestTaskBenchmark,
+  patchRootSuiteWithFullFilePath,
+} from "../common";
 import { extractBenchmarkResults } from "./utils";
 
 /**
@@ -15,6 +23,8 @@ export class WalltimeRunner extends NodeBenchmarkRunner {
   async runSuite(suite: RunnerTestSuite): Promise<void> {
     patchRootSuiteWithFullFilePath(suite);
 
+    this.populateBenchmarkUris(suite);
+
     setupCore();
 
     await super.runSuite(suite);
@@ -32,6 +42,53 @@ export class WalltimeRunner extends NodeBenchmarkRunner {
       );
     }
   }
+
+  private populateBenchmarkUris(suite: RunnerTestSuite, parentPath = ""): void {
+    const currentPath = parentPath
+      ? `${parentPath}::${suite.name}`
+      : suite.name;
+
+    for (const task of suite.tasks) {
+      if (isVitestTaskBenchmark(task)) {
+        const uri = `${currentPath}::${task.name}`;
+        this.benchmarkUris.set(task.name, uri);
+      } else if (task.type === "suite") {
+        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 runner = this;
+    tinybench.Task.prototype.run = async function () {
+      const { fn } = this as { fn: Fn };
+
+      function __codspeed_root_frame__() {
+        return fn();
+      }
+      (this as any).fn = __codspeed_root_frame__;
+
+      InstrumentHooks.startBenchmark();
+      await originalRun.call(this);
+      InstrumentHooks.stopBenchmark();
+
+      // Look up the URI by task name
+      const uri = runner.benchmarkUris.get(this.name) || this.name;
+      InstrumentHooks.setExecutedBenchmark(process.pid, uri);
+
+      return this;
+    };
+
+    return tinybench;
+  }
 }
 
 export default WalltimeRunner;