fix: replace plimit with own promises limiter implementation to plug mem leak (#386)

Signed-off-by: Jérôme Benoit <[email protected]>
Co-authored-by: Aras Abbasi <[email protected]>

Author: jerome-benoit

package.json

@@ -60,7 +60,6 @@
     "eslint-plugin-perfectionist": "^4.15.1",
     "nano-staged": "^0.8.0",
     "neostandard": "^0.12.2",
-    "p-limit": "^7.2.0",
     "simple-git-hooks": "^2.13.1",
     "size-limit": "^11.2.0",
     "tsdown": "^0.16.0",

pnpm-lock.yaml

@@ -1,5 +1,3 @@
-import pLimit from 'p-limit'
-
 import type {
   AddEventListenerOptionsArgument,
   BenchEvents,
@@ -26,7 +24,7 @@ import {
   defaultConvertTaskResultForConsoleTable,
   invariant,
   type JSRuntime,
-  now,
+  performanceNow,
   runtime,
   runtimeVersion,
 } from './utils'
@@ -109,10 +107,13 @@ export class Bench extends EventTarget {
     this.name = name
     this.runtime = runtime
     this.runtimeVersion = runtimeVersion
+    this.concurrency = restOptions.concurrency ?? null
+    this.threshold = restOptions.threshold ?? Infinity
+
     this.opts = {
       ...{
         iterations: defaultMinimumIterations,
-        now,
+        now: performanceNow,
         setup: emptyFunction,
         teardown: emptyFunction,
         throws: false,
@@ -122,6 +123,8 @@ export class Bench extends EventTarget {
         warmupTime: defaultMinimumWarmupTime,
       },
       ...restOptions,
+      ...(restOptions.iterations !== undefined && restOptions.time === undefined && { time: 0 }),
+      ...(restOptions.warmupIterations !== undefined && restOptions.warmupTime === undefined && { warmupTime: 0 }),
     }
 
     if (this.opts.signal) {
@@ -195,15 +198,23 @@ export class Bench extends EventTarget {
     if (this.opts.warmup) {
       await this.#warmupTasks()
     }
-    let values: Task[] = []
+
     this.dispatchEvent(new BenchEvent('start'))
+
+    let values: Task[] = []
+
     if (this.concurrency === 'bench') {
-      values = await this.#mapTasksConcurrently(task => task.run())
+      const taskPromises = []
+      for (const task of this.#tasks.values()) {
+        taskPromises.push(task.run())
+      }
+      values = await Promise.all(taskPromises)
     } else {
       for (const task of this.#tasks.values()) {
         values.push(await task.run())
       }
     }
+
     this.dispatchEvent(new BenchEvent('complete'))
     return values
   }
@@ -240,39 +251,18 @@ export class Bench extends EventTarget {
     return this.tasks.map(convert)
   }
 
-  /**
-   * Applies a worker function to all registered tasks using the concurrency limit.
-   *
-   * Scheduling is handled via p-limit with the current threshold. The returned array preserves
-   * the iteration order of the tasks. If any scheduled worker function rejects, the returned promise
-   * rejects with the first error after the scheduled worker functions settle, as per Promise.all semantics.
-   *
-   * Notes:
-   * - Concurrency is controlled by Bench.threshold (Infinity means unlimited).
-   * - No measurements are performed here; measurements happen inside Task.
-   * - Used internally by run() and warmupTasks() when concurrency === 'bench'.
-   * @template R The resolved type produced by the worker function for each task.
-   * @param workerFn A function invoked for each Task; it must return a Promise<R>.
-   * @returns Promise that resolves to an array of results in the same order as task iteration.
-   */
-  async #mapTasksConcurrently<R>(
-    workerFn: (task: Task) => Promise<R>
-  ): Promise<R[]> {
-    const limit = pLimit(Math.max(1, Math.floor(this.threshold)))
-    const promises: Promise<R>[] = []
-    for (const task of this.#tasks.values()) {
-      promises.push(limit(() => workerFn(task)))
-    }
-    return Promise.all(promises)
-  }
-
   /**
    * warmup the benchmark tasks.
    */
   async #warmupTasks (): Promise<void> {
     this.dispatchEvent(new BenchEvent('warmup'))
+
     if (this.concurrency === 'bench') {
-      await this.#mapTasksConcurrently(task => task.warmup())
+      const taskPromises = []
+      for (const task of this.#tasks.values()) {
+        taskPromises.push(task.warmup())
+      }
+      await Promise.all(taskPromises)
     } else {
       for (const task of this.#tasks.values()) {
         await task.warmup()

src/bench.ts

@@ -16,4 +16,4 @@ export type {
   TaskResult,
 } from './types'
 export type { JSRuntime } from './utils'
-export { hrtimeNow, now, nToMs } from './utils'
+export { hrtimeNow, performanceNow as now, nToMs } from './utils'

src/index.ts

@@ -1,5 +1,3 @@
-import pLimit from 'p-limit'
-
 import type { Bench } from './bench'
 import type {
   AddEventListenerOptionsArgument,
@@ -14,6 +12,7 @@ import type {
 } from './types'
 
 import { BenchEvent } from './event'
+import { withConcurrency } from './utils'
 import {
   getStatisticsSorted,
   invariant,
@@ -181,11 +180,11 @@ export class Task extends EventTarget {
     this.#result = { state: 'started' }
     this.dispatchEvent(new BenchEvent('start', this))
     await this.#bench.opts.setup(this, 'run')
-    const { error, samples: latencySamples } = (await this.#benchmark(
+    const { error, samples: latencySamples } = await this.#benchmark(
       'run',
       this.#bench.opts.time,
       this.#bench.opts.iterations
-    ))
+    )
     await this.#bench.opts.teardown(this, 'run')
 
     this.#processRunResult({ error, latencySamples })
@@ -325,37 +324,33 @@ export class Task extends EventTarget {
         }
       }
     }
-
-    try {
-      const promises: Promise<void>[] = [] // only for task level concurrency
-      let limit: ReturnType<typeof pLimit> | undefined // only for task level concurrency
-
-      if (this.#bench.concurrency === 'task') {
-        limit = pLimit(Math.max(1, Math.floor(this.#bench.threshold)))
+    if (this.#bench.concurrency === 'task') {
+      try {
+        await withConcurrency({
+          fn: benchmarkTask,
+          iterations,
+          limit: Math.max(1, Math.floor(this.#bench.threshold)),
+          now: this.#bench.opts.now,
+          signal: this.#signal ?? this.#bench.opts.signal,
+          time,
+        })
+      } catch (error) {
+        return { error: toError(error) }
       }
-
-      while (
-        // eslint-disable-next-line no-unmodified-loop-condition
-        (totalTime < time ||
-          samples.length + (limit?.activeCount ?? 0) + (limit?.pendingCount ?? 0) < iterations) &&
-        !this.#aborted
-      ) {
-        if (this.#bench.concurrency === 'task') {
-          // eslint-disable-next-line @typescript-eslint/no-non-null-assertion
-          promises.push((limit!)(benchmarkTask))
-        } else {
+      this.#runs = samples.length
+    } else {
+      try {
+        while (
+          // eslint-disable-next-line no-unmodified-loop-condition
+          (totalTime < time ||
+            samples.length < iterations) &&
+          !this.#aborted
+        ) {
           await benchmarkTask()
         }
+      } catch (error) {
+        return { error: toError(error) }
       }
-      if (!this.#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: toError(error) }
     }
 
     if (this.#fnOpts.afterAll != null) {

src/task.ts

@@ -29,6 +29,15 @@ export type BenchEventsWithTask = Extract<BenchEvents, 'add' | 'cycle' | 'error'
  * Bench options
  */
 export interface BenchOptions {
+  /**
+   * Executes tasks concurrently based on the specified concurrency mode.
+   *
+   * - When `mode` is set to `null` (default), concurrency is disabled.
+   * - When `mode` is set to 'task', each task's iterations (calls of a task function) run concurrently.
+   * - When `mode` is set to 'bench', different tasks within the bench run concurrently.
+   */
+  concurrency?: 'bench' | 'task' | null
+
   /**
    * number of times that a task should run if even the time option is finished
    * @default 64
@@ -60,6 +69,12 @@ export interface BenchOptions {
    */
   teardown?: Hook
 
+  /**
+   * The maximum number of concurrent tasks to run
+   * @default Infinity
+   */
+  threshold?: number
+
   /**
    * Throws if a task fails
    * @default false

src/types.ts

@@ -2,7 +2,11 @@
 // Portions copyright QuiiBz. 2023-2024. All Rights Reserved.
 
 import type { Task } from './task'
-import type { ConsoleTableConverter, Fn, Statistics } from './types'
+import type {
+  ConsoleTableConverter,
+  Fn,
+  Statistics,
+} from './types'
 
 import { emptyFunction, tTable } from './constants'
 
@@ -194,7 +198,7 @@ const hrtimeBigint: () => bigint = typeof (globalThis as { process?: { hrtime?:
  */
 export const hrtimeNow = () => nToMs(Number(hrtimeBigint()))
 
-export const now = performance.now.bind(performance)
+export const performanceNow = performance.now.bind(performance)
 
 /**
  * Checks if a value is a promise-like object.
@@ -506,3 +510,74 @@ export const defaultConvertTaskResultForConsoleTable: ConsoleTableConverter = (
   }
   /* eslint-enable perfectionist/sort-objects */
 }
+
+interface WithConcurrencyOptions<R> {
+  fn: () => Promise<R>
+  iterations: number
+  limit: number
+  now?: () => number
+  signal?: AbortSignal
+  time?: number
+}
+
+/**
+ * Creates a concurrency limiter that can execute functions with a maximum concurrency limit.
+ * @param options - The resource containing the function to execute and other options
+ * @returns A promise that resolves to an array of results.
+ * @throws {Error} if a single error occurs during execution
+ * @throws {AggregateError} if multiple errors occur during execution
+ */
+export const withConcurrency = async <R>(options: WithConcurrencyOptions<R>): Promise<R[]> => {
+  const { fn, iterations, limit, now = performanceNow, signal, time = 0 } = options
+
+  const maxWorkers = iterations === 0 ? limit : Math.max(0, Math.min(limit, iterations))
+
+  const errors: Error[] = []
+  const results: R[] = []
+
+  let isRunning = true
+  let nextIndex = 0
+
+  const hasTimeLimit = Number.isFinite(time) && time > 0
+  const hasIterationsLimit = iterations > 0
+  let targetTime = 0
+
+  // Reduce checks based on provided limits to avoid tainting the benchmark results
+  const doNext: () => boolean = hasIterationsLimit
+    ? hasTimeLimit
+      ? () => isRunning && (nextIndex++ < iterations) && ((now() < targetTime) || (isRunning = false))
+      : () => isRunning && (nextIndex++ < iterations)
+    : hasTimeLimit
+      ? () => isRunning && ((now() < targetTime) || (isRunning = false))
+      : () => isRunning
+
+  const pushResult = (r: R) => { isRunning && results.push(r) }
+  const pushError = (e: unknown) => { errors.push(toError(e)) }
+
+  const onAbort = () => (isRunning = false)
+
+  if (signal) {
+    if (signal.aborted) return []
+    signal.addEventListener('abort', onAbort)
+  }
+
+  const worker = async () => {
+    while (doNext()) {
+      try {
+        pushResult(await fn())
+      } catch (err) {
+        isRunning = false
+        pushError(err)
+        break
+      }
+    }
+  }
+
+  if (hasTimeLimit) targetTime = now() + time
+  const promises = Array.from({ length: maxWorkers }, () => worker())
+  await Promise.allSettled(promises)
+
+  if (errors.length === 0) return results
+  if (errors.length === 1) throw toError(errors[0])
+  throw new AggregateError(errors, 'Multiple errors occurred during concurrent execution')
+}