Python asyncio超时控制与优雅关闭方法

异步任务阻塞问题解析

在Python的 asyncio 编程中，我们经常使用 asyncio.gather 来并发执行多个异步任务。然而，当某些任务内部包含长时间阻塞的 await 调用（例如等待网络数据或消息队列事件），并且这些调用可能长时间不返回时，即使外部设置了停止标志，整个 gather 操作也可能无法按预期在指定时间内终止。

考虑以下场景：

import asyncio

stop = False

async def watch_task1(client):
    while not stop:
        print("watch_task1: Waiting for data...")
        await client.ws.get_data() # 可能会长时间阻塞
        print("watch_task1: Data received.")

async def watch_task2(client):
    while not stop:
        print("watch_task2: Waiting for news...")
        await client.ws.get_news() # 可能会长时间阻塞
        print("watch_task2: News received.")

async def stop_after(delay):
    global stop
    print(f"stop_after: Will stop after {delay} seconds.")
    await asyncio.sleep(delay)
    stop = True
    print("stop_after: Stop flag set to True.")

async def main_gather(client):
    tasks = [
        watch_task1(client),
        watch_task2(client),
        stop_after(60),
    ]

    try:
        # 使用 gather，如果 watch_task1/2 内部的 await 阻塞，则无法按时停止
        await asyncio.gather(*tasks, return_exceptions=True)
    except Exception as e:
        print(f"main_gather: An exception occurred: {e}")
    finally:
        print("main_gather: All tasks finished or gathered.")

# 模拟一个简化的客户端
class MockClient:
    def __init__(self):
        self.ws = self.MockWebSocket()

    class MockWebSocket:
        async def get_data(self):
            # 模拟长时间阻塞，除非外部取消
            await asyncio.sleep(3600) # 模拟24小时，或直到被取消
            return "some_data"

        async def get_news(self):
            await asyncio.sleep(3600) # 模拟24小时，或直到被取消
            return "some_news"

# 运行示例 (不会真正运行，因为需要一个 client 实例)
# asyncio.run(main_gather(MockClient()))

在这个例子中，stop_after 函数会在60秒后将 stop 标志设置为 True。然而，watch_task1 和 watch_task2 内部的 await client.ws.get_data() 和 await client.ws.get_news() 调用可能会无限期地阻塞，直到有数据到来。这意味着即使 stop 标志为 True，这些任务也无法退出其 while 循环，导致 asyncio.gather 无法在60秒后完成。

解决方案：使用 asyncio.wait 进行超时控制

为了解决这个问题，我们可以使用 asyncio.wait 函数，它提供了强大的超时管理能力。asyncio.wait 允许我们指定一个 timeout 参数，在达到指定时间后，它会返回已完成的任务和未完成的任务。

asyncio.wait 的基本签名如下：

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

• aws: 一个可迭代对象，包含要等待的 awaitable 对象（通常是 Task 或 Future）。
• timeout: 可选参数，指定等待的最大秒数。如果在此时间内所有任务都未完成，wait 会提前返回。
• return_when: 一个常量，指定何时返回。常见的有：
  • asyncio.ALL_COMPLETED (默认): 所有任务都完成。
  • asyncio.FIRST_COMPLETED: 任意一个任务完成。
  • asyncio.FIRST_EXCEPTION: 任意一个任务抛出异常。

当 timeout 参数被设置时，asyncio.wait 会返回两个集合：done 和 pending。

• done: 包含在超时时间内完成（或抛出异常）的任务。
• pending: 包含在超时时间内未完成的任务。

以下是使用 asyncio.wait 改进后的 main 函数示例：

import asyncio

# 假设 watch_task1, watch_task2, stop_after 和 MockClient 的定义与上文相同
# ...

async def main_wait(client):
    tasks = [
        asyncio.create_task(watch_task1(client)), # 显式创建任务
        asyncio.create_task(watch_task2(client)),
        asyncio.create_task(stop_after(60)),
    ]

    print("main_wait: Starting tasks with a 60-second timeout...")
    done, pending = await asyncio.wait(tasks, timeout=60)

    print(f"main_wait: Wait completed. Done tasks: {len(done)}, Pending tasks: {len(pending)}")

    # 处理已完成的任务
    for task in done:
        try:
            # 获取任务结果，如果任务抛出异常，这里会重新抛出
            result = task.result()
            print(f"main_wait: Task {task.get_name()} completed with result: {result if result is not None else 'None'}")
        except asyncio.CancelledError:
            print(f"main_wait: Task {task.get_name()} was cancelled.")
        except Exception as e:
            print(f"main_wait: Task {task.get_name()} raised an exception: {e}")

    # 处理未完成的任务：通常需要取消它们以释放资源
    if pending:
        print("main_wait: Cancelling pending tasks...")
        for task in pending:
            task.cancel()

        # 等待所有取消操作完成，或者设置一个短的超时
        # 注意：task.cancel() 只是请求取消，任务需要自行处理 CancelledError
        await asyncio.gather(*pending, return_exceptions=True) # 等待取消请求生效

    print("main_wait: All tasks processed.")

