多客户端连接FastAPI WebSocket出现ConnectionClosedError问题求助

错误现象 #

多客户端连接/ws端点时，抛出如下错误：

await self.close()
File "/usr/local/lib/python3.10/site-packages/websockets/legacy/protocol.py", line 766, in close
    await self.write_close_frame(Close(code, reason))
File "/usr/local/lib/python3.10/site-packages/websockets/legacy/protocol.py", line 1232, in write_close_frame
    await self.write_frame(True, OP_CLOSE, data, _state=State.CLOSING)
File "/usr/local/lib/python3.10/site-packages/websockets/legacy/protocol.py", line 1205, in write_frame
    await self.drain()
File "/usr/local/lib/python3.10/site-packages/websockets/legacy/protocol.py", line 1194, in drain
    await self.ensure_open()
File "/usr/local/lib/python3.10/site-packages/websockets/legacy/protocol.py", line 935, in ensure_open
    raise self.connection_closed_exc()
websockets.exceptions.ConnectionClosedError: sent 1000 (OK); no close frame received

错误根源分析 #

1. WebSocket消费者重复实例化：每个客户端连接/ws时都会创建新的KafkaConsumer实例，每个实例启动独立的AIOKafkaConsumer监听同一topic和group id。Kafka会将消息分发给group内的消费者，导致多客户端时消息被分散消费，同时重复实例引发连接冲突，导致WebSocket连接异常。
2. send_via_websocket方法逻辑缺陷：
   • 嵌套的while True和async for msg循环导致客户端断开后无法正确清理Kafka消费者连接
   • 未处理WebSocketDisconnect外的其他异常，导致资源泄漏
3. ConnectionManager的disconnect方法未处理异常：当WebSocket连接已被移除或不存在时，调用remove会抛出错误，中断后续流程
4. Kafka消费者未正确关闭：客户端断开后，对应的websocketconsumer未调用stop方法，导致Kafka连接资源泄漏
5. 代码拼写错误：KAFKA_WEEBSOCKET_GROUP_ID、paylaod、country_codeal等拼写错误会引发隐性逻辑问题

修复方案 #

1. 全局复用WebSocket Kafka消费者实例 #

启动时创建一个全局WebSocket消费者，统一接收消息后通过ConnectionManager广播给所有客户端，避免每个客户端创建独立消费者。

2. 重构WebSocket端点逻辑 #

客户端连接时仅处理连接管理，不再启动新的Kafka消费者，由全局消费者统一推送消息。

3. 完善异常处理与资源清理 #

• 给Kafka消费者方法添加finally块确保连接停止
• 给ConnectionManager的disconnect方法添加异常捕获
• 广播消息时遍历连接副本，避免列表变化引发错误
• 处理WebSocket连接中的各类异常

4. 修复代码拼写错误 #

修改后的完整代码 #

KafkaConsumer类 #

class KafkaConsumer:
    """
    Kafka Consumer class to consume messages from Kafka topic
    """

    def __init__(self, loop):
        """
        :param loop: asyncio event loop

        Initialize the consumer with the event loop and Kafka settings
        having two consumers with different group id
        one for normal database query and another for websocket"""
        self.consumer = AIOKafkaConsumer(
            "whatsapp-bot",
            loop=loop,
            bootstrap_servers=settings.KAFKA_BOOTSTRAP_SERVERS,
            group_id=settings.KAFKA_GROUP_ID,
        )
        self.websocket_consumer = AIOKafkaConsumer(
            "whatsapp-bot",
            loop=loop,
            bootstrap_servers=settings.KAFKA_BOOTSTRAP_SERVERS,
            group_id=settings.KAFKA_WEBSOCKET_GROUP_ID,  # 修复拼写错误
        )

    async def consume(self):
        """
        Consume messages from Kafka topic and create users based on the received data.
        """
        try:
            await self.consumer.start()
        except Exception as e:
            logging.error(f"Database consumer start failed: {str(e)}")
            return

        try:
            async for msg in self.consumer:
                try:
                    string_data = msg.value.decode("utf-8")
                    data = json.loads(string_data)
                    phone, country_code = extract_phone_number_and_country_code(data["phone_number"])
                    payload = UserCreate(phone=phone, country_code=country_code)  # 修复拼写错误
                    if not await user_repository.get_by_phone(phone):
                        await user_repository.create(payload)
                        logging.info("New user created Successfully")
                except Exception as e:
                    logging.error(f"Error processing database message: {str(e)}")
        finally:
            await self.consumer.stop()
            logging.info("Database consumer stopped")

    async def websocket_broadcast_worker(self):
        """
        Global worker to consume Kafka messages and broadcast to all connected WebSocket clients
        """
        try:
            await self.websocket_consumer.start()
            logging.info("WebSocket Kafka consumer started")
        except Exception as e:
            logging.error(f"WebSocket consumer start failed: {str(e)}")
            return

        try:
            async for msg in self.websocket_consumer:
                try:
                    string_data = msg.value.decode("utf-8")
                    logging.info(f"Message received from Kafka for WebSocket: {string_data}")
                    await websocket_manager.broadcast(string_data)
                    logging.info("Message broadcast to all clients")
                except Exception as e:
                    logging.error(f"Error processing WebSocket message: {str(e)}")
        finally:
            await self.websocket_consumer.stop()
            logging.info("WebSocket consumer stopped")

