30 行代码带你实现 p-limit 并发控制

前言

最近笔者在做大文件的分片上传，为了提高文件分片的传输效率，通常需要使用浏览器的请求并发能力。但是浏览器的请求并发量是有限制的，像 Chrome、Firefox的并发量是 6，为了解决分片上传的性能问题同时兼顾其他请求正常执行，需要对文件分片上传做并发控制。

当前最受欢迎的并发控制器即 p-limit，下面是笔者通过借鉴 p-limit 实现的并发控制器。

p-limit原理

p-limit 使用非常简单：

import pLimit from 'p-limit';

const limit = pLimit(2);

const input = [
    limit(() => fetchSomething('foo')),
    limit(() => fetchSomething('bar')),
    limit(() => doSomething())
];

const result = await Promise.all(input);

开发者通过 pLimit 创建了一个并发量是 2 的控制器，后面只需要通过控制器添加异步任务即可。下面我们来手写一个。

创建任务控制器

首先需要创建一个并发控制器并传入并发数量：

const pLimit = (concurrency) => {
    const generator = (fn, ...args) => new Promise(resolve => {
		//...
	});

    return generator;
}

这里的 generator 就是最终的并发任务控制器，支持导入异步任务，最终返回一个 Promise。

管理异步任务

通过 generator 导入的异步任务需要统一管理，这里我们使用队列来维护，并创建一个计数器来管理当前正在运行的异步任务数量：

const queue = [];
let activeCount = 0;

const enqueue = (fn, resolve, args) => {
    queue.push(run.bind(null, fn, resolve, args));

    if (activeCount < concurrency && queue.length > 0) {
        queue.shift()();
    }
}

const generator = (fn, ...args) => new Promise(resolve => {
	enqueue(fn, resolve, args);
});

在 enqueue 中会将异步任务添加到队列中，然后判断当前正在执行的任务是否到达上限，如果没有则从队列中取出任务执行，具体的运行逻辑如下：

const run = async (n, resolve, args) => {
	activeCount++;
	// 包装为异步任务
	const result = (async () => fn(...args))();

	resolve(result);

	try {
		await result;
	} catch {}

	next();
}

运行任务数量加一，包装并执行任务，将结果传递给 generator 创建的 Promise，等待任务执行完成之后执行下一个任务。下一步的的运行逻辑非常简单，将计数器减一并运行一个任务：

const next = () => {
	activeCount--;
	if (activeCount < concurrency && queue.length > 0) {
		queue.shift()();
    }
}

这样就能实现并发量的控制，整个代码只有 30 行，完整代码如下：

const pLimit = (concurrency) => {
    const queue = [];
    let activeCount = 0;
    const next = () => {
		activeCount--;
		if (activeCount < concurrency && queue.length > 0) {
			queue.shift()();
		}
    }
    const run = async (n, resolve, args) => {
        activeCount++;
		const result = (async () => fn(...args))();
		resolve(result);
		try {
			await result;
		} catch {}
		next();
    }
    const enqueue = (fn, resolve, args) => {
        queue.push(run.bind(null, fn, resolve, args));
        if (activeCount < concurrency && queue.length > 0) {
            queue.shift()();
        }
    };
    const generator = (fn, ...args) => new Promise(resolve => {
        enqueue(fn, resolve, args);
    });
    return generator;
}

下面用一个测试用例来测试效果，这里设置并发量是 3，使用 Promise + setTimeout 模拟异步任务：

const plimit = pLimit(3);

const asyncTask = (name, time) => new Promise(resolve => {
    console.log(`${name}开始执行`);
    setTimeout(() => {
        resolve(name);
    }, time);
});

(async () => {
    const res = await Promise.all([
        plimit(() => asyncTask('a', 2000)),
        plimit(() => asyncTask('b', 3000)),
        plimit(() => asyncTask('c', 3000)),
        plimit(() => asyncTask('d', 1000)),
        plimit(() => asyncTask('e', 1000)),
        plimit(() => asyncTask('f', 1000)),
        plimit(() => asyncTask('g', 1000)),
    ]);

    console.log(res);
})()

执行结果如下：

暴露任务管理方法

创建并发任务管理器后外部还没有任何渠道去获取任务管理器内部的执行状态，比如当前正在执行的任务个数、剩余任务个数以及终止并发任务，由于我们的并发任务管理器是一个方法，我们可以通过 Object.defineProperties 为这个方法添加一些属性：

Object.defineProperties(generator, {
	  activeCount: {
		    get: () => activeCount
	  },
	  pendingCount: {
		    get: () => queue.length
	  },
	  clearQueue: {
		    value: () => {
			      queue.length = 0;
		    }
	  }
});

Challenge

在 p-limit 的源码中，官方是这么实现任务入队的逻辑：

const enqueue = (function_, resolve, arguments_) => {
		// Queue `internalResolve` instead of the `run` function
		// to preserve asynchronous context.
		new Promise(internalResolve => {
			queue.enqueue(internalResolve);
		}).then(
			run.bind(undefined, function_, resolve, arguments_),
		);

		(async () => {
			// This function needs to wait until the next microtask before comparing
			// `activeCount` to `concurrency`, because `activeCount` is updated asynchronously
			// after the `internalResolve` function is dequeued and called. The comparison in the if-statement
			// needs to happen asynchronously as well to get an up-to-date value for `activeCount`.
			await Promise.resolve();

			if (activeCount < concurrency) {
				resumeNext();
			}
		})();
};

前面一段添加方法至队列的逻辑和我们实现相似，官方在队列中保存的是一个 Promise 的 resolve 方法，相当于异步任务执行的开关，只有执行 resolve 才会执行对应的异步任务，从官方注释也可以看到这是为了保留异步任务的上下文。

这里笔者需要 challenge 的点是下面的 await Promise.resolve() 按照官方的注释理解，这一段代码是为了解决 activeCount 不同步的问题。

但是笔者认为这行代码是多余的，如果去掉，那么 await Promise.resolve() 后面的代码作为同步代码在添加完异步任务之后立马执行，下面是我该写 p-limit 源码之后使用上面的测试用例的测试结果：

	const enqueue = (function_, resolve, arguments_) => {
		new Promise(internalResolve => {
			console.log(`添加任务${arguments_}`);
			queue.push(internalResolve);
		}).then(
			run.bind(undefined, function_, resolve, arguments_),
		);

		(async () => {
			// await Promise.resolve();
			console.log(`调度任务${arguments_}`);
			
			if (activeCount < concurrency) {
				resumeNext();
			}
		})();
	};

结论和我们一致，如果不删除，那么调度任务会作为一个微任务，在所有的异步任务添加完成之后再去执行：

这两种执行顺序都不影响最终的执行结果，因为这只是影响首次调度异步任务执行的时机。

总结

在实现实中除了文件分片上传还有很多需要使用并发控制的场景，而并发控制本身的实现逻辑非常简单，其控制的核心就是通过一个队列保存所有的任务，首次会调度 n 个任务执行，然后每异步任务执行完之后会自动调度下一个任务。