Asyncio TaskGroup 超时聚合

示例说明

gather 的痛点是孤儿任务：某协程超时后 gather 自己抛 TimeoutError，但其他协程还在后台跑。TaskGroup 把并发任务的生命周期圈成一个整体——组内任一任务异常会取消其余任务；外层再套 wait_for，超时同样触发全组取消。所有分片要么都贡献结果，要么一起被掐断，没有漏网孤儿。

技术要点

• TaskGroup 圈定并发任务共同生命周期
• wait_for 超时触发全组取消
• result 逐任务收集展平结果

代码解读

gather_with_timeout 把 TaskGroup 的创建与结果收集封装成 runner；正常路径在组退出后逐任务取结果展平。超时路径里 wait_for 取消 runner，TaskGroup 的清理机制随之取消所有子任务。fast 场景四分片全部按时完成，timeout 场景慢分片被取消返回空表。

使用建议

Python 3.11 起才有 TaskGroup，旧版本用 gather 加手动取消组合。被取消的任务若持锁要保证释放。超时预算按最快路径的下游要求设定。

代码示例

import asyncio


async def fetch_shard(shard: int) -> list[int]:
    await asyncio.sleep(0.2 * shard)  # 各分片耗时不一
    return [shard * 10, shard * 10 + 1]


async def fetch_slow_shard(shard: int) -> list[int]:
    await asyncio.sleep(10)
    return []


async def gather_with_timeout(coros, timeout: float) -> list[int]:
    async def runner():
        results: list[int] = []
        async with asyncio.TaskGroup() as tg:
            tasks = [tg.create_task(c) for c in coros]
        results = [item for task in tasks for item in task.result()]
        return results

    try:
        return await asyncio.wait_for(runner(), timeout=timeout)
    except TimeoutError:
        print('deadline hit, all pending tasks cancelled')
        return []


async def main() -> None:
    fast = await gather_with_timeout(
        [fetch_shard(i) for i in range(4)], timeout=2.0
    )
    print('fast path collected:', fast)

    interrupted = await gather_with_timeout(
        [fetch_shard(1), fetch_slow_shard(9)], timeout=1.0
    )
    print('timeout path result:', interrupted)


asyncio.run(main())