erpnew/sse/services.py

import queue
import logging

logger = logging.getLogger(__name__)

# 存储所有活动的 SSE 连接队列（按商户ID组织）
_connections = {}  # {merchant_id: set([conn_queue1, conn_queue2, ...])}


def get_all_connections_count():
    """获取所有商户的连接总数"""
    total = 0
    for merchant_id, connections in _connections.items():
        total += len(connections)
    return total


def get_merchant_connections(merchant_id):
    """获取指定商户的所有连接队列"""
    return _connections.get(merchant_id, set())


def get_all_merchant_ids():
    """获取所有有活跃连接的商户ID列表"""
    return list(_connections.keys())


def push_connection(merchant_id, conn_queue):
    """
    将一个新的连接队列添加到指定商户的连接集合中
    
    参数:
    - merchant_id: 商户ID
    - conn_queue: queue.Queue 实例
    """
    if merchant_id not in _connections:
        _connections[merchant_id] = set()
    
    _connections[merchant_id].add(conn_queue)
    total_connections = sum(len(conns) for conns in _connections.values())
    logger.info(f"商户 {merchant_id} 的新 SSE 连接建立，该商户连接数: {len(_connections[merchant_id])}，总连接数: {total_connections}")


def remove_connection(merchant_id, conn_queue):
    """
    从指定商户的连接集合中移除一个连接队列
    
    参数:
    - merchant_id: 商户ID
    - conn_queue: queue.Queue 实例
    """
    if merchant_id in _connections:
        _connections[merchant_id].discard(conn_queue)
        
        # 如果该商户没有其他连接，则移除整个商户记录
        if not _connections[merchant_id]:
            del _connections[merchant_id]
        
        total_connections = sum(len(conns) for conns in _connections.values())
        logger.info(f"商户 {merchant_id} 的 SSE 连接断开，该商户连接数: {len(_connections.get(merchant_id, set()))}，总连接数: {total_connections}")
    else:
        logger.warning(f"尝试从未注册的商户 {merchant_id} 移除连接")


def cleanup_all_connections():
    """
    清理所有连接（用于服务器关闭时）
    """
    all_connections = {}
    for merchant_id, connections in list(_connections.items()):
        all_connections[merchant_id] = list(connections)
        for conn_queue in connections:
            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} for merchant {merchant_id}')
            except Exception:
                pass  # 忽略错误，因为连接可能已经断开
    
    _connections.clear()
    logger.info("所有SSE连接已清理")


def shutdown_merchant_connections(merchant_id):
    """
    关闭指定商户的所有连接
    
    参数:
    - merchant_id: 商户ID
    
    返回:
    - int: 关闭的连接数
    """
    if merchant_id not in _connections:
        return 0
    
    connections = _connections[merchant_id]
    count = len(connections)
    
    # 向每个连接发送关闭信号
    for conn_queue in list(connections):
        try:
            conn_queue.put_nowait({
                'type': 'server_shutdown',
                'message': 'Server shutting down your connections, please reconnect later'
            })
        except queue.Full:
            # 队列满了，跳过
            pass
        except Exception:
            # 其他异常，也跳过
            pass
    
    # 删除所有连接
    del _connections[merchant_id]
    
    logger.info(f"关闭商户 {merchant_id} 的 {count} 个 SSE 连接")
    return count


def push_sse_event_to_all(event_data: dict):
    """
    向所有连接的客户端广播一个 SSE 事件（同步队列版本）
    
    参数:
    - event_data: 要发送的事件数据（字典）
    """
    disconnected = {}
    
    for merchant_id, connections in list(_connections.items()):
        disconnected[merchant_id] = []
        
        for conn_queue in connections:
            try:
                # 使用put_nowait非阻塞发送消息
                conn_queue.put_nowait(event_data)
            except queue.Full:
                # 队列满了，说明客户端处理消息太慢，跳过这条消息
                logger.warning(f"商户 {merchant_id} 队列已满，跳过消息: {event_data.get('type', 'unknown')}")
            except Exception as e:
                # 其他异常可能表示连接已断开
                logger.error(f"向商户 {merchant_id} 的队列发送消息失败: {e}")
                disconnected[merchant_id].append(conn_queue)
    
    # 清理断开的连接
    for merchant_id, conn_queues in disconnected.items():
        for conn_queue in conn_queues:
            remove_connection(merchant_id, conn_queue)


def push_sse_event_to_merchant(merchant_id, event_data: dict):
    """
    向指定商户的所有连接客户端广播一个 SSE 事件（同步队列版本）
    
    参数:
    - merchant_id: 目标商户ID
    - event_data: 要发送的事件数据（字典）
    """
    if merchant_id not in _connections:
        logger.info(f"商户 {merchant_id} 没有活跃连接，跳过消息发送")
        return
    
    disconnected = []
    
    for conn_queue in list(_connections[merchant_id]):
        try:
            # 使用put_nowait非阻塞发送消息
            conn_queue.put_nowait(event_data)
        except queue.Full:
            # 队列满了，说明客户端处理消息太慢，跳过这条消息
            logger.warning(f"商户 {merchant_id} 队列已满，跳过消息: {event_data.get('type', 'unknown')}")
        except Exception as e:
            # 其他异常可能表示连接已断开
            logger.error(f"向商户 {merchant_id} 的队列发送消息失败: {e}")
            disconnected.append(conn_queue)
    
    # 清理断开的连接
    for conn_queue in disconnected:
        remove_connection(merchant_id, conn_queue)
    
    logger.info(f"向商户 {merchant_id} 广播 SSE 事件: type={event_data.get('type', 'unknown')}, 接收客户端数={len(_connections.get(merchant_id, set()))}")


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)
    total_connections = sum(len(conns) for conns in _connections.values())
    logger.info(f"广播 SSE 事件: type={event_type}, object_id={object_id}, 接收客户端数={total_connections}")


def push_simple_message_with_object_id_to_merchant(merchant_id, event_type: str, message: str, object_id):
    """
    向指定商户的所有连接客户端广播一个简单消息事件，包含关联对象ID
    
    参数:
    - merchant_id: 目标商户ID
    - event_type: 事件类型字符串
    - message: 消息内容字符串
    - object_id: 关联对象的ID（整数或字符串）
    """
    event_data = {
        'mode': 'simple_message',
        'type': event_type,
        'message': message,
        'object_id': object_id,
        'merchant_id': merchant_id,  # 添加商户ID，便于客户端验证
    }
    push_sse_event_to_merchant(merchant_id, event_data)