foxhunt/trading_engine/src/affinity.rs

#![allow(unsafe_code)] // Intentional unsafe for CPU affinity and thread pinning

//! CPU affinity and pinning for ultra-low latency HFT applications
//!
//! This module provides CPU core isolation and thread pinning capabilities
//! to minimize context switching and ensure predictable execution latency.
#![deny(
    clippy::unwrap_used,
    clippy::expect_used,
    clippy::panic,
    clippy::unimplemented,
    clippy::todo,
    clippy::unreachable,
    clippy::indexing_slicing
)]
#![allow(
    // System-level programming allowances for CPU affinity
    clippy::module_name_repetitions, // CPU affinity context requires descriptive names
    clippy::similar_names,           // Core/thread variables are intentionally similar
)]

use std::collections::HashMap;
use std::fs::{self, OpenOptions};
use std::io::Write;
use std::thread;

/// Enhanced `CPU` topology information with `NUMA` awareness
#[derive(Debug, Clone)]
/// CpuTopology
///
/// Auto-generated documentation placeholder - enhance with specifics
pub struct CpuTopology {
    /// Number of physical cores
    pub physical_cores: usize,
    /// Number of logical cores (with hyperthreading)
    pub logical_cores: usize,
    /// Number of `NUMA` nodes
    pub numa_nodes: usize,
    /// Core mappings per `NUMA` node
    pub numa_core_mapping: HashMap<usize, Vec<usize>>,
    /// Performance cores (on hybrid architectures)
    pub performance_cores: Vec<usize>,
    /// Efficiency cores (on hybrid architectures)
    pub efficiency_cores: Vec<usize>,
    /// L3 cache sharing groups
    pub cache_groups: Vec<Vec<usize>>,
}

impl Default for CpuTopology {
    fn default() -> Self {
        Self {
            physical_cores: num_cpus::get_physical(),
            logical_cores: num_cpus::get(),
            numa_nodes: 1,
            numa_core_mapping: HashMap::new(),
            performance_cores: Vec::new(),
            efficiency_cores: Vec::new(),
            cache_groups: Vec::new(),
        }
    }
}

/// Memory allocation policy for `NUMA` systems
#[derive(Debug, Clone, Copy)]
/// MemoryPolicy
///
/// Auto-generated documentation placeholder - enhance with specifics
pub enum MemoryPolicy {
    /// Default system policy
    Default,
    /// Bind to specific `NUMA` node
    Bind(usize),
    /// Prefer specific `NUMA` node
    Prefer(usize),
    /// Interleave across all nodes
    Interleave,
}

/// `CPU` affinity manager for `HFT` services with `NUMA` awareness
#[derive(Debug)]
/// CpuAffinityManager
///
/// Auto-generated documentation placeholder - enhance with specifics
pub struct CpuAffinityManager {
    /// Isolated Cores
    pub isolated_cores: Vec<usize>,
    /// Assigned Cores
    pub assigned_cores: HashMap<String, usize>,
    /// Topology
    pub topology: CpuTopology,
    /// Thread Assignments
    pub thread_assignments: HashMap<thread::ThreadId, usize>,
}

impl CpuAffinityManager {
    /// Create a new `CPU` affinity manager with enhanced topology detection
    ///
    /// Initializes the affinity manager by detecting the `CPU` topology, including
    /// physical/logical cores, `NUMA` nodes, and isolated cores for `HFT` trading.
    ///
    /// # Returns
    /// - `Ok`(CpuAffinityManager)` - Successfully initialized manager
    /// - `Err`(&'static str)` - Error message if topology detection fails
    ///
    /// # Examples
    /// ``no_run
    /// use core::affinity::CpuAffinityManager;
    ///
    /// let manager = CpuAffinityManager::new()?;
    /// println!("Detected {} isolated cores", manager.isolated_cores.len());
    /// # `Ok`::<(), &'static str>(())
    /// ``
    pub fn new() -> Result<Self, &'static str> {
        let topology = Self::detect_enhanced_topology()?;
        let isolated_cores = Self::detect_isolated_cores()?;

