124 lines
4.3 KiB
Python
124 lines
4.3 KiB
Python
# -*- coding = utf-8 -*-
|
|
# @Time : 2025/3/1 下午8:47
|
|
# @File : managers.py
|
|
# @Software : PyCharm
|
|
# @Author : xingc
|
|
# @Desc :
|
|
import asyncio
|
|
import logging
|
|
|
|
import redis.asyncio as aioredis
|
|
import ujson
|
|
from fastapi import WebSocket
|
|
from fastapi.websockets import WebSocketState
|
|
from simple_spider_tool import jsonpath
|
|
|
|
from core.db.engine import get_cache
|
|
|
|
logger = logging.getLogger(__name__)
|
|
|
|
|
|
class RedisPubSubManager:
|
|
def __init__(self):
|
|
self.redis_connection = None
|
|
self.pubsub = None
|
|
|
|
async def _get_redis_connection(self) -> aioredis:
|
|
return await get_cache()
|
|
|
|
async def connect(self):
|
|
self.redis_connection = await self._get_redis_connection()
|
|
self.pubsub = self.redis_connection.pubsub()
|
|
|
|
async def subscribe(self, chat_guid: str) -> aioredis:
|
|
await self.pubsub.subscribe(chat_guid)
|
|
return self.pubsub
|
|
|
|
async def unsubscribe(self, chat_guid: str):
|
|
await self.pubsub.unsubscribe(chat_guid)
|
|
|
|
async def publish(self, chat_guid: str, message: str):
|
|
await self.redis_connection.publish(chat_guid, message)
|
|
|
|
async def disconnect(self):
|
|
if self.redis_connection is not None:
|
|
await self.redis_connection.close()
|
|
|
|
|
|
class SocketManager:
|
|
|
|
def __init__(self):
|
|
self.handlers: dict = {}
|
|
self.chats: dict = {}
|
|
self.pubsub_client = RedisPubSubManager()
|
|
self.user_guid_to_websocket: dict = {}
|
|
self._task_cancel = True
|
|
|
|
def handler(self, message_type):
|
|
def decorator(func):
|
|
self.handlers[message_type] = func
|
|
return func
|
|
|
|
return decorator
|
|
|
|
async def connect_socket(self, websocket: WebSocket):
|
|
await websocket.accept()
|
|
|
|
async def add_user_socket_connection(self, user_guid: str, websocket: WebSocket):
|
|
self.user_guid_to_websocket.setdefault(user_guid, set()).add(websocket)
|
|
|
|
async def add_user_to_chat(self, chat_guid: str, websocket: WebSocket):
|
|
if chat_guid in self.chats:
|
|
self.chats[chat_guid].add(websocket)
|
|
else:
|
|
self.chats[chat_guid] = {websocket}
|
|
await self.pubsub_client.connect()
|
|
pubsub_subscriber = await self.pubsub_client.subscribe(chat_guid)
|
|
asyncio.create_task(self._pubsub_data_reader(pubsub_subscriber))
|
|
|
|
async def broadcast_to_chat(self, chat_guid: str, message: str | dict) -> None:
|
|
if isinstance(message, dict):
|
|
message = ujson.dumps(message)
|
|
await self.pubsub_client.publish(chat_guid, message)
|
|
|
|
async def remove_user_from_chat(self, chat_guid: str, websocket: WebSocket) -> None:
|
|
self.chats[chat_guid].remove(websocket)
|
|
if len(self.chats[chat_guid]) == 0:
|
|
del self.chats[chat_guid]
|
|
logger.info(f'Removing user from PubSub channel {chat_guid}')
|
|
await self.pubsub_client.unsubscribe(chat_guid)
|
|
|
|
async def remove_user_guid_to_websocket(self, user_guid: str, websocket: WebSocket):
|
|
if user_guid in self.user_guid_to_websocket:
|
|
self.user_guid_to_websocket.get(user_guid).remove(websocket)
|
|
|
|
# https://github.com/redis/redis-py/issues/2523
|
|
async def _pubsub_data_reader(self, pubsub_subscriber):
|
|
try:
|
|
while True:
|
|
message = await pubsub_subscriber.get_message(ignore_subscribe_messages=True)
|
|
if message is not None:
|
|
chat_guid = message['channel']
|
|
sockets = self.chats.get(chat_guid)
|
|
if sockets:
|
|
for socket in sockets:
|
|
if socket.client_state != WebSocketState.CONNECTED:
|
|
continue
|
|
data = message['data']
|
|
if f',"user_id":{socket.user_id}' in data:
|
|
continue
|
|
await socket.send_json(data)
|
|
else:
|
|
if self._task_cancel is False:
|
|
break
|
|
except Exception as exc:
|
|
logger.exception(f'Exception occurred: {exc}')
|
|
|
|
async def send_error(self, message: str, websocket: WebSocket):
|
|
if websocket.client_state == WebSocketState.CONNECTED:
|
|
await websocket.send_json({'status': 'error', 'content': message})
|
|
|
|
async def quit(self):
|
|
self._task_cancel = False
|
|
await self.pubsub_client.disconnect()
|