feat(vitest-plugin): add walltime support

Author: GuillaumeLagrange

packages/vitest-plugin/benches/timing.bench.ts

@@ -0,0 +1,19 @@
+import { bench, describe } from "vitest";
+
+const sleep = (ms: number): Promise<void> => {
+  return new Promise((resolve) => setTimeout(resolve, ms));
+};
+
+describe("timing tests", () => {
+  bench("wait 1ms", async () => {
+    await sleep(1);
+  });
+
+  bench("wait 500ms", async () => {
+    await sleep(1);
+  });
+
+  bench("wait 1sec", async () => {
+    await sleep(1);
+  });
+});

packages/vitest-plugin/moon.yml

@@ -6,6 +6,8 @@ tasks:
     local: true
     options:
       cache: false
+    deps:
+      - build
 
   test:
     command: vitest --run

packages/vitest-plugin/rollup.config.ts

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

packages/vitest-plugin/src/__tests__/index.test.ts

@@ -1,5 +1,5 @@
 import { fromPartial } from "@total-typescript/shoehorn";
-import { describe, expect, it, vi } from "vitest";
+import { afterAll, beforeAll, describe, expect, it, vi } from "vitest";
 import codspeedPlugin from "../index";
 
 const coreMocks = vi.hoisted(() => {
@@ -25,6 +25,18 @@ vi.mock("@codspeed/core", async (importOriginal) => {
 console.warn = vi.fn();
 
 describe("codSpeedPlugin", () => {
+  beforeAll(() => {
+    // Set environment variables to trigger instrumented mode
+    process.env.CODSPEED_ENV = "1";
+    process.env.CODSPEED_RUNNER_MODE = "instrumentation";
+  });
+
+  afterAll(() => {
+    // Clean up environment variables
+    delete process.env.CODSPEED_ENV;
+    delete process.env.CODSPEED_RUNNER_MODE;
+  });
+
   it("should have a name", async () => {
     expect(resolvedCodSpeedPlugin.name).toBe("codspeed:vitest");
   });

packages/vitest-plugin/src/__tests__/runner.test.ts renamed to packages/vitest-plugin/src/__tests__/instrumented.test.ts

@@ -1,7 +1,7 @@
 import { fromPartial } from "@total-typescript/shoehorn";
 import { describe, expect, it, Suite, vi } from "vitest";
 import { getBenchFn } from "vitest/suite";
-import CodSpeedRunner from "../runner";
+import { InstrumentedRunner as CodSpeedRunner } from "../instrumented";
 
 const coreMocks = vi.hoisted(() => {
   return {
@@ -53,24 +53,24 @@ describe("CodSpeedRunner", () => {
     // setup
     expect(coreMocks.setupCore).toHaveBeenCalledTimes(1);
     expect(console.log).toHaveBeenCalledWith(
-      "[CodSpeed] running suite packages/vitest-plugin/src/__tests__/runner.test.ts"
+      "[CodSpeed] running suite packages/vitest-plugin/src/__tests__/instrumented.test.ts"
     );
 
     // run
     expect(coreMocks.mongoMeasurement.start).toHaveBeenCalledWith(
-      "packages/vitest-plugin/src/__tests__/runner.test.ts::test bench"
+      "packages/vitest-plugin/src/__tests__/instrumented.test.ts::test bench"
     );
     expect(coreMocks.Measurement.startInstrumentation).toHaveBeenCalledTimes(1);
     expect(benchFn).toHaveBeenCalledTimes(8);
     expect(coreMocks.Measurement.stopInstrumentation).toHaveBeenCalledTimes(1);
     expect(coreMocks.mongoMeasurement.stop).toHaveBeenCalledTimes(1);
     expect(console.log).toHaveBeenCalledWith(
-      "[CodSpeed] packages/vitest-plugin/src/__tests__/runner.test.ts::test bench done"
+      "[CodSpeed] packages/vitest-plugin/src/__tests__/instrumented.test.ts::test bench done"
     );
 
     // teardown
     expect(console.log).toHaveBeenCalledWith(
-      "[CodSpeed] running suite packages/vitest-plugin/src/__tests__/runner.test.ts done"
+      "[CodSpeed] running suite packages/vitest-plugin/src/__tests__/instrumented.test.ts done"
     );
     expect(coreMocks.teardownCore).toHaveBeenCalledTimes(1);
   });
@@ -107,24 +107,24 @@ describe("CodSpeedRunner", () => {
     // setup
     expect(coreMocks.setupCore).toHaveBeenCalledTimes(1);
     expect(console.log).toHaveBeenCalledWith(
-      "[CodSpeed] running suite packages/vitest-plugin/src/__tests__/runner.test.ts"
+      "[CodSpeed] running suite packages/vitest-plugin/src/__tests__/instrumented.test.ts"
     );
 
     // run
     expect(coreMocks.mongoMeasurement.start).toHaveBeenCalledWith(
-      "packages/vitest-plugin/src/__tests__/runner.test.ts::nested suite::test bench"
+      "packages/vitest-plugin/src/__tests__/instrumented.test.ts::nested suite::test bench"
     );
     expect(coreMocks.Measurement.startInstrumentation).toHaveBeenCalledTimes(1);
     expect(benchFn).toHaveBeenCalledTimes(8);
     expect(coreMocks.Measurement.stopInstrumentation).toHaveBeenCalledTimes(1);
     expect(coreMocks.mongoMeasurement.stop).toHaveBeenCalledTimes(1);
     expect(console.log).toHaveBeenCalledWith(
-      "[CodSpeed] packages/vitest-plugin/src/__tests__/runner.test.ts::nested suite::test bench done"
+      "[CodSpeed] packages/vitest-plugin/src/__tests__/instrumented.test.ts::nested suite::test bench done"
     );
 
     // teardown
     expect(console.log).toHaveBeenCalledWith(
-      "[CodSpeed] running suite packages/vitest-plugin/src/__tests__/runner.test.ts done"
+      "[CodSpeed] running suite packages/vitest-plugin/src/__tests__/instrumented.test.ts done"
     );
     expect(coreMocks.teardownCore).toHaveBeenCalledTimes(1);
   });

packages/vitest-plugin/src/common.ts

@@ -36,4 +36,4 @@ export function patchRootSuiteWithFullFilePath(suite: Suite) {
     throw new Error("Could not find a git repository");
   }
   suite.name = path.relative(gitDir, suite.filepath);
-}
+}

packages/vitest-plugin/src/index.ts

@@ -1,4 +1,5 @@
 import {
+  getCodspeedRunnerMode,
   getV8Flags,
   Measurement,
   mongoMeasurement,
@@ -19,32 +20,51 @@ function getCodSpeedFileFromName(name: string) {
   return join(__dirname, `${name}.${fileExtension}`);
 }
 
+function getRunnerFile(): string | undefined {
+  const codspeedRunnerMode = getCodspeedRunnerMode();
+  if (codspeedRunnerMode === "disabled") {
+    return undefined;
+  }
+
+  return getCodSpeedFileFromName(codspeedRunnerMode);
+}
+
 export default function codspeedPlugin(): Plugin {
   return {
     name: "codspeed:vitest",
     apply(_, { mode }) {
       if (mode !== "benchmark") {
         return false;
       }
-      if (!Measurement.isInstrumented()) {
+      if (
+        getCodspeedRunnerMode() == "instrumented" &&
+        !Measurement.isInstrumented()
+      ) {
         console.warn("[CodSpeed] bench detected but no instrumentation found");
       }
       return true;
     },
     enforce: "post",
     config(): UserConfig {
-      return {
+      const runnerFile = getRunnerFile();
+      const config: UserConfig = {
         test: {
           pool: "forks",
           poolOptions: {
             forks: {
               execArgv: getV8Flags(),
             },
           },
-          runner: getCodSpeedFileFromName("runner"),
           globalSetup: [getCodSpeedFileFromName("globalSetup")],
         },
       };
+
+      // Only set custom runner when CODSPEED_ENV is set
+      if (runnerFile) {
+        config.test!.runner = runnerFile;
+      }
+
+      return config;
     },
   };
 }

packages/vitest-plugin/src/instrumented.ts

@@ -106,3 +106,4 @@ export class InstrumentedRunner extends NodeBenchmarkRunner {
   }
 }
 
+export default InstrumentedRunner;

packages/vitest-plugin/src/walltime.ts

@@ -0,0 +1,178 @@
+import {
+  calculateQuantiles,
+  mongoMeasurement,
+  writeWalltimeResults,
+  type Benchmark,
+  type BenchmarkStats,
+} from "@codspeed/core";
+import { BenchTaskResult, Suite, Task } from "vitest";
+import { NodeBenchmarkRunner } from "vitest/runners";
+import { patchRootSuiteWithFullFilePath } from "./common";
+
+declare const __VERSION__: string;
+
+// Vitest's task result structure extends Tinybench with additional properties
+interface VitestTaskResult {
+  state: "pass" | "fail" | "skip" | "todo";
+  startTime: number;
+  duration?: number;
+  benchmark?: BenchTaskResult;
+}
+
+/**
+ * WalltimeRunner uses Vitest's default benchmark execution
+ * and extracts results from the suite after completion
+ */
+export class WalltimeRunner extends NodeBenchmarkRunner {
+  async runSuite(suite: Suite): Promise<void> {
+    patchRootSuiteWithFullFilePath(suite);
+
+    console.log(
+      `[CodSpeed] running with @codspeed/vitest-plugin v${__VERSION__} (walltime mode)`
+    );
+
+    // Let Vitest's default benchmark runner handle execution
+    await super.runSuite(suite);
+
+    // Extract benchmark results from the completed suite
+    const benchmarks = await this.extractBenchmarkResults(suite);
+
+    // Write results to JSON file using core function
+    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 async extractBenchmarkResults(
+    suite: Suite,
+    parentPath = ""
+  ): Promise<Benchmark[]> {
+    const benchmarks: Benchmark[] = [];
+    const currentPath = parentPath
+      ? `${parentPath}::${suite.name}`
+      : suite.name;
+
+    for (const task of suite.tasks) {
+      if (task.meta?.benchmark && task.result?.state === "pass") {
+        const benchmark = await this.processBenchmarkTask(task, currentPath);
+        if (benchmark) {
+          benchmarks.push(benchmark);
+        }
+      } else if (task.type === "suite") {
+        const nestedBenchmarks = await this.extractBenchmarkResults(
+          task,
+          currentPath
+        );
+        benchmarks.push(...nestedBenchmarks);
+      }
+    }
+
+    return benchmarks;
+  }
+
+  private async processBenchmarkTask(
+    task: Task,
+    suitePath: string
+  ): Promise<Benchmark | null> {
+    const uri = `${suitePath}::${task.name}`;
+
+    // Notify mongoMeasurement about the benchmark
+    await mongoMeasurement.start(uri);
+    await mongoMeasurement.stop(uri);
+
+    const result = task.result;
+    if (!result) {
+      console.warn(`    ⚠ No result data available for ${uri}`);
+      return null;
+    }
+
+    try {
+      const stats = this.convertVitestResultToBenchmarkStats(
+        result as VitestTaskResult
+      );
+
+      if (stats === null) {
+        console.log(`    ✔ No walltime data to collect for ${uri}`);
+        return null;
+      }
+
+      const coreBenchmark: Benchmark = {
+        name: task.name,
+        uri,
+        config: {
+          warmup_time_ns: null, // Vitest doesn't expose this in task.result
+          min_round_time_ns: null, // Vitest doesn't expose this in task.result
+        },
+        stats,
+      };
+
+      console.log(`    ✔ Collected walltime data for ${uri}`);
+      return coreBenchmark;
+    } catch (error) {
+      console.warn(
+        `    ⚠ Failed to process benchmark result for ${uri}:`,
+        error
+      );
+      return null;
+    }
+  }
+
+  private convertVitestResultToBenchmarkStats(
+    result: VitestTaskResult
+  ): BenchmarkStats | null {
+    const benchmark = result.benchmark;
+
+    if (!benchmark) {
+      throw new Error("No benchmark data available in result");
+    }
+
+    // All tinybench times are in milliseconds, convert to nanoseconds
+    const ms_to_ns = (ms: number) => ms * 1_000_000;
+
+    const { totalTime, min, max, mean, sd, samples } = benchmark;
+
+    // Get individual sample times in nanoseconds and sort them
+    const sortedTimesNs = samples.map(ms_to_ns).sort((a, b) => a - b);
+    const meanNs = ms_to_ns(mean);
+    const stdevNs = ms_to_ns(sd);
+
+    if (sortedTimesNs.length == 0) {
+      // Sometimes the benchmarks can be completely optimized out and not even run, but its beforeEach and afterEach hooks are still executed, and the task is still considered a success.
+      // This is the case for the hooks.bench.ts example in this package
+      return null;
+    }
+
+    const {
+      q1_ns,
+      q3_ns,
+      median_ns,
+      iqr_outlier_rounds,
+      stdev_outlier_rounds,
+    } = calculateQuantiles({ meanNs, stdevNs, sortedTimesNs });
+
+    return {
+      min_ns: ms_to_ns(min),
+      max_ns: ms_to_ns(max),
+      mean_ns: meanNs,
+      stdev_ns: stdevNs,
+      q1_ns,
+      median_ns,
+      q3_ns,
+      total_time: totalTime / 1_000, // convert from ms to seconds
+      iter_per_round: sortedTimesNs.length,
+      rounds: 1, // Tinybench only runs one round
+      iqr_outlier_rounds,
+      stdev_outlier_rounds,
+      warmup_iters: 0, // TODO: get warmup iters here
+    };
+  }
+}
+
+export default WalltimeRunner;