使用Python协程并发测试cdn响应速度

代码干净清爽才能看着赏心悦目：

#!/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"
    )


async def find_fastest_host(loop_interval=0.1) -> str:
    urls = list(CdnHost)
    results = [None] * len(urls)

    async def fetch(httpx_client, url, index):
        try:
            r = await httpx_client.get(url)
        except (httpx.ConnectError, httpx.ReadError):
            ...
        else:
            print(f"{url = }\n{r.elapsed = }")
            if r.status_code < 300:
                results[index] = r.content

    async with (anyio.create_task_group() as tg, httpx.AsyncClient() as client):
        for i, url in enumerate(urls):
            tg.start_soon(fetch, client, url, i)
        for _ in range(int(5 / loop_interval) + 1):
            if any(r is not None for r in results):
                tg.cancel_scope.cancel()
                break
            await anyio.sleep(0.1)
    for url, res in zip(urls, results):
        if res is not None:
            return url
    return urls[0]


async def main():
    with timeit("Sniff hosts"):
        url = await find_fastest_host()
    print("result:", url)


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

执行结果：

url = <CdnHost.unpkg: 'https://unpkg.com/swagger-ui-dist@5.9.0/swagger-ui.css'>
r.elapsed = datetime.timedelta(seconds=1, microseconds=418428)
Sniff hosts cost = 1.4775612354278564
result: https://unpkg.com/swagger-ui-dist@5.9.0/swagger-ui.css

2023-12-14 后记：

在此基础上开发了fastapi的文档cdn竞速插件fastapi-cdn-host用于解决/docs文档页面空白或加载不了的问题，它会自动查找本地目录是否有swagger-ui的js和css文件，有的话直接用它们来渲染页面；没有的话，对比几个免费CDN的响应速度，自动采用最快的那个。
示例代码如下：

from fastapi import FastAPI
from fastapi_cdn_host import monkey_patch_for_docs_ui  # pip install fastapi_cdn_host

app = FastAPI()
monkey_patch_for_docs_ui(app)  # 会自动替换/docs页面的CDN地址