Python asyncio 任务管理实战：取消、超时与结构化并发

async def wait_until_cancelled(released):
    try:
        await asyncio.Event().wait()
    finally:
        released.set()

class RequestDeadlineExceeded(TimeoutError):
    pass


async def call_with_deadline(operation, seconds):
    deadline = asyncio.timeout(seconds)
    try:
        async with deadline:
            return await operation()
    except TimeoutError as error:
        if not deadline.expired():
            raise
        raise RequestDeadlineExceeded("请求超过总耗时限制") from error

async def bounded_map(items, limit, process):
    if limit < 1:
        raise ValueError("并发数必须大于 0")

    items = list(items)
    queue = asyncio.Queue()
    for index, item in enumerate(items):
        queue.put_nowait((index, item))

    results = {}

    async def worker():
        while True:
            try:
                index, item = queue.get_nowait()
            except asyncio.QueueEmpty:
                return
            results[index] = await process(item)

    async with asyncio.TaskGroup() as group:
        for _ in range(min(limit, len(items))):
            group.create_task(worker())

    return [results[index] for index in range(len(items))]

async def task_group_cancels_siblings():
    cleaned = asyncio.Event()

    async def slow_task():
        try:
            await asyncio.Event().wait()
        finally:
            cleaned.set()

    async def failing_task():
        await asyncio.sleep(0)
        raise ValueError("example failure")

    try:
        async with asyncio.TaskGroup() as group:
            group.create_task(slow_task())
            group.create_task(failing_task())
    except ExceptionGroup as errors:
        contains_expected_error = any(isinstance(error, ValueError) for error in errors.exceptions)
        return contains_expected_error and cleaned.is_set()
    return False

python -m unittest -v test_asyncio_guide.py

import asyncio
import unittest

from asyncio_guide import (
    RequestDeadlineExceeded,
    bounded_map,
    call_with_deadline,
    task_group_cancels_siblings,
    wait_until_cancelled,
)


class AsyncioGuideTests(unittest.IsolatedAsyncioTestCase):
    async def test_cancellation_runs_finally_and_propagates(self):
        released = asyncio.Event()
        task = asyncio.create_task(wait_until_cancelled(released))
        await asyncio.sleep(0)
        task.cancel()

        with self.assertRaises(asyncio.CancelledError):
            await task
        self.assertTrue(released.is_set())

    async def test_deadline_is_distinct_from_operation_timeout(self):
        async def slow_operation():
            await asyncio.sleep(0.05)
            return "late"

        with self.assertRaises(RequestDeadlineExceeded):
            await call_with_deadline(slow_operation, 0.005)

        async def operation_with_its_own_timeout():
            raise TimeoutError("operation failure")

        with self.assertRaisesRegex(TimeoutError, "operation failure"):
            await call_with_deadline(operation_with_its_own_timeout, 1)

    async def test_worker_count_is_bounded_and_results_keep_input_order(self):
        active = 0
        peak = 0

        async def process(item):
            nonlocal active, peak
            active += 1
            peak = max(peak, active)
            try:
                await asyncio.sleep(0)
                return item * 2
            finally:
                active -= 1

        result = await bounded_map([3, 1, 4, 2, 5], 2, process)
        self.assertEqual(result, [6, 2, 8, 4, 10])
        self.assertLessEqual(peak, 2)

    async def test_task_group_cancels_siblings_and_waits_for_cleanup(self):
        self.assertTrue(await task_group_cancels_siblings())


if __name__ == "__main__":
    unittest.main()