带超时的线程池映射

功能说明

标准库的 map 是串行执行，逐个手动提交 future 又容易在异常路径丢失结果。pool_map 封装了常见的批量并发模式：给定函数和输入序列，限制并发线程数，按原始顺序返回结果，并支持整体超时。

ThreadPoolExecutor 适合 I/O 密集型工作，例如批量请求接口、读取多个文件、调用阻塞式 SDK。它不会绕过 GIL 做 CPU 密集计算，那类场景应改用进程池。

关键代码

提交时为每个 future 记录原始下标，完成顺序与输入顺序无关，但写回字典时使用下标，最后按零到长度减一取出，即可恢复顺序。as_completed 的 timeout 参数作用于整体等待时长，超时后触发 FutureTimeout，逐个调用 cancel 取消尚未开始的任务，然后重新抛出异常。上下文管理器负责关闭线程池并等待已启动的线程结束。

使用注意

超时只能取消尚未开始的任务，正在运行的线程无法强制中断，因此总耗时可能略超设定值。函数内部应尊重超时约定，阻塞调用尽量带自己的超时。任意任务抛出异常时，with 块退出会取消剩余任务并传播异常，已经完成的结果不会返回。不要在任务中修改共享可变对象而不加锁。线程数应结合下游连接池上限设置，避免把压力转嫁给服务端。

验证方式

六个任务、三个线程、每个休眠五十毫秒时，squares 输出一到六的平方，耗时约零点一秒，说明并发生效且顺序与输入一致。第二个调用把超时设为五十毫秒而任务需要三百毫秒，会捕获超时并打印取消提示，程序不会卡住。可以把超时调大，确认结果正常返回。

代码示例

import time
from collections.abc import Callable, Sequence
from concurrent.futures import ThreadPoolExecutor, TimeoutError as FutureTimeout
from concurrent.futures import as_completed
from typing import TypeVar

T = TypeVar("T")
R = TypeVar("R")


def pool_map(
    func: Callable[[T], R],
    items: Sequence[T],
    *,
    workers: int = 4,
    timeout: float | None = None,
) -> list[R]:
    if workers < 1:
        raise ValueError("workers must be positive")
    if not items:
        return []
    results: dict[int, R] = {}
    with ThreadPoolExecutor(max_workers=workers, thread_name_prefix="map") as pool:
        futures = {pool.submit(func, item): index for index, item in enumerate(items)}
        try:
            for future in as_completed(futures, timeout=timeout):
                results[futures[future]] = future.result()
        except FutureTimeout:
            for future in futures:
                future.cancel()
            raise
    return [results[index] for index in range(len(items))]


def main() -> None:
    def slow_square(value: int) -> int:
        time.sleep(0.05)
        return value * value

    numbers = list(range(1, 7))
    start = time.perf_counter()
    squares = pool_map(slow_square, numbers, workers=3)
    print("squares:", squares)
    print(f"elapsed: {time.perf_counter() - start:.2f}s")

    try:
        pool_map(lambda value: time.sleep(0.3) or value, [1, 2], timeout=0.05)
    except FutureTimeout:
        print("timed out and cancelled pending work")


if __name__ == "__main__":
    main()