WebSocket未正常关闭：FastAPI日志流场景下Uvicorn退出阻塞

问题说明 #

通过FastAPI实现后端到前端的单向日志流（模拟tail -f功能），已完成WebSocket端点和测试客户端代码，能实现实时日志推送，但存在以下问题：

• 客户端断开连接后，执行Ctrl+C关闭Uvicorn服务时，控制台会提示等待后台任务完成
• 尽管Uvicorn日志显示connection closed，但WebSocket任务并未正常终止，导致服务无法快速退出

原控制器代码 #

"""
WebSocket endpoint to stream PM2 logs.
"""
try:
    await websocket.accept()
    filepath = "mylogs.txt"  # Adjust to your PM2 log file path
    with open(filepath) as f:
        while websocket.application_state ==  WebSocketState.CONNECTED and  websocket.client_state ==  WebSocketState.CONNECTED:
            if content:=f.read(4096):
                await websocket.send_text(content)
            else:
                await asyncio.sleep(0.5)
except Exception as e:
    print(e)

原测试客户端代码 #

import sys
import asyncio
import websockets

async def connect_websocket():
    url = "ws://localhost:8000/logs/ws"
    try:
        async with websockets.connect(url) as wsoc:
            print(wsoc.local_address)
            while True:
                content = await wsoc.recv()
                sys.stdout.write(str(content))
                sys.stdout.flush()
    except websockets.exceptions.ConnectionClosed as e:
        print("连接意外关闭")
        print(e)
    except Exception as e:
        print(e)

if __name__ == "__main__":
    asyncio.run(connect_websocket())

服务退出时的日志 #

INFO   ('127.0.0.1', 55536) - "WebSocket /logs/ws"
      INFO   connection open
      INFO   connection closed
      INFO   Shutting down
      INFO   Waiting for background tasks to complete. (CTRL+C to force quit)

问题根源 #

1. 连接状态判断不准确：仅通过websocket.application_state和client_state判断连接状态，无法及时感知客户端主动断开的情况
2. 未处理WebSocket断开异常：客户端断开后，await websocket.send_text()会抛出WebSocketDisconnect异常，但原代码仅捕获通用Exception，且未在异常后终止任务
3. 同步文件读取+睡眠阻塞任务取消：asyncio.sleep(0.5)会阻塞任务，导致Uvicorn的取消信号无法及时被处理

修复方案 #

修改后的WebSocket控制器代码 #

from fastapi import WebSocket, WebSocketDisconnect
import asyncio
from starlette.websockets import WebSocketState

async def log_stream_websocket(websocket: WebSocket):
    await websocket.accept()
    filepath = "mylogs.txt"
    try:
        with open(filepath, "r") as f:
            # 可选：跳转到文件末尾，避免推送历史日志（需要历史日志可删除此行）
            f.seek(0, 2)
            while True:
                content = f.read(4096)
                if content:
                    await websocket.send_text(content)
                # 使用带超时的等待，让任务能响应取消信号
                try:
                    await asyncio.wait_for(asyncio.sleep(0.5), timeout=1.0)
                except asyncio.TimeoutError:
                    pass
                # 双重检查连接状态，确保及时终止任务
                if not websocket.client_state == WebSocketState.CONNECTED:
                    break
    except WebSocketDisconnect:
        # 捕获客户端主动断开异常，直接终止任务
        pass
    except asyncio.CancelledError:
        # 捕获服务终止的取消信号，主动关闭连接
        await websocket.close()
    except Exception as e:
        print(f"日志流异常: {e}")
        await websocket.close()

关键修改点 #

• 精准捕获异常：新增WebSocketDisconnect和asyncio.CancelledError的捕获逻辑，分别处理客户端断开、服务终止两种场景
• 响应取消信号：将asyncio.sleep(0.5)替换为带超时的asyncio.wait_for，确保任务能及时感知Uvicorn的退出请求
• 双重连接校验：在循环中定期检查客户端连接状态，结合异常处理形成双重保障，确保任务及时终止
• 可选优化：添加f.seek(0, 2)跳转到文件末尾，避免客户端连接时推送历史冗余日志

验证方法 #

1. 启动服务：uv run fastapi dev --no-reload
2. 运行测试客户端，建立WebSocket连接
3. 关闭客户端，等待Uvicorn日志显示connection closed
4. 执行Ctrl+C关闭服务，此时服务会直接退出，不再提示等待后台任务完成

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