Python Asyncio：优雅地管理与终止长时间运行的任务

import asyncio

stop = False

async def watch_task1(client):
    while not stop:
        # 假设 client.ws.get_data() 可能长时间没有数据返回
        await client.ws.get_data()
        print("Task 1 received data")

async def watch_task2(client):
    while not stop:
        # 假设 client.ws.get_news() 也可能长时间没有数据返回
        await client.ws.get_news()
        print("Task 2 received news")

async def stop_after(delay_seconds):
    global stop
    await asyncio.sleep(delay_seconds)
    print(f"Stopping after {delay_seconds} seconds...")
    stop = True

class MockClient:
    async def get_data(self):
        await asyncio.sleep(100) # 模拟长时间阻塞
    async def get_news(self):
        await asyncio.sleep(100) # 模拟长时间阻塞
    async def sleep(self, delay):
        await asyncio.sleep(delay)

async def main_gather():
    client = MockClient()
    tasks = [
        watch_task1(client),
        watch_task2(client),
        stop_after(5), # 尝试在5秒后停止
    ]

    try:
        # 使用 gather，即使 stop 变为 True，阻塞的 get_data/get_news 仍会阻止 gather 完成
        await asyncio.gather(*tasks, return_exceptions=True)
    except Exception as e:
        print(f"An exception occurred: {e}")
    print("Main gather finished.")

# 运行 main_gather() 会发现程序在5秒后并不会立即停止，而是会等待 get_data/get_news 结束

import asyncio

stop = False # 这是一个共享状态，用于控制协程的内部循环

async def watch_task1(client):
    try:
        while not stop:
            print("Task 1: Waiting for data...")
            await client.ws.get_data() # 可能会阻塞
            print("Task 1: Data received.")
    except asyncio.CancelledError:
        print("Task 1: Cancelled.")
    finally:
        print("Task 1: Exiting.")

async def watch_task2(client):
    try:
        while not stop:
            print("Task 2: Waiting for news...")
            await client.ws.get_news() # 可能会阻塞
            print("Task 2: News received.")
    except asyncio.CancelledError:
        print("Task 2: Cancelled.")
    finally:
        print("Task 2: Exiting.")

# MockClient 保持不变
class MockClient:
    async def get_data(self):
        # 模拟长时间阻塞，但为了演示，将其缩短
        await asyncio.sleep(5)
    async def get_news(self):
        # 模拟长时间阻塞
        await asyncio.sleep(5)
    async def sleep(self, delay):
        await asyncio.sleep(delay)

async def main_wait_timeout():
    client = MockClient()
    tasks_to_run = [
        asyncio.create_task(watch_task1(client)), # 显式创建任务
        asyncio.create_task(watch_task2(client)),
    ]

    print("Starting tasks with a 3-second timeout...")
    # 设置一个全局超时，例如3秒
    done, pending = await asyncio.wait(tasks_to_run, timeout=3, return_when=asyncio.ALL_COMPLETED)

    print("\n--- After asyncio.wait ---")
    print(f"Completed tasks ({len(done)}):")
    for task in done:
        try:
            # 获取任务结果，如果任务抛出异常，这里会重新抛出
            result = task.result()
            print(f"  Task finished successfully: {task.get_name()}, Result: {result}")
        except asyncio.CancelledError:
            print(f"  Task {task.get_name()} was cancelled (expected for pending tasks).")
        except Exception as e:
            print(f"  Task {task.get_name()} raised an exception: {e}")

    print(f"\nPending tasks ({len(pending)}):")
    for task in pending:
        print(f"  Task {task.get_name()} is still pending. Cancelling...")
        task.cancel() # 取消未完成的任务
        try:
            await task # 等待任务真正结束，以便其处理 CancelledError
        except asyncio.CancelledError:
            print(f"  Task {task.get_name()} successfully cancelled and cleaned up.")
        except Exception as e:
            print(f"  Task {task.get_name()} raised an exception during cancellation cleanup: {e}")

    print("Main wait_timeout finished.")

# 运行主函数
if __name__ == "__main__":
    asyncio.run(main_wait_timeout())

async def example_wait_for():
    client = MockClient()
    try:
        print("Attempting to get data with a 2-second timeout...")
        # watch_task1 内部的循环会因为超时而中断
        await asyncio.wait_for(watch_task1(client), timeout=2)
        print("watch_task1 finished within timeout.")
    except asyncio.TimeoutError:
        print("watch_task1 timed out!")
    except Exception as e:
        print(f"watch_task1 raised an unexpected exception: {e}")
    print("Example wait_for finished.")

# 如果在 main_wait_timeout 之后运行
# asyncio.run(example_wait_for())