Skip to content

Commit eec7ec7

Browse files
feat(vitest-plugin): use the vitest 5 benchmark provider API
Vitest 5 exposes a `benchmark.provider` option that hands a plugin the raw benchmark functions and their hooks directly. Wire the V5 backend to register a provider through `test.benchmark.provider` and add the provider itself, which runs analysis and walltime modes off the registrations it receives. Drop the previous V5 seam that patched the static `TestRunner.runBenchmarks` and `Bench.prototype.add` from a setup file, along with the `WeakMap` fn-capture it needed to work around tinybench v6's private `fn` field. The provider gets the fn and options as data, so none of that interception remains. The legacy (Vitest 3/4) runner path is untouched. Point the rollup build and the config-injection test at the new provider entry. Refs COD-2931
1 parent c3a356b commit eec7ec7

8 files changed

Lines changed: 284 additions & 357 deletions

File tree

packages/vitest-plugin/rollup.config.ts

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -35,8 +35,8 @@ export default defineConfig([
3535
external: ["@codspeed/core", /^vitest/],
3636
},
3737
{
38-
input: "src/v5/setup.ts",
39-
output: { file: "dist/v5/setup.mjs", format: "es" },
38+
input: "src/v5/provider.ts",
39+
output: { file: "dist/v5/provider.mjs", format: "es" },
4040
plugins: jsPlugins(pkg.version),
4141
external: ["@codspeed/core", /^vitest/, "tinybench"],
4242
},

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

Lines changed: 12 additions & 8 deletions
Original file line numberDiff line numberDiff line change
@@ -209,7 +209,7 @@ describe("codSpeedPlugin", () => {
209209
fsMocks.setMockVersion("4.0.18");
210210
});
211211

212-
it("should inject the v5 setup file (not a runner) when benchmarks are enabled", () => {
212+
it("should wire the v5 benchmark provider (not a runner or setup file) when benchmarks are enabled", () => {
213213
fsMocks.setMockVersion("5.0.0-beta.5");
214214
const v5Plugin = codspeedPlugin();
215215
const config = v5Plugin.config;
@@ -233,15 +233,19 @@ describe("codSpeedPlugin", () => {
233233
],
234234
pool: "forks",
235235
execArgv: EXPECTED_EXEC_ARGV,
236-
setupFiles: [
237-
expect.stringContaining("packages/vitest-plugin/src/v5/setup.ts"),
238-
],
236+
benchmark: {
237+
provider: expect.stringContaining(
238+
"packages/vitest-plugin/src/v5/provider.ts",
239+
),
240+
},
239241
},
240242
});
241-
// The v5 path must not set a custom runner.
242-
expect(
243-
(result as { test?: { runner?: unknown } })?.test?.runner,
244-
).toBeUndefined();
243+
// The v5 path must not set a custom runner or a setup file.
244+
const test = (
245+
result as { test?: { runner?: unknown; setupFiles?: unknown } }
246+
)?.test;
247+
expect(test?.runner).toBeUndefined();
248+
expect(test?.setupFiles).toBeUndefined();
245249
fsMocks.setMockVersion("4.0.18");
246250
});
247251
});

packages/vitest-plugin/src/index.ts

Lines changed: 0 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -1,5 +1,4 @@
11
import {
2-
getCodspeedRunnerMode,
32
getInstrumentMode,
43
getV8Flags,
54
InstrumentHooks,
@@ -56,9 +55,6 @@ export default function codspeedPlugin(): Plugin {
5655
pool: "forks",
5756
globalSetup: [resolveFile("globalSetup")],
5857
...backend.getBenchmarkTestConfig(getV8Flags(), resolveFile),
59-
...(getCodspeedRunnerMode() === "walltime" && {
60-
benchmark: backend.getWalltimeBenchmarkConfig(),
61-
}),
6258
},
6359
};
6460

packages/vitest-plugin/src/instrument.ts

Lines changed: 3 additions & 74 deletions
Original file line numberDiff line numberDiff line change
@@ -52,7 +52,6 @@ export interface TinybenchBench {
5252
teardown: TinybenchHook;
5353
}
5454

55-
5655
/**
5756
* The tinybench statistics shape (latency/throughput) shared across the v2 and
5857
* v6 lines. Only the fields the conversion needs are modeled.
@@ -84,78 +83,13 @@ interface InstrumentWindow {
8483
runStart: bigint | null;
8584
}
8685

87-
let isBenchAddPatched = false;
88-
8986
/**
9087
* The window bracketing the currently running task's measured loop, driven by
9188
* the setup/teardown hooks below. Tasks run strictly sequentially within a
9289
* worker, so one shared value suffices.
9390
*/
9491
const instrumentWindow: InstrumentWindow = { runStart: null };
9592

