Skip to content

Commit 8410705

Browse files
GuillaumeLagrangeGuillaumeLagrange
committed
refactor(vitest-plugin): move instrumented behavior in dedicated file
1 parent 022c470 commit 8410705

3 files changed

Lines changed: 149 additions & 139 deletions

File tree

‎packages/vitest-plugin/src/common.ts‎

Lines changed: 39 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,39 @@
1+
import { getGitDir } from "@codspeed/core";
2+
import path from "path";
3+
import { Suite, Task } from "vitest";
4+
import { getHooks } from "vitest/suite";
5+
6+
type SuiteHooks = ReturnType<typeof getHooks>;
7+
8+
function getSuiteHooks(suite: Suite, name: keyof SuiteHooks) {
9+
return getHooks(suite)?.[name] ?? [];
10+
}
11+
12+
export async function callSuiteHook<T extends keyof SuiteHooks>(
13+
suite: Suite,
14+
currentTask: Task,
15+
name: T
16+
): Promise<void> {
17+
if (name === "beforeEach" && suite?.suite) {
18+
await callSuiteHook(suite.suite, currentTask, name);
19+
}
20+
21+
const hooks = getSuiteHooks(suite, name);
22+
23+
await Promise.all(hooks.map((fn) => fn()));
24+
25+
if (name === "afterEach" && suite?.suite) {
26+
await callSuiteHook(suite.suite, currentTask, name);
27+
}
28+
}
29+
30+
export function patchRootSuiteWithFullFilePath(suite: Suite) {
31+
if (suite.filepath === undefined) {
32+
throw new Error("filepath is undefined is the root suite");
33+
}
34+
const gitDir = getGitDir(suite.filepath);
35+
if (gitDir === undefined) {
36+
throw new Error("Could not find a git repository");
37+
}
38+
suite.name = path.relative(gitDir, suite.filepath);
39+
}

‎packages/vitest-plugin/src/instrumented.ts‎

