Asyncio任务超时与优雅取消技巧

Asyncio异步编程中，任务超时与优雅终止是构建健壮应用的关键。当任务包含长时间I/O等待且可能无限期阻塞时，传统的全局停止标志往往失效。本文聚焦于解决这一难题，深入探讨了两种核心Asyncio超时管理方案。首先，介绍了`asyncio.wait_for`，它能为单个任务设置超时，当任务在指定时间内未完成时抛出`asyncio.TimeoutError`，实现精细化控制。其次，详细阐述了`asyncio.wait`，它可为一组任务设定整体超时，并返回已完成和未完成的任务集合，便于程序按预期退出，避免资源悬挂。通过合理运用这两个工具，并结合异常处理和任务取消策略，开发者能够构建更健壮、可控的Asyncio应用，确保即使面对长时间阻塞的任务，程序也能在预设时间内优雅地完成或终止。

在Asyncio应用中，我们经常需要运行多个并发任务。然而，当某些任务内部包含长时间的I/O等待（如client.ws.get_data()或client.ws.get_news()），并且这些等待可能在很长时间内都没有数据返回时，即使设置了全局的停止标志，这些任务也无法及时响应，导致程序无法在预期时间内退出。传统的asyncio.gather会等待所有任务完成，这在任务可能无限期阻塞的情况下并不适用。为了解决这一问题，Asyncio提供了更精细的任务管理和超时机制。

1. 问题分析：全局停止标志的局限性

在上述场景中，尽管我们尝试通过stop_after任务设置stop标志来通知watch_tasks停止，但如果await client.ws.get_data()或await client.ws.get_news()处于阻塞状态，并且在stop标志设置为True之后仍未收到任何数据，那么这些await调用将继续等待，导致任务无法检查stop标志并退出循环。这使得程序无法在设定的超时时间（例如60秒）内停止。

2. 解决方案一：使用 asyncio.wait_for 为单个任务设置超时

asyncio.wait_for(fut, timeout) 函数可以用于为单个可等待对象（如协程或Future）设置一个超时时间。如果在指定时间内协程没有完成，它将抛出asyncio.TimeoutError异常。这使得我们可以为每个可能长时间阻塞的I/O操作设置独立的超时。

示例：

import asyncio

async def watch_task1():
    try:
        while True:
            # 假设 client.ws.get_data() 是一个可能阻塞的IO操作
            # 为此操作设置一个短期的超时，例如5秒
            data = await asyncio.wait_for(client.ws.get_data(), timeout=5)
            print(f"Received data: {data}")
    except asyncio.TimeoutError:
        print("watch_task1: get_data timed out, retrying...")
        # 可以在这里选择重试、记录日志或退出
    except asyncio.CancelledError:
        print("watch_task1 cancelled.")
    except Exception as e:
        print(f"watch_task1 error: {e}")

async def main():
    # 启动任务，并在外部处理其可能的超时
    task1 = asyncio.create_task(watch_task1())
    try:
        await asyncio.sleep(60) # 主程序运行60秒
    finally:
        task1.cancel() # 取消任务
        await task1 # 等待任务真正结束

注意事项：

• asyncio.wait_for只对它所包装的awaitable有效。如果watch_task1内部有多个await点，每个点都需要单独考虑超时。
• 当超时发生时，asyncio.TimeoutError会被抛出，任务需要捕获并处理此异常。

3. 解决方案二：使用 asyncio.wait 进行任务组超时管理

asyncio.wait(aws, *, timeout=None, return_when=ALL_COMPLETED) 是一个更强大的工具，它允许我们等待一组可等待对象，并提供一个全局的超时机制。它返回两个集合：done（已完成的任务）和pending（未完成的任务）。

关键参数：

• aws: 可等待对象的集合（列表、元组或集合）。
• timeout: 可选的超时时间（秒）。如果指定，wait将在该时间后返回，无论所有任务是否完成。
• return_when: 决定何时返回，默认为ALL_COMPLETED。其他选项包括FIRST_COMPLETED（任一任务完成即返回）和FIRST_EXCEPTION（任一任务抛出异常即返回）。

示例：

import asyncio
import time

# 模拟一个可能长时间阻塞的WebSocket客户端
class MockWebSocketClient:
    async def get_data(self):
        print("get_data: Waiting for data...")
        await asyncio.sleep(100) # 模拟长时间阻塞
        return "Some data"

    async def get_news(self):
        print("get_news: Waiting for news...")
        await asyncio.sleep(120) # 模拟更长时间阻塞
        return "Some news"

# 模拟全局停止标志
stop = False
client = MockWebSocketClient()

async def watch_task1():
    global stop
    while not stop:
        try:
            # 这里的阻塞由 asyncio.wait 的外部超时控制
            data = await client.ws.get_data()
            print(f"Watch Task 1 received: {data}")
        except asyncio.CancelledError:
            print("Watch Task 1 was cancelled.")
            break
        except Exception as e:
            print(f"Watch Task 1 error: {e}")
            break
    print("Watch Task 1 exiting.")