96-
// tinybench keeps a task's fn and options as `#private` fields (v6+), so we
97-
// capture them ourselves when `Bench.add` runs, keyed by bench then task name.
98-
// The analysis seam needs the raw fn to run it under its own tight window
99-
// instead of tinybench's timing loop.
100-
const capturedTasks = new WeakMap<object, Map<string, CapturedTask>>();
101-
102-
/** The minimal tinybench Bench prototype we patch to capture registrations. */
103-
interface TinybenchBenchClass {
104-
prototype: {
105-
// eslint-disable-next-line @typescript-eslint/no-explicit-any
106-
add: (...args: any[]) => unknown;
107-
};
108-
}
109-
110-
/**
111-
* Patch `Bench.prototype.add` to record each task's fn and options, keyed by
112-
* bench then task name. Idempotent, and applied to the prototype so it captures
113-
* registrations on every Bench the host constructs.
114-
*
115-
* `BenchClass` must be the exact class the host instantiates. In tinybench v6 a
116-
* task's fn is a true `#private` field — it cannot be read or replaced on the
117-
* task afterwards — so capturing (and, for walltime, root-frame-wrapping) has to
118-
* happen here, as the fn is registered.
119-
*
120-
* `registerFn` transforms the fn actually handed to tinybench: identity for
121-
* analysis (which runs the captured fn itself), or a root-frame wrap for
122-
* walltime (where tinybench drives the fn and the frame must already be baked
123-
* in).
124-
*/
125-
export function captureBenchAddOnce(
126-
BenchClass: TinybenchBenchClass,
127-
registerFn: (fn: CapturedTask["fn"]) => CapturedTask["fn"],
128-
): void {
129-
if (isBenchAddPatched) {
130-
return;
131-
}
132-
isBenchAddPatched = true;
133-
134-
const originalAdd = BenchClass.prototype.add;
135-
BenchClass.prototype.add = function (
136-
this: object,
137-
name: string,
138-
fn: CapturedTask["fn"],
139-
fnOpts?: TinybenchFnOptions,
140-
) {
141-
let byName = capturedTasks.get(this);
142-
if (!byName) {
143-
byName = new Map<string, CapturedTask>();
144-
capturedTasks.set(this, byName);
145-
}
146-
byName.set(name, { fn, fnOpts });
147-
return originalAdd.call(this, name, registerFn(fn), fnOpts);
148-
};
149-
}
150-
151-
/** Retrieve the fn/options captured for a task on a given bench, if any. */
152-
export function getCapturedTask(
153-
bench: object,
154-
taskName: string,
155-
): CapturedTask | undefined {
156-
return capturedTasks.get(bench)?.get(taskName);
157-
}
158-
15993
/** The tinybench Task prototype whose `run` the legacy seam wraps. */
16094
interface TinybenchTaskClass {
16195
// eslint-disable-next-line @typescript-eslint/no-explicit-any
@@ -193,21 +127,16 @@ export function patchTaskRunOnce(TaskClass: TinybenchTaskClass): void {
193127
}
194128

195129
/**
196-
* The root-frame wrap to hand tinybench at registration time (walltime, v5).
197-
* Post-hoc assignment to a task's `fn` is a no-op on tinybench v6 (private
198-
* field), so the frame must be baked into the registered fn instead.
130+
* The root-frame wrap to hand tinybench at registration time. Post-hoc
131+
* assignment to a task's `fn` is a no-op on tinybench v6 (private field), so the
132+
* frame must be baked into the registered fn instead.
199133
*/
200134
export function rootFrameRegisterFn(
201135
fn: CapturedTask["fn"],
202136
): CapturedTask["fn"] {
203137
return wrapWithRootFrame(() => fn());
204138
}
205139

206-
/** Identity registration: analysis runs the captured fn itself, unwrapped. */
207-
export function identityRegisterFn(fn: CapturedTask["fn"]): CapturedTask["fn"] {
208-
return fn;
209-
}
210-
211140
/**
212141
* Run one benchmark under instrumentation, matching the analysis window the
213142
* Vitest 3/4 runner uses exactly: warm the JIT with `optimizeFunction` outside

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

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -22,8 +22,8 @@ type Tinybench = typeof tinybench;
2222
/**
2323
* Lets tinybench run the benches through Vitest's default benchmark execution,
2424
* instrumenting each measured loop, then extracts the results from the suite
25-
* tree afterwards. (The v5 seam instruments the same way but reads results off
26-
* the live tinybench tasks instead — see `v5/setup.ts`.)
25+
* tree afterwards. (The v5 provider instruments the same way but reads results
26+
* off the live tinybench tasks instead — see `v5/provider.ts`.)
2727
*/
2828
export class WalltimeRunner extends NodeBenchmarkRunner {
2929
private suiteUris = new Map<string, string>();

0 commit comments

Comments
 (0)