异步任务阻塞问题解析
在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())
代码解释：
- asyncio.create_task(): 在将 awaitable 对象传递给 asyncio.wait 之前，最好使用 asyncio.create_task() 将它们包装成 Task 对象。这使得我们可以更好地管理和取消这些任务。
- await asyncio.wait(tasks, timeout=60): 这是核心部分。它会等待 tasks 列表中的所有任务，但最长等待60秒。一旦60秒过去，或者所有任务都已完成，它就会返回 done 和 pending 两个集合。
-
处理 done 任务:
- 遍历 done 集合中的每个任务。
- 使用 task.result() 来获取任务的返回值或重新抛出任务内部发生的任何异常。这是处理 asyncio.wait 返回的已完成任务的关键。
- 通过 task.get_name() 可以获取任务的名称（如果设置了）。
-
处理 pending 任务:
- 遍历 pending 集合中的每个任务。这些任务在超时时间内未能完成。
- 调用 task.cancel() 方法来请求取消这些任务。cancel() 方法会向任务内部注入一个 asyncio.CancelledError 异常。任务需要自行捕获并处理这个异常，以执行必要的清理工作。
- await asyncio.gather(*pending, return_exceptions=True): 在取消所有 pending 任务后，我们通常需要等待这些取消操作真正生效，即等待这些任务处理完 CancelledError 并最终退出。使用 gather 并设置 return_exceptions=True 可以确保即使有任务在取消过程中抛出其他异常，也不会中断整个流程。
注意事项与最佳实践
-
任务取消的响应: 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.") 资源清理: 确保你的异步任务在被取消或正常完成时，能够正确地关闭文件句柄、网络连接或其他系统资源。try...finally 块是实现这一点的常用模式。
异常处理: 当从 task.result() 中获取结果时，如果任务内部发生了未捕获的异常，task.result() 会重新抛出该异常。务必在处理 done 任务时捕获这些潜在的异常，以防止主程序崩溃。
-
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 应用程序的关键。
Python免费学习笔记（深入）：立即使用
在学习笔记中，你将探索 Python 的核心概念和高级技巧！
