230 lines
7.2 KiB
Python
230 lines
7.2 KiB
Python
"""WebSocket 端点模块
|
||
|
||
提供 WebSocket 连接端点，用于实时推送提醒消息到 Web 客户端。
|
||
|
||
主要功能：
|
||
- /ws/reminders: 提醒消息的 WebSocket 端点
|
||
- 连接管理：维护活跃连接，处理断线重连
|
||
- 消息分发：从 Redis Pub/Sub 接收消息并推送到客户端
|
||
|
||
依赖：
|
||
- backend.app.core.pubsub: Redis Pub/Sub 消息通道
|
||
- backend.app.core.security: Token 验证
|
||
"""
|
||
|
||
import asyncio
|
||
import json
|
||
import logging
|
||
from typing import Dict, Set
|
||
|
||
from fastapi import APIRouter, WebSocket, WebSocketDisconnect, Query
|
||
|
||
from backend.app.core.pubsub import reminder_pubsub
|
||
from backend.app.core.security import decode_access_token
|
||
|
||
logger = logging.getLogger(__name__)
|
||
|
||
router = APIRouter(tags=["websocket"])
|
||
|
||
|
||
class ConnectionManager:
|
||
"""WebSocket 连接管理器"""
|
||
|
||
def __init__(self):
|
||
# user_id -> Set[WebSocket]
|
||
self.active_connections: Dict[int, Set[WebSocket]] = {}
|
||
|
||
async def connect(self, websocket: WebSocket, user_id: int) -> None:
|
||
"""接受 WebSocket 连接"""
|
||
await websocket.accept()
|
||
if user_id not in self.active_connections:
|
||
self.active_connections[user_id] = set()
|
||
self.active_connections[user_id].add(websocket)
|
||
logger.info(f"用户 {user_id} 建立 WebSocket 连接，当前连接数: {len(self.active_connections[user_id])}")
|
||
|
||
def disconnect(self, websocket: WebSocket, user_id: int) -> None:
|
||
"""断开 WebSocket 连接"""
|
||
if user_id in self.active_connections:
|
||
self.active_connections[user_id].discard(websocket)
|
||
if not self.active_connections[user_id]:
|
||
del self.active_connections[user_id]
|
||
logger.info(f"用户 {user_id} 断开 WebSocket 连接")
|
||
|
||
async def send_message(self, websocket: WebSocket, message: dict) -> None:
|
||
"""发送消息到单个连接"""
|
||
try:
|
||
await websocket.send_json(message)
|
||
except Exception as e:
|
||
logger.error(f"发送消息失败: {e}")
|
||
|
||
async def broadcast_to_user(self, user_id: int, message: dict) -> None:
|
||
"""广播消息到指定用户的所有连接"""
|
||
if user_id in self.active_connections:
|
||
for websocket in self.active_connections[user_id]:
|
||
await self.send_message(websocket, message)
|
||
|
||
def get_connection_count(self, user_id: int) -> int:
|
||
"""获取用户的连接数"""
|
||
return len(self.active_connections.get(user_id, set()))
|
||
|
||
|
||
# 全局连接管理器
|
||
manager = ConnectionManager()
|
||
|
||
|
||
async def _redis_listener(pubsub, websocket: WebSocket, user_id: int):
|
||
"""Redis Pub/Sub 消息监听协程
|
||
|
||
持续从 Redis 订阅频道读取消息，收到后推送到 WebSocket 客户端。
|
||
当 WebSocket 断开时，通过 CancelledError 退出。
|
||
"""
|
||
try:
|
||
while True:
|
||
message = await pubsub.get_message(
|
||
ignore_subscribe_messages=True,
|
||
timeout=1.0,
|
||
)
|
||
if message and message["type"] == "message":
|
||
try:
|
||
data = json.loads(message["data"])
|
||
await manager.send_message(websocket, {
|
||
"type": "reminder",
|
||
"data": data,
|
||
})
|
||
except json.JSONDecodeError:
|
||
logger.error(f"消息解析失败: {message['data']}")
|
||
except asyncio.CancelledError:
|
||
pass
|
||
except Exception as e:
|
||
logger.error(f"Redis 监听异常: {e}")
|
||
|
||
|
||
async def _websocket_listener(websocket: WebSocket, user_id: int):
|
||
"""WebSocket 客户端消息监听协程
|
||
|
||
持续读取客户端消息（心跳 ping），当客户端断开时抛出 WebSocketDisconnect。
|
||
"""
|
||
try:
|
||
while True:
|
||
data = await websocket.receive_text()
|
||
if data == "ping":
|
||
await websocket.send_text("pong")
|
||
except WebSocketDisconnect:
|
||
raise
|
||
except Exception:
|
||
pass
|
||
|
||
|
||
@router.websocket("/ws/reminders")
|
||
async def websocket_reminders(
|
||
websocket: WebSocket,
|
||
token: str = Query(..., description="认证 Token"),
|
||
):
|
||
"""提醒消息的 WebSocket 端点
|
||
|
||
连接参数：
|
||
token: JWT 认证 Token（通过 query 参数传递）
|
||
|
||
消息格式：
|
||
{
|
||
"type": "reminder",
|
||
"data": {
|
||
"reminder_id": 123,
|
||
"title": "提醒标题",
|
||
"content": "提醒内容",
|
||
"type": "arrears",
|
||
"biz_type": "customer_arrears",
|
||
"biz_id": 456,
|
||
"status": "pending",
|
||
"created_at": "2024-01-01 12:00:00"
|
||
}
|
||
}
|
||
"""
|
||
# 验证 Token
|
||
try:
|
||
payload = decode_access_token(token)
|
||
if not payload:
|
||
await websocket.close(code=4001, reason="无效的 Token")
|
||
return
|
||
|
||
user_id = payload.get("user_id")
|
||
if not user_id:
|
||
await websocket.close(code=4001, reason="无效的用户信息")
|
||
return
|
||
except Exception as e:
|
||
logger.error(f"Token 验证失败: {e}")
|
||
await websocket.close(code=4001, reason="Token 验证失败")
|
||
return
|
||
|
||
# 建立连接
|
||
await manager.connect(websocket, user_id)
|
||
|
||
# 订阅 Redis 频道
|
||
pubsub = None
|
||
redis_task = None
|
||
ws_task = None
|
||
try:
|
||
pubsub = await reminder_pubsub.subscribe(user_id)
|
||
|
||
# 发送连接成功消息
|
||
await websocket.send_json({
|
||
"type": "connected",
|
||
"message": "连接成功",
|
||
})
|
||
|
||
# 并发运行 Redis 监听和 WebSocket 监听
|
||
redis_task = asyncio.create_task(
|
||
_redis_listener(pubsub, websocket, user_id),
|
||
)
|
||
ws_task = asyncio.create_task(
|
||
_websocket_listener(websocket, user_id),
|
||
)
|
||
|
||
# 等待任一任务结束（WebSocket 断开 或 Redis 异常）
|
||
done, pending = await asyncio.wait(
|
||
[redis_task, ws_task],
|
||
return_when=asyncio.FIRST_COMPLETED,
|
||
)
|
||
|
||
# 取消未完成的任务
|
||
for task in pending:
|
||
task.cancel()
|
||
try:
|
||
await task
|
||
except asyncio.CancelledError:
|
||
pass
|
||
|
||
# 检查是否有异常
|
||
for task in done:
|
||
if task.exception() and not isinstance(task.exception(), WebSocketDisconnect):
|
||
logger.error(f"WebSocket 任务异常: {task.exception()}")
|
||
|
||
except WebSocketDisconnect:
|
||
logger.info(f"用户 {user_id} 主动断开连接")
|
||
except Exception as e:
|
||
logger.error(f"WebSocket 错误: {e}")
|
||
finally:
|
||
# 清理资源
|
||
manager.disconnect(websocket, user_id)
|
||
if redis_task and not redis_task.done():
|
||
redis_task.cancel()
|
||
if ws_task and not ws_task.done():
|
||
ws_task.cancel()
|
||
if pubsub:
|
||
try:
|
||
await reminder_pubsub.unsubscribe(pubsub, user_id)
|
||
except Exception:
|
||
pass
|
||
|
||
|
||
@router.get("/ws/status")
|
||
async def websocket_status():
|
||
"""查看 WebSocket 连接状态（调试用）"""
|
||
return {
|
||
"active_users": len(manager.active_connections),
|
||
"connections": {
|
||
user_id: len(conns)
|
||
for user_id, conns in manager.active_connections.items()
|
||
},
|
||
}
|