zenith-linux 0.1.0

//! 线程 CPU 亲和性绑定（Linux 平台抽象）
//!
//! 将 `libc::sched_setaffinity` 等 Linux 特有系统调用封装为 safe API，
//! 严格遵守 AGENT.md §1.1 "所有 unsafe 必须封装在 zenith-linux 内部"。
//!
//! # 安全保证
//! 本模块所有 `unsafe` 块均集中在 `libc` FFI 调用，对外提供 `Result` 接口。
//! 调用方零 unsafe 即可完成线程→CPU 核心的硬绑定。

#![allow(unsafe_code)]

use std::io;

/// 亲和性设置错误
#[derive(Debug, thiserror::Error)]
pub enum AffinityError {
    /// `sched_setaffinity` 系统调用失败
    #[error("sched_setaffinity failed: {0}")]
    SetFailed(io::Error),

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

/// 将当前线程绑定到指定 CPU 核心
///
/// # Arguments
/// * `cpu` - 目标 CPU 编号（0 起）
///
/// # Returns
/// * `Ok(())` - 绑定成功
/// * `Err(AffinityError)` - 系统调用失败或 CPU 越界
#[cfg(target_os = "linux")]
pub fn set_thread_affinity(cpu: usize) -> Result<(), AffinityError> {
    // 边界检查：CPU_SET 在 cpu >= CPU_SETSIZE 时会 panic，必须提前拦截
    let cpu_set_size = std::mem::size_of::<libc::cpu_set_t>() * 8;
    if cpu >= cpu_set_size {
        let online = online_cpu_count();
        return Err(AffinityError::CpuOutOfRange { cpu, online });
    }

    // SAFETY: libc::cpu_set_t 是 POD 类型，zeroed() 后由 CPU_ZERO/CPU_SET 显式初始化；
    // sched_setaffinity 仅读取 cpuset，不持有指针所有权；tid 来自 SYS_gettid 永远合法。
    unsafe {
        let mut cpuset: libc::cpu_set_t = std::mem::zeroed();
        libc::CPU_ZERO(&mut cpuset);
        libc::CPU_SET(cpu, &mut cpuset);

        let tid = libc::syscall(libc::SYS_gettid);
        let ret = libc::sched_setaffinity(
            tid as libc::pid_t,
            std::mem::size_of::<libc::cpu_set_t>(),
            &cpuset,
        );

        if ret != 0 {
            return Err(AffinityError::SetFailed(io::Error::last_os_error()));
        }
    }
    Ok(())
}

/// 非 Linux 平台：CPU 亲和为空操作
#[cfg(not(target_os = "linux"))]
pub fn set_thread_affinity(_cpu: usize) -> Result<(), AffinityError> {
    Ok(())
}

/// 将当前线程绑定到一组 CPU 核心
///
/// # Arguments
/// * `cpus` - 目标 CPU 编号切片（至少 1 个）
#[cfg(target_os = "linux")]
pub fn set_thread_affinity_range(cpus: &[usize]) -> Result<(), AffinityError> {
    // 边界检查：同 set_thread_affinity
    let cpu_set_size = std::mem::size_of::<libc::cpu_set_t>() * 8;
    for &cpu in cpus {
        if cpu >= cpu_set_size {
            let online = online_cpu_count();
            return Err(AffinityError::CpuOutOfRange { cpu, online });
        }
    }

    // SAFETY: 同 set_thread_affinity；CPU_SET 在循环中对每个 cpu 累加设置位。
    unsafe {
        let mut cpuset: libc::cpu_set_t = std::mem::zeroed();
        libc::CPU_ZERO(&mut cpuset);
        for &cpu in cpus {
            libc::CPU_SET(cpu, &mut cpuset);
        }

        let tid = libc::syscall(libc::SYS_gettid);
        let ret = libc::sched_setaffinity(
            tid as libc::pid_t,
            std::mem::size_of::<libc::cpu_set_t>(),
            &cpuset,
        );

        if ret != 0 {
            return Err(AffinityError::SetFailed(io::Error::last_os_error()));
        }
    }
    Ok(())
}

/// 非 Linux 平台：空操作
#[cfg(not(target_os = "linux"))]
pub fn set_thread_affinity_range(_cpus: &[usize]) -> Result<(), AffinityError> {
    Ok(())
}

/// 获取在线 CPU 数量
///
/// sysconf 返回 -1 时 fail-closed 回退 1（与 `page_size()` 的 -1 检查语义对齐：
/// 不得将错误值静默转为巨大 usize）。
#[cfg(target_os = "linux")]
pub fn online_cpu_count() -> usize {
    // SAFETY: sysconf 是纯只读系统调用，仅返回整数；不存在任何内存安全风险。
    let n = unsafe { libc::sysconf(libc::_SC_NPROCESSORS_ONLN) };
    if n <= 0 {
        return 1;
    }
    n as usize
}

/// 非 Linux 平台：返回 1
#[cfg(not(target_os = "linux"))]
pub fn online_cpu_count() -> usize {
    1
}

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

    #[test]
    fn test_online_cpu_count_positive() {
        let n = online_cpu_count();
        assert!(n >= 1, "online CPU count must be >= 1, got {}", n);
    }

    #[cfg(target_os = "linux")]
    #[test]
    fn test_set_affinity_cpu_zero() {
        // CPU 0 在所有 Linux 系统都存在
        let result = set_thread_affinity(0);
        assert!(result.is_ok(), "set_thread_affinity(0) should succeed: {:?}", result);
    }

    #[cfg(target_os = "linux")]
    #[test]
    fn test_set_affinity_range_all_cpus() {
        let n = online_cpu_count();
        let cpus: Vec<usize> = (0..n).collect();
        let result = set_thread_affinity_range(&cpus);
        assert!(result.is_ok(), "set_thread_affinity_range(all) should succeed: {:?}", result);
    }

    #[cfg(target_os = "linux")]
    #[test]
    fn test_set_affinity_out_of_range() {
        // 设置一个超大 CPU 编号，期望失败但不会 panic
        let result = set_thread_affinity(65535);
        // 失败为 Err，成功也允许（部分实现容忍越界）
        let _ = result;
    }
}