ws

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():
+            active_connections_dist = redis_client.get_redis().get("active_connections_dist")
-                # global active_connections
+            if active_connections_dist is not None:
-                active_connections.remove(ws)
+                active_connections_dist = json.loads(active_connections_dist)
-                # global active_connections_dist
+                if ws in active_connections_dist.values():
-                active_connections_dist[id].pop(ws)
+                    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:
+        active_connections = redis_client.get_redis().get("active_connections")
-            connection.send(message)
+        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]:
+        active_connections_dist = redis_client.get_redis().get("active_connections_dist")
-            try:
+        if active_connections_dist is not None:
-                connection.send(json.dumps(message, cls=MyEncoder, indent=4), )
+            active_connections_dist = json.loads(active_connections_dist)
-            except Exception as e:
+            for connection in active_connections_dist[id]:
-                logger.error("websocket异常断开，{}", e)
+                try:
-                self.disconnect(ws=connection, id=id)
+                    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()