diff --git a/app/utils/websocket_tool.py b/app/utils/websocket_tool.py index c1ae7cc..13db4a2 100644 --- a/app/utils/websocket_tool.py +++ b/app/utils/websocket_tool.py @@ -14,53 +14,38 @@ 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 +from multiprocessing import Process, Queue class WebsocketUtil: - # def __init__(self): - # self.active_connections = multiprocessing.Manager().list() - # self.active_connections_dist = multiprocessing.Manager().dict() + def __init__(self): + # self.active_connections = multiprocessing.Manager().list() + active_connections_dist = {} + self.queue = Queue(maxsize=0) + self.queue.put(active_connections_dist) def connect(self, ws, id: str): # 等待连接 msg = ws.receive() # 存储ws连接对象 - 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) - 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) + active_connections_dist = Queue.get() 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)) + self.queue.put(active_connections_dist) def disconnect(self, ws, id): # ws关闭时 移除ws对象 if ws.closed: - active_connections_dist = redis_client.get_redis().get("active_connections_dist") + active_connections_dist = Queue.get() 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 @@ -68,27 +53,26 @@ class WebsocketUtil: # 发送个人消息 await ws.send(message) - def broadcast(self, message: str): - # 广播消息 - # global active_connections - 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 broadcast(self, message: str): + # # 广播消息 + # # global active_connections + # 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 - 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) + active_connections_dist = self.queue.get() + self.queue.put(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