core/aggregation_protocol/protocol.py

#!/usr/bin/env python3
"""
虫群聚合协议 — 主入口

AggregationProtocol 是整个聚合协议的统一入口，整合：
- 节点发现(discovery)
- 通信层(transport)
- 任务调度(scheduler)
- 权限管理(permission)

使用方式：
    protocol = AggregationProtocol(
        node_id="queen_zhangming",
        role=NodeRole.QUEEN,
        address="http://localhost:8080",
    )
    protocol.start()
    
    # 提交聚合任务
    result = protocol.submit_task("分析这段代码的性能瓶颈", ...)
    
    # 关闭
    protocol.stop()
"""

import hashlib
import logging
import threading
import time
from datetime import datetime
from typing import Callable, Dict, List, Optional

from .types import (
    NodeInfo, NodeCapability, NodeRole, NodeStatus,
    PermissionLevel, AggregationStrategy, AggregationTask,
    TaskForce, ProtocolMessage,
)
from .discovery import NodeRegistry
from .transport import MessageBus, LocalTransport
from .scheduler import TaskForceManager
from .permission import PermissionManager

logger = logging.getLogger(__name__)


class AggregationProtocol:
    """
    虫群聚合协议 — 核心入口
    
    设计理念：
    - 像GPU集群一样，将分散的设备临时组合成服务器
    - 每个设备可同时参与多个临时服务器
    - 权限控制调用范围和数量
    - 任务完成自动解散，设备回归空闲池
    """

    def __init__(
        self,
        node_id: str,
        role: NodeRole = NodeRole.QUEEN,
        name: str = "",
        address: str = "",
        permission_level: PermissionLevel = PermissionLevel.QUEEN,
        capabilities: List[str] = None,
        compute_score: float = 1.0,
        max_concurrent: int = 3,
        heartbeat_timeout: float = 30.0,
    ):
        # 自身节点信息
        self.node_info = NodeInfo(
            node_id=node_id,
            name=name or node_id,
            role=role,
            status=NodeStatus.ONLINE,
            address=address,
            capability=NodeCapability(
                specializations=capabilities or ["general"],
                compute_score=compute_score,
                max_concurrent=max_concurrent,
            ),
            permission_level=permission_level,
        )
        
        # 核心组件
        self.registry = NodeRegistry(node_id, heartbeat_timeout)
        self.bus = MessageBus(self.node_info)
        self.scheduler = TaskForceManager(self.registry, self.bus)
        self.permission = PermissionManager(permission_level)
        
        # 本地传输（同一进程内通信）
        self._local_transport = LocalTransport()
        
        # 运行状态
        self._running = False
        self._heartbeat_thread = None
        
        # 消息处理
        self._setup_message_handlers()
        
        # 统计
        self._stats = {
            "start_time": None,
            "tasks_submitted": 0,
            "tasks_completed": 0,
            "messages_sent": 0,
        }

    # ============================================================
    # 生命周期
    # ============================================================

    def start(self):
        """启动聚合协议"""
        if self._running:
            return
        
        self._running = True
        self.node_info.status = NodeStatus.ONLINE
        
        # 注册自身
        self.registry.register(self.node_info)
        self.bus.register_node_address(self.node_info.node_id, self.node_info.address)
        self._local_transport.register_bus(self.node_info.node_id, self.bus)
        self.permission.set_node_permission(self.node_info.node_id, self.node_info.permission_level)
        
        # 启动消息总线
        self.bus.start()
        
        # 启动心跳线程
        self._heartbeat_thread = threading.Thread(
            target=self._heartbeat_loop, daemon=True
        )
        self._heartbeat_thread.start()
        
        self._stats["start_time"] = datetime.now().isoformat()
        logger.info(f"聚合协议启动: {self.node_info.node_id} ({self.node_info.role.value})")

    def stop(self):
        """停止聚合协议"""
        self._running = False
        self.node_info.status = NodeStatus.OFFLINE
        self.bus.stop()
        logger.info(f"聚合协议停止: {self.node_info.node_id}")

    # ============================================================
    # 节点管理
    # ============================================================

    def add_remote_node(self, node_id: str, name: str = "",
                        role: NodeRole = NodeRole.QUEEN,
                        address: str = "",
                        capabilities: List[str] = None,
                        compute_score: float = 1.0,
                        permission_level: PermissionLevel = PermissionLevel.QUEEN):
        """添加远程节点"""
        node = NodeInfo(
            node_id=node_id,
            name=name or node_id,
            role=role,
            address=address,
            capability=NodeCapability(
                specializations=capabilities or ["general"],
                compute_score=compute_score,
            ),
            permission_level=permission_level,
        )
        
        self.registry.register(node)
        if address:
            self.bus.register_node_address(node_id, address)
        self.permission.set_node_permission(node_id, permission_level)

    def add_local_node(self, protocol: 'AggregationProtocol'):
        """添加本地节点（同一进程）"""
        self.registry.register(protocol.node_info)
        self._local_transport.register_bus(protocol.node_info.node_id, protocol.bus)
        self.permission.set_node_permission(
            protocol.node_info.node_id, 
            protocol.node_info.permission_level
        )

    def remove_node(self, node_id: str):
        """移除节点"""
        self.registry.unregister(node_id)
        self.bus.unregister_node_address(node_id)

    # ============================================================
    # 聚合任务
    # ============================================================

    def submit_task(
        self,
        query: str,
        required_capabilities: List[str] = None,
        min_nodes: int = 2,
        max_nodes: int = 5,
        strategy: AggregationStrategy = AggregationStrategy.ADAPTIVE_MIX,
        priority: int = 5,
        timeout_sec: float = 60.0,
    ) -> Dict:
        """
        提交聚合任务
        
        流程：
        1. 创建任务
        2. 权限检查（可调用多少节点）
        3. 组建临时服务器
        4. 分发任务到各节点
        5. 聚合结果
        6. 解散临时服务器
        
        Returns:
            {"taskforce_id": str, "status": str, "result": Dict}
        """
        self._stats["tasks_submitted"] += 1
        
        # 1. 创建任务
        task_id = self._gen_id("task")
        task = AggregationTask(
            task_id=task_id,
            query=query,
            requester=self.node_info.node_id,
            required_capabilities=required_capabilities or [],
            min_nodes=min_nodes,
            max_nodes=max_nodes,
            strategy=strategy,
            priority=priority,
            timeout_sec=timeout_sec,
        )
        
        # 权限检查
        max_allowed = self.permission.get_max_callable_count(self.node_info.node_id)
        task.max_nodes = min(task.max_nodes, max_allowed)
        
        # 虫皇权限：至少可以调用min_nodes个
        if self.node_info.permission_level == PermissionLevel.OVERMIND:
            task.max_nodes = max(task.max_nodes, max(min_nodes, 1))
        
        if task.max_nodes < 1:
            return {
                "task_id": task_id,
                "status": "denied",
                "error": "权限不足，无法调用其他节点",
            }
        
        # 3. 组建临时服务器
        tf = self.scheduler.create_taskforce(task)
        if not tf:
            return {
                "task_id": task_id,
                "status": "failed",
                "error": "无法组建临时服务器，可用节点不足",
            }
        
        # 4. 分发任务到各节点
        node_results = {}
        for member_id in tf.members:
            # 权限检查（虫皇自动通过）
            if (self.node_info.permission_level != PermissionLevel.OVERMIND and 
                not self.permission.can_call_node(self.node_info.node_id, member_id)):
                continue
            
            # 本地模拟：生成占位结果
            # 实际场景中通过消息总线发送任务并等待回复
            node = self.registry.get_node(member_id)
            node_results[member_id] = {
                "response": f"[来自{node.name if node else member_id}针对\"{query[:20]}\"的回答]",
                "confidence": 0.7 + (node.capability.compute_score * 0.05 if node else 0),
            }
        
        # 5. 聚合结果
        final_result = self.scheduler.aggregate_results(tf.taskforce_id, node_results)
        
        # 6. 解散临时服务器
        self.scheduler.complete_taskforce(tf.taskforce_id, final_result)
        
        self._stats["tasks_completed"] += 1
        
        return {
            "task_id": task_id,
            "taskforce_id": tf.taskforce_id,
            "status": "completed",
            "result": final_result,
            "members": tf.members,
            "strategy": strategy.value,
        }

    # ============================================================
    # 消息处理
    # ============================================================

    def _setup_message_handlers(self):
        """设置消息处理器"""
        self.bus.register_handler("discover", self._handle_discover)
        self.bus.register_handler("join_taskforce", self._handle_join)
        self.bus.register_handler("leave_taskforce", self._handle_leave)
        self.bus.register_handler("task_assign", self._handle_task_assign)
        self.bus.register_handler("task_result", self._handle_task_result)
        self.bus.register_handler("heartbeat", self._handle_heartbeat)

    def _handle_discover(self, msg: ProtocolMessage):
        """处理发现消息"""
        node_data = msg.payload.get("node_info", {})
        if node_data:
            node = NodeInfo(
                node_id=node_data.get("node_id", msg.sender),
                name=node_data.get("name", ""),
                role=NodeRole(node_data.get("role", "queen")),
                address=node_data.get("address", ""),
                capability=NodeCapability(
                    specializations=node_data.get("capabilities", ["general"]),
                    compute_score=node_data.get("compute_score", 1.0),
                ),
            )
            self.registry.register(node)

    def _handle_join(self, msg: ProtocolMessage):
        """处理加入临时服务器"""
        tf_id = msg.payload.get("taskforce_id", "")
        if tf_id and self.node_info.node_id in msg.payload.get("members", []):
            self.node_info.current_taskforces.append(tf_id)

    def _handle_leave(self, msg: ProtocolMessage):
        """处理离开临时服务器"""
        tf_id = msg.payload.get("taskforce_id", "")
        if tf_id in self.node_info.current_taskforces:
            self.node_info.current_taskforces.remove(tf_id)

    def _handle_task_assign(self, msg: ProtocolMessage):
        """处理任务分配"""
        # 子类可覆写此方法实现具体的任务执行逻辑
        pass

    def _handle_task_result(self, msg: ProtocolMessage):
        """处理任务结果"""
        pass

    def _handle_heartbeat(self, msg: ProtocolMessage):
        """处理心跳"""
        self.registry.update_heartbeat(msg.sender)

    # ============================================================
    # 心跳
    # ============================================================

    def _heartbeat_loop(self):
        """心跳循环"""
        while self._running:
            try:
                # 广播心跳
                self.bus.broadcast("heartbeat", {
                    "node_id": self.node_info.node_id,
                    "status": self.node_info.status.value,
                    "active_taskforces": len(self.node_info.current_taskforces),
                })
                
                # 检查其他节点心跳
                timeout_nodes = self.registry.check_heartbeats()
                if timeout_nodes:
                    logger.warning(f"心跳超时节点: {timeout_nodes}")
                
            except Exception as e:
                logger.error(f"心跳异常: {e}")
            
            time.sleep(10)  # 每10秒一次

    # ============================================================
    # 工具方法
    # ============================================================

    def _gen_id(self, prefix: str) -> str:
        return f"{prefix}_{hashlib.md5(f'{time.time()}{prefix}'.encode()).hexdigest()[:8]}"

    def get_status(self) -> Dict:
        """获取协议状态"""
        return {
            "node": self.node_info.to_dict(),
            "registry": self.registry.get_status(),
            "scheduler": self.scheduler.get_stats(),
            "permission": self.permission.get_stats(),
            "transport": self.bus.get_stats(),
            "protocol_stats": self._stats,
            "running": self._running,
        }

    def discover_nodes(self) -> List[Dict]:
        """发现所有可用节点"""
        nodes = self.registry.get_all_nodes()
        return [n.to_dict() for n in nodes if n.is_available()]