indexbus_platform_ops/
cpu.rs

1//! CPU-related operational helpers.
2//!
3//! **Platform support**
4//! - Linux: uses `sched_getaffinity` / `sched_setaffinity`.
5//! - Other OSes: functions return an error.
6//!
7//! **Contracts**
8//! - Pinning affects the *current thread* only.
9//! - CPU ids are OS CPU indices as used by the kernel.
10
11use crate::errors::{Error, Result};
12
13/// Returns the list of CPUs the current thread is allowed to run on (Linux).
14///
15/// On non-Linux targets this returns an error.
16pub fn current_thread_allowed_cpus() -> Result<Vec<usize>> {
17    #[cfg(target_os = "linux")]
18    {
19        let mut set: libc::cpu_set_t = unsafe { std::mem::zeroed() };
20        let rc =
21            unsafe { libc::sched_getaffinity(0, std::mem::size_of::<libc::cpu_set_t>(), &mut set) };
22        if rc != 0 {
23            return Err(Error::msg(format!(
24                "sched_getaffinity failed: {}",
25                std::io::Error::last_os_error()
26            )));
27        }
28
29        let max_cpus = std::mem::size_of::<libc::cpu_set_t>() * 8;
30        let mut out = Vec::new();
31        for cpu in 0..max_cpus {
32            if unsafe { libc::CPU_ISSET(cpu, &set) } {
33                out.push(cpu);
34            }
35        }
36        Ok(out)
37    }
38
39    #[cfg(not(target_os = "linux"))]
40    {
41        Err(Error::msg("CPU affinity is only supported on Linux"))
42    }
43}
44
45/// Format a list of CPU ids as a comma-separated string.
46///
47/// Returns `"<none>"` when `cpus` is empty.
48pub fn fmt_cpu_list(cpus: &[usize]) -> String {
49    if cpus.is_empty() {
50        return "<none>".to_string();
51    }
52
53    let mut out = String::new();
54    for (i, cpu) in cpus.iter().enumerate() {
55        if i > 0 {
56            out.push(',');
57        }
58        out.push_str(&cpu.to_string());
59    }
60    out
61}
62
63/// Pins the current thread to the given CPU (Linux).
64///
65/// On non-Linux targets this returns an error.
66///
67/// # Errors
68/// Returns an error when the requested CPU is out of range, when affinity syscalls fail, or when
69/// CPU pinning is not supported on this platform.
70pub fn pin_current_thread_to_cpu(cpu: usize) -> Result<()> {
71    #[cfg(target_os = "linux")]
72    {
73        let max_cpus = std::mem::size_of::<libc::cpu_set_t>() * 8;
74        if cpu >= max_cpus {
75            return Err(Error::msg(format!(
76                "cpu pin out of range for cpu_set_t: requested cpu={cpu}, max_supported={max_cpus}"
77            )));
78        }
79
80        let mut set: libc::cpu_set_t = unsafe { std::mem::zeroed() };
81        unsafe {
82            libc::CPU_ZERO(&mut set);
83            libc::CPU_SET(cpu, &mut set);
84        }
85
86        let rc = unsafe {
87            libc::sched_setaffinity(
88                0,
89                std::mem::size_of::<libc::cpu_set_t>(),
90                &set as *const libc::cpu_set_t,
91            )
92        };
93
94        if rc != 0 {
95            let os_err = std::io::Error::last_os_error();
96            let allowed = current_thread_allowed_cpus().unwrap_or_default();
97            let allowed_list = fmt_cpu_list(&allowed);
98            return Err(Error::msg(format!(
99                "sched_setaffinity(cpu={cpu}) failed: {os_err}; allowed_cpus={allowed_list}"
100            )));
101        }
102
103        Ok(())
104    }
105
106    #[cfg(not(target_os = "linux"))]
107    {
108        Err(Error::msg("CPU pinning is only supported on Linux"))
109    }
110}
111
112#[cfg(test)]
113mod tests {
114    use super::*;
115
116    #[test]
117    fn fmt_cpu_list_formats_empty() {
118        assert_eq!(fmt_cpu_list(&[]), "<none>");
119    }
120
121    #[test]
122    fn fmt_cpu_list_formats_values() {
123        assert_eq!(fmt_cpu_list(&[0]), "0");
124        assert_eq!(fmt_cpu_list(&[0, 2, 7]), "0,2,7");
125    }
126}
indexbus_platform_ops/cpu.rs

indexbus_platform_ops/
cpu.rs
