Skip to content

Commit 3e49eb1

Browse files
--wip-- [skip ci]
1 parent 3cfa4e9 commit 3e49eb1

File tree

5 files changed

+244
-16
lines changed

5 files changed

+244
-16
lines changed

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

Lines changed: 14 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -1,5 +1,5 @@
11
import { fromPartial } from "@total-typescript/shoehorn";
2-
import { describe, expect, it, vi } from "vitest";
2+
import { afterAll, beforeAll, describe, expect, it, vi } from "vitest";
33
import codspeedPlugin from "../index";
44

55
const coreMocks = vi.hoisted(() => {
@@ -25,6 +25,18 @@ vi.mock("@codspeed/core", async (importOriginal) => {
2525
console.warn = vi.fn();
2626

2727
describe("codSpeedPlugin", () => {
28+
beforeAll(() => {
29+
// Set environment variables to trigger instrumented mode
30+
process.env.CODSPEED_ENV = "1";
31+
process.env.CODSPEED_RUNNER_MODE = "instrumentation";
32+
});
33+
34+
afterAll(() => {
35+
// Clean up environment variables
36+
delete process.env.CODSPEED_ENV;
37+
delete process.env.CODSPEED_RUNNER_MODE;
38+
});
39+
2840
it("should have a name", async () => {
2941
expect(resolvedCodSpeedPlugin.name).toBe("codspeed:vitest");
3042
});
@@ -96,7 +108,7 @@ describe("codSpeedPlugin", () => {
96108
],
97109
},
98110
},
99-
runner: expect.stringContaining("packages/vitest-plugin/src/runner.ts"),
111+
runner: expect.stringContaining("packages/vitest-plugin/src/instrumented.ts"),
100112
},
101113
});
102114
});

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

Lines changed: 9 additions & 9 deletions
Original file line numberDiff line numberDiff line change
@@ -1,7 +1,7 @@
11
import { fromPartial } from "@total-typescript/shoehorn";
22
import { describe, expect, it, Suite, vi } from "vitest";
33
import { getBenchFn } from "vitest/suite";
4-
import CodSpeedRunner from "../runner";
4+
import { InstrumentedRunner as CodSpeedRunner } from "../instrumented";
55

