asyncio.wait 更精细化并发

关于并发，前面我提到了asyncio.gather和asyncio.as_completed。但是多少这个两个进行并发任务的都有点小遗憾，就是无法更精细化进行并发任务的监控和处理。本小结我们针对asyncio.wait展开说明。

首先我们回顾一下asyncio.gather, asyncio.as_completed, 和 asyncio.wait 是它们都是 asyncio 库中用于处理并发任务的三种不同方法。每种方法都有其自身存在特定的用途和局限性。在并处理上：

• asyncio.gather asyncio.gather(*aws, loop=None, return_exceptions=False) 函数用于并发运行多个协程并收集它们的结果。它存在问题有：

• 同步返回：asyncio.gather() 会等待所有协程完成后才返回结果，这意味着如果有一个协程运行时间很长，其他已经完成的协程的结果也会被延迟返回。
• 异常处理：如果不设置 return_exceptions=True，第一个抛出异常的协程会导致 asyncio.gather() 抛出该异常，而其他协程的结果将不会被返回。
• 资源等待：如果其中一个协程被阻塞或无法正常完成处理的话，那么它会阻塞整个 asyncio.gather() 调用，导致其他协程的结果也无法正常进行获取。

• asyncio.as_completed(aws, loop=None, timeout=None) 是一个异步生成器，它按照协程完成的顺序返回协程的结果。它存在问题有：

• 无序返回：并发处理的结果按照协程完成的顺序返回，而不是按照启动的顺序，这与可能会导致结果的顺序和我们最初所预期的不同。

• 并发处理：我们需要用户自己处理每个协程的结果，这可能会使代码更加复杂，尤其是在处理大量协程时。

• 超时处理：虽然可以设置超时，但超时只会影响 as_completed() 的迭代过程，而不是协程本身的执行。如果需要取消超时的协程，需要手动实现。

而关于 asyncio.wait在前面的中，我们也有大概梳理过，其中一些梳理内容如下：asyncio.wait的参数说明如下：

asyncio.wait(aws, *, loop=None, timeout=None, return_when=ALL_COMPLETED)

• aws：它表示了一个可迭代的协程对象列表（或任何可迭代对象），表示要等待的任务列表。
• loop：可选参数，它表示了指定要使用的事件循环。如果未提供，则使用默认的事件循环。
• timeout：可选参数，它表示了指定等待的最长时间（以秒为单位）。如果超时时间到达而任务还未完成，则等待函数会返回。
• return_when：它表示了可选参数，指定返回的条件。可以选择以下值之一：
• FIRST_COMPLETED：如果有任何一个任务完成，则立即返回。
• FIRST_EXCEPTION：如果有任何一个任务抛出异常，则立即返回。
• ALL_COMPLETED：如果所有任务都完成，则返回。

最简单示例，默认是如果所有任务都完成，则返回：

import asyncio

async def coro(i):
    print(f"开始执行协程 {i}")
    await asyncio.sleep(i)
    print(f"协程 {i} 执行完毕")
    return i

async def main():
    tasks = [asyncio.create_task(coro(1)), asyncio.create_task(coro(2)), asyncio.create_task(coro(3))]
    done, pending = await asyncio.wait(tasks)
    for task in done:
        print(f"协程 {task.result()} 完成")

asyncio.run(main())

这里的wait默认的使用了条件是：ALL_COMPLETED。

需要注意上面的代码中我们需要的list列表对象包含的是一个asyncio.create_task的任务的对象，并将任务列表传递给 asyncio.wait()。这个函数会挂起当前任务，等待所有任务完成。

上面的代码执行的结果会等待我们的所有的任务的执行完成后会再统一的进行输出。上面输出的结果为(这里也是等待任务组的结果返回)：

开始执行协程 1
开始执行协程 2
开始执行协程 3
协程 1 执行完毕
协程 2 执行完毕
协程 3 执行完毕
协程 1 完成
协程 2 完成
协程 3 完成

超时示例：

import asyncio

async def coro(i):
    print(f"开始执行协程 {i}")
    await asyncio.sleep(i)
    print(f"协程 {i} 执行完毕")
    return i