# 实际运行示例
async def run_example():
    client = MockClient()
    await main_wait(client)

if __name__ == "__main__":
    asyncio.run(run_example())

代码解释：

1. asyncio.create_task(): 在将 awaitable 对象传递给 asyncio.wait 之前，最好使用 asyncio.create_task() 将它们包装成 Task 对象。这使得我们可以更好地管理和取消这些任务。
2. await asyncio.wait(tasks, timeout=60): 这是核心部分。它会等待 tasks 列表中的所有任务，但最长等待60秒。一旦60秒过去，或者所有任务都已完成，它就会返回 done 和 pending 两个集合。
3. 处理 done 任务:
   • 遍历 done 集合中的每个任务。
   • 使用 task.result() 来获取任务的返回值或重新抛出任务内部发生的任何异常。这是处理 asyncio.wait 返回的已完成任务的关键。
   • 通过 task.get_name() 可以获取任务的名称（如果设置了）。
4. 处理 pending 任务:
   • 遍历 pending 集合中的每个任务。这些任务在超时时间内未能完成。
   • 调用 task.cancel() 方法来请求取消这些任务。cancel() 方法会向任务内部注入一个 asyncio.CancelledError 异常。任务需要自行捕获并处理这个异常，以执行必要的清理工作。
   • await asyncio.gather(*pending, return_exceptions=True): 在取消所有 pending 任务后，我们通常需要等待这些取消操作真正生效，即等待这些任务处理完 CancelledError 并最终退出。使用 gather 并设置 return_exceptions=True 可以确保即使有任务在取消过程中抛出其他异常，也不会中断整个流程。

注意事项与最佳实践

1. 任务取消的响应: task.cancel() 仅仅是发送一个取消请求。任务本身必须在适当的位置检查 CancelledError 并进行响应。如果任务内部有 await 调用，当 CancelledError 抛出时，await 表达式会立即抛出该异常。如果任务内部没有 await 调用或不处理异常，它将继续运行直到完成。

   async def my_cancellable_task():
       try:
           while True:
               # 模拟工作
               await asyncio.sleep(1)
               print("Task working...")
       except asyncio.CancelledError:
           print("Task was cancelled, performing cleanup...")
           # 执行清理操作
           await asyncio.sleep(0.1)
           print("Cleanup complete.")
       finally:
           print("Task finished.")

2. 资源清理: 确保你的异步任务在被取消或正常完成时，能够正确地关闭文件句柄、网络连接或其他系统资源。try...finally 块是实现这一点的常用模式。

3. 异常处理: 当从 task.result() 中获取结果时，如果任务内部发生了未捕获的异常，task.result() 会重新抛出该异常。务必在处理 done 任务时捕获这些潜在的异常，以防止主程序崩溃。

4. asyncio.wait_for 的替代: 另一种方法是使用 asyncio.wait_for 为每个单独的长时间运行任务设置超时。

   async def main_wait_for(client):
       try:
           await asyncio.wait_for(watch_task1(client), timeout=60)
       except asyncio.TimeoutError:
           print("watch_task1 timed out!")
   
       try:
           await asyncio.wait_for(watch_task2(client), timeout=60)
       except asyncio.TimeoutError:
           print("watch_task2 timed out!")
   
       # 对于多个任务，这种方式可能不如 asyncio.wait 灵活，
       # 因为它会串行处理超时，而不是全局并行。
       # 但如果只需要对单个任务施加超时，则非常适用。

   asyncio.wait_for 更适合对单个 awaitable 设置超时，如果超时，它会取消该 awaitable 并抛出 asyncio.TimeoutError。对于需要全局管理一组任务并在超时时统一处理的场景，asyncio.wait 提供了更强大的控制。

总结

通过 asyncio.wait 及其 timeout 参数，我们可以精确地控制一组并发异步任务的最大执行时间。这种方法不仅能够确保应用程序在预设时间内响应，还提供了清晰的机制来识别已完成和未完成的任务，并允许我们优雅地取消那些未能按时完成的任务，从而实现健壮且可控的异步编程实践。理解并正确应用 asyncio.wait 是构建高性能、可靠的 asyncio 应用程序的关键。