|
| 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