SenSu/bridges/core_bridge.py

#!/usr/bin/env python3
# -*- coding: utf-8 -*-

import logging
import asyncio
from typing import Dict, List, Callable, Any
from enum import Enum
import json

logger = logging.getLogger(__name__)

class MessageType(Enum):
    """消息类型枚举"""
    EVENT = "event"
    COMMAND = "command"
    DATA = "data"
    STATUS = "status"
    ERROR = "error"

class CoreBridge:
    """核心桥接服务 - 用于框架模块间通信"""
    
    def __init__(self):
        self.subscribers: Dict[str, List[Callable]] = {}
        self.message_queue: asyncio.Queue = asyncio.Queue()
        self.is_running = False
        self.processing_task = None
        logger.debug("CoreBridge初始化开始")
    
    async def start(self):
        """启动桥接服务"""
        try:
            logger.info("启动核心桥接服务")
            self.is_running = True
            self.processing_task = asyncio.create_task(self._process_messages())
            logger.debug("核心桥接服务启动完成")
        except Exception as e:
            logger.error(f"启动核心桥接服务时出错: {str(e)}", exc_info=True)
            raise
    
    def subscribe(self, topic: str, callback: Callable):
        """订阅主题"""
        try:
            if topic not in self.subscribers:
                self.subscribers[topic] = []
            
            self.subscribers[topic].append(callback)
            logger.debug(f"订阅主题: {topic}, 当前订阅者数: {len(self.subscribers[topic])}")
        except Exception as e:
            logger.error(f"订阅主题 {topic} 时出错: {str(e)}", exc_info=True)
            raise
    
    def unsubscribe(self, topic: str, callback: Callable):
        """取消订阅"""
        try:
            if topic in self.subscribers and callback in self.subscribers[topic]:
                self.subscribers[topic].remove(callback)
                logger.debug(f"取消订阅主题: {topic}, 剩余订阅者数: {len(self.subscribers[topic])}")
        except Exception as e:
            logger.error(f"取消订阅主题 {topic} 时出错: {str(e)}", exc_info=True)
    
    async def publish(self, topic: str, message: Dict, msg_type: MessageType = MessageType.DATA):
        """发布消息"""
        try:
            message_data = {
                "topic": topic,
                "type": msg_type.value,
                "data": message,
                "timestamp": asyncio.get_event_loop().time()
            }
            
            await self.message_queue.put(message_data)
            logger.debug(f"发布消息到主题: {topic}, 类型: {msg_type.value}")
            
        except Exception as e:
            logger.error(f"发布消息到主题 {topic} 时出错: {str(e)}", exc_info=True)
            raise
    
    async def _process_messages(self):
        """处理消息队列"""
        try:
            logger.debug("开始处理消息队列")
            
            while self.is_running:
                try:
                    # 等待消息，带超时以便检查运行状态
                    message = await asyncio.wait_for(self.message_queue.get(), timeout=1.0)
                    
                    # 分发消息给订阅者
                    await self._dispatch_message(message)
                    
                    # 标记任务完成
                    self.message_queue.task_done()
                    
                except asyncio.TimeoutError:
                    continue
                except Exception as e:
                    logger.error(f"处理消息时出错: {str(e)}", exc_info=True)
                    continue
            
            logger.debug("消息队列处理结束")
            
        except Exception as e:
            logger.error(f"消息队列处理循环出错: {str(e)}", exc_info=True)
    
    async def _dispatch_message(self, message: Dict):
        """分发消息给订阅者"""
        try:
            topic = message["topic"]
            
            if topic not in self.subscribers:
                logger.debug(f"主题 {topic} 没有订阅者")
                return
            
            subscribers = self.subscribers[topic][:]  # 复制列表避免在迭代时修改
            
            # 并行调用所有订阅者
            tasks = []
            for callback in subscribers:
                task = asyncio.create_task(self._call_subscriber(callback, message))
                tasks.append(task)
            
            if tasks:
                await asyncio.gather(*tasks, return_exceptions=True)
            
            logger.debug(f"消息分发完成，主题: {topic}, 订阅者数: {len(subscribers)}")
            
        except Exception as e:
            logger.error(f"分发消息时出错: {str(e)}", exc_info=True)
    
    async def _call_subscriber(self, callback: Callable, message: Dict):
        """调用订阅者回调"""
        try:
            if asyncio.iscoroutinefunction(callback):
                await callback(message)
            else:
                callback(message)
                
        except Exception as e:
            logger.error(f"调用订阅者回调时出错: {str(e)}", exc_info=True)
    
    def get_subscriber_count(self, topic: str = None) -> int:
        """获取订阅者数量"""
        try:
            if topic:
                count = len(self.subscribers.get(topic, []))
                logger.debug(f"主题 {topic} 的订阅者数量: {count}")
                return count
            else:
                total = sum(len(subs) for subs in self.subscribers.values())
                logger.debug(f"总订阅者数量: {total}")
                return total
        except Exception as e:
            logger.error(f"获取订阅者数量时出错: {str(e)}", exc_info=True)
            return 0
    
    async def shutdown(self):
        """关闭桥接服务"""
        try:
            logger.info("关闭核心桥接服务")
            self.is_running = False
            
            # 等待处理任务结束
            if self.processing_task:
                await asyncio.wait_for(self.processing_task, timeout=5.0)
            
            # 清空队列和订阅者
            self.subscribers.clear()
            while not self.message_queue.empty():
                try:
                    self.message_queue.get_nowait()
                    self.message_queue.task_done()
                except asyncio.QueueEmpty:
                    break
            
            logger.debug("核心桥接服务关闭完成")
            
        except Exception as e:
            logger.error(f"关闭核心桥接服务时出错: {str(e)}", exc_info=True)