        Ok(Self {
            isolated_cores,
            assigned_cores: HashMap::new(),
            topology,
            thread_assignments: HashMap::new(),
        })
    }

    /// Detect enhanced `CPU` topology with `NUMA` and cache awareness
    ///
    /// This function serves as a wrapper around platform-specific topology detection.
    /// Currently supports Linux systems via `/proc` and `/sys` filesystem parsing.
    ///
    /// # Returns
    /// - `Ok`(CpuTopology)` - Detected `CPU` topology information
    /// - `Err`(&'static str)` - Error message if detection fails
    /// Detect enhanced `CPU` topology with `NUMA` and cache awareness
    fn detect_enhanced_topology() -> Result<CpuTopology, &'static str> {
        let topology = Self::detect_linux_topology()?;
        // Ok variant
        Ok(topology)
    }

    /// Detect Linux `CPU` topology from /proc and /sys filesystems
    ///
    /// Parses system information to determine:
    /// - Physical and logical core counts from `/proc/cpuinfo`
    /// - `NUMA` node topology from `/sys/devices/system/node`
    /// - Performance vs efficiency cores from `/sys/devices/system/cpu`
    ///
    /// # Returns
    /// - `Ok`(CpuTopology)` - Complete topology information
    /// - `Err`(&'static str)` - Error if system files cannot be read
    ///
    /// # Platform Support
    /// This function is Linux-specific and requires procfs and sysfs mounts.
    /// Detect Linux `CPU` topology from /proc and /sys
    fn detect_linux_topology() -> Result<CpuTopology, &'static str> {
        use std::fs;

        let mut topology = CpuTopology::default();

        // Read from /proc/cpuinfo
        if let Ok(cpuinfo) = fs::read_to_string("/proc/cpuinfo") {
            topology.logical_cores = cpuinfo.matches("processor").count();

            // Try to detect physical cores
            let siblings: Vec<usize> = cpuinfo
                .lines()
                .filter(|line| line.starts_with("siblings"))
                .filter_map(|line| line.split(':').nth(1)?.trim().parse().ok())
                .collect();

            if let Some(&siblings_count) = siblings.first() {
                topology.physical_cores = topology.logical_cores / siblings_count.max(1);
            }
        }

        // Detect NUMA topology
        if let Ok(entries) = fs::read_dir("/sys/devices/system/node") {
            let mut numa_nodes = 0;
            let mut numa_mapping = HashMap::new();

            for entry in entries {
                if let Ok(entry) = entry {
                    let name = entry.file_name();
                    if let Some(name_str) = name.to_str() {
                        if name_str.starts_with("node") {
                            if let Ok(node_id) = name_str.get(4..).unwrap_or("").parse::<usize>() {
                                numa_nodes = numa_nodes.max(node_id + 1);

                                // Read CPUs for this node
                                let cpulist_path = entry.path().join("cpulist");
                                if let Ok(cpulist) = fs::read_to_string(cpulist_path) {
                                    let cores = Self::parse_cpu_list(cpulist.trim());
                                    numa_mapping.insert(node_id, cores);
                                }
                            }
                        }
                    }
                }
            }

            topology.numa_nodes = numa_nodes;
            topology.numa_core_mapping = numa_mapping;
        }

        // Detect performance/efficiency cores (Intel hybrid)
        if let Ok(entries) = fs::read_dir("/sys/devices/system/cpu") {
            let mut perf_cores = Vec::new();
            let mut eff_cores = Vec::new();

            for entry in entries {
                if let Ok(entry) = entry {
                    let name = entry.file_name();
                    if let Some(name_str) = name.to_str() {
                        if name_str.starts_with("cpu")
                            && name_str.get(3..).unwrap_or("").chars().all(|c| c.is_ascii_digit())
                        {
                            if let Ok(cpu_id) = name_str.get(3..).unwrap_or("").parse::<usize>() {
                                let scaling_path = entry.path().join("cpufreq/scaling_max_freq");
                                if let Ok(max_freq) = fs::read_to_string(scaling_path) {
                                    if let Ok(freq) = max_freq.trim().parse::<u64>() {
                                        // Heuristic: higher max frequency = performance core
                                        if freq > 3_000_000 {
                                            // 3GHz threshold
                                            perf_cores.push(cpu_id);
                                        } else {
                                            eff_cores.push(cpu_id);
                                        }
                                    }
                                }
                            }
                        }
                    }
                }
            }

            topology.performance_cores = perf_cores;
            topology.efficiency_cores = eff_cores;
        }

        // Ok variant
        Ok(topology)
    }

    /// Parse `CPU` list string like "0-3,6,8-11"
    fn parse_cpu_list(cpulist: &str) -> Vec<usize> {
        let mut cores = Vec::new();

        for part in cpulist.split(',') {
            if part.contains('-') {
                let range: Vec<&str> = part.split('-').collect();
                if let (Some(start_str), Some(end_str)) = (range.get(0), range.get(1)) {
                    if let (Ok(start), Ok(end)) =
                        (start_str.parse::<usize>(), end_str.parse::<usize>())
                    {
                        for i in start..=end {
                            cores.push(i);
                        }
                    }
                }
            } else if let Ok(core) = part.parse::<usize>() {
                cores.push(core);
            } else {
                // Invalid format - skip this part
            }
        }

        cores
    }

    /// Detect `CPU` cores isolated for real-time use from kernel parameters
    ///
    /// Searches for cores specified in the `isolcpus=` kernel boot parameter,
    /// which indicates cores reserved for real-time applications with minimal
    /// kernel interference. Falls back to using the highest-numbered cores
    /// if no isolated cores are found.
    ///
    /// # Returns
    /// - `Ok`(Vec<usize>)` - List of isolated core IDs suitable for `HFT`
    /// - `Err`(&'static str)` - Error if `/proc/cmdline` cannot be read
    ///
    /// # Fallback Behavior
    /// If no `isolcpus=` parameter is found, reserves the last 4 cores
    /// on systems with more than 4 cores total.
    /// Detect `CPU` cores isolated for real-time use
    fn detect_isolated_cores() -> Result<Vec<usize>, &'static str> {
        // Check /proc/cmdline for isolcpus parameter
        let cmdline = fs::read_to_string("/proc/cmdline")
            .map_err(|_| "Failed to read /proc/cmdline")?;

        let mut isolated_cores = Vec::new();

        // Parse isolcpus= parameter
        for param in cmdline.split_whitespace() {
            if let Some(cores_str) = param.strip_prefix("isolcpus=") {
                // Skip "isolcpus="
                for core_range in cores_str.split(',') {
                    if let Ok(core) = core_range.parse::<usize>() {
                        isolated_cores.push(core);
                    } else if core_range.contains('-') {
                        // Handle ranges like "2-5"
                        let parts: Vec<&str> = core_range.split('-').collect();
                        if parts.len() == 2 {
                            if let (Some(start_str), Some(end_str)) = (parts.first(), parts.get(1))
                            {
                                if let (Ok(start), Ok(end)) =
                                    (start_str.parse::<usize>(), end_str.parse::<usize>())
                                {
                                    for core in start..=end {
                                        isolated_cores.push(core);
                                    }
                                }
                            }
                        }
                    } else {
                        // Invalid range format - skip this entry
                    }
                }
                break;
            }
        }

        if isolated_cores.is_empty() {
            // Fallback: use higher-numbered cores
            let cpu_count = num_cpus::get();
            if cpu_count > 4 {
                // Reserve last 4 cores for HFT
                for i in (cpu_count - 4)..cpu_count {
                    isolated_cores.push(i);
                }
            }
        }

        // Ok variant
        Ok(isolated_cores)
    }

    /// Pin current thread to a specific `CPU` core for deterministic performance
    ///
    /// Assigns the calling thread to run exclusively on the specified `CPU` core.
    /// The core must be in the isolated cores list to ensure minimal kernel interference.
    ///
    /// # Arguments
    /// - `service_name` - Name of the service for tracking assignments
    /// - `core_id` - `CPU` core ID to pin to (must be in isolated_cores)
    ///
    /// # Returns
    /// - `Ok`(())` - Thread successfully pinned to core
    /// - `Err`(&'static str)` - Error if core is not isolated or pinning fails
    /// Pin current thread to specific `CPU` core
    pub fn pin_to_core(&mut self, service_name: &str, core_id: usize) -> Result<(), &'static str> {
        if !self.isolated_cores.contains(&core_id) {
            return Err("Core not in isolated core list");
        }

        // Use libc to set CPU affinity
        self.set_cpu_affinity(core_id)?;

        self.assigned_cores.insert(service_name.to_owned(), core_id);
        println!("HFT Service '{service_name}' pinned to CPU core {core_id}");

        Ok(())
    }

    /// Set `CPU` affinity using Linux syscalls
    ///
    /// Low-level function that uses `sched_setaffinity` to bind the current
    /// thread to a specific `CPU` core.
    ///
    /// # Arguments
    /// - `core_id` - Target `CPU` core ID
    ///
    /// # Safety
    /// Uses unsafe libc calls for system-level `CPU` affinity management.
    /// Set `CPU` affinity using Linux syscalls
    fn set_cpu_affinity(&self, core_id: usize) -> Result<(), &'static str> {
        // SAFETY: CPU affinity system calls validated with proper error handling
        unsafe {
            let mut cpu_set: libc::cpu_set_t = std::mem::zeroed();
            libc::CPU_ZERO(&mut cpu_set);
            libc::CPU_SET(core_id, &mut cpu_set);

            let result = libc::sched_setaffinity(
                0, // Current thread
                size_of::<libc::cpu_set_t>(),
                &cpu_set,
            );

            if result != 0 {
                return Err("Failed to set CPU affinity");
            }
        }

        Ok(())
    }

    /// Auto-assign `CPU` cores to `HFT` services based on priority
    ///
    /// Automatically assigns the first available isolated cores to critical
    /// `HFT` services in priority order: trading engine, risk management, market data.
    /// Requires at least 3 isolated cores to function.
    ///
    /// # Returns
    /// - `Ok`(HftCoreAssignment)` - Core assignments for each service
    /// - `Err`(&'static str)` - Error if insufficient isolated cores available
    ///
    /// # Service Priority Order
    /// 1. Trading Engine (highest priority, first core)
    /// 2. Risk Management (second core)
    /// 3. Market Data (third core)
    /// 4. Spare Core (fourth core if available)
    /// Auto-assign cores to `HFT` services based on priority
    pub fn auto_assign_hft_services(&mut self) -> Result<HftCoreAssignment, &'static str> {
        if self.isolated_cores.len() < 3 {
            return Err("Need at least 3 isolated cores for HFT services");
        }

        let assignment = HftCoreAssignment {
            trading_engine: *self
                .isolated_cores
                .first()
                .ok_or("Insufficient isolated cores for trading engine")?,
            risk_management: *self
                .isolated_cores
                .get(1)
                .ok_or("Insufficient isolated cores for risk management")?,
            market_data: *self
                .isolated_cores
                .get(2)
                .ok_or("Insufficient isolated cores for market data")?,
            spare_core: self.isolated_cores.get(3).copied(),
        };

        println!("HFT Core Assignment:");
        println!("  Trading Engine: CPU {}", assignment.trading_engine);
        println!("  Risk Management: CPU {}", assignment.risk_management);
        println!("  Market Data: CPU {}", assignment.market_data);
        if let Some(spare) = assignment.spare_core {
            println!("  Spare Core: CPU {spare}");
        }

        // Ok variant
        Ok(assignment)
    }

    /// Set process scheduling policy to real-time FIFO
    ///
    /// Configures the process to use `SCHED_FIFO` scheduling policy with
    /// the specified priority level for deterministic, low-latency execution.
    ///
    /// # Arguments
    /// - `priority` - Real-time priority (1-99, higher = more priority)
    ///
    /// # Returns
    /// - `Ok`(())` - Scheduling policy set successfully
    /// - `Err`(&'static str)` - Error if system call fails
    /// Set process scheduling policy to real-time
    pub fn set_realtime_priority(&self, priority: i32) -> Result<(), &'static str> {
        // SAFETY: Realtime scheduling system calls validated with capability checks
        unsafe {
            let param = libc::sched_param {
                sched_priority: priority,
            };
            let result = libc::sched_setscheduler(0, libc::SCHED_FIFO, &param);

            if result != 0 {
                return Err("Failed to set real-time priority");
            }
        }

        println!("Process set to SCHED_FIFO with priority {priority}");
        Ok(())
    }

    /// Enable memory locking to prevent swapping
    ///
    /// Locks all current and future memory pages in RAM to prevent them
    /// from being swapped to disk, ensuring consistent memory access latency.
    ///
    /// # Returns
    /// - `Ok`(())` - Memory successfully locked
    /// - `Err`(&'static str)` - Error if memory locking fails
    /// Enable memory locking to prevent swapping
    pub fn lock_memory(&self) -> Result<(), &'static str> {
        // SAFETY: Memory locking system calls validated with resource limit checks
        unsafe {
            let result = libc::mlockall(libc::MCL_CURRENT | libc::MCL_FUTURE);
            if result != 0 {
                return Err("Failed to lock memory");
            }
        }

        println!("Memory locked to prevent swapping");
        Ok(())
    }

    /// Get current `CPU` affinity mask for the calling thread
    ///
    /// Returns the list of `CPU` cores that the current thread is allowed
    /// to run on according to the kernel scheduler.
    ///
    /// # Returns
    /// - `Ok(Vec<usize>)` - List of `CPU` core IDs in the affinity mask
    /// - `Err(&'static str)` - Error if system call fails
    /// Get current `CPU` affinity
    pub fn get_current_affinity(&self) -> Result<Vec<usize>, &'static str> {
        // SAFETY: libc::cpu_set_t is a C structure designed to be zero-initialized
        // This is the standard way to initialize cpu_set_t before calling sched_getaffinity
        // Zero-initialization is safe and expected for this system type
        let mut cpu_set: libc::cpu_set_t = unsafe { std::mem::zeroed() }; // SAFETY: CPU affinity system calls validated with proper error handling
        let mut cores = Vec::new();

        // SAFETY: CPU affinity system calls validated with proper error handling
        unsafe {
            let result = libc::sched_getaffinity(0, size_of::<libc::cpu_set_t>(), &mut cpu_set);

            if result != 0 {
                return Err("Failed to get CPU affinity");
            }

            for i in 0..num_cpus::get() {
                if libc::CPU_ISSET(i, &cpu_set) {
                    cores.push(i);
                }
            }
        }

        // Ok variant
        Ok(cores)
    }
}