async def main():
    tasks = [asyncio.create_task(coro(1)), asyncio.create_task(coro(2)), asyncio.create_task(coro(3))]
    done, pending = await asyncio.wait(tasks, timeout=2, return_when=asyncio.FIRST_COMPLETED)
    for task in done:
        print(f"协程 {task.result()} 完成")

asyncio.run(main())

再上面的代码加入某个任务的执行超过等待的时间。

等待所有任务完成，如果有任务抛出异常，则立即返回

import asyncio

async def coro(i):
    if i == 2:
        raise ValueError("协程任务错粗错误~~~~~任务抛出异常")
    print(f"开始执行协程 {i}")
    await asyncio.sleep(i)
    print(f"协程 {i} 执行完毕")
    return i

async def main():
    tasks = [asyncio.create_task(coro(1)), asyncio.create_task(coro(2)), asyncio.create_task(coro(3))]
    done, pending = await asyncio.wait(tasks, return_when=asyncio.FIRST_EXCEPTION)
    for task in done:
        if task.exception() is not None:
            print(f"协程 {task.get_coro().__name__} 抛出异常: {task.exception()}")
        else:
            print(f"协程 {task.result()} 完成")

asyncio.run(main())

上述代码执行输出的结果如下：

开始执行协程 1
开始执行协程 3
协程 coro 抛出异常: 协程任务错误~~~~~任务抛出异常

超时和错误结合示例

import asyncio

async def coro(i):
    if i == 2:
        raise ValueError("协程错误~~任务抛出异常")
    print(f"开始执行协程 {i}")
    await asyncio.sleep(i)
    print(f"协程 {i} 执行完毕")
    return i

async def main():
    tasks = [asyncio.create_task(coro(1)), asyncio.create_task(coro(2)), asyncio.create_task(coro(3))]
    done, pending = await asyncio.wait(tasks, timeout=2, return_when=asyncio.FIRST_EXCEPTION)
    for task in done:
        if task.exception() is not None:
            print(f"协程 {task.get_coro().__name__} 抛出异常: {task.exception()}")
        else:
            print(f"协程 {task.result()} 完成")

    for task in pending:
        print(f"协程 {task.get_coro().__name__} 未完成")
        # task.cancel() 这里还可以发生异常的时候把未完成的任务进行取消

asyncio.run(main())

输出结果如下：

开始执行协程 1
开始执行协程 3
协程 coro 抛出异常: 任务抛出异常
协程 coro 未完成
协程 coro 未完成

真实业务逻辑中，上面的代码会存在以下的可能的几种情况，

• 如果在超时时间内有任务完成（即有任务变为已完成状态），那么done集合将包含这些已完成的任务，而pending集合将包含未完成的任务。

• 如果在超时时间内没有任务完成，但至少有一个任务抛出了异常，那么done集合将包含这个抛出异常的任务，而pending集合将包含未完成的任务。

• 如果在超时时间内没有任务完成，也没有任务抛出异常，那么done集合将为空，而pending集合将包含所有的任务。

归纳总结

总体上我们使用asyncio.wait的好处在在于：

• 返回了成功和还在处理中的任务，我们可以有针对性的对done集合和pending集合分别的处理。
• wait() 本身不会取消任何任务。如果需要取消任务，你需要在 wait() 返回后对 pending 集合中的任务调用 cancel() 方法。
• 如果设置了超时，wait() 会在超时后立即返回，不管是否有任务完成。如果超时发生，pending 集合将包含尚未完成的任务。

番外篇之 区别于 asyncio.wait_for()。

• asyncio.wait_for() 主要是用于等待单个协程完成的函数。它接受一个协程或任务以及一个超时时间，如果在指定的时间内协程没有完成，则会抛出一个 asyncio.TimeoutError 异常。asyncio.wait_for() 会在超时后尝试取消任务。也就是用于自动进行任务取消的处理。

• asyncio.wait() 是一个用于并发运行多个协程的函数。它接受一个包含协程或任务的列表，并返回两个集合：done 和 pending。done 集合包含已完成的任务，而 pending 集合包含未完成的任务。asyncio.wait() 的行为取决于 return_when 参数的值，该参数决定了何时返回结果。