Lines changed: 108 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,108 @@
1+
import {
2+
logDebug,
3+
Measurement,
4+
mongoMeasurement,
5+
optimizeFunction,
6+
setupCore,
7+
teardownCore,
8+
} from "@codspeed/core";
9+
import { Benchmark, chai, Suite } from "vitest";
10+
import { NodeBenchmarkRunner } from "vitest/runners";
11+
import { getBenchFn } from "vitest/suite";
12+
import { callSuiteHook, patchRootSuiteWithFullFilePath } from "./common";
13+
14+
const currentFileName =
15+
typeof __filename === "string"
16+
? __filename
17+
: new URL("instrumented.mjs", import.meta.url).pathname;
18+
19+
/**
20+
* @deprecated
21+
* TODO: try to use something like `updateTask` from `@vitest/runner` instead to use the output
22+
* of vitest instead console.log but at the moment, `updateTask` is not exposed
23+
*/
24+
function logCodSpeed(message: string) {
25+
console.log(`[CodSpeed] ${message}`);
26+
}
27+
28+
async function runInstrumentedBench(
29+
benchmark: Benchmark,
30+
currentSuiteName: string
31+
) {
32+
const uri = `${currentSuiteName}::${benchmark.name}`;
33+
const fn = getBenchFn(benchmark);
34+
35+
await callSuiteHook(benchmark.suite, benchmark, "beforeEach");
36+
try {
37+
await optimizeFunction(fn);
38+
} catch (e) {
39+
// if the error is not an assertion error, we want to fail the run
40+
// we allow assertion errors because we want to be able to use `expect` in the benchmark to allow for better authoring
41+
// assertions are allowed to fail in the optimization phase since it might be linked to stateful code
42+
if (!(e instanceof chai.AssertionError)) {
43+
throw e;
44+
}
45+
}
46+
await callSuiteHook(benchmark.suite, benchmark, "afterEach");
47+
48+
await callSuiteHook(benchmark.suite, benchmark, "beforeEach");
49+
await mongoMeasurement.start(uri);
50+
global.gc?.();
51+
await (async function __codspeed_root_frame__() {
52+
Measurement.startInstrumentation();
53+
// @ts-expect-error we do not need to bind the function to an instance of tinybench's Bench
54+
await fn();
55+
Measurement.stopInstrumentation(uri);
56+
})();
57+
await mongoMeasurement.stop(uri);
58+
await callSuiteHook(benchmark.suite, benchmark, "afterEach");
59+
60+
logCodSpeed(`${uri} done`);
61+
}
62+
63+
async function runInstrumentedBenchmarkSuite(
64+
suite: Suite,
65+
parentSuiteName?: string
66+
) {
67+
const currentSuiteName = parentSuiteName
68+
? parentSuiteName + "::" + suite.name
69+
: suite.name;
70+
71+
// do not call `beforeAll` if we are in the root suite, since it is already called by vitest
72+
// see https://github.com/vitest-dev/vitest/blob/1fee63f2598edc228017f18eca325f85ee54aee0/packages/runner/src/run.ts#L293
73+
if (parentSuiteName !== undefined) {
74+
await callSuiteHook(suite, suite, "beforeAll");
75+
}
76+
77+
for (const task of suite.tasks) {
78+
if (task.mode !== "run") continue;
79+
80+
if (task.meta?.benchmark) {
81+
await runInstrumentedBench(task as Benchmark, currentSuiteName);
82+
} else if (task.type === "suite") {
83+
await runInstrumentedBenchmarkSuite(task, currentSuiteName);
84+
}
85+
}
86+
87+
// do not call `afterAll` if we are in the root suite, since it is already called by vitest
88+
// see https://github.com/vitest-dev/vitest/blob/1fee63f2598edc228017f18eca325f85ee54aee0/packages/runner/src/run.ts#L324
89+
if (parentSuiteName !== undefined) {
90+
await callSuiteHook(suite, suite, "afterAll");
91+
}
92+
}
93+
94+
export class InstrumentedRunner extends NodeBenchmarkRunner {
95+
async runSuite(suite: Suite): Promise<void> {
96+
logDebug(`PROCESS PID: ${process.pid} in ${currentFileName}`);
97+
setupCore();
98+
99+
patchRootSuiteWithFullFilePath(suite);
100+
101+
logCodSpeed(`running suite ${suite.name}`);
102+
await runInstrumentedBenchmarkSuite(suite);
103+
logCodSpeed(`running suite ${suite.name} done`);
104+
105+
teardownCore();
106+
}
107+
}
108+

‎packages/vitest-plugin/src/runner.ts‎

Lines changed: 2 additions & 139 deletions
Original file line numberDiff line numberDiff line change
@@ -1,140 +1,3 @@
1-
import {
2-
getGitDir,
3-
logDebug,
4-
Measurement,
5-
mongoMeasurement,
6-
optimizeFunction,
7-
setupCore,
8-
teardownCore,
9-
} from "@codspeed/core";
10-
import path from "path";
11-
import { Benchmark, chai, Suite, Task } from "vitest";
12-
import { NodeBenchmarkRunner } from "vitest/runners";
13-
import { getBenchFn, getHooks } from "vitest/suite";
1+
import { InstrumentedRunner } from "./instrumented";
142

