如何使用ThreadPoolExecutor实现生产者消费者模式及线程启停控制

优化后符合规范的实现方案 #

import logging
import random
import concurrent.futures
from queue import Queue, Empty
from threading import Event
from time import sleep


class Updater:
    def __init__(self, max_workers: int = 5, queue_size: int = 5):
        self.update_queue = Queue(queue_size)
        self._logger = logging.getLogger(__name__)
        # 线程池统一在初始化时创建
        self.executor = concurrent.futures.ThreadPoolExecutor(max_workers=max_workers)
        self._producer_future = None
        self._consumer_future = None
        self.producer_running = Event()
        self.consumer_running = Event()

    def producer(self):
        while self.producer_running.is_set():
            try:
                item = random.randint(0, 10)
                # 阻塞放入队列，超时1s避免停止时永久阻塞
                self.update_queue.put(item, timeout=1)
                self._logger.info(f"生产元素: {item}")
                sleep(0.1 * float(random.randint(0, 10)))
            except Exception as e:
                self._logger.warning(f"生产者运行异常: {str(e)}", exc_info=True)
                break

    def consumer(self):
        while self.consumer_running.is_set() or not self.update_queue.empty():
            try:
                # 阻塞获取队列元素，超时1s避免停止时永久阻塞
                item = self.update_queue.get(timeout=1)
                self._logger.info(f"消费元素: {item}")
                sleep(0.2 * float(random.randint(0, 10)))
                # 标记队列任务处理完成
                self.update_queue.task_done()
            except Exception as e:
                # 队列空超时属于正常情况，只打印其他异常
                if not isinstance(e, Empty):
                    self._logger.warning(f"消费者运行异常: {str(e)}", exc_info=True)
                continue

    def start(self):
        if not self.producer_running.is_set():
            self.producer_running.set()
            self._producer_future = self.executor.submit(self.producer)

            self.consumer_running.set()
            self._consumer_future = self.executor.submit(self.consumer)
            self._logger.info("生产者消费者已启动")

    def stop(self):
        # 先停止生产者，不再生产新元素
        if self.producer_running.is_set():
            self.producer_running.clear()
            # 等待生产者任务完全退出
            concurrent.futures.wait([self._producer_future])
            self._logger.info("生产者已停止")

        # 等待队列中剩余元素全部消费完成
        self.update_queue.join()
        # 再停止消费者
        if self.consumer_running.is_set():
            self.consumer_running.clear()
            concurrent.futures.wait([self._consumer_future])
            self._logger.info("消费者已停止")
        
        # 关闭线程池，等待所有任务退出
        self.executor.shutdown(wait=True)
        self._logger.info("线程池已关闭，所有服务停止完成")

if __name__ == "__main__":
    logging.basicConfig(level=logging.INFO,format='%(asctime)s - %(name)s - %(levelname)s - %(message)s')

    foo = Updater()
    foo.start()
    sleep(5)
    foo.stop()

核心优化点说明 #

• 线程池初始化逻辑调整：将ThreadPoolExecutor的创建放到类的__init__方法中，符合对象初始化的设计规范，避免重复创建线程池的风险
• 启停逻辑优化：停止时先停生产者，再等待队列所有元素消费完成，最后停消费者，保证所有已生产的元素都会被处理，不会出现数据丢失
• 阻塞操作加超时：队列的put和get操作都增加了超时时间，避免服务停止时因为队列满/空导致线程永久阻塞无法退出
• 异常捕获：对生产消费逻辑增加异常捕获，避免单个任务出错直接导致整个服务崩溃，同时保留错误日志方便排查
• 队列任务标记：增加task_done调用和队列join等待，准确感知队列中所有任务的处理状态
• 支持参数配置：开放最大线程数、队列大小的配置入口，适配不同的业务性能需求

扩展建议 #

• 如果需要支持多生产者多消费者，只需要在start方法中循环submit多个生产/消费任务即可，停止逻辑不需要修改
• 如果生产消费逻辑是CPU密集型，可以替换ThreadPoolExecutor为ProcessPoolExecutor，仅需要修改初始化部分的代码即可，整体架构不用调整
• 可以增加指标监控，比如队列长度、生产消费速度的统计，方便做性能调优

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