p-limit 源码学习

岭子笑笑
193 阅读5分钟

前言

p-limit简介

p-limit是一个限制promise并发量的函数，可用于限制HTTP请求数等

先来一个简单的案例了解一下，如下pLimit的限制传参为3
代码执行2秒后会立刻按照顺序在控制台输出如下文案
“执行我啦 foo” “执行我啦 bar” ”执行我啦 hi“ "['foo', 'bar', 'hi']"

pLimit的并发限制设为3时，若有3个Promise同时执行，由于达到但未超过限制，因此不会触发限流控制！！！

import pLimit from 'p-limit';

const delayDo = (arg) => {
	return new Promise((resolve) => {
		setTimeout(() => {
			console.log('执行我啦', arg)
			resolve(arg)
		}, 2000);
	});
}

const limit = pLimit(3);

const input = [
	limit(() => delayDo('foo')),
	limit(() => delayDo('bar')),
	limit(() => delayDo('hi'))
];

const result = await Promise.all(input);
console.log(result); // ['foo', 'bar', 'hi']

我们更改pLimit函数的传参为1，如下
2秒后输出”执行我啦 foo“
过2秒后输出”执行我拉 bar“
再过2秒后输出”执行我拉hi“”['foo', 'bar', 'hi']“

触发了限流控制，每次执行一个promise

const limit = pLimit(1);

const input = [
	limit(() => delayDo('foo')),
	limit(() => delayDo('bar')),
	limit(() => delayDo('hi'))
];

// Only one promise is run at once
const result = await Promise.all(input);
console.log(result);

源码分析

1. validateConcurrency

validateConcurrency 校验promise最大并发数concurrency，
concurrency需要是一个正整数，否则throw new TypeError

function validateConcurrency(concurrency) {
	if (!((Number.isInteger(concurrency) || concurrency === Number.POSITIVE_INFINITY) && concurrency > 0)) {
		throw new TypeError('Expected `concurrency` to be a number from 1 and up');
	}
}

2. generator

generator是pLimit返回的函数，将当前的执行过程中的activeCount、pendingCount、clearQueue和concurrency注入成generator函数的属性，方便用户使用和更改
generator函数之间调用了enqueue函数，并将当前function_和resolve传入了进入

	const generator = (function_, ...arguments_) => new Promise(resolve => {
		enqueue(function_, resolve, arguments_);
	});

	Object.defineProperties(generator, {
		activeCount: {
			get: () => activeCount,
		},
		pendingCount: {
			get: () => queue.size,
		},
		clearQueue: {
			value() {
				queue.clear();
			},
		},
		concurrency: {
			get: () => concurrency,

			set(newConcurrency) {
				validateConcurrency(newConcurrency);
				concurrency = newConcurrency;

				queueMicrotask(() => {
					// eslint-disable-next-line no-unmodified-loop-condition
					while (activeCount < concurrency && queue.size > 0) {
						resumeNext();
					}
				});
			},
		},
	});

3.enqueue 函数

enqueue 函数通过queue和activeCount来控制run函数执行时机

为啥要用2个promise ？

  • 第一个promise是为了把run注入promise的回调函数中并且将internalResolve函数入队缓存起来，注意internalResolve函数未被执行时，run是不会被注入到微任务队列中等待执行的，参照第4点resumeNext函数一起理解，重要的事情说三遍！！！
  • 第二个promise是为了控制resumeNext中internalResolve函数出队并执行
	const enqueue = (function_, resolve, arguments_) => {
		// Queue `internalResolve` instead of the `run` function
		// to preserve asynchronous context.
		new Promise(internalResolve => {
                        // 将internalResolve入队，目的是控制run的执行
			queue.enqueue(internalResolve);
		}).then(
                        // .then中注入run函数，注意internalResolve未执行时，run函数是不会进入微任务队列的
			run.bind(undefined, function_, resolve, arguments_),
		);

		(async () => {
			// 保证上面internalResolve已经如队了
			await Promise.resolve();
                        // 此时在微任务中判断activeCount < concurrency，条件成立则执行resumeNext
			if (activeCount < concurrency) {
				resumeNext();
			}
		})();
	};

