Fix AsyncResource propagation issue (#71)

Author: hyperair

async-hooks-stub.js

@@ -0,0 +1,15 @@
+export const AsyncResource = {
+	bind(fn, _type, thisArg) {
+		return fn.bind(thisArg);
+	},
+};
+
+export class AsyncLocalStorage {
+	getStore() {
+		return undefined;
+	}
+
+	run(_store, callback) {
+		return callback();
+	}
+}

index.js

@@ -1,4 +1,5 @@
 import Queue from 'yocto-queue';
+import {AsyncResource} from '#async_hooks';
 
 export default function pLimit(concurrency) {
 	if (!((Number.isInteger(concurrency) || concurrency === Number.POSITIVE_INFINITY) && concurrency > 0)) {
@@ -31,7 +32,9 @@ export default function pLimit(concurrency) {
 	};
 
 	const enqueue = (fn, resolve, args) => {
-		queue.enqueue(run.bind(undefined, fn, resolve, args));
+		queue.enqueue(
+			AsyncResource.bind(run.bind(undefined, fn, resolve, args)),
+		);
 
 		(async () => {
 			// This function needs to wait until the next microtask before comparing

package.json

@@ -12,6 +12,12 @@
 	},
 	"type": "module",
 	"exports": "./index.js",
+	"imports": {
+		"#async_hooks": {
+			"node": "async_hooks",
+			"default": "./async-hooks-stub.js"
+		}
+	},
 	"engines": {
 		"node": "^12.20.0 || ^14.13.1 || >=16.0.0"
 	},
@@ -20,7 +26,8 @@
 	},
 	"files": [
 		"index.js",
-		"index.d.ts"
+		"index.d.ts",
+		"async-hooks-stub.js"
 	],
 	"keywords": [
 		"promise",

test.js

@@ -3,6 +3,7 @@ import delay from 'delay';
 import inRange from 'in-range';
 import timeSpan from 'time-span';
 import randomInt from 'random-int';
+import {AsyncLocalStorage} from '#async_hooks';
 import pLimit from './index.js';
 
 test('concurrency: 1', async t => {
@@ -40,6 +41,23 @@ test('concurrency: 4', async t => {
 	await Promise.all(input);
 });
 
+test('propagates async execution context properly', async t => {
+	const concurrency = 2;
+	const limit = pLimit(concurrency);
+	const store = new AsyncLocalStorage();
+
+	const checkId = async id => {
+		await Promise.resolve();
+		t.is(id, store.getStore()?.id);
+	};
+
+	const startContext = async id => store.run({id}, () => limit(checkId, id));
+
+	await Promise.all(
+		Array.from({length: 100}, (_, id) => startContext(id)),
+	);
+});
+
 test('non-promise returning function', async t => {
 	await t.notThrowsAsync(async () => {
 		const limit = pLimit(1);