/
benchmark.ts
146 lines (122 loc) · 4.15 KB
/
benchmark.ts
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
import type { Suite, Task, VitestRunner, VitestRunnerImportSource } from '@vitest/runner'
import { updateTask as updateRunnerTask } from '@vitest/runner'
import { createDefer, getSafeTimers } from '@vitest/utils'
import { getBenchFn, getBenchOptions } from '../benchmark'
import { getWorkerState } from '../../utils'
import type { BenchTask, Benchmark, BenchmarkResult } from '../../types/benchmark'
import type { ResolvedConfig } from '../../types/config'
import type { VitestExecutor } from '../execute'
function createBenchmarkResult(name: string): BenchmarkResult {
return {
name,
rank: 0,
rme: 0,
samples: [] as number[],
} as BenchmarkResult
}
const benchmarkTasks = new WeakMap<Benchmark, import('tinybench').Task>()
async function runBenchmarkSuite(suite: Suite, runner: NodeBenchmarkRunner) {
const { Task, Bench } = await runner.importTinybench()
const start = performance.now()
const benchmarkGroup: Benchmark[] = []
const benchmarkSuiteGroup = []
for (const task of suite.tasks) {
if (task.mode !== 'run')
continue
if (task.meta?.benchmark)
benchmarkGroup.push(task as Benchmark)
else if (task.type === 'suite')
benchmarkSuiteGroup.push(task)
}
// run sub suites sequentially
for (const subSuite of benchmarkSuiteGroup)
await runBenchmarkSuite(subSuite, runner)
if (benchmarkGroup.length) {
const defer = createDefer()
suite.result = {
state: 'run',
startTime: start,
benchmark: createBenchmarkResult(suite.name),
}
updateTask(suite)
const addBenchTaskListener = (task: InstanceType<typeof Task>, benchmark: Benchmark) => {
task.addEventListener('complete', (e) => {
const task = e.task
const taskRes = task.result!
const result = benchmark.result!.benchmark!
Object.assign(result, taskRes)
updateTask(benchmark)
}, {
once: true,
})
task.addEventListener('error', (e) => {
const task = e.task
defer.reject(benchmark ? task.result!.error : e)
}, {
once: true,
})
}
benchmarkGroup.forEach((benchmark) => {
const options = getBenchOptions(benchmark)
const benchmarkInstance = new Bench(options)
const benchmarkFn = getBenchFn(benchmark)
benchmark.result = {
state: 'run',
startTime: start,
benchmark: createBenchmarkResult(benchmark.name),
}
const task = new Task(benchmarkInstance, benchmark.name, benchmarkFn)
benchmarkTasks.set(benchmark, task)
addBenchTaskListener(task, benchmark)
updateTask(benchmark)
})
const { setTimeout } = getSafeTimers()
const tasks: [BenchTask, Benchmark][] = []
for (const benchmark of benchmarkGroup) {
const task = benchmarkTasks.get(benchmark)!
await task.warmup()
tasks.push([
await new Promise<BenchTask>(resolve => setTimeout(async () => {
resolve(await task.run())
})),
benchmark,
])
}
suite.result!.duration = performance.now() - start
suite.result!.state = 'pass'
tasks
.sort(([taskA], [taskB]) => taskA.result!.mean - taskB.result!.mean)
.forEach(([, benchmark], idx) => {
benchmark.result!.state = 'pass'
if (benchmark) {
const result = benchmark.result!.benchmark!
result.rank = Number(idx) + 1
updateTask(benchmark)
}
})
updateTask(suite)
defer.resolve(null)
await defer
}
function updateTask(task: Task) {
updateRunnerTask(task, runner)
}
}
export class NodeBenchmarkRunner implements VitestRunner {
private __vitest_executor!: VitestExecutor
constructor(public config: ResolvedConfig) {}
async importTinybench() {
return await import('tinybench')
}
importFile(filepath: string, source: VitestRunnerImportSource): unknown {
if (source === 'setup')
getWorkerState().moduleCache.delete(filepath)
return this.__vitest_executor.executeId(filepath)
}
async runSuite(suite: Suite): Promise<void> {
await runBenchmarkSuite(suite, this)
}
async runTask(): Promise<void> {
throw new Error('`test()` and `it()` is only available in test mode.')
}
}