diff --git a/app/utils/websocket_tool.py b/app/utils/websocket_tool.py index 2704546..c1ae7cc 100644 --- a/app/utils/websocket_tool.py +++ b/app/utils/websocket_tool.py @@ -13,12 +13,8 @@ from typing import Union, List, Dict from app.core.common_utils import logger from app.utils.JSONEncodeTools import MyEncoder from configs.global_var import * +from app.utils.redis_config import redis_client -global active_connections -active_connections = multiprocessing.Manager().list() - -global active_connections_dist -active_connections_dist = multiprocessing.Manager().dict() class WebsocketUtil: @@ -30,23 +26,42 @@ class WebsocketUtil: # 等待连接 msg = ws.receive() # 存储ws连接对象 - global active_connections + active_connections = redis_client.get_redis().get("active_connections") + if active_connections == None: + active_connections = [] + else: + active_connections = json.loads(active_connections) active_connections.append(ws) - global active_connections_dist + active_connections_dist = redis_client.get_redis().get("active_connections_dist") + if active_connections_dist == None: + active_connections_dist = {} + else: + active_connections_dist = json.loads(active_connections_dist) if id in active_connections_dist: active_connections_dist[id].append(ws) else: ws_list = [ws, ] active_connections_dist[id] = ws_list + redis_client.get_redis().set("active_connections", json.dumps(active_connections)) + redis_client.get_redis().set("active_connections_dist", json.dumps(active_connections_dist)) + def disconnect(self, ws, id): # ws关闭时 移除ws对象 if ws.closed: - if ws in active_connections_dist.values(): - # global active_connections - active_connections.remove(ws) - # global active_connections_dist - active_connections_dist[id].pop(ws) + active_connections_dist = redis_client.get_redis().get("active_connections_dist") + if active_connections_dist is not None: + active_connections_dist = json.loads(active_connections_dist) + if ws in active_connections_dist.values(): + active_connections_dist[id].pop(ws) + redis_client.get_redis().set("active_connections", json.dumps(active_connections)) + active_connections = redis_client.get_redis().get("active_connections") + if active_connections is not None: + active_connections = json.loads(active_connections) + if ws in active_connections.values(): + active_connections.remove(ws) + redis_client.get_redis().set("active_connections_dist", json.dumps(active_connections_dist)) + @staticmethod async def send_personal_message(message: str, ws): @@ -56,18 +71,24 @@ class WebsocketUtil: def broadcast(self, message: str): # 广播消息 # global active_connections - for connection in active_connections: - connection.send(message) + active_connections = redis_client.get_redis().get("active_connections") + if active_connections is not None: + active_connections = json.loads(active_connections) + for connection in active_connections: + connection.send(message) def send_message_proj_json(self, message: Union[str, int, List, Dict], id: str): # 广播该项目的消息 # global active_connections_dist - for connection in active_connections_dist[id]: - try: - connection.send(json.dumps(message, cls=MyEncoder, indent=4), ) - except Exception as e: - logger.error("websocket异常断开,{}", e) - self.disconnect(ws=connection, id=id) + active_connections_dist = redis_client.get_redis().get("active_connections_dist") + if active_connections_dist is not None: + active_connections_dist = json.loads(active_connections_dist) + for connection in active_connections_dist[id]: + try: + connection.send(json.dumps(message, cls=MyEncoder, indent=4), ) + except Exception as e: + logger.error("websocket异常断开,{}", e) + self.disconnect(ws=connection, id=id) manager = WebsocketUtil() \ No newline at end of file