Skip to content

Commit e55aee6

Browse files
GuillaumeLagrangeGuillaumeLagrange
committed
refactor(tinybench): share structure across walltime and instrumented
1 parent 7577855 commit e55aee6

File tree

5 files changed

+241
-198
lines changed

5 files changed

+241
-198
lines changed

‎packages/tinybench-plugin/src/index.ts‎

Lines changed: 0 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -3,7 +3,6 @@ import {
33
getGitDir,
44
InstrumentHooks,
55
mongoMeasurement,
6-
setupCore,
76
SetupInstrumentsRequestBody,
87
SetupInstrumentsResponse,
98
tryIntrospect,
@@ -23,7 +22,6 @@ export function withCodSpeed(bench: Bench): Bench {
2322
if (codspeedRunnerMode === "disabled") {
2423
return bench;
2524
}
26-
setupCore();
2725

2826
const rootCallingFile = getCallingFile();
2927

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

Lines changed: 37 additions & 66 deletions
Original file line numberDiff line numberDiff line change
@@ -2,18 +2,20 @@ import {
22
InstrumentHooks,
33
mongoMeasurement,
44
optimizeFunction,
5-
teardownCore,
65
} from "@codspeed/core";
76
import { Bench, Fn, FnOptions, Task } from "tinybench";
8-
import { getTaskUri } from "./uri";
7+
import { BaseBenchRunner } from "./shared";
98

10-
declare const __VERSION__: string;
9+
class InstrumentedBenchRunner extends BaseBenchRunner {
10+
protected getModeName(): string {
11+
return "instrumented mode";
12+
}
1113

12-
export function runInstrumentedBench(
13-
bench: Bench,
14-
rootCallingFile: string
15-
): void {
16-
const runTaskAsync = async (task: Task, uri: string): Promise<void> => {
14+
private taskCompletionMessage() {
15+
return InstrumentHooks.isInstrumented() ? "Measured" : "Checked";
16+
}
17+
18+
protected async runTaskAsync(task: Task, uri: string): Promise<void> {
1719
const { fnOpts, fn } = task as unknown as { fnOpts?: FnOptions; fn: Fn };
1820

1921
await fnOpts?.beforeAll?.call(task, "run");
@@ -25,79 +27,48 @@ export function runInstrumentedBench(
2527
await fnOpts?.beforeEach?.call(task, "run");
2628
await mongoMeasurement.start(uri);
2729

28-
await (async function __codspeed_root_frame__() {
29-
global.gc?.();
30-
InstrumentHooks.startBenchmark();
31-
await fn();
32-
InstrumentHooks.stopBenchmark();
33-
InstrumentHooks.setExecutedBenchmark(process.pid, uri);
34-
})();
30+
await this.wrapWithHooksAsync(
31+
this.wrapFunctionWithFrame(fn, true) as () => Promise<void>,
32+
uri
33+
);
3534

3635
await mongoMeasurement.stop(uri);
3736
await fnOpts?.afterEach?.call(task, "run");
3837
await fnOpts?.afterAll?.call(task, "run");
39-
};
4038

41-
// Sync task runner
42-
const runTaskSync = (task: Task, uri: string): void => {
39+
this.logTaskCompletion(uri, this.taskCompletionMessage());
40+
}
41+
42+
protected runTaskSync(task: Task, uri: string): void {
4343
const { fnOpts, fn } = task as unknown as { fnOpts?: FnOptions; fn: Fn };
4444

4545
fnOpts?.beforeAll?.call(task, "run");
4646
fnOpts?.beforeEach?.call(task, "run");
4747

48-
(function __codspeed_root_frame__() {
49-
global.gc?.();
50-
InstrumentHooks.startBenchmark();
51-
fn();
52-
InstrumentHooks.stopBenchmark();
53-
InstrumentHooks.setExecutedBenchmark(process.pid, uri);
54-
})();
48+
this.wrapWithHooks(
49+
this.wrapFunctionWithFrame(fn, false) as () => void,
50+
uri
51+
);
5552

5653
fnOpts?.afterEach?.call(task, "run");
5754
fnOpts?.afterAll?.call(task, "run");
58-
};
59-
60-
bench.run = async () => {
61-
logStart();
62-
63-
for (const task of bench.tasks) {
64-
const uri = getTaskUri(bench, task.name, rootCallingFile);
65-
await runTaskAsync(task, uri);
66-
logTaskCompletion(uri);
67-
}
6855

69-
return logEnd();
70-
};
56+
this.logTaskCompletion(uri, this.taskCompletionMessage());
57+
}
7158

72-
bench.runSync = () => {
73-
logStart();
59+
protected finalizeAsyncRun(): Task[] {
60+
return this.finalizeBenchRun();
61+
}
7462

75-
for (const task of bench.tasks) {
76-
const uri = getTaskUri(bench, task.name, rootCallingFile);
77-
runTaskSync(task, uri);
78-
logTaskCompletion(uri);
79-
}
80-
81-
return logEnd();
82-
};
83-
84-
const logStart = () => {
85-
console.log(
86-
`[CodSpeed] running with @codspeed/tinybench v${__VERSION__} (instrumented mode)`
87-
);
88-
};
89-
90-
const logTaskCompletion = (uri: string) => {
91-
console.log(
92-
` ✔ ${
93-
InstrumentHooks.isInstrumented() ? "Measured" : "Checked"
94-
} ${uri}`
95-
);
96-
};
63+
protected finalizeSyncRun(): Task[] {
64+
return this.finalizeBenchRun();
65+
}
66+
}
9767

98-
const logEnd = () => {
99-
teardownCore();
100-
console.log(`[CodSpeed] Done running ${bench.tasks.length} benches.`);
101-
return bench.tasks;
102-
};
68+
export function runInstrumentedBench(
69+
bench: Bench,
70+
rootCallingFile: string
71+
): void {
72+
const runner = new InstrumentedBenchRunner(bench, rootCallingFile);
73+
runner.setupBenchMethods();
10374
}

‎packages/tinybench-plugin/src/shared.ts‎

Lines changed: 105 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,105 @@
1+
import { InstrumentHooks, setupCore, teardownCore } from "@codspeed/core";
2+
import { Bench, Task } from "tinybench";
3+
import { getTaskUri } from "./uri";
4+
5+
declare const __VERSION__: string;
6+
7+
export abstract class BaseBenchRunner {
8+
protected bench: Bench;
9+
protected rootCallingFile: string;
10+
11+
constructor(bench: Bench, rootCallingFile: string) {
12+
this.bench = bench;
13+
this.rootCallingFile = rootCallingFile;
14+
}
15+
16+
protected setupBenchRun(): void {
17+
setupCore();
18+
this.logStart();
19+
}
20+
21+
protected abstract getModeName(): string;
22+
23+
protected logStart(): void {
24+
console.log(
25+
`[CodSpeed] running with @codspeed/tinybench v${__VERSION__} (${this.getModeName()})`
26+
);
27+
}
28+
29+
protected logTaskCompletion(uri: string, status: string): void {
30+
console.log(`[CodSpeed] ${status} ${uri}`);
31+
}
32+
33+
protected finalizeBenchRun(): Task[] {
34+
teardownCore();
35+
console.log(`[CodSpeed] Done running ${this.bench.tasks.length} benches.`);
36+
return this.bench.tasks;
37+
}
38+
39+
protected getTaskUri(task: Task): string {
40+
return getTaskUri(this.bench, task.name, this.rootCallingFile);
41+
}
42+
43+
protected wrapWithHooks<T>(fn: () => T, uri: string): T {
44+
InstrumentHooks.startBenchmark();
45+
const result = fn();
46+
InstrumentHooks.stopBenchmark();
47+
InstrumentHooks.setExecutedBenchmark(process.pid, uri);
48+
return result;
49+
}
50+
51+
protected async wrapWithHooksAsync<T>(
52+
fn: () => Promise<T>,
53+
uri: string
54+
): Promise<T> {
55+
InstrumentHooks.startBenchmark();
56+
const result = await fn();
57+
InstrumentHooks.stopBenchmark();
58+
InstrumentHooks.setExecutedBenchmark(process.pid, uri);
59+
return result;
60+
}
61+
62+
protected wrapFunctionWithFrame(fn: Function, isAsync: boolean): Function {
63+
if (isAsync) {
64+
return async function __codspeed_root_frame__() {
65+
global.gc?.();
66+
await fn();
67+
};
68+
} else {
69+
return function __codspeed_root_frame__() {
70+
global.gc?.();
71+
fn();
72+
};
73+
}
74+
}
75+
76+
// Template methods to be implemented by subclasses
77+
protected abstract runTaskAsync(task: Task, uri: string): Promise<void>;
78+
protected abstract runTaskSync(task: Task, uri: string): void;
79+
protected abstract finalizeAsyncRun(): Task[];
80+
protected abstract finalizeSyncRun(): Task[];
81+
82+
public setupBenchMethods(): void {
83+
// Call setupCore immediately when bench methods are set up to maintain compatibility
84+
this.setupBenchRun();
85+
86+
this.bench.run = async () => {
87+
for (const task of this.bench.tasks) {
88+
const uri = this.getTaskUri(task);
89+
await this.runTaskAsync(task, uri);
90+
}
91+
92+
return this.finalizeAsyncRun();
93+
};
94+
95+
this.bench.runSync = () => {
96+
for (const task of this.bench.tasks) {
97+
const uri = this.getTaskUri(task);
98+
this.runTaskSync(task, uri);
99+
}
100+
101+
return this.finalizeSyncRun();
102+
};
103+
}
104+
}
105+

0 commit comments

Comments
 (0)