Starlette项目中WebSocket连接断开异常处理指南

在基于Starlette框架开发WebSocket应用时，开发者可能会遇到一个常见的运行时错误："Cannot call 'receive' once a disconnect message has been received"。这个问题通常发生在WebSocket连接断开时，但开发者没有正确处理断开逻辑的情况下。

问题现象

当客户端突然断开WebSocket连接（例如用户关闭浏览器标签或刷新页面），服务器端可能会出现以下错误堆栈：

RuntimeError: Cannot call "receive" once a disconnect message has been received.

这个错误表明在已经接收到断开消息后，代码仍然尝试调用WebSocket的receive方法。

问题根源

该问题通常由以下两种开发模式引起：

1. 错误使用WebSocketEndpoint类：开发者直接在on_connect方法中使用websocket.iter_json()或类似方法，而没有使用框架提供的on_receive回调。

2. 任务取消处理不当：在使用任务组(TaskGroup)或异步任务时，没有正确处理连接断开时的任务取消逻辑。

最佳实践解决方案

方案一：正确使用WebSocketEndpoint类

Starlette框架为WebSocket提供了专门的WebSocketEndpoint基类，它已经封装了正确的消息处理循环。开发者应该重写其三个关键方法：

class ChatRoomWebsocket(WebSocketEndpoint):
    async def on_connect(self, websocket: WebSocket):
        await websocket.accept()
        # 初始化工作

    async def on_receive(self, websocket: WebSocket, data):
        # 处理接收到的消息
        pass

    async def on_disconnect(self, websocket: WebSocket, close_code: int):
        # 清理资源
        pass

方案二：结合Broadcaster的正确实现

当需要与Broadcaster等消息广播系统集成时，可以这样实现：

class ChatRoomWebsocket(WebSocketEndpoint):
    async def on_connect(self, websocket: WebSocket):
        await websocket.accept()
        self._listener_task = asyncio.create_task(self.chatroom_ws_sender(websocket))

    async def on_receive(self, websocket: WebSocket, data):
        await broadcast.publish(channel="chatroom", message=data)

    async def chatroom_ws_sender(self, websocket) -> None:
        async with broadcast.subscribe(channel="chatroom") as subscriber:
            async for event in subscriber:
                await websocket.send_text(event.message)

    async def on_disconnect(self, websocket: WebSocket, close_code: int) -> None:
        if hasattr(self, '_listener_task'):
            self._listener_task.cancel()

关键注意事项

1. 不要在on_connect中直接使用iter_json：这会绕过框架的消息循环机制，导致连接状态管理混乱。

2. 正确处理异步任务：任何在连接期间创建的长期运行任务都必须在on_disconnect中妥善取消。

3. 资源清理：确保在断开连接时释放所有相关资源，如数据库连接、订阅等。

4. 异常处理：为可能出现的网络异常添加适当的错误处理逻辑。

总结

通过遵循Starlette框架的设计模式，特别是正确使用WebSocketEndpoint类及其生命周期方法，可以避免大多数WebSocket连接管理问题。对于需要集成消息广播系统的场景，确保在断开连接时正确取消订阅和清理资源是关键。这些最佳实践不仅能解决当前的运行时错误，还能提高WebSocket应用的健壮性和可维护性。