feat(vitest-plugin): add 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,5 +1,14 @@
-import { setupCore, writeWalltimeResults } from "@codspeed/core";
-import { type RunnerTestSuite } from "vitest";
+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";
@@ -10,10 +19,13 @@ import { extractBenchmarkResults } from "./utils";
  */
 export class WalltimeRunner extends NodeBenchmarkRunner {
   private isTinybenchHookedWithCodspeed = false;
-  private benchmarkUris = new Map<string, string>();
+  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();
 
@@ -32,6 +44,73 @@ 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 (task.type === "suite") {
+        this.populateBenchmarkUris(task, currentPath);
+        this.suiteUris.set(task.id, 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;