zenith-runtime 0.1.0

//! CPU 亲和绑定模块
//!
//! 提供线程到指定 CPU 核心的硬绑定，避免缓存抖动和 NUMA 远程访问。
//! 集成 `ExtremePlanner` 的拓扑规划结果，实现自动最优绑定。
//!
//! # Safety
//! 本模块**不含任何 unsafe**。所有 Linux 系统调用 (`sched_setaffinity` 等)
//! 已封装在 `zenith_linux::affinity` 模块内，遵循 AGENT.md §1.1
//! "所有 unsafe 必须封装在 zenith-linux 内部" 的硬性约束。

use thiserror::Error;

/// 亲和性操作错误
///
/// 包装 `zenith_linux::AffinityError` 并扩展业务级错误（如 CPU 越界、NUMA 节点未找到）。
#[derive(Debug, Error)]
pub enum CpuAffinityError {
    /// `sched_setaffinity` 系统调用失败
    #[error("sched_setaffinity failed: {0}")]
    SetAffinity(#[from] zenith_linux::AffinityError),

    /// CPU 编号超出范围
    #[error("CPU {cpu} out of range (max: {max})")]
    CpuOutOfRange {
        /// 非法的 CPU 编号
        cpu: usize,
        /// 当前在线 CPU 总数
        max: usize,
    },

    /// NUMA 节点不存在
    #[error("NUMA node {node} not found")]
    NumaNodeNotFound {
        /// 不存在的 NUMA 节点编号
        node: u32,
    },

    /// Worker 配置未找到（worker_id 不在 AffinityManager 配置列表中）
    #[error("worker {worker_id} not found in affinity configs (count: {config_count})")]
    WorkerNotFound {
        /// 未找到的 worker ID
        worker_id: u32,
        /// 当前已注册的配置总数
        config_count: usize,
    },
}

/// CPU 标识（包含核号与 NUMA 节点号）
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub struct CpuId {
    /// CPU 编号（0 起）
    pub id: usize,
    /// 所属 NUMA 节点
    pub node: u32,
}

impl CpuId {
    /// 创建新的 CpuId
    pub fn new(id: usize, node: u32) -> Self {
        Self { id, node }
    }
}

/// 将当前线程绑定到指定 CPU 核心
///
/// 直接委托给 `zenith_linux::set_thread_affinity`，本函数零 unsafe。
pub fn set_thread_affinity(cpu: usize) -> Result<(), CpuAffinityError> {
    zenith_linux::set_thread_affinity(cpu).map_err(CpuAffinityError::SetAffinity)
}

/// 将当前线程绑定到一组 CPU 核心
///
/// 直接委托给 `zenith_linux::set_thread_affinity_range`，本函数零 unsafe。
pub fn set_thread_affinity_range(cpus: &[usize]) -> Result<(), CpuAffinityError> {
    zenith_linux::set_thread_affinity_range(cpus).map_err(CpuAffinityError::SetAffinity)
}

/// 获取在线 CPU 数量
///
/// 直接委托给 `zenith_linux::online_cpu_count`，本函数零 unsafe。
pub fn get_online_cpus() -> usize {
    zenith_linux::online_cpu_count()
}

/// CPU 亲和配置：将一个 Worker 绑定到一组 CPU
#[derive(Debug, Clone)]
pub struct CpuAffinityConfig {
    /// Worker 标识
    pub worker_id: u32,
    /// 目标 CPU 编号列表
    pub cpu_ids: Vec<usize>,
    /// 所属 NUMA 节点
    pub numa_node: u32,
}

impl CpuAffinityConfig {
    /// 创建新的亲和配置
    pub fn new(worker_id: u32, cpu_ids: Vec<usize>, numa_node: u32) -> Self {
        Self { worker_id, cpu_ids, numa_node }
    }

    /// 应用配置到当前线程
    ///
    /// SYS-016：应用前校验每个 CPU 编号在在线范围内（越界 fail-closed，
    /// 返回 `CpuOutOfRange` 而非传给底层造成未定义行为）。
    pub fn apply(&self) -> Result<(), CpuAffinityError> {
        let online = get_online_cpus();
        for &cpu in &self.cpu_ids {
            if cpu >= online {
                return Err(CpuAffinityError::CpuOutOfRange { cpu, max: online });
            }
        }
        if self.cpu_ids.len() == 1 {
            set_thread_affinity(self.cpu_ids[0])
        } else {
            set_thread_affinity_range(&self.cpu_ids)
        }
    }
}

/// 亲和性管理器：管理多个 Worker 的亲和配置
#[derive(Debug)]
pub struct AffinityManager {
    configs: Vec<CpuAffinityConfig>,
}

impl AffinityManager {
    /// 创建空的管理器
    pub fn new() -> Self {
        Self { configs: Vec::new() }
    }

    /// 添加一个 Worker 的亲和配置
    pub fn add_config(&mut self, config: CpuAffinityConfig) {
        self.configs.push(config);
    }

    /// 应用所有配置并返回每个 Worker 的结果
    pub fn apply_all(&self) -> Vec<(u32, Result<(), CpuAffinityError>)> {
        self.configs
            .iter()
            .map(|c| {
                let result = c.apply();
                (c.worker_id, result)
            })
            .collect()
    }

    /// 应用指定 Worker 的亲和配置
    pub fn apply_for_worker(&self, worker_id: u32) -> Result<(), CpuAffinityError> {
        let config = self
            .configs
            .iter()
            .find(|c| c.worker_id == worker_id)
            .ok_or(CpuAffinityError::WorkerNotFound {
                worker_id,
                config_count: self.configs.len(),
            })?;
        config.apply()
    }

    /// 获取所有配置
    pub fn configs(&self) -> &[CpuAffinityConfig] {
        &self.configs
    }
}

impl Default for AffinityManager {
    fn default() -> Self {
        Self::new()
    }
}

#[cfg(test)]
mod tests {
    use super::*;

    #[test]
    fn test_cpu_id() {
        let cpu = CpuId::new(3, 0);
        assert_eq!(cpu.id, 3);
        assert_eq!(cpu.node, 0);
    }

    #[test]
    fn test_affinity_config() {
        let config = CpuAffinityConfig::new(1, vec![2, 3], 0);
        assert_eq!(config.worker_id, 1);
        assert_eq!(config.cpu_ids, vec![2, 3]);
        assert_eq!(config.numa_node, 0);
    }

    #[test]
    fn test_affinity_manager() {
        let mut mgr = AffinityManager::new();
        mgr.add_config(CpuAffinityConfig::new(0, vec![0], 0));
        mgr.add_config(CpuAffinityConfig::new(1, vec![1], 0));
        assert_eq!(mgr.configs().len(), 2);
    }

    #[test]
    fn test_apply_validates_cpu_bounds() {
        // SYS-016 回归：越界 CPU 编号在应用前 fail-closed 拒绝（返回 CpuOutOfRange），
        // 不得传给底层 set_affinity（此前缺越界校验）。
        let max = get_online_cpus().max(1);
        let config = CpuAffinityConfig::new(1, vec![max + 100], 0);
        let err = config.apply().unwrap_err();
        assert!(matches!(err, CpuAffinityError::CpuOutOfRange { .. }));
    }

    #[test]
    fn test_get_online_cpus() {
        let cpus = get_online_cpus();
        assert!(cpus >= 1);
    }

    #[cfg(target_os = "linux")]
    #[test]
    fn test_set_affinity_valid_cpu() {
        let max = get_online_cpus();
        if max > 0 {
            let result = set_thread_affinity(0);
            assert!(result.is_ok());
        }
    }
}