ConnectionManager类 #

from fastapi import WebSocket
import logging

class ConnectionManager:
    """
    Handling websocket connections for multiple clients"""
    def __init__(self):
        self.active_connections: list[WebSocket] = []

    async def connect(self, websocket: WebSocket):
        await websocket.accept()
        self.active_connections.append(websocket)
        logging.info(f"Client connected: {id(websocket)}")

    def disconnect(self, websocket: WebSocket):
        try:
            self.active_connections.remove(websocket)
            logging.info(f"Client disconnected: {id(websocket)}")
        except ValueError:
            logging.warning(f"Trying to disconnect non-existent client: {id(websocket)}")

    async def send_personal_message(self, message: str, websocket: WebSocket):
        try:
            await websocket.send_text(message)
        except Exception as e:
            logging.error(f"Failed to send personal message: {str(e)}")
            self.disconnect(websocket)

    async def broadcast(self, message: str):
        # 遍历副本，避免遍历过程中列表变化导致错误
        for connection in list(self.active_connections):
            try:
                await connection.send_text(message)
            except Exception as e:
                logging.error(f"Failed to broadcast to client {id(connection)}: {str(e)}")
                self.disconnect(connection)

websocket_manager = ConnectionManager()

main.py #

import asyncio
from fastapi import FastAPI, WebSocket
from fastapi.middleware.cors import CORSMiddleware
import logging

# 导入其他依赖（database, settings, KafkaConsumer, websocket_manager等）

loop = asyncio.get_event_loop()
metadata.create_all(engine)
app = FastAPI(title=settings.PROJECT_NAME, openapi_url=f"{settings.API_V1_STR}/openapi.json")

kafka_consumer = None

@app.on_event("startup")
async def startup():
    await database.connect()
    await consume_kafka()

@app.on_event("shutdown")
async def shutdown():
    await database.disconnect()
    # 停止Kafka消费者
    if kafka_consumer:
        pass

# Set all CORS enabled origins
app.add_middleware(
    CORSMiddleware,
    allow_origins=["*"],
    allow_credentials=True,
    allow_methods=["*"],
    allow_headers=["*"],
)

app.include_router(api_router, prefix=settings.API_V1_STR)

@app.websocket("/ws")
async def websocket_endpoint(websocket: WebSocket):
    """
    Websocket endpoint to handle client connections"""
    await websocket_manager.connect(websocket)
    try:
        # 保持连接存活，监听客户端断开事件
        while True:
            await websocket.receive_text()
    except WebSocketDisconnect:
        websocket_manager.disconnect(websocket)
    except Exception as e:
        logging.error(f"WebSocket error: {str(e)}")
        websocket_manager.disconnect(websocket)

async def consume_kafka():
    global kafka_consumer
    kafka_consumer = KafkaConsumer(loop)
    # 启动数据库消费任务
    asyncio.create_task(kafka_consumer.consume())
    # 启动WebSocket广播任务
    asyncio.create_task(kafka_consumer.websocket_broadcast_worker())

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