Skip to content

Commit f372161

Browse files
boneskullboneskull
authored
feat: add AbortSignal support at the task level (#364)
1 parent 68d61ce commit f372161

File tree

4 files changed

+371
-7
lines changed

4 files changed

+371
-7
lines changed

‎README.md‎

Lines changed: 101 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -123,6 +123,107 @@ bench.concurrency = 'task' // The concurrency mode to determine how tasks are ru
123123
await bench.run()
124124
```
125125

126+
## Aborting Benchmarks
127+
128+
Tinybench supports aborting benchmarks using `AbortSignal` at both the bench and task levels:
129+
130+
### Bench-level Abort
131+
132+
Abort all tasks in a benchmark by passing a signal to the `Bench` constructor:
133+
134+
```ts
135+
const controller = new AbortController()
136+
137+
const bench = new Bench({ signal: controller.signal })
138+
139+
bench
140+
.add('task1', () => {
141+
// This will be aborted
142+
})
143+
.add('task2', () => {
144+
// This will also be aborted
145+
})
146+
147+
// Abort all tasks
148+
controller.abort()
149+
150+
await bench.run()
151+
// Both tasks will be aborted
152+
```
153+
154+
### Task-level Abort
155+
156+
Abort individual tasks without affecting other tasks by passing a signal to the task options:
157+
158+
```ts
159+
const controller = new AbortController()
160+
161+
const bench = new Bench()
162+
163+
bench
164+
.add('abortable task', () => {
165+
// This task can be aborted independently
166+
}, { signal: controller.signal })
167+
.add('normal task', () => {
168+
// This task will continue normally
169+
})
170+
171+
// Abort only the first task
172+
controller.abort()
173+
174+
await bench.run()
175+
// Only 'abortable task' will be aborted, 'normal task' continues
176+
```
177+
178+
### Abort During Execution
179+
180+
You can abort benchmarks while they're running:
181+
182+
```ts
183+
const controller = new AbortController()
184+
185+
const bench = new Bench({ time: 10000 }) // Long-running benchmark
186+
187+
bench.add('long task', async () => {
188+
await new Promise(resolve => setTimeout(resolve, 100))
189+
}, { signal: controller.signal })
190+
191+
// Abort after 1 second
192+
setTimeout(() => controller.abort(), 1000)
193+
194+
await bench.run()
195+
// Task will stop after ~1 second instead of running for 10 seconds
196+
```
197+
198+
### Abort Events
199+
200+
Both `Bench` and `Task` emit `abort` events when aborted:
201+
202+
```ts
203+
const controller = new AbortController()
204+
const bench = new Bench()
205+
206+
bench.add('task', () => {
207+
// Task function
208+
}, { signal: controller.signal })
209+
210+
const task = bench.getTask('task')
211+
212+
// Listen for abort events
213+
task.addEventListener('abort', () => {
214+
console.log('Task aborted!')
215+
})
216+
217+
bench.addEventListener('abort', () => {
218+
console.log('Bench received abort event!')
219+
})
220+
221+
controller.abort()
222+
await bench.run()
223+
```
224+
225+
**Note:** When a task is aborted, `task.result.aborted` will be `true`, and the task will have completed any iterations that were running when the abort signal was received.
226+
126227
## Prior art
127228

128229
- [Benchmark.js](https://github.com/bestiejs/benchmark.js)

‎src/task.ts‎

Lines changed: 37 additions & 7 deletions
Original file line numberDiff line numberDiff line change
@@ -59,14 +59,30 @@ export class Task extends EventTarget {
5959
*/
6060
private readonly fnOpts: Readonly<FnOptions>
6161

62+
/**
63+
* The task-level abort signal
64+
*/
65+
private readonly signal: AbortSignal | undefined
66+
6267
constructor (bench: Bench, name: string, fn: Fn, fnOpts: FnOptions = {}) {
6368
super()
6469
this.bench = bench
6570
this.name = name
6671
this.fn = fn
6772
this.fnOpts = fnOpts
6873
this.async = isFnAsyncResource(fn)
69-
// TODO: support signal in Tasks
74+
this.signal = fnOpts.signal
75+
76+
if (this.signal) {
77+
this.signal.addEventListener(
78+
'abort',
79+
() => {
80+
this.dispatchEvent(createBenchEvent('abort', this))
81+
this.bench.dispatchEvent(createBenchEvent('abort', this))
82+
},
83+
{ once: true }
84+
)
85+
}
7086
}
7187

7288
addEventListener<K extends TaskEvents>(
@@ -225,7 +241,7 @@ export class Task extends EventTarget {
225241
let totalTime = 0 // ms
226242
const samples: number[] = []
227243
const benchmarkTask = async () => {
228-
if (this.bench.opts.signal?.aborted) {
244+
if (this.isAborted()) {
229245
return
230246
}
231247
try {
@@ -259,7 +275,7 @@ export class Task extends EventTarget {
259275
// eslint-disable-next-line no-unmodified-loop-condition
260276
(totalTime < time ||
261277
samples.length + (limit?.activeCount ?? 0) + (limit?.pendingCount ?? 0) < iterations) &&
262-
!this.bench.opts.signal?.aborted
278+
!this.isAborted()
263279
) {
264280
if (this.bench.concurrency === 'task') {
265281
// eslint-disable-next-line @typescript-eslint/no-non-null-assertion
@@ -268,7 +284,7 @@ export class Task extends EventTarget {
268284
await benchmarkTask()
269285
}
270286
}
271-
if (!this.bench.opts.signal?.aborted && promises.length > 0) {
287+
if (!this.isAborted() && promises.length > 0) {
272288
await Promise.all(promises)
273289
} else if (promises.length > 0) {
274290
// Abort path
@@ -309,7 +325,7 @@ export class Task extends EventTarget {
309325
let totalTime = 0
310326
const samples: number[] = []
311327
const benchmarkTask = () => {
312-
if (this.bench.opts.signal?.aborted) {
328+
if (this.isAborted()) {
313329
return
314330
}
315331
try {
@@ -341,7 +357,7 @@ export class Task extends EventTarget {
341357
// eslint-disable-next-line no-unmodified-loop-condition
342358
(totalTime < time ||
343359
samples.length < iterations) &&
344-
!this.bench.opts.signal?.aborted
360+
!this.isAborted()
345361
) {
346362
benchmarkTask()
347363
}
@@ -363,6 +379,14 @@ export class Task extends EventTarget {
363379
return { samples }
364380
}
365381

382+
/**
383+
* Check if either our signal or the bench-level signal is aborted
384+
* @returns `true` if either signal is aborted
385+
*/
386+
private isAborted (): boolean {
387+
return this.signal?.aborted === true || this.bench.opts.signal?.aborted === true
388+
}
389+
366390
private async measureOnce (): Promise<{ fnResult: ReturnType<Fn>; taskTime: number }> {
367391
const taskStart = this.bench.opts.now()
368392
// eslint-disable-next-line no-useless-call
@@ -420,6 +444,9 @@ export class Task extends EventTarget {
420444
error?: Error
421445
latencySamples?: number[]
422446
}): void {
447+
// Always set aborted status, even if no samples were collected
448+
const isAborted = this.isAborted()
449+
423450
if (latencySamples && latencySamples.length > 0) {
424451
this.runs = latencySamples.length
425452
const totalTime = latencySamples.reduce((a, b) => a + b, 0)
@@ -442,7 +469,7 @@ export class Task extends EventTarget {
442469
const throughputStatistics = getStatisticsSorted(throughputSamples)
443470

444471
this.mergeTaskResult({
445-
aborted: this.bench.opts.signal?.aborted ?? false,
472+
aborted: isAborted,
446473
critical: latencyStatistics.critical,
447474
df: latencyStatistics.df,
448475
hz: throughputStatistics.mean,
@@ -466,6 +493,9 @@ export class Task extends EventTarget {
466493
totalTime,
467494
variance: latencyStatistics.variance,
468495
})
496+
} else if (isAborted) {
497+
// If aborted with no samples, still set the aborted flag
498+
this.mergeTaskResult({ aborted: true })
469499
}
470500

471501
if (error) {

‎src/types.ts‎

Lines changed: 7 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -156,6 +156,13 @@ export interface FnOptions {
156156
* An optional function that is run before each iteration of this task
157157
*/
158158
beforeEach?: FnHook
159+
160+
/**
161+
* An AbortSignal for aborting this specific task
162+
*
163+
* If not provided, falls back to {@link BenchOptions.signal}
164+
*/
165+
signal?: AbortSignal
159166
}
160167

161168
/**

0 commit comments

Comments
 (0)