feat(chat): 聊天功能

core/routers/chat/managers.py

@@ -0,0 +1,116 @@
+# -*- 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 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)
+            await 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:
+                            data = message['data']
+                            await socket.send_text(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):
+        await websocket.send_json({'status': 'error', 'message': message})
+
+    async def quit(self):
+        self._task_cancel = False
+        await self.pubsub_client.disconnect()