如何在anyio.TaskGroup上下文内取消任务？CDN测速场景示例

你当前通过自定义StopNow异常终止任务的方式，可以替换为anyio原生的任务取消机制，无需自定义异常，同时还能简化轮询逻辑，让代码更简洁高效。

修改后的完整代码 #

#!/usr/bin/env python3.11
import time
from contextlib import contextmanager
from enum import StrEnum

import anyio
import httpx


@contextmanager
def timeit(msg: str):
    start = time.time()
    yield
    cost = time.time() - start
    print(msg, f"{cost = }")


class CdnHost(StrEnum):
    jsdelivr = "https://cdn.jsdelivr.net/npm/swagger-ui-dist@5.9.0/swagger-ui.css"
    unpkg = "https://unpkg.com/swagger-ui-dist@5.9.0/swagger-ui.css"
    cloudflare = (
        "https://cdnjs.cloudflare.com/ajax/libs/swagger-ui/5.9.0/swagger-ui.css"
    )


TIMEOUT = 5


async def find_fastest_host(timeout=TIMEOUT) -> str:
    urls = list(CdnHost)
    success_url = None
    success_lock = anyio.Lock()

    async def fetch(client, url):
        nonlocal success_url
        try:
            r = await client.get(url, timeout=timeout)
        except (httpx.ConnectError, httpx.ReadError, httpx.TimeoutException):
            return
        print(f"{url = }\n{r.elapsed = }")
        if r.status_code < 300:
            async with success_lock:
                if success_url is None:
                    success_url = url
                    # 触发任务组取消所有子任务
                    anyio.get_current_task_group().cancel()

    try:
        async with anyio.create_task_group() as tg:
            with anyio.move_on_after(timeout):
                async with httpx.AsyncClient() as client:
                    for url in urls:
                        tg.start_soon(fetch, client, url)
    except anyio.CancelledError:
        # 主动取消任务属于预期行为，忽略该异常
        pass

    # 返回第一个成功的URL，超时则返回默认第一个CDN
    return success_url or urls[0]


async def main():
    with timeit("Sniff hosts"):
        url = await find_fastest_host()
    print("cdn host options:", list(CdnHost))
    print("fastest result:", url)


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

关键改动说明 #

• 移除自定义异常与轮询逻辑：删掉了StopNow异常和循环检查结果的代码，改为让每个fetch任务在成功获取有效响应后直接触发任务取消，避免了轮询带来的延迟。
• 使用anyio原生任务取消机制：调用anyio.get_current_task_group().cancel()可以直接终止任务组内的所有子任务，这是框架原生支持的方式，无需自定义异常。
• 线程安全的结果记录：通过anyio.Lock确保多个任务不会同时修改success_url，保证我们拿到的是第一个成功响应的CDN地址（也就是最快的那个）。
• 简化异常处理：任务被主动取消时会抛出CancelledError，我们只需捕获并忽略这个异常即可，因为这是预期内的流程。

内容的提问来源于stack exchange，提问作者Waket Zheng