Celery任务完成后WebSocket广播消息异常问题求助

核心问题梳理 #

你的实现中导致WebSocket连接不稳定、消息推送失败的主要原因包括：

• 无节制的Ping消息发送（/connect端点的无限循环无延迟Ping）导致连接过载
• Redis PubSub连接重复初始化，导致旧订阅失效
• 消息发送时未正确处理无效/已关闭的WebSocket连接
• 消息序列化重复处理，导致客户端解析异常
• 连接断开后未清理管理器中的无效连接记录

修复后的代码实现 #

import asyncio
import redis.asyncio as aioredis
import json
from fastapi import WebSocket, WebSocketDisconnect, APIRouter, FastAPI, HTTPException
from typing import Dict, Set, Any
import enum

from app.api.utils.verifyJWTToken import verify_jwt_token

app = FastAPI()
web_socket_router = APIRouter()

class WebSocketState(enum.Enum):
    CONNECTING = 0
    CONNECTED = 1
    DISCONNECTED = 2

class RedisPubSubManager:
    """Redis发布订阅管理器"""
    def __init__(self, host="localhost", port=6379):
        self.redis_host = host
        self.redis_port = port
        self.redis_connection = None
        self.pubsub = None
        self.subscribed_channels: Set[str] = set()

    async def _get_redis_connection(self) -> aioredis.Redis:
        """获取Redis连接（单例模式）"""
        if not self.redis_connection or self.redis_connection.connection_pool.closed:
            self.redis_connection = aioredis.Redis(
                host=self.redis_host, port=self.redis_port, auto_close_connection_pool=False
            )
        return self.redis_connection

    async def connect(self) -> None:
        """初始化PubSub客户端（仅执行一次）"""
        if not self.pubsub:
            await self._get_redis_connection()
            self.pubsub = self.redis_connection.pubsub()

    async def publish(self, channel: str, message: str) -> None:
        """发布消息到指定频道"""
        conn = await self._get_redis_connection()
        await conn.publish(channel, message)

    async def subscribe(self, channel: str) -> None:
        """订阅指定频道（避免重复订阅）"""
        if channel not in self.subscribed_channels:
            await self.connect()
            await self.pubsub.subscribe(channel)
            self.subscribed_channels.add(channel)

    async def unsubscribe(self, channel: str) -> None:
        """取消订阅指定频道"""
        if channel in self.subscribed_channels:
            await self.pubsub.unsubscribe(channel)
            self.subscribed_channels.remove(channel)
            # 如果没有订阅任何频道，关闭PubSub
            if not self.subscribed_channels:
                await self.pubsub.close()
                self.pubsub = None

    async def get_message(self, ignore_subscribe_messages: bool = True):
        """获取PubSub消息"""
        if self.pubsub:
            return await self.pubsub.get_message(ignore_subscribe_messages=ignore_subscribe_messages)
        return None

