Python线程池任务提交时如何在池满时阻塞以限制消费速率？

不需要单独使用常规线程从队列拉取记录，通过Python标准库封装或第三方库就能实现需求，避免线程池任务队列无界增长。以下是具体实现方式：

一、Python标准库实现 #

1. 结合queue.Queue做任务限流 #

用有界queue.Queue控制待提交任务的数量，当队列满时，消费逻辑会自动阻塞，直到线程池处理完任务腾出位置。

import concurrent.futures
import queue
import time

def process_data(data):
    # 模拟任务处理耗时
    time.sleep(0.5)
    print(f"Processed: {data}")

def main():
    pool_size = 4
    # 有界队列，大小可根据业务调整（示例设为线程池的2倍）
    task_queue = queue.Queue(maxsize=pool_size * 2)
    
    with concurrent.futures.ThreadPoolExecutor(max_workers=pool_size) as executor:
        # 模拟外部数据队列的消费源
        data_source = range(20)
        for data in data_source:
            # 队列满时自动阻塞，直到有空闲位置
            task_queue.put(data)
            # 提交任务到线程池，同时从队列取数据
            executor.submit(process_data, task_queue.get())
            # 标记任务完成，确保队列join()能正常等待所有任务结束
            task_queue.task_done()
        
        # 等待所有任务处理完成
        task_queue.join()

if __name__ == "__main__":
    main()

2. 自定义有界ThreadPoolExecutor #

继承标准库ThreadPoolExecutor，将其底层无界队列替换为有界queue.Queue，让submit()方法在队列满时自动阻塞：

from concurrent.futures import ThreadPoolExecutor
from queue import Queue

class BoundedThreadPoolExecutor(ThreadPoolExecutor):
    def __init__(self, max_workers=None, thread_name_prefix='', max_queue_size=10):
        super().__init__(max_workers, thread_name_prefix)
        # 替换无界队列为有界队列
        self._work_queue = Queue(maxsize=max_queue_size)

def process_data(data):
    import time
    time.sleep(0.5)
    print(f"Processed: {data}")

def main():
    with BoundedThreadPoolExecutor(max_workers=4, max_queue_size=8) as executor:
        data_source = range(20)
        for data in data_source:
            # 队列满时submit自动阻塞，直到线程池有空闲
            executor.submit(process_data, data)
            print(f"Submitted: {data}")

if __name__ == "__main__":
    main()

二、第三方库方案 #

使用billiard（Celery底层线程/进程库）的Pool，它原生支持设置任务队列大小，队列满时提交任务会自动阻塞：

from billiard import Pool
import time

def process_data(data):
    time.sleep(0.5)
    print(f"Processed: {data}")

def main():
    # 配置线程池大小和任务队列上限
    pool = Pool(processes=4, max_queue_size=8)
    data_source = range(20)
    for data in data_source:
        # 队列满时自动阻塞
        pool.apply_async(process_data, args=(data,))
        print(f"Submitted: {data}")
    pool.close()
    pool.join()

if __name__ == "__main__":
    main()

总结 #

通过上述方法，无需额外编写单独的拉取线程，就能让线程池自动调控队列的消费速率，避免任务积压导致的内存问题。

内容的提问来源于stack exchange，提问作者Ivan Voras