feat: allow overridding task function durations (#329)

rubennorte · web-flow · commit 67b46e7cb037 · 2025-08-13T12:43:38.000+02:00
diff --git a/src/bench.ts b/src/bench.ts
@@ -8,6 +8,7 @@ import type {
   Fn,
   FnOptions,
   RemoveEventListenerOptionsArgument,
+  ResolvedBenchOptions,
   TaskResult,
 } from './types'
 
@@ -51,7 +52,7 @@ export class Bench extends EventTarget {
   /**
    * The options.
    */
-  readonly opts: Readonly<BenchOptions>
+  readonly opts: Readonly<ResolvedBenchOptions>
 
   /**
    * The JavaScript runtime environment.
diff --git a/src/index.ts b/src/index.ts
@@ -9,7 +9,9 @@ export type {
   Fn,
   FnHook,
   FnOptions,
+  FnReturnedObject,
   Hook,
+  ResolvedBenchOptions,
   Statistics,
   TaskEvents,
   TaskEventsMap,
diff --git a/src/task.ts b/src/task.ts
@@ -105,15 +105,15 @@ export class Task extends EventTarget {
       return this
     }
     this.dispatchEvent(createBenchEvent('start', this))
-    await this.bench.opts.setup?.(this, 'run')
+    await this.bench.opts.setup(this, 'run')
     const { error, samples: latencySamples } = (await this.benchmark(
       'run',
-      // eslint-disable-next-line @typescript-eslint/no-non-null-assertion
-      this.bench.opts.time!,
-      // eslint-disable-next-line @typescript-eslint/no-non-null-assertion
-      this.bench.opts.iterations!
+
+      this.bench.opts.time,
+
+      this.bench.opts.iterations
     )) as { error?: Error; samples?: number[] }
-    await this.bench.opts.teardown?.(this, 'run')
+    await this.bench.opts.teardown(this, 'run')
 
     this.processRunResult({ error, latencySamples })
 
@@ -136,21 +136,21 @@ export class Task extends EventTarget {
     )
     this.dispatchEvent(createBenchEvent('start', this))
 
-    const setupResult = this.bench.opts.setup?.(this, 'run')
+    const setupResult = this.bench.opts.setup(this, 'run')
     invariant(
       !isPromiseLike(setupResult),
       '`setup` function must be sync when using `runSync()`'
     )
 
     const { error, samples: latencySamples } = this.benchmarkSync(
       'run',
-      // eslint-disable-next-line @typescript-eslint/no-non-null-assertion
-      this.bench.opts.time!,
-      // eslint-disable-next-line @typescript-eslint/no-non-null-assertion
-      this.bench.opts.iterations!
+
+      this.bench.opts.time,
+
+      this.bench.opts.iterations
     ) as { error?: Error; samples?: number[] }
 
-    const teardownResult = this.bench.opts.teardown?.(this, 'run')
+    const teardownResult = this.bench.opts.teardown(this, 'run')
     invariant(
       !isPromiseLike(teardownResult),
       '`teardown` function must be sync when using `runSync()`'
@@ -170,15 +170,15 @@ export class Task extends EventTarget {
       return
     }
     this.dispatchEvent(createBenchEvent('warmup', this))
-    await this.bench.opts.setup?.(this, 'warmup')
+    await this.bench.opts.setup(this, 'warmup')
     const { error } = (await this.benchmark(
       'warmup',
-      // eslint-disable-next-line @typescript-eslint/no-non-null-assertion
-      this.bench.opts.warmupTime!,
-      // eslint-disable-next-line @typescript-eslint/no-non-null-assertion
-      this.bench.opts.warmupIterations!
+
+      this.bench.opts.warmupTime,
+
+      this.bench.opts.warmupIterations
     )) as { error?: Error }
-    await this.bench.opts.teardown?.(this, 'warmup')
+    await this.bench.opts.teardown(this, 'warmup')
 
     this.postWarmup(error)
   }
@@ -194,21 +194,21 @@ export class Task extends EventTarget {
 
     this.dispatchEvent(createBenchEvent('warmup', this))
 
-    const setupResult = this.bench.opts.setup?.(this, 'warmup')
+    const setupResult = this.bench.opts.setup(this, 'warmup')
     invariant(
       !isPromiseLike(setupResult),
       '`setup` function must be sync when using `runSync()`'
     )
 
     const { error } = this.benchmarkSync(
       'warmup',
-      // eslint-disable-next-line @typescript-eslint/no-non-null-assertion
-      this.bench.opts.warmupTime!,
-      // eslint-disable-next-line @typescript-eslint/no-non-null-assertion
-      this.bench.opts.warmupIterations!
+
+      this.bench.opts.warmupTime,
+
+      this.bench.opts.warmupIterations
     ) as { error?: Error }
 
-    const teardownResult = this.bench.opts.teardown?.(this, 'warmup')
+    const teardownResult = this.bench.opts.teardown(this, 'warmup')
     invariant(
       !isPromiseLike(teardownResult),
       '`teardown` function must be sync when using `runSync()`'
@@ -240,19 +240,25 @@ export class Task extends EventTarget {
 
       let taskTime = 0 // ms;
       if (this.async) {
-        // eslint-disable-next-line @typescript-eslint/no-non-null-assertion
-        const taskStart = this.bench.opts.now!()
+        const taskStart = this.bench.opts.now()
         // eslint-disable-next-line no-useless-call
-        await this.fn.call(this)
-        // eslint-disable-next-line @typescript-eslint/no-non-null-assertion
-        taskTime = this.bench.opts.now!() - taskStart
+        const fnResult = await this.fn.call(this)
+        taskTime = this.bench.opts.now() - taskStart
+
+        const overriddenDuration = getOverriddenDurationFromFnResult(fnResult)
+        if (overriddenDuration != null) {
+          taskTime = overriddenDuration
+        }
       } else {
-        // eslint-disable-next-line @typescript-eslint/no-non-null-assertion
-        const taskStart = this.bench.opts.now!()
+        const taskStart = this.bench.opts.now()
         // eslint-disable-next-line no-useless-call
-        this.fn.call(this)
-        // eslint-disable-next-line @typescript-eslint/no-non-null-assertion
-        taskTime = this.bench.opts.now!() - taskStart
+        const fnResult = this.fn.call(this)
+        taskTime = this.bench.opts.now() - taskStart
+
+        const overriddenDuration = getOverriddenDurationFromFnResult(fnResult)
+        if (overriddenDuration != null) {
+          taskTime = overriddenDuration
+        }
       }
 
       samples.push(taskTime)
@@ -326,18 +332,21 @@ export class Task extends EventTarget {
 
       let taskTime = 0 // ms;
 
-      // eslint-disable-next-line @typescript-eslint/no-non-null-assertion
-      const taskStart = this.bench.opts.now!()
+      const taskStart = this.bench.opts.now()
       // eslint-disable-next-line no-useless-call
-      const result = this.fn.call(this)
-      // eslint-disable-next-line @typescript-eslint/no-non-null-assertion
-      taskTime = this.bench.opts.now!() - taskStart
+      const fnResult = this.fn.call(this)
+      taskTime = this.bench.opts.now() - taskStart
 
       invariant(
-        !isPromiseLike(result),
+        !isPromiseLike(fnResult),
         'task function must be sync when using `runSync()`'
       )
 
+      const overriddenDuration = getOverriddenDurationFromFnResult(fnResult)
+      if (overriddenDuration != null) {
+        taskTime = overriddenDuration
+      }
+
       samples.push(taskTime)
       totalTime += taskTime
 
@@ -467,3 +476,14 @@ export class Task extends EventTarget {
     this.dispatchEvent(createBenchEvent('complete', this))
   }
 }
+
+/**
+ *
+ * @param fnResult - the result of the task function.
+ * @returns the overridden duration if defined by the function.
+ */
+function getOverriddenDurationFromFnResult (fnResult: ReturnType<Fn>): number | undefined {
+  if (fnResult != null && typeof fnResult === 'object' && 'overriddenDuration' in fnResult && typeof fnResult.overriddenDuration === 'number') {
+    return fnResult.overriddenDuration
+  }
+}
diff --git a/src/types.ts b/src/types.ts
@@ -108,10 +108,15 @@ export interface BenchOptions {
 export type EventListener = (evt: BenchEvent) => void
 
 /**
- * the task function
+ * The task function.
+ *
+ * If you need to provide a custom duration for the task (e.g.: because
+ * you want to measure a specific part of its execution), you can return an
+ * object with a `overriddenDuration` field. You should still use
+ * `bench.opts.now()` to measure that duration.
  */
 // eslint-disable-next-line @typescript-eslint/no-redundant-type-constituents
-export type Fn = () => Promise<unknown> | unknown
+export type Fn = () => FnReturnedObject | Promise<FnReturnedObject | unknown> | unknown
 
 /**
  * The task hook function signature.
@@ -148,6 +153,21 @@ export interface FnOptions {
   beforeEach?: FnHook
 }
 
+/**
+ * A possible object returned by task functions to override default behaviors,
+ * like the duration of the function itself.
+ */
+export interface FnReturnedObject {
+  /**
+   * An overridden duration for the task function, to be used instead of the
+   * duration measured by tinybench when running the benchmark.
+   *
+   * This can be useful to measure parts of the execution of a function that are
+   * hard to execute independently.
+   */
+  overriddenDuration?: number;
+}
+
 /**
  * The hook function signature.
  * If warmup is enabled, the hook will be called twice, once for the warmup and once for the run.
@@ -164,6 +184,18 @@ export type RemoveEventListenerOptionsArgument = Parameters<
   typeof EventTarget.prototype.removeEventListener
 >[2]
 
+export interface ResolvedBenchOptions extends BenchOptions {
+  iterations: NonNullable<BenchOptions['iterations']>,
+  now: NonNullable<BenchOptions['now']>,
+  setup: NonNullable<BenchOptions['setup']>,
+  teardown: NonNullable<BenchOptions['teardown']>,
+  throws: NonNullable<BenchOptions['throws']>,
+  time: NonNullable<BenchOptions['time']>,
+  warmup: NonNullable<BenchOptions['warmup']>,
+  warmupIterations: NonNullable<BenchOptions['warmupIterations']>,
+  warmupTime: NonNullable<BenchOptions['warmupTime']>,
+}
+
 /**
  * the statistics object
  */
diff --git a/test/index.test.ts b/test/index.test.ts
@@ -205,10 +205,9 @@ test('bench task runs and time consistency (async)', async () => {
 
   const fooTask = bench.getTask('foo')
 
-  // eslint-disable-next-line @typescript-eslint/no-non-null-assertion
-  expect(fooTask?.runs).toBeGreaterThanOrEqual(bench.opts.iterations!)
-  // eslint-disable-next-line @typescript-eslint/no-non-null-assertion
-  expect(fooTask?.result?.totalTime).toBeGreaterThanOrEqual(bench.opts.time!)
+  expect(fooTask?.runs).toBeGreaterThanOrEqual(bench.opts.iterations)
+
+  expect(fooTask?.result?.totalTime).toBeGreaterThanOrEqual(bench.opts.time)
 })
 
 test('bench task runs and time consistency (sync)', () => {
@@ -221,10 +220,9 @@ test('bench task runs and time consistency (sync)', () => {
 
   const fooTask = bench.getTask('foo')
 
-  // eslint-disable-next-line @typescript-eslint/no-non-null-assertion
-  expect(fooTask?.runs).toBeGreaterThanOrEqual(bench.opts.iterations!)
-  // eslint-disable-next-line @typescript-eslint/no-non-null-assertion
-  expect(fooTask?.result?.totalTime).toBeGreaterThanOrEqual(bench.opts.time!)
+  expect(fooTask?.runs).toBeGreaterThanOrEqual(bench.opts.iterations)
+
+  expect(fooTask?.result?.totalTime).toBeGreaterThanOrEqual(bench.opts.time)
 })
 
 test('events order (async)', async () => {
@@ -1309,3 +1307,43 @@ test('using concurrency should throw (sync)', () => {
     bench.runSync()
   }).toThrowError('Cannot use `concurrency` option when using `runSync`')
 })
+
+test('uses overridden task durations (async)', async () => {
+  const bench = new Bench({
+    iterations: 16,
+    now: () => 100,
+    throws: true,
+  })
+
+  bench.add('foo', () => {
+    return {
+      overriddenDuration: bench.opts.now() + 50,
+    }
+  })
+
+  await bench.run()
+
+  expect(bench.getTask('foo')?.result?.latency.mean).toBe(150)
+  expect(bench.getTask('foo')?.result?.latency.min).toBe(150)
+  expect(bench.getTask('foo')?.result?.latency.max).toBe(150)
+})
+
+test('uses overridden task durations (sync)', () => {
+  const bench = new Bench({
+    iterations: 16,
+    now: () => 100,
+    throws: true,
+  })
+
+  bench.add('foo', () => {
+    return {
+      overriddenDuration: bench.opts.now() + 50,
+    }
+  })
+
+  bench.runSync()
+
+  expect(bench.getTask('foo')?.result?.latency.mean).toBe(150)
+  expect(bench.getTask('foo')?.result?.latency.min).toBe(150)
+  expect(bench.getTask('foo')?.result?.latency.max).toBe(150)
+})
