xkrs_manan/RODY
1
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

ws

Browse Source
This commit is contained in:
JIAKUNHAO 2022-11-24 15:21:00 +08:00
parent 4842024f99
commit fcb39c8863
1 changed files with 25 additions and 41 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

66
app/utils/websocket_tool.py
View File

@ -14,53 +14,38 @@ from app.core.common_utils import logger
from app.utils.JSONEncodeTools import MyEncoder from app.utils.JSONEncodeTools import MyEncoder
from configs.global_var import * from configs.global_var import *
from app.utils.redis_config import redis_client from app.utils.redis_config import redis_client
from multiprocessing import Process, Queue
class WebsocketUtil: class WebsocketUtil:
# def __init__(self): def __init__(self):
# self.active_connections = multiprocessing.Manager().list() # self.active_connections = multiprocessing.Manager().list()
# self.active_connections_dist = multiprocessing.Manager().dict() active_connections_dist = {}
self.queue = Queue(maxsize=0)
self.queue.put(active_connections_dist)
def connect(self, ws, id: str): def connect(self, ws, id: str):
# 等待连接 # 等待连接
msg = ws.receive() msg = ws.receive()
# 存储ws连接对象 # 存储ws连接对象
active_connections = redis_client.get_redis().get("active_connections") active_connections_dist = Queue.get()
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)
if id in active_connections_dist: if id in active_connections_dist:
active_connections_dist[id].append(ws) active_connections_dist[id].append(ws)
else: else:
ws_list = [ws, ] ws_list = [ws, ]
active_connections_dist[id] = ws_list active_connections_dist[id] = ws_list
redis_client.get_redis().set("active_connections", json.dumps(active_connections)) self.queue.put(active_connections_dist)
redis_client.get_redis().set("active_connections_dist", json.dumps(active_connections_dist))
def disconnect(self, ws, id): def disconnect(self, ws, id):
# ws关闭时 移除ws对象 # ws关闭时 移除ws对象
if ws.closed: 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: if active_connections_dist is not None:
active_connections_dist = json.loads(active_connections_dist)
if ws in active_connections_dist.values(): 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)) 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 @staticmethod
@ -68,27 +53,26 @@ class WebsocketUtil:
# 发送个人消息 # 发送个人消息
await ws.send(message) await ws.send(message)
def broadcast(self, message: str): # def broadcast(self, message: str):
# 广播消息 # # 广播消息
# global active_connections # # global active_connections
active_connections = redis_client.get_redis().get("active_connections") # active_connections = redis_client.get_redis().get("active_connections")
if active_connections is not None: # if active_connections is not None:
active_connections = json.loads(active_connections) # active_connections = json.loads(active_connections)
for connection in active_connections: # for connection in active_connections:
connection.send(message) # connection.send(message)
def send_message_proj_json(self, message: Union[str, int, List, Dict], id: str): def send_message_proj_json(self, message: Union[str, int, List, Dict], id: str):
# 广播该项目的消息 # 广播该项目的消息
# global active_connections_dist # global active_connections_dist
active_connections_dist = redis_client.get_redis().get("active_connections_dist") active_connections_dist = self.queue.get()
if active_connections_dist is not None: self.queue.put(active_connections_dist)
active_connections_dist = json.loads(active_connections_dist) for connection in active_connections_dist[id]:
for connection in active_connections_dist[id]: try:
try: connection.send(json.dumps(message, cls=MyEncoder, indent=4), )
connection.send(json.dumps(message, cls=MyEncoder, indent=4), ) except Exception as e:
except Exception as e: logger.error("websocket异常断开，{}", e)
logger.error("websocket异常断开，{}", e) self.disconnect(ws=connection, id=id)
self.disconnect(ws=connection, id=id)
manager = WebsocketUtil() manager = WebsocketUtil()