fix: properly handle abortion at task level and benchmark fn error

jerome-benoit · jerome-benoit · commit fbf91ae9ef14 · 2025-08-15T17:37:44.000+02:00
Signed-off-by: Jérôme Benoit &lt;jerome.benoit@piment-noir.org&gt;
diff --git a/src/task.ts b/src/task.ts
@@ -40,7 +40,7 @@ export class Task extends EventTarget {
   runs = 0
 
   /**
-   * The task synchronous status
+   * The task asynchronous status
    */
   private readonly async: boolean
 
@@ -222,62 +222,58 @@ export class Task extends EventTarget {
       }
     }
 
-    // TODO: factor out
     let totalTime = 0 // ms
     const samples: number[] = []
     const benchmarkTask = async () => {
-      if (this.fnOpts.beforeEach != null) {
-        await this.fnOpts.beforeEach.call(this, mode)
+      if (this.bench.opts.signal?.aborted) {
+        return
       }
-
-      let taskTime = 0 // ms;
-      if (this.async) {
-        const taskStart = this.bench.opts.now()
-        // eslint-disable-next-line no-useless-call
-        const fnResult = await this.fn.call(this)
-        taskTime = this.bench.opts.now() - taskStart
-
-        const overriddenDuration = getOverriddenDurationFromFnResult(fnResult)
-        if (overriddenDuration != null) {
-          taskTime = overriddenDuration
-        }
-      } else {
-        const taskStart = this.bench.opts.now()
-        // eslint-disable-next-line no-useless-call
-        const fnResult = this.fn.call(this)
-        taskTime = this.bench.opts.now() - taskStart
-
-        const overriddenDuration = getOverriddenDurationFromFnResult(fnResult)
-        if (overriddenDuration != null) {
-          taskTime = overriddenDuration
+      try {
+        if (this.fnOpts.beforeEach != null) {
+          await this.fnOpts.beforeEach.call(this, mode)
         }
-      }
 
-      samples.push(taskTime)
-      totalTime += taskTime
+        let taskTime: number
+        if (this.async) {
+          ({ taskTime } = await this.measureOnce())
+        } else {
+          ({ taskTime } = this.measureOnceSync())
+        }
 
-      if (this.fnOpts.afterEach != null) {
-        await this.fnOpts.afterEach.call(this, mode)
+        samples.push(taskTime)
+        totalTime += taskTime
+      } finally {
+        if (this.fnOpts.afterEach != null) {
+          await this.fnOpts.afterEach.call(this, mode)
+        }
       }
     }
 
     try {
-      const limit = pLimit(this.bench.threshold) // only for task level concurrency
+      let limit: ReturnType<typeof pLimit> | undefined // only for task level concurrency
+      if (this.bench.concurrency === 'task') {
+        limit = pLimit(this.bench.threshold)
+      }
       const promises: Promise<void>[] = [] // only for task level concurrency
       while (
         // eslint-disable-next-line no-unmodified-loop-condition
         (totalTime < time ||
-          samples.length + limit.activeCount + limit.pendingCount < iterations) &&
+          samples.length + (limit?.activeCount ?? 0) + (limit?.pendingCount ?? 0) < iterations) &&
         !this.bench.opts.signal?.aborted
       ) {
         if (this.bench.concurrency === 'task') {
-          promises.push(limit(benchmarkTask))
+          // eslint-disable-next-line @typescript-eslint/no-non-null-assertion
+          promises.push((limit!)(benchmarkTask))
         } else {
           await benchmarkTask()
         }
       }
       if (!this.bench.opts.signal?.aborted && promises.length > 0) {
         await Promise.all(promises)
+      } else if (promises.length > 0) {
+        // Abort path
+        // eslint-disable-next-line no-void
+        void Promise.allSettled(promises)
       }
     } catch (error) {
       return { error }
@@ -310,52 +306,42 @@ export class Task extends EventTarget {
       }
     }
 
-    // TODO: factor out
-    let totalTime = 0 // ms
+    let totalTime = 0
     const samples: number[] = []
     const benchmarkTask = () => {
-      if (this.fnOpts.beforeEach != null) {
-        const beforeEachResult = this.fnOpts.beforeEach.call(this, mode)
-        invariant(
-          !isPromiseLike(beforeEachResult),
-          '`beforeEach` function must be sync when using `runSync()`'
-        )
-      }
-
-      let taskTime = 0 // ms;
-
-      const taskStart = this.bench.opts.now()
-      // eslint-disable-next-line no-useless-call
-      const fnResult = this.fn.call(this)
-      taskTime = this.bench.opts.now() - taskStart
-
-      invariant(
-        !isPromiseLike(fnResult),
-        'task function must be sync when using `runSync()`'
-      )
-
-      const overriddenDuration = getOverriddenDurationFromFnResult(fnResult)
-      if (overriddenDuration != null) {
-        taskTime = overriddenDuration
+      if (this.bench.opts.signal?.aborted) {
+        return
       }
+      try {
+        if (this.fnOpts.beforeEach != null) {
+          const beforeEachResult = this.fnOpts.beforeEach.call(this, mode)
+          invariant(
+            !isPromiseLike(beforeEachResult),
+            '`beforeEach` function must be sync when using `runSync()`'
+          )
+        }
 
-      samples.push(taskTime)
-      totalTime += taskTime
-
-      if (this.fnOpts.afterEach != null) {
-        const afterEachResult = this.fnOpts.afterEach.call(this, mode)
-        invariant(
-          !isPromiseLike(afterEachResult),
-          '`afterEach` function must be sync when using `runSync()`'
-        )
+        const { taskTime } = this.measureOnceSync()
+
+        samples.push(taskTime)
+        totalTime += taskTime
+      } finally {
+        if (this.fnOpts.afterEach != null) {
+          const afterEachResult = this.fnOpts.afterEach.call(this, mode)
+          invariant(
+            !isPromiseLike(afterEachResult),
+            '`afterEach` function must be sync when using `runSync()`'
+          )
+        }
       }
     }
 
     try {
       while (
         // eslint-disable-next-line no-unmodified-loop-condition
-        totalTime < time ||
-        samples.length < iterations
+        (totalTime < time ||
+        samples.length < iterations) &&
+        !this.bench.opts.signal?.aborted
       ) {
         benchmarkTask()
       }
@@ -377,6 +363,34 @@ export class Task extends EventTarget {
     return { samples }
   }
 
+  private async measureOnce (): Promise<{ fnResult: ReturnType<Fn>; taskTime: number }> {
+    const taskStart = this.bench.opts.now()
+    // eslint-disable-next-line no-useless-call
+    const fnResult = await this.fn.call(this)
+    let taskTime = this.bench.opts.now() - taskStart
+    const overriddenDuration = getOverriddenDurationFromFnResult(fnResult)
+    if (overriddenDuration != null) {
+      taskTime = overriddenDuration
+    }
+    return { fnResult, taskTime }
+  }
+
+  private measureOnceSync (): { fnResult: ReturnType<Fn>; taskTime: number } {
+    const taskStart = this.bench.opts.now()
+    // eslint-disable-next-line no-useless-call
+    const fnResult = this.fn.call(this)
+    let taskTime = this.bench.opts.now() - taskStart
+    invariant(
+      !isPromiseLike(fnResult),
+      'task function must be sync when using `runSync()`'
+    )
+    const overriddenDuration = getOverriddenDurationFromFnResult(fnResult)
+    if (overriddenDuration != null) {
+      taskTime = overriddenDuration
+    }
+    return { fnResult, taskTime }
+  }
+
   /**
    * merge into the result object values
    * @param result - the task result object to merge with the current result object values
@@ -406,7 +420,7 @@ export class Task extends EventTarget {
     error?: Error
     latencySamples?: number[]
   }): void {
-    if (latencySamples) {
+    if (latencySamples && latencySamples.length > 0) {
       this.runs = latencySamples.length
       const totalTime = latencySamples.reduce((a, b) => a + b, 0)
 
@@ -418,16 +432,17 @@ export class Task extends EventTarget {
       // Throughput statistics
       const throughputSamples = latencySamples
         .map(sample =>
-          sample !== 0 ? 1000 / sample : 1000 / latencyStatistics.mean
+          sample !== 0
+            ? 1000 / sample
+            : latencyStatistics.mean !== 0
+              ? 1000 / latencyStatistics.mean
+              : 0
         ) // Use latency average as imputed sample
         .sort((a, b) => a - b)
       const throughputStatistics = getStatisticsSorted(throughputSamples)
 
-      if (this.bench.opts.signal?.aborted) {
-        return
-      }
-
       this.mergeTaskResult({
+        aborted: this.bench.opts.signal?.aborted ?? false,
         critical: latencyStatistics.critical,
         df: latencyStatistics.df,
         hz: throughputStatistics.mean,
@@ -481,7 +496,9 @@ function getOverriddenDurationFromFnResult (
     fnResult != null &&
     typeof fnResult === 'object' &&
     'overriddenDuration' in fnResult &&
-    typeof fnResult.overriddenDuration === 'number'
+    typeof fnResult.overriddenDuration === 'number' &&
+    Number.isFinite(fnResult.overriddenDuration) &&
+    fnResult.overriddenDuration >= 0
   ) {
     return fnResult.overriddenDuration
   }
diff --git a/src/types.ts b/src/types.ts
@@ -321,6 +321,11 @@ export interface TaskEventsMap {
  * The task result object
  */
 export interface TaskResult {
+  /**
+   * whether the task was aborted
+   */
+  aborted: boolean
+
   /**
    * the latency samples critical value
    * @deprecated use `.latency.critical` instead
diff --git a/src/utils.ts b/src/utils.ts
@@ -261,8 +261,8 @@ const average = (samples: number[]) => {
  * @returns the variance of the sample
  */
 const variance = (samples: number[], avg = average(samples)) => {
-  const result = samples.reduce((sum, n) => sum + (n - avg) ** 2, 0)
-  return result / (samples.length - 1) || 0
+  const sumSq = samples.reduce((sum, n) => sum + (n - avg) ** 2, 0)
+  return samples.length <= 1 ? 0 : sumSq / (samples.length - 1)
 }
 
 /**
diff --git a/test/index.test.ts b/test/index.test.ts
@@ -332,6 +332,7 @@ test('events order (async)', async () => {
     'cycle',
     'error-complete',
     'abort',
+    'cycle',
     'complete',
     'reset',
   ])