4. resumeNext

resumeNext将promise的回调函数放入到微任务队列中等待事件循环机制执行

这里需要注意queue队列里面装的是什么数据？

  • 参照上文的enqueue函数，入队的是internalResolve函数， 是promise的resolve函数
  • queue.dequeue()()表示把resolve函数出队且立即执行该resolve函数

执行internalResolve函数有什么作用？

  • internalResolve执行之后上文enqueue函数中的Promise的回调run函数就能进入微任务队列等待事件循环机制执行拉，只有执行了run函数，用户从外部传入的promise函数才能被执行，这样就达到限流的目的了！！！重点
    const queue = new Queue(); // 队列
    let activeCount = 0; // 激活执行的promise

    const resumeNext = () => {
        // active的promise小于并发数限制数concurrency且queue里面有值
            if (activeCount < concurrency && queue.size > 0) {
                    // 出队列且执行**************重点
                    queue.dequeue()();
                    // 活动状态的promise加1
                    activeCount++;
            }
    };

5 run 和 next

next是继续将下一个promise回调函数注入到微任务队列中
run是执行用户传入的promise函数

	const next = () => {
                // 活动状态的promise执行减1
		activeCount--;
                // 继续查看queue中是否含有未出列的resolve待执行
		resumeNext();
	};
        // 执行用户传入的promise函数
	const run = async (function_, resolve, arguments_) => {
                // ()()立即执行function_函数
		const result = (async () => function_(...arguments_))();
                此时的result是一个pending状态的promise, resolve回去
		resolve(result);

		try {
                    // pending状态的promise resolve 或者 reject之后， 在微任务中调用next()
                     await result;
		} catch {}
                
		next();
	};

完整源码如下

import Queue from 'yocto-queue';

export default function pLimit(concurrency) {
	validateConcurrency(concurrency);

	const queue = new Queue();
	let activeCount = 0;

	const resumeNext = () => {
		if (activeCount < concurrency && queue.size > 0) {
			queue.dequeue()();
			// Since `pendingCount` has been decreased by one, increase `activeCount` by one.
			activeCount++;
		}
	};

	const next = () => {
		activeCount--;

		resumeNext();
	};

	const run = async (function_, resolve, arguments_) => {
		const result = (async () => function_(...arguments_))();

		resolve(result);

		try {
			await result;
		} catch {}

		next();
	};

	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();
			}
		})();
	};

	const generator = (function_, ...arguments_) => new Promise(resolve => {
		enqueue(function_, resolve, arguments_);
	});

	Object.defineProperties(generator, {
		activeCount: {
			get: () => activeCount,
		},
		pendingCount: {
			get: () => queue.size,
		},
		clearQueue: {
			value() {
				queue.clear();
			},
		},
		concurrency: {
			get: () => concurrency,

			set(newConcurrency) {
				validateConcurrency(newConcurrency);
				concurrency = newConcurrency;

				queueMicrotask(() => {
					// eslint-disable-next-line no-unmodified-loop-condition
					while (activeCount < concurrency && queue.size > 0) {
						resumeNext();
					}
				});
			},
		},
	});
        
        // 返回了一个generator函数
	return generator;
}

export function limitFunction(function_, option) {
	const {concurrency} = option;
	const limit = pLimit(concurrency);

	return (...arguments_) => limit(() => function_(...arguments_));
}

function validateConcurrency(concurrency) {
	if (!((Number.isInteger(concurrency) || concurrency === Number.POSITIVE_INFINITY) && concurrency > 0)) {
		throw new TypeError('Expected `concurrency` to be a number from 1 and up');
	}
}

总结

  1. enqueue函数纠结了很久，为啥run函数比resumeNext后执行，查资料才知道promise必须执行resolve函数之后then的回调函数才会注入微任务队列等待执行
  2. 对promise的了解又加深了一些
avatar
文章
阅读
粉丝