forked from erp-dev/erp
97 lines
2.9 KiB
Python
97 lines
2.9 KiB
Python
import queue
|
||
import logging
|
||
|
||
logger = logging.getLogger(__name__)
|
||
|
||
# 存储所有活动的 SSE 连接队列（同步队列）
|
||
_connections = set()
|
||
|
||
|
||
def get_active_connections():
|
||
"""获取当前所有活动的 SSE 连接队列"""
|
||
return _connections
|
||
|
||
|
||
def push_connection(conn_queue):
|
||
"""
|
||
将一个新的连接队列添加到活动连接集合中
|
||
|
||
参数:
|
||
- conn_queue: queue.Queue 实例
|
||
"""
|
||
_connections.add(conn_queue)
|
||
logger.info(f"新的 SSE 连接建立，当前连接数: {len(_connections)}")
|
||
|
||
|
||
def remove_connection(conn_queue):
|
||
"""
|
||
从活动连接集合中移除一个连接队列
|
||
|
||
参数:
|
||
- conn_queue: queue.Queue 实例
|
||
"""
|
||
_connections.discard(conn_queue)
|
||
logger.info(f"SSE 连接断开，当前连接数: {len(_connections)}")
|
||
|
||
|
||
def cleanup_all_connections():
|
||
"""
|
||
清理所有连接（用于服务器关闭时）
|
||
"""
|
||
connections_list = list(_connections)
|
||
for conn_queue in connections_list:
|
||
try:
|
||
# 发送关闭信号到队列
|
||
if hasattr(conn_queue, 'put_nowait'):
|
||
conn_queue.put_nowait({'type': 'server_shutdown', 'message': 'Server shutting down'})
|
||
logger.debug(f'shutdown connection {conn_queue}')
|
||
except Exception:
|
||
pass # 忽略错误，因为连接可能已经断开
|
||
_connections.clear()
|
||
logger.info("所有SSE连接已清理")
|
||
|
||
|
||
def push_sse_event_to_all(event_data: dict):
|
||
"""
|
||
向所有连接的客户端广播一个 SSE 事件（同步队列版本）
|
||
|
||
参数:
|
||
- event_data: 要发送的事件数据（字典）
|
||
"""
|
||
disconnected = []
|
||
|
||
for conn_queue in list(_connections):
|
||
try:
|
||
# 使用put_nowait非阻塞发送消息
|
||
conn_queue.put_nowait(event_data)
|
||
except queue.Full:
|
||
# 队列满了，说明客户端处理消息太慢，跳过这条消息
|
||
logger.warning(f"队列已满，跳过消息: {event_data.get('type', 'unknown')}")
|
||
except Exception as e:
|
||
# 其他异常可能表示连接已断开
|
||
logger.error(f"向队列发送消息失败: {e}")
|
||
disconnected.append(conn_queue)
|
||
|
||
# 清理断开的连接
|
||
for conn_queue in disconnected:
|
||
remove_connection(conn_queue)
|
||
|
||
|
||
def push_simple_message_with_object_id(event_type: str, message: str, object_id):
|
||
"""
|
||
向所有连接的客户端广播一个简单消息事件，包含关联对象ID（WSGI同步版本）
|
||
|
||
参数:
|
||
- event_type: 事件类型字符串
|
||
- message: 消息内容字符串
|
||
- object_id: 关联对象的ID（整数或字符串）
|
||
"""
|
||
event_data = {
|
||
'mode': 'simple_message',
|
||
'type': event_type,
|
||
'message': message,
|
||
'object_id': object_id,
|
||
}
|
||
push_sse_event_to_all(event_data)
|
||
logger.info(f"广播 SSE 事件: type={event_type}, object_id={object_id}, 接收客户端数={len(_connections)}")
|