RODY/app/utils/websocket_tool.py

79 lines
2.5 KiB
Python
Raw Normal View History

2022-11-04 17:37:08 +08:00
"""
@Time 2022/10/12 17:55
@Auth
2022-11-24 10:58:58 +08:00
@File websocket_tool.py
@IDE PyCharm
2022-11-04 17:37:08 +08:00
@MottoABC(Always Be Coding)
@Desc
"""
import json
from typing import Union, List, Dict
from app.core.common_utils import logger
from app.utils.JSONEncodeTools import MyEncoder
2022-11-24 10:45:05 +08:00
from configs.global_var import *
2022-11-24 14:15:23 +08:00
from app.utils.redis_config import redis_client
2022-11-24 15:21:00 +08:00
from multiprocessing import Process, Queue
2022-11-04 17:37:08 +08:00
2022-11-24 11:15:07 +08:00
class WebsocketUtil:
2022-11-24 15:21:00 +08:00
def __init__(self):
# self.active_connections = multiprocessing.Manager().list()
active_connections_dist = {}
self.queue = Queue(maxsize=0)
self.queue.put(active_connections_dist)
2022-11-24 10:58:58 +08:00
2022-11-04 17:37:08 +08:00
def connect(self, ws, id: str):
# 等待连接
msg = ws.receive()
# 存储ws连接对象
2022-11-24 15:22:49 +08:00
active_connections_dist = self.queue.get()
2022-11-24 11:18:07 +08:00
if id in active_connections_dist:
active_connections_dist[id].append(ws)
2022-11-04 17:37:08 +08:00
else:
ws_list = [ws, ]
2022-11-24 11:18:07 +08:00
active_connections_dist[id] = ws_list
2022-11-24 15:21:00 +08:00
self.queue.put(active_connections_dist)
2022-11-24 14:15:23 +08:00
2022-11-04 17:37:08 +08:00
def disconnect(self, ws, id):
# ws关闭时 移除ws对象
if ws.closed:
2022-11-24 15:30:06 +08:00
active_connections_dist = self.queue.get()
2022-11-24 14:15:23 +08:00
if active_connections_dist is not None:
if ws in active_connections_dist.values():
active_connections_dist[id].pop(ws)
2022-11-24 15:30:06 +08:00
self.queue.put(active_connections_dist)
2022-11-04 17:37:08 +08:00
@staticmethod
async def send_personal_message(message: str, ws):
# 发送个人消息
await ws.send(message)
2022-11-24 15:21:00 +08:00
# 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)
2022-11-04 17:37:08 +08:00
def send_message_proj_json(self, message: Union[str, int, List, Dict], id: str):
# 广播该项目的消息
2022-11-24 11:19:34 +08:00
# global active_connections_dist
2022-11-24 15:30:06 +08:00
print(self.queue.qsize())
2022-11-24 15:21:00 +08:00
active_connections_dist = self.queue.get()
self.queue.put(active_connections_dist)
2022-11-24 15:25:44 +08:00
print(active_connections_dist)
2022-11-24 15:21:00 +08:00
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)
2022-11-04 17:37:08 +08:00
2022-11-24 10:58:58 +08:00
manager = WebsocketUtil()