feat: add pLimit & createLimitedVersionOfFn; refactor Dispatcher idleState

Author: magicdawn

package.json

@@ -34,23 +34,10 @@
     "gen-readme": "swig render ./readme/readme.md > README.md && prettier --write README.md",
     "dev": "tsup --watch",
     "build": "tsup",
+    "typecheck": "tsc --noEmit",
     "prepublishOnly": "npm run build"
   },
-  "repository": {
-    "type": "git",
-    "url": "git+https://github.com/magicdawn/promise.map.git"
-  },
-  "keywords": [
-    "promise",
-    "map",
-    "parallel"
-  ],
-  "author": "magicdawn",
-  "license": "MIT",
-  "bugs": {
-    "url": "https://github.com/magicdawn/promise.map/issues"
-  },
-  "homepage": "https://github.com/magicdawn/promise.map#readme",
+  "packageManager": "[email protected]+sha512.76e2379760a4328ec4415815bcd6628dee727af3779aaa4c914e3944156c4299921a89f976381ee107d41f12cfa4b66681ca9c718f0668fa0831ed4c6d8ba56c",
   "devDependencies": {
     "@magicdawn/eslint-config": "^0.1.0",
     "@magicdawn/prettier-config": "^0.0.4",
@@ -68,6 +55,21 @@
     "typescript": "^5.7.2",
     "vitest": "^2.1.8"
   },
+  "repository": {
+    "type": "git",
+    "url": "git+https://github.com/magicdawn/promise.map.git"
+  },
+  "keywords": [
+    "promise",
+    "map",
+    "parallel"
+  ],
+  "author": "magicdawn",
+  "license": "MIT",
+  "bugs": {
+    "url": "https://github.com/magicdawn/promise.map/issues"
+  },
+  "homepage": "https://github.com/magicdawn/promise.map#readme",
   "husky": {
     "hooks": {
       "pre-commit": "lint-staged"
@@ -78,7 +80,6 @@
       "prettier --write"
     ]
   },
-  "packageManager": "[email protected]+sha512.76e2379760a4328ec4415815bcd6628dee727af3779aaa4c914e3944156c4299921a89f976381ee107d41f12cfa4b66681ca9c718f0668fa0831ed4c6d8ba56c",
   "publishConfig": {
     "registry": "https://registry.npmjs.org/"
   }

src/dispatcher.ts