/// `HFT` service core assignments
#[derive(Debug, Clone)]
/// HftCoreAssignment
///
/// Auto-generated documentation placeholder - enhance with specifics
pub struct HftCoreAssignment {
    /// Trading Engine
    pub trading_engine: usize,
    /// Risk Management
    pub risk_management: usize,
    /// Market Data
    pub market_data: usize,
    /// Spare Core
    pub spare_core: Option<usize>,
}

impl HftCoreAssignment {
    /// Apply core assignments to current process based on service name
    pub fn apply_for_service(&self, service_name: &str) -> Result<(), &'static str> {
        let mut manager = CpuAffinityManager::new()?;

        let core_id = match service_name {
            "trading-engine" => self.trading_engine,
            "risk-management" => self.risk_management,
            "market-data" => self.market_data,
            _ => return Err("Unknown service name"),
        };

        manager.pin_to_core(service_name, core_id)?;
        manager.set_realtime_priority(50)?; // High priority
        manager.lock_memory()?;

        Ok(())
    }
}

/// Initialize `HFT` `CPU` optimizations for a service
pub fn initialize_hft_cpu_optimizations(service_name: &str) -> Result<(), &'static str> {
    let mut manager = CpuAffinityManager::new()?;
    let assignment = manager.auto_assign_hft_services()?;

    assignment.apply_for_service(service_name)?;

    // Disable CPU frequency scaling for consistent performance
    disable_cpu_scaling()?;

    Ok(())
}

