如何在FastAPI中实现非阻塞的WebSocket端点？

问题描述 #

我刚接触Python异步编程，现在用FastAPI实现了一个WebSocket端点来执行长任务，但发现当这个任务运行时，其他HTTP端点会被阻塞。比如下面的WebSocket代码在执行长任务时，/health.check端点要等长任务结束才能响应。

我的WebSocket端点代码：

@router.websocket("/ws/catalog.check")
async def websocket_endpoint(ws: WebSocket, user: Annotated[UserWithGroupEntity, Depends(user_middleware_ws)]) -> None:
    manager = WSManager()
    await manager.connect(ws)

    try:
        raw_req = await ws.receive_json()
        req = CatalogCheckRequest.model_validate(raw_req)

        await manager.send_personal_json(CatalogCheckResponse(message=CatalogCheckMessage(data='Process started')),
                                         ws)

        #long-running task
        check_result = await create_task(CatalogHandler.catalog_check(user.group_name, req.source, req.csv))

        await manager.send_personal_json(CatalogCheckResponse(result=check_result), ws)
    except Exception as e:
        await manager.send_personal_json(CatalogCheckResponse(message=CatalogCheckMessage(data=repr(e), warning=True)),
                                         ws)

    await manager.disconnect(ws)

被阻塞的/health.check端点代码：

@router.get("/health.check")
def health_check(config: Annotated[Config, Depends(get_config)],
                 log: Annotated[Log, Depends(get_logger)]) -> Dict[str, Any]:
    try:
        conf = config.model_dump()

        log.info(conf)

        resp = {"status": "ok"}
    except Exception as e:
        resp = {"status": repr(e)}

    return resp

问题原因 #

出现阻塞的核心原因是FastAPI的异步事件循环是单线程的：

• 如果你的长任务CatalogHandler.catalog_check是同步实现（不管是CPU密集型计算，还是用了同步IO库如requests、同步文件操作），那么即使你用await create_task(...)包装，任务仍然会在事件循环线程中同步执行，完全占满线程，导致事件循环无法处理其他请求（包括路由分发、其他异步/同步端点的执行）。
• 你的/health.check是同步def端点，FastAPI会把它放到线程池执行，但事件循环被阻塞时，无法将新的同步请求提交到线程池，所以健康检查也会被卡住。

解决方案 #

根据长任务的类型，选择对应的方案让它不占用事件循环线程：

1. CPU密集型任务（如大量计算、数据处理） #

适合用进程池将任务放到独立进程执行，避免阻塞事件循环，还能利用多核CPU：

import asyncio
import concurrent.futures
from fastapi import Depends, WebSocket
from your_modules import CatalogHandler, CatalogCheckRequest, CatalogCheckResponse, WSManager, UserWithGroupEntity, user_middleware_ws

# 初始化全局进程池（或用依赖注入管理，更优雅）
process_executor = concurrent.futures.ProcessPoolExecutor()

@router.websocket("/ws/catalog.check")
async def websocket_endpoint(ws: WebSocket, user: Annotated[UserWithGroupEntity, Depends(user_middleware_ws)]) -> None:
    manager = WSManager()
    await manager.connect(ws)

    try:
        raw_req = await ws.receive_json()
        req = CatalogCheckRequest.model_validate(raw_req)

        await manager.send_personal_json(CatalogCheckResponse(message=CatalogCheckMessage(data='Process started')), ws)

        # 用进程池执行长任务，不阻塞事件循环
        loop = asyncio.get_running_loop()
        check_result = await loop.run_in_executor(
            process_executor,
            CatalogHandler.catalog_check,
            user.group_name,
            req.source,
            req.csv
        )

        await manager.send_personal_json(CatalogCheckResponse(result=check_result), ws)
    except Exception as e:
        await manager.send_personal_json(CatalogCheckResponse(message=CatalogCheckMessage(data=repr(e), warning=True)), ws)
    finally:
        await manager.disconnect(ws)

2. IO密集型任务（如调用外部API、读写文件、数据库操作） #

有两种选择：

方案A：用线程池执行同步IO任务

Python 3.9+提供了asyncio.to_thread，可以方便地把同步函数放到线程池执行，不阻塞事件循环：

import asyncio
from fastapi import Depends, WebSocket
from your_modules import CatalogHandler, CatalogCheckRequest, CatalogCheckResponse, WSManager, UserWithGroupEntity, user_middleware_ws

@router.websocket("/ws/catalog.check")
async def websocket_endpoint(ws: WebSocket, user: Annotated[UserWithGroupEntity, Depends(user_middleware_ws)]) -> None:
    manager = WSManager()
    await manager.connect(ws)

    try:
        raw_req = await ws.receive_json()
        req = CatalogCheckRequest.model_validate(raw_req)

        await manager.send_personal_json(CatalogCheckResponse(message=CatalogCheckMessage(data='Process started')), ws)

        # 用线程池执行同步IO任务
        check_result = await asyncio.to_thread(
            CatalogHandler.catalog_check,
            user.group_name,
            req.source,
            req.csv
        )

        await manager.send_personal_json(CatalogCheckResponse(result=check_result), ws)
    except Exception as e:
        await manager.send_personal_json(CatalogCheckResponse(message=CatalogCheckMessage(data=repr(e), warning=True)), ws)
    finally:
        await manager.disconnect(ws)

方案B：将任务改造成异步实现

如果可以修改CatalogHandler.catalog_check的代码，把同步IO操作替换成异步库（比如用aiohttp替代requests，asyncpg替代psycopg2），这样任务就能真正异步执行，await时会自动让出事件循环：

# 修改后的异步catalog_check示例
import aiohttp

class CatalogHandler:
    @staticmethod
    async def catalog_check(group_name, source, csv):
        # 异步请求外部API
        async with aiohttp.ClientSession() as session:
            async with session.get(source) as resp:
                data = await resp.text()
                # 其他异步操作...
        return check_result

然后WebSocket端点里直接用：

check_result = await CatalogHandler.catalog_check(user.group_name, req.source, req.csv)

额外优化建议 #

1. WSManager单例化：你现在每次WebSocket请求都创建新的WSManager()，这会导致每个连接的管理器独立，无法实现广播等功能。建议做成单例：

from fastapi import Depends
from functools import lru_cache

@lru_cache()
def get_ws_manager():
    return WSManager()  # 每次依赖注入返回同一个实例

@router.websocket("/ws/catalog.check")
async def websocket_endpoint(
    ws: WebSocket, 
    user: Annotated[UserWithGroupEntity, Depends(user_middleware_ws)],
    manager: WSManager = Depends(get_ws_manager)
) -> None:
    # 去掉manager = WSManager()
    await manager.connect(ws)
    # 后续代码不变

2. 避免创建无用的asyncio.Task：如果你的create_task是asyncio.create_task，那么await create_task(sync_func())等价于直接await sync_func()，因为create_task只是把任务加入事件循环，但同步函数仍然会阻塞事件循环，所以不需要额外包装。

备注：内容来源于stack exchange，提问作者Lev Marder