@@ -1,18 +1,29 @@
-export type UniversalMap<E, V> = E extends object ? WeakMap<E, V> : Map<E, V>
+type Ref<T> = { val: T }
 
 export class Dispatcher<E> {
-  private executors: E[]
-  private idleMap: UniversalMap<E, boolean>
+  static fromConcurrency(concurrency: number, label = '') {
+    if (!isFinite(concurrency)) throw new Error('concurrency must be finite')
+    if (!(concurrency > 0)) throw new Error('concurrency must be greater than 0')
+    const executors = new Array(concurrency).fill(0).map((_, index) =>
+      [label, `executors(${concurrency})`, index]
+        .map((x) => x && x.toString().trim())
+        .filter(Boolean)
+        .join('.'),
+    )
+    return new Dispatcher(executors)
+  }
+
+  private unwrapRef<T>(ref: Ref<T>): T {
+    return ref.val
+  }
+
+  private executorRefs: Ref<E>[]
+  private idleState: WeakMap<Ref<E>, boolean>
   constructor(executors: E[]) {
     if (!executors.length) throw new Error('executors can not be empty')
-    this.executors = executors
-
-    this.idleMap = (
-      executors[0] && typeof executors[0] === 'object'
-        ? new WeakMap<E & object, boolean>()
-        : new Map<E, boolean>()
-    ) as UniversalMap<E, boolean>
-    executors.forEach((x) => this.idleMap.set(x, true))
+    this.executorRefs = executors.map((x) => ({ val: x })) // weakmap 需要 reference type 作为 key
+    this.idleState = new WeakMap()
+    this.executorRefs.forEach((x) => this.idleState.set(x, true))
   }
 
   private aborted = false
@@ -22,35 +33,36 @@ export class Dispatcher<E> {
   }
 
   pendingResolves: Array<() => void> = []
-  replenish = (executor: E) => {
+  replenish = (executor: Ref<E>) => {
     if (this.aborted) return
     if (!this.pendingResolves.length) return
     this.pendingResolves.shift()?.()
   }
 
-  private async getExecutor() {
-    const find = () => this.executors.find((x) => this.idleMap.get(x))
+  private async getExecutorRef() {
+    const find = () => this.executorRefs.find((x) => this.idleState.get(x))
 
-    let executor = find()
-    while (!executor) {
+    let executorRef = find()
+    while (!executorRef) {
       const { promise, resolve } = Promise.withResolvers<void>()
       this.pendingResolves.push(resolve)
       await promise
-      executor = find()
+      executorRef = find()
     }
 
-    this.idleMap.set(executor, false) // mark used
-    return executor
+    this.idleState.set(executorRef, false) // mark used
+    return executorRef
   }
 
-  async dispatch<T>(action: (executor: E) => T) {
-    const executor = await this.getExecutor()
+  async dispatch<R>(action: (executor: E) => R): Promise<Awaited<R>> {
+    const executorRef = await this.getExecutorRef()
+    const executor = this.unwrapRef(executorRef)
     try {
       return await action(executor)
     } finally {
-      this.idleMap.set(executor, true)
+      this.idleState.set(executorRef, true)
       // replenish run as a `macro task`, before this macro task, `abort` can be called
-      setTimeout(() => this.replenish(executor))
+      setTimeout(() => this.replenish(executorRef))
     }
   }
 }

src/index.ts

@@ -1,7 +1,4 @@
 export { Dispatcher } from './dispatcher'
 export { pmapWorker } from './worker'
-export { pmap } from './pmap'
-
-// pmap default
-import { pmap } from './pmap'
-export default pmap
+export { pLimit, createLimitedVersionOfFn } from './limit'
+export { pmap, pmap as default } from './pmap'

src/limit.ts

@@ -0,0 +1,19 @@
+import { Dispatcher } from './dispatcher'
+
+// api same as `p-limit` package
+export function pLimit(concurrency: number) {
+  const dispatcher = Dispatcher.fromConcurrency(concurrency, 'pLimit')
+  return function limit<T extends unknown[], R>(fn: (...args: T) => R, ...args: NoInfer<T>) {
+    return dispatcher.dispatch(() => fn(...args))
+  }
+}
+
+export function createLimitedVersionOfFn<T extends unknown[], R>(
+  fn: (...args: T) => R,
+  concurrency: number,
+) {
+  const dispatcher = Dispatcher.fromConcurrency(concurrency, 'createLimitedVersionOfFn')
+  return function limitedVersionOfFn(...args: NoInfer<T>) {
+    return dispatcher.dispatch(() => fn(...args))
+  }
+}

src/pmap.ts

@@ -6,8 +6,7 @@ export async function pmap<T, R>(
   concurrency: number,
 ) {
   concurrency = Math.min(concurrency, arr.length)
-  const executors = new Array(concurrency).fill(0).map((_, index) => `pmap.executor.${index}`)
-  const dispatcher = new Dispatcher(executors)
+  const dispatcher = Dispatcher.fromConcurrency(concurrency, 'pmap')
   try {
     return await Promise.all(
       arr.map((item, index) => dispatcher.dispatch(() => fn(item, index, arr))),

test/.eslintrc.yml

@@ -0,0 +1,5 @@
+export function approximateCostTime(time: number, tolerance: number) {
+  return function satisfy(val: number) {
+    return val >= time && val - time <= Math.abs(tolerance)
+  }
+}

test/_shared.ts

@@ -0,0 +1,8 @@
+import { describe, expect, it } from 'vitest'
+import { Dispatcher } from '../src'
+
+describe('constructor check', () => {
+  it('should throws when executors is empty', () => {
+    expect(() => new Dispatcher([])).toThrow('executors can not be empty')
+  })
+})

test/dispatcher.test.ts

@@ -0,0 +1,98 @@
+import { delay, noop, range } from 'es-toolkit'
+import { expect, describe, it } from 'vitest'
+import { createLimitedVersionOfFn, pLimit } from '../src'
+import { approximateCostTime } from './_shared'
+
+describe('concurrency invalid check', () => {
+  describe('pLimit(concurrency)', () => {
+    it('should throws when concurrency is invalid', () => {
+      expect(() => pLimit(0)).toThrowError('concurrency must be greater than 0')
+      expect(() => pLimit(Infinity)).toThrowError('concurrency must be finite')
+    })
+  })
+
+  describe('createLimitedVersionOfFn(fn, concurrency)', () => {
+    it('should throws when concurrency is invalid', () => {
+      expect(() => createLimitedVersionOfFn(noop, 0)).toThrowError(
+        'concurrency must be greater than 0',
+      )
+      expect(() => createLimitedVersionOfFn(noop, Infinity)).toThrowError(
+        'concurrency must be finite',
+      )
+    })
+  })
+})
+
+describe('concurrency control is correct', () => {
+  it('pLimit', async function () {
+    let arr = range(5) // [0 .. 4]
+    async function measureCostTime(concurrency: number) {
+      let start = performance.now()
+
+      async function work(x: number) {
+        await delay(x * 10)
+        return x * 10
+      }
+
+      const limit = pLimit(concurrency)
+      const result = await Promise.all(arr.map((value) => limit(() => work(value))))
+      expect(result).toEqual(arr.map((value) => value * 10))
+
+      const cost = performance.now() - start
+      return cost
+    }
+
+    function approximateCostTime(time: number, tolerance: number) {
+      return function satisfy(val: number) {
+        return Math.abs(val - time) <= Math.abs(tolerance)
+      }
+    }
+
+    const tolerance = 10
+
+    // Infinity
+    expect(await measureCostTime(arr.length)).toSatisfy(approximateCostTime(40, tolerance))
+
+    // executor-0: 0ms	 20ms  40ms
+    // executor-1: 10ms  30ms
+    expect(await measureCostTime(2)).toSatisfy(approximateCostTime(60, tolerance))
+
+    // executor-0: 0ms	 30ms
+    // executor-1: 10ms  40ms
+    // executor-2: 20ms
+    expect(await measureCostTime(3)).toSatisfy(approximateCostTime(50, tolerance))
+  })
+
+  it('createLimitedVersionOfFn', async function () {
+    let arr = range(5) // [0 .. 4]
+    async function measureCostTime(concurrency: number) {
+      let start = performance.now()
+
+      async function work(x: number) {
+        await delay(x * 10)
+        return x * 10
+      }
+
+      const workWithLimit = createLimitedVersionOfFn(work, concurrency)
+      const result = await Promise.all(arr.map((x) => workWithLimit(x)))
+      expect(result).toEqual(arr.map((value) => value * 10))
+
+      const cost = performance.now() - start
+      return cost
+    }
+
+    const tolerance = 10
+
+    // Infinity
+    expect(await measureCostTime(arr.length)).toSatisfy(approximateCostTime(40, tolerance))
+
+    // executor-0: 0ms	 20ms  40ms
+    // executor-1: 10ms  30ms
+    expect(await measureCostTime(2)).toSatisfy(approximateCostTime(60, tolerance))
+
+    // executor-0: 0ms	 30ms
+    // executor-1: 10ms  40ms
+    // executor-2: 20ms
+    expect(await measureCostTime(3)).toSatisfy(approximateCostTime(50, tolerance))
+  })
+})

test/limit.test.ts

@@ -1,6 +1,7 @@
 import { delay, range } from 'es-toolkit'
 import { describe, it, expect } from 'vitest'
 import pmap from '../src'
+import { approximateCostTime } from './_shared'
 
 describe('pmap', function () {
   it('it works', async function () {
@@ -32,12 +33,6 @@ describe('pmap', function () {
       return performance.now() - start
     }
 
-    function approximateCostTime(time: number, tolerance: number) {
-      return function satisfy(val: number) {
-        return Math.abs(val - time) <= Math.abs(tolerance)
-      }
-    }
-
     const tolerance = 10
 
     // Infinity
@@ -56,7 +51,7 @@ describe('pmap', function () {
   it('should not start new after errored', async () => {
     const arr = [1, 2, 3, 4, 5, 6]
     const called: number[] = []
-    await expect(
+    await expect(() =>
       pmap(
         arr,
         async (x: number) => {
@@ -68,7 +63,7 @@ describe('pmap', function () {
         },
         2,
       ),
-    ).rejects.toThrow('test error')
+    ).rejects.toThrowError('test error')
     expect(called).toEqual([1, 2, 3])
   })
 })