15-
type SuiteHooks = ReturnType<typeof getHooks>;
16-
17-
function getSuiteHooks(suite: Suite, name: keyof SuiteHooks) {
18-
return getHooks(suite)?.[name] ?? [];
19-
}
20-
21-
export async function callSuiteHook<T extends keyof SuiteHooks>(
22-
suite: Suite,
23-
currentTask: Task,
24-
name: T
25-
): Promise<void> {
26-
if (name === "beforeEach" && suite?.suite) {
27-
await callSuiteHook(suite.suite, currentTask, name);
28-
}
29-
30-
const hooks = getSuiteHooks(suite, name);
31-
32-
await Promise.all(hooks.map((fn) => fn()));
33-
34-
if (name === "afterEach" && suite?.suite) {
35-
await callSuiteHook(suite.suite, currentTask, name);
36-
}
37-
}
38-
39-
const currentFileName =
40-
typeof __filename === "string"
41-
? __filename
42-
: new URL("runner.mjs", import.meta.url).pathname;
43-
44-
/**
45-
* @deprecated
46-
* TODO: try to use something like `updateTask` from `@vitest/runner` instead to use the output
47-
* of vitest instead console.log but at the moment, `updateTask` is not exposed
48-
*/
49-
function logCodSpeed(message: string) {
50-
console.log(`[CodSpeed] ${message}`);
51-
}
52-
53-
async function runBench(benchmark: Benchmark, currentSuiteName: string) {
54-
const uri = `${currentSuiteName}::${benchmark.name}`;
55-
const fn = getBenchFn(benchmark);
56-
57-
await callSuiteHook(benchmark.suite, benchmark, "beforeEach");
58-
try {
59-
await optimizeFunction(fn);
60-
} catch (e) {
61-
// if the error is not an assertion error, we want to fail the run
62-
// we allow assertion errors because we want to be able to use `expect` in the benchmark to allow for better authoring
63-
// assertions are allowed to fail in the optimization phase since it might be linked to stateful code
64-
if (!(e instanceof chai.AssertionError)) {
65-
throw e;
66-
}
67-
}
68-
await callSuiteHook(benchmark.suite, benchmark, "afterEach");
69-
70-
await callSuiteHook(benchmark.suite, benchmark, "beforeEach");
71-
await mongoMeasurement.start(uri);
72-
global.gc?.();
73-
await (async function __codspeed_root_frame__() {
74-
Measurement.startInstrumentation();
75-
// @ts-expect-error we do not need to bind the function to an instance of tinybench's Bench
76-
await fn();
77-
Measurement.stopInstrumentation(uri);
78-
})();
79-
await mongoMeasurement.stop(uri);
80-
await callSuiteHook(benchmark.suite, benchmark, "afterEach");
81-
82-
logCodSpeed(`${uri} done`);
83-
}
84-
85-
async function runBenchmarkSuite(suite: Suite, parentSuiteName?: string) {
86-
const currentSuiteName = parentSuiteName
87-
? parentSuiteName + "::" + suite.name
88-
: suite.name;
89-
90-
// do not call `beforeAll` if we are in the root suite, since it is already called by vitest
91-
// see https://github.com/vitest-dev/vitest/blob/1fee63f2598edc228017f18eca325f85ee54aee0/packages/runner/src/run.ts#L293
92-
if (parentSuiteName !== undefined) {
93-
await callSuiteHook(suite, suite, "beforeAll");
94-
}
95-
96-
for (const task of suite.tasks) {
97-
if (task.mode !== "run") continue;
98-
99-
if (task.meta?.benchmark) {
100-
await runBench(task as Benchmark, currentSuiteName);
101-
} else if (task.type === "suite") {
102-
await runBenchmarkSuite(task, currentSuiteName);
103-
}
104-
}
105-
106-
// do not call `afterAll` if we are in the root suite, since it is already called by vitest
107-
// see https://github.com/vitest-dev/vitest/blob/1fee63f2598edc228017f18eca325f85ee54aee0/packages/runner/src/run.ts#L324
108-
if (parentSuiteName !== undefined) {
109-
await callSuiteHook(suite, suite, "afterAll");
110-
}
111-
}
112-
113-
function patchRootSuiteWithFullFilePath(suite: Suite) {
114-
if (suite.filepath === undefined) {
115-
throw new Error("filepath is undefined is the root suite");
116-
}
117-
const gitDir = getGitDir(suite.filepath);
118-
if (gitDir === undefined) {
119-
throw new Error("Could not find a git repository");
120-
}
121-
suite.name = path.relative(gitDir, suite.filepath);
122-
}
123-
124-
class CodSpeedRunner extends NodeBenchmarkRunner {
125-
async runSuite(suite: Suite): Promise<void> {
126-
logDebug(`PROCESS PID: ${process.pid} in ${currentFileName}`);
127-
setupCore();
128-
129-
patchRootSuiteWithFullFilePath(suite);
130-
131-
logCodSpeed(`running suite ${suite.name}`);
132-
133-
await runBenchmarkSuite(suite);
134-
logCodSpeed(`running suite ${suite.name} done`);
135-
136-
teardownCore();
137-
}
138-
}
139-
140-
export default CodSpeedRunner;
3+
export default InstrumentedRunner;

0 commit comments

Comments
 (0)