66
const coreMocks = vi.hoisted(() => {
77
return {
@@ -53,24 +53,24 @@ describe("CodSpeedRunner", () => {
5353
// setup
5454
expect(coreMocks.setupCore).toHaveBeenCalledTimes(1);
5555
expect(console.log).toHaveBeenCalledWith(
56-
"[CodSpeed] running suite packages/vitest-plugin/src/__tests__/runner.test.ts"
56+
"[CodSpeed] running suite packages/vitest-plugin/src/__tests__/instrumented.test.ts"
5757
);
5858

5959
// run
6060
expect(coreMocks.mongoMeasurement.start).toHaveBeenCalledWith(
61-
"packages/vitest-plugin/src/__tests__/runner.test.ts::test bench"
61+
"packages/vitest-plugin/src/__tests__/instrumented.test.ts::test bench"
6262
);
6363
expect(coreMocks.Measurement.startInstrumentation).toHaveBeenCalledTimes(1);
6464
expect(benchFn).toHaveBeenCalledTimes(8);
6565
expect(coreMocks.Measurement.stopInstrumentation).toHaveBeenCalledTimes(1);
6666
expect(coreMocks.mongoMeasurement.stop).toHaveBeenCalledTimes(1);
6767
expect(console.log).toHaveBeenCalledWith(
68-
"[CodSpeed] packages/vitest-plugin/src/__tests__/runner.test.ts::test bench done"
68+
"[CodSpeed] packages/vitest-plugin/src/__tests__/instrumented.test.ts::test bench done"
6969
);
7070

7171
// teardown
7272
expect(console.log).toHaveBeenCalledWith(
73-
"[CodSpeed] running suite packages/vitest-plugin/src/__tests__/runner.test.ts done"
73+
"[CodSpeed] running suite packages/vitest-plugin/src/__tests__/instrumented.test.ts done"
7474
);
7575
expect(coreMocks.teardownCore).toHaveBeenCalledTimes(1);
7676
});
@@ -107,24 +107,24 @@ describe("CodSpeedRunner", () => {
107107
// setup
108108
expect(coreMocks.setupCore).toHaveBeenCalledTimes(1);
109109
expect(console.log).toHaveBeenCalledWith(
110-
"[CodSpeed] running suite packages/vitest-plugin/src/__tests__/runner.test.ts"
110+
"[CodSpeed] running suite packages/vitest-plugin/src/__tests__/instrumented.test.ts"
111111
);
112112

113113
// run
114114
expect(coreMocks.mongoMeasurement.start).toHaveBeenCalledWith(
115-
"packages/vitest-plugin/src/__tests__/runner.test.ts::nested suite::test bench"
115+
"packages/vitest-plugin/src/__tests__/instrumented.test.ts::nested suite::test bench"
116116
);
117117
expect(coreMocks.Measurement.startInstrumentation).toHaveBeenCalledTimes(1);
118118
expect(benchFn).toHaveBeenCalledTimes(8);
119119
expect(coreMocks.Measurement.stopInstrumentation).toHaveBeenCalledTimes(1);
120120
expect(coreMocks.mongoMeasurement.stop).toHaveBeenCalledTimes(1);
121121
expect(console.log).toHaveBeenCalledWith(
122-
"[CodSpeed] packages/vitest-plugin/src/__tests__/runner.test.ts::nested suite::test bench done"
122+
"[CodSpeed] packages/vitest-plugin/src/__tests__/instrumented.test.ts::nested suite::test bench done"
123123
);
124124

125125
// teardown
126126
expect(console.log).toHaveBeenCalledWith(
127-
"[CodSpeed] running suite packages/vitest-plugin/src/__tests__/runner.test.ts done"
127+
"[CodSpeed] running suite packages/vitest-plugin/src/__tests__/instrumented.test.ts done"
128128
);
129129
expect(coreMocks.teardownCore).toHaveBeenCalledTimes(1);
130130
});

packages/vitest-plugin/src/index.ts

Lines changed: 25 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -19,6 +19,22 @@ function getCodSpeedFileFromName(name: string) {
1919
return join(__dirname, `${name}.${fileExtension}`);
2020
}
2121

22+
function getRunnerFile(): string | undefined {
23+
// If CODSPEED_ENV is not set, use default vitest behavior
24+
if (!process.env.CODSPEED_ENV) {
25+
return undefined;
26+
}
27+
28+
const runnerMode = process.env.CODSPEED_RUNNER_MODE;
29+
30+
if (runnerMode === "instrumentation") {
31+
return getCodSpeedFileFromName("instrumented");
32+
} else {
33+
// Default to walltime mode (includes "walltime" mode and unrecognized modes)
34+
return getCodSpeedFileFromName("walltime");
35+
}
36+
}
37+
2238
export default function codspeedPlugin(): Plugin {
2339
return {
2440
name: "codspeed:vitest",
@@ -33,18 +49,25 @@ export default function codspeedPlugin(): Plugin {
3349
},
3450
enforce: "post",
3551
config(): UserConfig {
36-
return {
52+
const runnerFile = getRunnerFile();
53+
const config: UserConfig = {
3754
test: {
3855
pool: "forks",
3956
poolOptions: {
4057
forks: {
4158
execArgv: getV8Flags(),
4259
},
4360
},
44-
runner: getCodSpeedFileFromName("runner"),
4561
globalSetup: [getCodSpeedFileFromName("globalSetup")],
4662
},
4763
};
64+
65+
// Only set custom runner when CODSPEED_ENV is set
66+
if (runnerFile) {
67+
config.test!.runner = runnerFile;
68+
}
69+
70+
return config;
4871
},
4972
};
5073
}

packages/vitest-plugin/src/runner.ts

Lines changed: 0 additions & 3 deletions
This file was deleted.
Lines changed: 196 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,196 @@
1+
import {
2+
calculateQuantiles,
3+
logDebug,
4+
mongoMeasurement,
5+
setupCore,
6+
teardownCore,
7+
writeWalltimeResults,
8+
type Benchmark as CoreBenchmark,
9+
type BenchmarkStats,
10+
} from "@codspeed/core";
11+
import { Benchmark, Suite } from "vitest";
12+
import { NodeBenchmarkRunner } from "vitest/runners";
13+
import { getBenchFn } from "vitest/suite";
14+
import { callSuiteHook, patchRootSuiteWithFullFilePath } from "./common";
15+
16+
declare const __VERSION__: string;
17+
18+
const currentFileName =
19+
typeof __filename === "string"
20+
? __filename
21+
: new URL("walltime.mjs", import.meta.url).pathname;
22+
23+
/**
24+
* @deprecated
25+
* TODO: try to use something like `updateTask` from `@vitest/runner` instead to use the output
26+
* of vitest instead console.log but at the moment, `updateTask` is not exposed
27+
*/
28+
function logCodSpeed(message: string) {
29+
console.log(`[CodSpeed] ${message}`);
30+
}
31+
32+
async function runWalltimeBench(
33+
benchmark: Benchmark,
34+
currentSuiteName: string
35+
): Promise<CoreBenchmark | null> {
36+
const uri = `${currentSuiteName}::${benchmark.name}`;
37+
const fn = getBenchFn(benchmark);
38+
39+
// Run the benchmark function multiple times to collect timing data
40+
const samples: number[] = [];
41+
const iterations = 100; // Similar to tinybench default
42+
43+
// Warmup phase
44+
await callSuiteHook(benchmark.suite, benchmark, "beforeEach");
45+
for (let i = 0; i < 10; i++) {
46+
await fn();
47+
}
48+
await callSuiteHook(benchmark.suite, benchmark, "afterEach");
49+
50+
// Measurement phase
51+
await mongoMeasurement.start(uri);
52+
53+
for (let i = 0; i < iterations; i++) {
54+
await callSuiteHook(benchmark.suite, benchmark, "beforeEach");
55+
56+
const startTime = process.hrtime.bigint();
57+
await fn();
58+
const endTime = process.hrtime.bigint();
59+
60+
await callSuiteHook(benchmark.suite, benchmark, "afterEach");
61+
62+
// Convert nanoseconds to milliseconds for consistency with tinybench
63+
const durationMs = Number(endTime - startTime) / 1_000_000;
64+
samples.push(durationMs);
65+
}
66+
67+
await mongoMeasurement.stop(uri);
68+
69+
if (samples.length === 0) {
70+
console.warn(` ⚠ No timing data available for ${uri}`);
71+
return null;
72+
}
73+
74+
// Convert to BenchmarkStats format
75+
const stats = convertSamplesToBenchmarkStats(samples);
76+
77+
const coreBenchmark: CoreBenchmark = {
78+
name: benchmark.name,
79+
uri,
80+
config: {
81+
warmup_time_ns: null, // Not configurable in Vitest
82+
min_round_time_ns: null, // Not configurable in Vitest
83+
},
84+
stats,
85+
};
86+
87+
console.log(` ✔ Collected walltime data for ${uri}`);
88+
return coreBenchmark;
89+
}
90+
91+
function convertSamplesToBenchmarkStats(samples: number[]): BenchmarkStats {
92+
// All samples are in milliseconds, convert to nanoseconds
93+
const ms_to_ns = (ms: number) => ms * 1_000_000;
94+
95+
const sortedTimesNs = samples.map(ms_to_ns).sort((a, b) => a - b);
96+
const meanMs = samples.reduce((a, b) => a + b, 0) / samples.length;
97+
const meanNs = ms_to_ns(meanMs);
98+
99+
// Calculate standard deviation
100+
const variance =
101+
samples.reduce((acc, val) => acc + Math.pow(val - meanMs, 2), 0) /
102+
samples.length;
103+
const stdevMs = Math.sqrt(variance);
104+
const stdevNs = ms_to_ns(stdevMs);
105+
106+
const { q1_ns, q3_ns, median_ns, iqr_outlier_rounds, stdev_outlier_rounds } =
107+
calculateQuantiles({ meanNs, stdevNs, sortedTimesNs });
108+
109+
return {
110+
min_ns: sortedTimesNs[0],
111+
max_ns: sortedTimesNs[sortedTimesNs.length - 1],
112+
mean_ns: meanNs,
113+
stdev_ns: stdevNs,
114+
q1_ns,
115+
median_ns,
116+
q3_ns,
117+
total_time: samples.reduce((a, b) => a + b, 0) / 1000, // convert from ms to seconds
118+
iter_per_round: sortedTimesNs.length,
119+
rounds: 1, // Single round like tinybench
120+
iqr_outlier_rounds,
121+
stdev_outlier_rounds,
122+
warmup_iters: 10, // Fixed warmup iterations
123+
};
124+
}
125+
126+
async function runWalltimeBenchmarkSuite(
127+
suite: Suite,
128+
parentSuiteName?: string
129+
): Promise<CoreBenchmark[]> {
130+
const currentSuiteName = parentSuiteName
131+
? parentSuiteName + "::" + suite.name
132+
: suite.name;
133+
134+
const benchmarks: CoreBenchmark[] = [];
135+
136+
// do not call `beforeAll` if we are in the root suite, since it is already called by vitest
137+
if (parentSuiteName !== undefined) {
138+
await callSuiteHook(suite, suite, "beforeAll");
139+
}
140+
141+
for (const task of suite.tasks) {
142+
if (task.mode !== "run") continue;
143+
144+
if (task.meta?.benchmark) {
145+
const benchmark = await runWalltimeBench(
146+
task as Benchmark,
147+
currentSuiteName
148+
);
149+
if (benchmark) {
150+
benchmarks.push(benchmark);
151+
}
152+
} else if (task.type === "suite") {
153+
const nestedBenchmarks = await runWalltimeBenchmarkSuite(
154+
task,
155+
currentSuiteName
156+
);
157+
benchmarks.push(...nestedBenchmarks);
158+
}
159+
}
160+
161+
// do not call `afterAll` if we are in the root suite, since it is already called by vitest
162+
if (parentSuiteName !== undefined) {
163+
await callSuiteHook(suite, suite, "afterAll");
164+
}
165+
166+
return benchmarks;
167+
}
168+
169+
export class WalltimeRunner extends NodeBenchmarkRunner {
170+
async runSuite(suite: Suite): Promise<void> {
171+
logDebug(`PROCESS PID: ${process.pid} in ${currentFileName}`);
172+
setupCore();
173+
174+
patchRootSuiteWithFullFilePath(suite);
175+
176+
console.log(
177+
`[CodSpeed] running with @codspeed/vitest-plugin v${__VERSION__} (walltime mode)`
178+
);
179+
logCodSpeed(`running suite ${suite.name}`);
180+
181+
const benchmarks = await runWalltimeBenchmarkSuite(suite);
182+
183+
// Write results to JSON file using core function
184+
if (benchmarks.length > 0) {
185+
writeWalltimeResults(benchmarks);
186+
}
187+
188+
console.log(
189+
`[CodSpeed] Done collecting walltime data for ${benchmarks.length} benches.`
190+
);
191+
logCodSpeed(`running suite ${suite.name} done`);
192+
193+
teardownCore();
194+
}
195+
}
196+

0 commit comments

Comments
 (0)