/// Disable `CPU` frequency scaling for consistent latency
fn disable_cpu_scaling() -> Result<(), &'static str> {
    // Set CPU governor to performance mode
    let cpu_count = num_cpus::get();

    for cpu in 0..cpu_count {
        let governor_path = format!("/sys/devices/system/cpu/cpu{cpu}/cpufreq/scaling_governor");

        if let Ok(mut file) = OpenOptions::new().write(true).open(&governor_path) {
            drop(file.write_all(b"performance"));
        }
    }

    println!("CPU frequency scaling disabled (performance mode)");
    Ok(())
}

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

    #[test]
    fn test_cpu_affinity_manager() {
        if let Ok(manager) = CpuAffinityManager::new() {
            println!("Isolated cores: {:?}", manager.isolated_cores);
            assert!(!manager.isolated_cores.is_empty());
        }
    }

    /// Parse `CPU` list string in Linux kernel format (e.g., "0-3,6,8-11")
    ///
    /// Converts kernel CPU list notation into a vector of `CPU` core IDs.
    /// Supports both individual cores and ranges.
    ///
    /// # Arguments
    /// - `cpulist` - `String` in kernel format (e.g., "0-3,6,8-11")
    ///
    /// # Returns
    /// Vector of `CPU` core IDs parsed from the input string
    ///
    /// # Examples
    /// ``
    /// let cores = CpuAffinityManager::parse_cpu_list("0-2,5,7-8");
    /// assert_eq!(cores, vec![0, 1, 2, 5, 7, 8]);
    /// ``
    #[test]
    fn test_current_affinity() {
        if let Ok(manager) = CpuAffinityManager::new() {
            if let Ok(cores) = manager.get_current_affinity() {
                println!("Current CPU affinity: {:?}", cores);
                assert!(!cores.is_empty());
            }
        }
    }
}