async def watch_task2():
    global stop
    while not stop:
        try:
            news = await client.ws.get_news()
            print(f"Watch Task 2 received: {news}")
        except asyncio.CancelledError:
            print("Watch Task 2 was cancelled.")
            break
        except Exception as e:
            print(f"Watch Task 2 error: {e}")
            break
    print("Watch Task 2 exiting.")

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

async def main():
    global stop
    stop = False # 重置停止标志
    client.ws = MockWebSocketClient() # 实例化模拟客户端

    tasks_to_run = [
        asyncio.create_task(watch_task1(), name="WatchTask1"),
        asyncio.create_task(watch_task2(), name="WatchTask2"),
        asyncio.create_task(stop_after_task(60), name="StopAfterTask"),
    ]

    print("Main: Starting tasks with a 60-second timeout...")
    # 使用 asyncio.wait 设置整体超时
    done, waiting = await asyncio.wait(tasks_to_run, timeout=60, return_when=asyncio.ALL_COMPLETED)
    # 注意：ALL_COMPLETED在这里的意思是，要么所有任务完成，要么超时时间到。
    # 如果超时时间到，则返回当前已完成的任务和未完成的任务。

    print(f"\nMain: Operation completed after {60} seconds.")
    print(f"Main: Number of done tasks: {len(done)}")
    print(f"Main: Number of waiting tasks: {len(waiting)}")

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

    # 处理未完成的任务：通常需要取消它们以释放资源
    if waiting:
        print("\nMain: Cancelling remaining waiting tasks...")
        for task in waiting:
            task.cancel()
            try:
                # 等待任务真正被取消并结束，捕获 CancelledError
                await task
            except asyncio.CancelledError:
                print(f"Main: Task '{task.get_name()}' successfully cancelled.")
            except Exception as e:
                print(f"Main: Task '{task.get_name()}' ended with error after cancellation: {e}")
    else:
        print("\nMain: All tasks completed within the timeout.")

    print("\nMain: All tasks processed. Exiting.")

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

代码解析与注意事项：

1. asyncio.create_task(): 推荐使用asyncio.create_task()来创建任务，而不是直接将协程对象放入asyncio.wait。create_task会立即调度协程运行，并返回一个Task对象，这更符合任务管理的语义。
2. done 和 waiting 集合: asyncio.wait返回的done集合包含了在超时时间内完成的任务，而waiting集合则包含了那些在超时时间到达时仍未完成的任务。
3. 处理 done 任务: 遍历done集合中的任务，通过task.result()获取其结果。如果任务在执行过程中抛出异常，task.result()也会重新抛出该异常，因此需要进行异常捕获。
4. 处理 waiting 任务: 这是关键一步。对于waiting集合中的任务，它们在超时时并未完成，通常需要显式地取消它们以确保资源被释放。调用task.cancel()会向任务发送一个取消请求，任务内部如果正确处理asyncio.CancelledError，就会停止执行。
5. 等待取消完成: 在调用task.cancel()之后，最好await task，以确保任务有机会捕获CancelledError并执行清理工作，然后再真正退出。
6. asyncio.CancelledError: 任务在被取消时会抛出asyncio.CancelledError。在任务内部，应捕获此异常以执行必要的清理工作（如关闭文件句柄、网络连接等），然后通常重新抛出或直接返回，以便外部能够感知到任务已被取消。

4. 总结与最佳实践

• 选择合适的工具:
  • 如果需要为单个可能阻塞的I/O操作设置超时，并且希望在超时后进行重试或特定处理，使用asyncio.wait_for。
  • 如果需要管理一组任务的整体生命周期，并在指定时间后检查哪些任务已完成、哪些未完成，并对未完成任务进行统一处理（如取消），则asyncio.wait是更强大的选择。
• 优雅地取消任务: 当使用asyncio.wait进行超时管理时，对waiting集合中的任务调用task.cancel()是至关重要的。同时，任务内部必须正确地捕获和处理asyncio.CancelledError，以确保资源得到释放，避免内存泄漏或连接悬挂。
• 错误处理: 无论是asyncio.wait_for还是asyncio.wait，都需要考虑任务内部可能抛出的异常。对于asyncio.wait，task.result()会重新抛出任务内部的异常，而对于asyncio.wait_for，超时会抛出asyncio.TimeoutError。
• 避免全局标志的陷阱: 对于依赖于外部I/O事件才能解除阻塞的任务，仅仅依赖全局停止标志是不足够的。必须结合Asyncio提供的超时和取消机制来强制任务在预期时间内停止。

通过合理运用asyncio.wait_for和asyncio.wait，我们可以构建出更加健壮和可控的Asyncio应用程序，确保即使面对长时间阻塞的任务，程序也能在预设的时间内优雅地完成或终止。

好了，本文到此结束，带大家了解了《Asyncio任务超时与优雅取消技巧》，希望本文对你有所帮助！关注golang学习网公众号，给大家分享更多文章知识！