五千字拆解 p-limit 源码,如何实现控制请求并发的数量,如何实现队列管理

190 阅读5分钟

最近在多个面试的过程中,频繁被问到关于如何实现控制请求并发的数量,如何实现队列管理的。

在学习这个过程中,慢慢了解到 p-limit,控制 p-limit 是 一个已经整合好相关功能的工具库。看实例使用起来很简单。

import pLimit from 'p-limit';

const limit = pLimit(1);

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

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

看了源码以后,其实内容很简单,只有短短的 104 行 代码

image.png

源码解析

第一步准备工作

// 引入队列库,用于管理待执行的任务
import Queue from 'yocto-queue';

这里引入了 目的是 创建一个队列用于存储待执行的任务。


/**
 * 验证并发数是否合法
 * @param {number} concurrency - 要验证的并发数
 */
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');
    }
}

**
 * 创建一个并发控制器,限制同时执行的 Promise 数量
 * @param {number} concurrency - 允许的最大并发数
 * @returns {Function} - 返回一个用于包装任务的函数
 */
export default function pLimit(concurrency) {
    // 验证并发数是否合法(必须为正整数或正无穷)
    validateConcurrency(concurrency);

    // 创建一个队列用于存储待执行的任务
    const queue = new Queue();
    // 当前活跃的任务数量
    let activeCount = 0;
    
    
    ...
}

第二步 Step by Step

先完成一个任务,然后在队列中取出一个任务,并执行它,这个部分很好理解。

执行任务,活跃数+1,执行下一个任务

/**
 * 尝试从队列中取出并执行下一个任务
 */
 const resumeNext = () => {
    // 当活跃任务数少于并发限制且队列中有等待任务时
    if (activeCount < concurrency && queue.size > 0) {
        // 从队列中取出任务并执行
        queue.dequeue()();
        // 增加活跃任务计数
        activeCount++;
    }
};

/**
 * 标记一个任务完成,并尝试执行下一个任务
 */
const next = () => {
    // 减少活跃任务计数
    activeCount--;
    // 尝试执行队列中的下一个任务
    resumeNext();
};

第三步 执行 Run

简单来说就是将传入的任务和参数等,执行任务。最后完成后,执行下一个任务。

/**
 * 执行实际的任务函数,并在完成后处理后续逻辑
 * @param {Function} function_ - 要执行的任务函数
 * @param {Function} resolve - 用于解析任务结果的函数
 * @param {Array} arguments_ - 传递给任务函数的参数
 */
const run = async (function_, resolve, arguments_) => {
    // 执行任务函数并获取结果
    const result = (async () => function_(...arguments_))();
    // 立即解析结果,让调用者可以获取Promise
    resolve(result);

    try {
        // 等待任务完成
        await result;
    } catch {}

    // 任务完成后,调用next函数处理队列中的下一个任务
    next();
};

第四步 enqueue 加入队列

这里利用到 Promise 的微任务机制,将任务包装到 Promise

/**
 * 将任务加入队列等待执行
 * @param {Function} function_ - 要执行的任务函数
 * @param {Function} resolve - 用于解析任务结果的函数
 * @param {Array} arguments_ - 传递给任务函数的参数
 */
const enqueue = (function_, resolve, arguments_) => {
    // 将任务包装在Promise中,确保任务按顺序执行
    new Promise(internalResolve => {
        // 将内部解析函数加入队列
        queue.enqueue(internalResolve);
    }).then(
        // 当内部Promise解析时,执行实际的任务
        run.bind(undefined, function_, resolve, arguments_),
    );

    // 异步检查是否可以立即执行任务
    (async () => {
        // 等待下一个微任务周期,确保activeCount已更新
        await Promise.resolve();

        // 如果活跃任务数少于并发限制,尝试执行下一个任务
        if (activeCount < concurrency) {
            resumeNext();
        }
    })();
};

最后一步 generator 函数

generator 函数 添加一下 queue 队列中的一些方法

  • activeCount 当前的活跃的任务数
  • pendingCount 队列中等待执行的任务数量
  • clearQueue 清空队列的方法
  • concurrency Getter Setter 并发限制数量
/**
 * 生成器函数,用于创建可限制并发的任务包装器
 * @param {Function} function_ - 要执行的任务函数
 * @param {...any} arguments_ - 传递给任务函数的参数
 * @returns {Promise} - 返回一个Promise,代表任务的执行结果
 */
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();
        },
    },
    // 并发限制数量的getter/setter
    concurrency: {
        get: () => concurrency,

        set(newConcurrency) {
            // 验证新的并发限制是否合法
            validateConcurrency(newConcurrency);
            // 更新并发限制
            concurrency = newConcurrency;

            // 在微任务队列中检查并执行更多任务
            queueMicrotask(() => {
                // 持续执行队列中的任务,直到达到新的并发限制
                while (activeCount < concurrency && queue.size > 0) {
                    resumeNext();
                }
            });
        },
    },
});

return generator;

其他 直接限制特定函数的并发调用

/**
 * 直接限制特定函数的并发调用
 * @param {Function} function_ - 要限制并发的函数
 * @param {Object} option - 配置选项,包含concurrency属性
 * @returns {Function} - 返回一个包装后的函数,具有并发限制
 */
export function limitFunction(function_, option) {
    const {concurrency} = option;
    const limit = pLimit(concurrency);

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

流程图

image.png

总结

文字版的流程

  1. 初始化并发控制器,设置最大并发数和创建任务队列
  2. 当有新任务添加时,检查当前活跃任务数量
  3. 如果活跃任务数小于并发限制,立即执行任务
  4. 如果达到并发限制,任务进入队列等待
  5. 当任务完成时,减少活跃计数并检查队列
  6. 如果队列中有等待任务,取出并执行
  7. 循环这个过程直到所有任务完成

这个流程确保了任何时候都不会有超过指定数量的任务同时执行,实现了对并发请求的有效控制。

关键机制解析

  1. 任务队列管理:使用 yocto-queue 库创建队列,将待执行的任务按顺序存储。

  2. 并发控制核心

    • activeCount 跟踪当前正在执行的任务数量
    • concurrency 限制最大并发数
    • resumeNext() 负责从队列中取出任务并执行
  3. 异步处理

    • 使用 Promise 和微任务确保任务按顺序执行
    • 通过 async/await 处理任务的异步执行和结果
  4. 动态调整

    • 支持动态修改并发限制 concurrency
    • 修改后会自动启动更多任务以达到新的限制
  5. 状态监控

    • 通过 activeCountpendingCount 提供实时状态
    • 支持清空队列操作 clearQueue()

这种实现方式确保了任务的有序执行,避免了过多并发导致的资源耗尽问题,非常适合需要控制并发数量的场景,如数据库操作、网络请求等。