Python中如何动态向线程池添加任务或选取线程分配任务

核心思路 #

• 放弃ThreadPoolExecutor的上下文管理器（with块），直接实例化并长期持有，这样就能随时提交新任务
• 给每个持续运行的监听器任务添加线程安全的停止信号，实现动态终止
• 自定义任务注册表，跟踪所有任务的状态，方便增删操作；如果需要随机选取空闲线程，可以扩展状态跟踪逻辑

实现代码 #

import concurrent.futures
import logging
from threading import Event
from typing import Dict

def run_listener(task_arg: dict, stop_event: Event):
    """带停止信号的事件监听器"""
    topic = task_arg["topic"]
    logging.info(f"Listener started for topic: {topic}")
    try:
        # 替换为你的事件总线监听逻辑
        while not stop_event.is_set():
            # 模拟监听等待，实际替换为事件总线的阻塞调用
            stop_event.wait(1)
    except Exception as e:
        logging.error(f"Listener failure for topic {topic}: {str(e)}")
    finally:
        logging.info(f"Listener stopped for topic: {topic}")

class DynamicThreadPoolManager:
    def __init__(self, max_workers: int = 20):
        self.executor = concurrent.futures.ThreadPoolExecutor(max_workers=max_workers)
        # 注册表：任务ID -> (停止信号, Future对象)
        self._task_registry: Dict[str, tuple[Event, concurrent.futures.Future]] = {}

    def add_task(self, task_id: str, task_arg: dict) -> bool:
        """新增任务：线程池自动分配空闲线程，若要随机选可扩展状态跟踪"""
        if task_id in self._task_registry:
            logging.warning(f"Task {task_id} already exists")
            return False
        
        stop_event = Event()
        # 提交任务到线程池，自动复用空闲线程
        future = self.executor.submit(run_listener, task_arg, stop_event)
        self._task_registry[task_id] = (stop_event, future)
        logging.info(f"Added task {task_id} for topic {task_arg['topic']}")
        return True

    def remove_task(self, task_id: str) -> bool:
        """终止指定任务"""
        if task_id not in self._task_registry:
            logging.warning(f"Task {task_id} not found")
            return False
        
        stop_event, future = self._task_registry.pop(task_id)
        stop_event.set()
        # 等待任务终止，超时则强制放弃等待
        try:
            future.result(timeout=5)
        except concurrent.futures.TimeoutError:
            logging.warning(f"Task {task_id} did not stop within 5 seconds")
        logging.info(f"Removed task {task_id}")
        return True

    def get_running_task_count(self) -> int:
        """获取当前运行的任务数"""
        return sum(1 for _, future in self._task_registry.values() if future.running())

    def shutdown(self, wait: bool = True):
        """关闭线程池，先终止所有任务"""
        for stop_event, _ in self._task_registry.values():
            stop_event.set()
        self.executor.shutdown(wait=wait)
        logging.info("ThreadPool shut down completely")

# 示例用法
if __name__ == "__main__":
    logging.basicConfig(level=logging.INFO)
    pool = DynamicThreadPoolManager(max_workers=3)

    # 初始添加任务
    pool.add_task("save_event_listener", {"topic": "/captureSaveEvent"})
    pool.add_task("upload_event_listener", {"topic": "/imageUploadEvent"})

    # 运行3秒后新增任务
    import time
    time.sleep(3)
    pool.add_task("stream_event_listener", {"topic": "/videoStreamEvent"})

    # 再运行3秒后删除一个任务
    time.sleep(3)
    pool.remove_task("upload_event_listener")

    # 继续运行5秒后关闭线程池
    time.sleep(5)
    pool.shutdown()

关键细节 #

• 动态添加任务：通过长期持有ThreadPoolExecutor实例，随时调用submit提交新任务，线程池会自动复用空闲线程，无需创建新线程
• 任务终止：用threading.Event作为优雅停止的信号，监听器在循环中检查该信号，收到停止指令后退出循环并清理资源
• 线程分配控制：如果需要随机选取空闲线程而非依赖线程池的默认调度，可以扩展DynamicThreadPoolManager，添加线程状态跟踪逻辑（比如记录每个线程的ID和是否空闲），提交任务时从空闲线程列表中随机选择并绑定任务（标准库的ThreadPoolExecutor内部调度是队列式，若要精确控制需自定义线程池实现）
• 进程内操作：所有任务都在同一个进程中运行，完全符合你避免重复创建进程/线程的需求

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