class WebSocketManager:
    """WebSocket连接管理器（单例）"""
    _instance = None

    def __new__(cls, *args, **kwargs):
        if not cls._instance:
            cls._instance = super().__new__(cls)
            cls._instance.chats: Dict[str, Set[WebSocket]] = {}  # 使用Set避免重复连接
            cls._instance.users: Dict[str, WebSocket] = {}
            cls._instance.pubsub_client = RedisPubSubManager()
            # 启动全局PubSub消息监听任务
            cls._instance._pubsub_listener_task = asyncio.create_task(
                cls._instance._global_pubsub_listener()
            )
        return cls._instance

    async def _global_pubsub_listener(self):
        """全局监听Redis PubSub消息，统一处理广播"""
        while True:
            message = await self.pubsub_client.get_message()
            if message:
                channel = message["channel"].decode("utf-8")
                if channel not in self.chats:
                    continue
                data = message["data"].decode("utf-8")
                # 批量发送消息，清理无效连接
                valid_sockets = set()
                for socket in self.chats[channel]:
                    try:
                        if socket.application_state == WebSocketState.CONNECTED.value:
                            await socket.send_text(data)
                            valid_sockets.add(socket)
                    except Exception:
                        # 移除无效连接
                        continue
                # 更新有效连接集合
                self.chats[channel] = valid_sockets
                if not self.chats[channel]:
                    del self.chats[channel]
                    await self.pubsub_client.unsubscribe(channel)
            await asyncio.sleep(0.1)  # 避免空循环占用资源

    async def connect_user(self, user_id: str, websocket: WebSocket):
        """建立用户与WebSocket的关联"""
        await websocket.accept()
        self.users[user_id] = websocket
        await self.send_to_user(user_id, json.dumps({"message": "connected"}))

    async def send_to_user(self, user_id: str, message: str):
        """向指定用户发送消息"""
        websocket = self.users.get(user_id)
        if not websocket:
            return
        try:
            if websocket.application_state == WebSocketState.CONNECTED.value:
                await websocket.send_text(message)
            else:
                del self.users[user_id]
        except Exception:
            del self.users[user_id]

    async def add_user_to_chat(self, chat_id: str, user_id: str):
        """将用户加入指定聊天频道"""
        websocket = self.users.get(user_id)
        if not websocket:
            return
        if chat_id not in self.chats:
            self.chats[chat_id] = set()
            await self.pubsub_client.subscribe(chat_id)
        self.chats[chat_id].add(websocket)
        # 发送连接通知
        notify_msg = json.dumps({
            "user_id": user_id,
            "chat_id": chat_id,
            "message": f"User {user_id} connected to chat - {chat_id}"
        })
        await self.pubsub_client.publish(chat_id, notify_msg)

    async def remove_user_from_chat(self, chat_id: str, websocket: WebSocket):
        """从聊天频道移除用户"""
        if chat_id in self.chats:
            self.chats[chat_id].discard(websocket)
            if not self.chats[chat_id]:
                del self.chats[chat_id]
                await self.pubsub_client.unsubscribe(chat_id)
            # 发送断开通知
            notify_msg = json.dumps({
                "user_id": next((uid for uid, ws in self.users.items() if ws == websocket), "unknown"),
                "chat_id": chat_id,
                "message": f"User disconnected from chat - {chat_id}"
            })
            await self.pubsub_client.publish(chat_id, notify_msg)

    async def disconnect_user(self, user_id: str, websocket: WebSocket, code: int = 1000, message: str = ""):
        """断开用户连接并清理资源"""
        try:
            if message:
                await websocket.send_text(message)
            await websocket.close(code)
        except Exception:
            pass
        # 从用户列表移除
        if user_id in self.users:
            del self.users[user_id]
        # 从所有聊天频道移除
        for chat_id in list(self.chats.keys()):
            self.chats[chat_id].discard(websocket)
            if not self.chats[chat_id]:
                del self.chats[chat_id]
                await self.pubsub_client.unsubscribe(chat_id)

socket_manager = WebSocketManager()

@web_socket_router.websocket("/connect")
async def websocket_connect(websocket: WebSocket):
    token = websocket.headers.get("Authorization")
    error_message = ""

    if token:
        try:
            token_data = await verify_jwt_token(token)
        except HTTPException as http_exc:
            error_message = http_exc.detail
    else:
        error_message = "Authorization token is missing"

    if error_message:
        await websocket.send_text(error_message)
        await websocket.close(4001)
        return

    user_id = token_data["user"]["_id"]
    try:
        await socket_manager.connect_user(user_id, websocket)
        # 合理的Ping间隔（例如30秒），避免频繁发送
        while True:
            await asyncio.sleep(30)
            await websocket.send_text(json.dumps({"type": "ping"}))
    except WebSocketDisconnect:
        await socket_manager.disconnect_user(user_id, websocket)

@web_socket_router.websocket("/api/v1/ws/{chat_id}/{user_id}")
async def websocket_chat(websocket: WebSocket, chat_id: str, user_id: str):
    # 先验证用户（如果需要）
    token = websocket.headers.get("Authorization")
    if token:
        try:
            await verify_jwt_token(token)
        except HTTPException:
            await websocket.close(4001)
            return
    else:
        await websocket.close(4001)
        return

    await websocket.accept()
    try:
        await socket_manager.add_user_to_chat(chat_id, user_id)
        # 监听客户端消息
        while True:
            data = await websocket.receive_text()
            message = json.dumps({
                "user_id": user_id,
                "chat_id": chat_id,
                "message": data
            })
            await socket_manager.pubsub_client.publish(chat_id, message)
    except WebSocketDisconnect:
        await socket_manager.remove_user_from_chat(chat_id, websocket)

# 注册路由
app.include_router(web_socket_router)

关键修复点说明 #

1. Redis PubSub优化

   • 采用单例Redis连接，避免重复初始化
   • 全局统一监听PubSub消息，避免每个频道创建独立监听任务
   • 自动管理订阅频道，无订阅时关闭PubSub节省资源

2. WebSocket连接管理

   • 使用Set存储聊天频道的WebSocket连接，避免重复添加
   • 发送消息时自动检测连接状态，清理无效连接
   • 连接断开后自动从用户列表和所有聊天频道中移除

3. Ping消息优化

   • 设置合理的Ping间隔（30秒），避免无节制发送导致连接过载
   • Ping消息携带明确类型标识，便于客户端区分业务消息

4. 消息序列化修复

   • 避免重复调用json.dumps，确保发送给客户端的是正确格式的JSON字符串

5. 代码结构简化

   • 合并重复的WebSocket端点逻辑，减少冗余代码
   • 单例管理器优化，避免线程锁与asyncio的冲突

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