core/aggregation_protocol/discovery.py

#!/usr/bin/env python3
"""
虫群聚合协议 — 节点发现模块

功能：
- 节点注册与注销
- 心跳检测（判断节点在线/离线）
- 能力查询（找到满足条件的节点）
- 广播发现（局域网/已知节点广播）
"""

import hashlib
import threading
import time
from collections import defaultdict
from datetime import datetime, timedelta
from typing import Callable, Dict, List, Optional

from .types import (
    NodeInfo, NodeCapability, NodeRole, NodeStatus,
    PermissionLevel, ProtocolMessage,
)


class NodeRegistry:
    """
    节点注册中心 — 维护所有已知节点的信息
    
    类似GPU集群的服务器列表，但这里是分布式的：
    - 每个节点维护自己看到的节点列表
    - 心跳机制检测节点存活
    - 按能力匹配找到合适的节点组建临时服务器
    """

    def __init__(self, self_node_id: str, heartbeat_timeout: float = 30.0):
        self.self_node_id = self_node_id
        self.heartbeat_timeout = heartbeat_timeout
        
        # 已知节点 node_id -> NodeInfo
        self._nodes: Dict[str, NodeInfo] = {}
        self._lock = threading.RLock()
        
        # 能力索引: capability_name -> set of node_ids
        self._cap_index: Dict[str, set] = defaultdict(set)
        
        # 角色索引: role -> set of node_ids
        self._role_index: Dict[NodeRole, set] = defaultdict(set)
        
        # 事件回调
        self._callbacks: Dict[str, List[Callable]] = defaultdict(list)

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

    def register(self, node: NodeInfo) -> bool:
        """注册节点"""
        with self._lock:
            self._nodes[node.node_id] = node
            node.last_heartbeat = datetime.now()
            
            # 更新索引
            for cap in node.capability.specializations:
                self._cap_index[cap].add(node.node_id)
            self._role_index[node.role].add(node.node_id)
        
        self._fire("node_joined", node)
        return True

    def unregister(self, node_id: str) -> bool:
        """注销节点"""
        with self._lock:
            node = self._nodes.pop(node_id, None)
            if not node:
                return False
            
            # 清理索引
            for cap in node.capability.specializations:
                self._cap_index[cap].discard(node_id)
            self._role_index[node.role].discard(node_id)
        
        self._fire("node_left", node)
        return True

    def update_heartbeat(self, node_id: str) -> bool:
        """更新心跳"""
        with self._lock:
            node = self._nodes.get(node_id)
            if node:
                node.last_heartbeat = datetime.now()
                node.status = NodeStatus.ONLINE
                return True
        return False

    def get_node(self, node_id: str) -> Optional[NodeInfo]:
        """获取节点信息"""
        with self._lock:
            return self._nodes.get(node_id)

    def get_all_nodes(self) -> List[NodeInfo]:
        """获取所有已知节点"""
        with self._lock:
            return list(self._nodes.values())

    # ============================================================
    # 节点发现 — 找到满足条件的节点
    # ============================================================

    def discover_by_capability(self, capability: str, 
                                min_compute: float = 0.0,
                                exclude: List[str] = None) -> List[NodeInfo]:
        """按能力发现节点"""
        with self._lock:
            candidates = self._cap_index.get(capability, set())
            exclude = exclude or []
            
            results = []
            for nid in candidates:
                if nid in exclude or nid == self.self_node_id:
                    continue
                node = self._nodes.get(nid)
                if node and node.is_available() and node.capability.compute_score >= min_compute:
                    results.append(node)
            
            # 按计算能力排序
            results.sort(key=lambda n: n.capability.compute_score, reverse=True)
            return results

    def discover_by_role(self, role: NodeRole,
                         exclude: List[str] = None) -> List[NodeInfo]:
        """按角色发现节点"""
        with self._lock:
            candidates = self._role_index.get(role, set())
            exclude = exclude or []
            
            results = []
            for nid in candidates:
                if nid in exclude or nid == self.self_node_id:
                    continue
                node = self._nodes.get(nid)
                if node and node.is_available():
                    results.append(node)
            return results

    def discover_for_task(self, required_caps: List[str],
                          min_nodes: int = 2,
                          max_nodes: int = 5,
                          exclude: List[str] = None) -> List[NodeInfo]:
        """
        为任务发现合适的节点
        
        核心逻辑：找到同时满足所有所需能力的节点集合
        类似于GPU集群中找有空闲显存的显卡
        """
        with self._lock:
            exclude = exclude or []
            
            # 计算每个可用节点的匹配度
            scored_nodes = []
            for node in self._nodes.values():
                if not node.can_accept_task():
                    continue
                if node.node_id in exclude or node.node_id == self.self_node_id:
                    continue
                
                # 计算能力匹配度
                matched = sum(1 for cap in required_caps 
                             if cap in node.capability.specializations)
                score = matched / max(len(required_caps), 1)
                
                if matched > 0 or not required_caps:  # 至少匹配一个或无特殊要求
                    scored_nodes.append((node, score, matched))
            
            # 先按匹配度，再按计算能力排序
            scored_nodes.sort(key=lambda x: (x[1], x[0].capability.compute_score), reverse=True)
            
            # 取top max_nodes
            selected = [n for n, s, m in scored_nodes[:max_nodes]]
            
            # 确保最少节点数
            if len(selected) < min_nodes:
                # 放宽条件，加入任何可用节点
                for node in self._nodes.values():
                    if (node.can_accept_task() and 
                        node.node_id not in exclude and
                        node.node_id != self.self_node_id and
                        node.node_id not in [n.node_id for n in selected]):
                        selected.append(node)
                        if len(selected) >= min_nodes:
                            break
            
            return selected[:max_nodes]

    # ============================================================
    # 心跳检测
    # ============================================================

    def check_heartbeats(self) -> List[str]:
        """检查心跳，标记超时节点"""
        now = datetime.now()
        timeout_ids = []
        
        with self._lock:
            for node_id, node in list(self._nodes.items()):
                if node_id == self.self_node_id:
                    continue
                if node.last_heartbeat:
                    elapsed = (now - node.last_heartbeat).total_seconds()
                    if elapsed > self.heartbeat_timeout:
                        node.status = NodeStatus.OFFLINE
                        timeout_ids.append(node_id)
                else:
                    # 没有心跳记录
                    node.status = NodeStatus.OFFLINE
        
        return timeout_ids

    def get_online_count(self) -> int:
        """获取在线节点数"""
        with self._lock:
            return sum(1 for n in self._nodes.values() if n.is_available())

    # ============================================================
    # 事件回调
    # ============================================================

    def on(self, event: str, callback: Callable):
        """注册事件回调"""
        self._callbacks[event].append(callback)

    def _fire(self, event: str, data=None):
        """触发事件"""
        for cb in self._callbacks.get(event, []):
            try:
                cb(data)
            except Exception:
                pass

    # ============================================================
    # 状态信息
    # ============================================================

    def get_status(self) -> Dict:
        """获取注册中心状态"""
        with self._lock:
            online = sum(1 for n in self._nodes.values() if n.is_available())
            total = len(self._nodes)
            return {
                "self_node_id": self.self_node_id,
                "total_nodes": total,
                "online_nodes": online,
                "offline_nodes": total - online,
                "capabilities_indexed": len(self._cap_index),
                "roles": {r.value: len(ids) for r, ids in self._role_index.items()},
            }