FastAPI中无法管理多用户WebSocket连接的问题排查

我想用字典存储并管理用户WebSocket连接，实现给每个用户发送自定义事件。下面是我的实现代码：

from fastapi import FastAPI, WebSocket
app = FastAPI()
con = {}
@app.websocket("/ws/{id}")
async def websocket_endpoint(websocket: WebSocket, id: str):
    try:
        print("connect", id)
        await websocket.accept()
        con[id] = websocket
        print("check", con.keys())
        for uid in con:
            print(uid)
            await con[uid].send_text("hello")
    except Exception as e:
        print(e)

运行日志如下：

INFO:     Uvicorn running on http://127.0.0.1:8000 (Press CTRL+C to quit)
INFO:     Started reloader process [26605] using StatReload
INFO:     Started server process [26610]
INFO:     Waiting for application startup.
INFO:     Application startup complete.
connect 1
INFO:     ('127.0.0.1', 48346) - "WebSocket /ws/1" [accepted]
check dict_keys(['1'])
1
INFO:     connection open
connect 2
INFO:     ('127.0.0.1', 48354) - "WebSocket /ws/2" [accepted]
check dict_keys(['1', '2'])
1
Unexpected ASGI message 'websocket.send', after sending 'websocket.close'.
INFO:     connection open

客户端1连接时一切正常，能收到"hello"；但客户端2连接时，服务器抛出异常Unexpected ASGI message 'websocket.send', after sending 'websocket.close'.，客户端2也收不到消息。我用uvicorn main:app --reload启动应用，请问哪里操作错了？

• 错误核心原因：你的WebSocket端点函数在完成消息发送后直接退出，FastAPI会默认在函数结束时关闭当前的WebSocket连接。当客户端2连接时，遍历字典给客户端1发消息，但此时客户端1的连接已经因端点函数执行完毕被关闭，给已关闭的连接发消息就会触发异常，同时客户端2的端点函数会因异常中断，导致它也收不到消息。

• 修复步骤：

  1. 保持WebSocket连接活跃，不要让端点函数直接退出。可以在函数末尾添加无限循环监听客户端消息，哪怕不处理消息，也能维持连接：

     from fastapi import FastAPI, WebSocket
     app = FastAPI()
     con = {}
     @app.websocket("/ws/{id}")
     async def websocket_endpoint(websocket: WebSocket, id: str):
         try:
             print("connect", id)
             await websocket.accept()
             con[id] = websocket
             print("check", con.keys())
             # 给所有在线连接发消息
             for uid in con:
                 print(uid)
                 await con[uid].send_text("hello")
             # 保持连接活跃，监听客户端消息
             while True:
                 data = await websocket.receive_text()
                 # 可在此添加消息处理逻辑，或仅维持连接
         except Exception as e:
             print(e)
             # 连接关闭时从字典移除对应条目
             if id in con:
                 del con[id]
     

  2. 处理连接关闭的情况，确保断开连接的用户从字典中移除，避免后续给已关闭的连接发消息。
  3. 生产环境注意：普通字典在多线程/多进程环境下存在线程安全问题，建议用asyncio.Lock保护字典的读写操作，或者使用线程安全的数据结构。

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