agave_cpu_utils/
affinity.rs

1//! Core CPU affinity operations.
2
3use std::{io, mem, ops::Deref};
4
5const CPU_SETSIZE: usize = libc::CPU_SETSIZE as usize;
6
7/// Identifies a logical CPU (hardware thread) by its kernel-assigned ID.
8#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)]
9pub struct CpuId(usize);
10
11impl CpuId {
12    pub fn new(cpu: usize) -> io::Result<Self> {
13        if cpu < CPU_SETSIZE {
14            Ok(Self(cpu))
15        } else {
16            Err(io::Error::from_raw_os_error(libc::EINVAL))
17        }
18    }
19}
20
21impl Deref for CpuId {
22    type Target = usize;
23
24    fn deref(&self) -> &Self::Target {
25        &self.0
26    }
27}
28
29/// Set CPU affinity for a thread.
30///
31/// Restricts the thread to run only on the specified CPUs.
32///
33/// # Arguments
34/// * `thread_id` - Thread ID to set affinity for. `None` means the calling thread.
35/// * `cpus` - CPU IDs to bind the thread to. Can be any iterable collection.
36///
37/// # Examples
38///
39/// ```no_run
40/// # use agave_cpu_utils::*;
41/// # fn main() -> std::io::Result<()> {
42/// // Pin current thread to CPU 0
43/// set_cpu_affinity(None, [CpuId::new(0)?])?;
44///
45/// // Pin current thread to multiple CPUs
46/// set_cpu_affinity(None, [CpuId::new(0)?, CpuId::new(1)?, CpuId::new(2)?])?;
47/// # Ok(())
48/// # }
49/// ```
50///
51/// # Errors
52///
53/// Returns [`io::ErrorKind::InvalidInput`] if any CPU ID is too large for `cpu_set_t`.
54/// Returns [`io::Error`] if the system call fails, including offline CPUs or CPUs
55/// disallowed by the task's hard cpuset/cgroup.
56pub fn set_cpu_affinity(
57    thread_id: Option<libc::pid_t>,
58    cpus: impl IntoIterator<Item = CpuId>,
59) -> io::Result<()> {
60    // safety: cpu_set_t is a POD type, zero-initialization is standard
61    let mut cpu_set: libc::cpu_set_t = unsafe { mem::zeroed() };
62
63    for cpu in cpus {
64        // safety: CpuId values are constructed only after validating cpu < CPU_SETSIZE.
65        unsafe { libc::CPU_SET(*cpu, &mut cpu_set) };
66    }
67
68    let tid = thread_id.unwrap_or(0);
69
70    // safety: sched_setaffinity is safe with valid parameters
71    let result =
72        unsafe { libc::sched_setaffinity(tid, mem::size_of::<libc::cpu_set_t>(), &cpu_set) };
73
74    if result != 0 {
75        return Err(io::Error::last_os_error());
76    }
77
78    Ok(())
79}
80
81/// Get the CPU affinity mask for a thread.
82///
83/// Returns a sorted vector of CPU IDs that the thread is allowed to run on.
84///
85/// # Arguments
86/// * `thread_id` - Thread ID to query. `None` means the calling thread.
87///
88/// # Examples
89///
90/// ```no_run
91/// # use agave_cpu_utils::*;
92/// # fn main() -> std::io::Result<()> {
93/// let cpus = cpu_affinity(None)?;
94/// println!("Thread can run on CPUs: {:?}", cpus);
95/// # Ok(())
96/// # }
97/// ```
98///
99/// # Errors
100///
101/// Returns [`io::Error`] if the system call fails.
102pub fn cpu_affinity(thread_id: Option<libc::pid_t>) -> io::Result<Vec<CpuId>> {
103    // safety: cpu_set_t is a POD type, zero-initialization is standard
104    let mut cpu_set: libc::cpu_set_t = unsafe { mem::zeroed() };
105
106    let tid = thread_id.unwrap_or(0);
107
108    // safety: sched_getaffinity is safe with valid parameters
109    let result =
110        unsafe { libc::sched_getaffinity(tid, mem::size_of::<libc::cpu_set_t>(), &mut cpu_set) };
111
112    if result != 0 {
113        return Err(io::Error::last_os_error());
114    }
115
116    let mut cpus = Vec::new();
117    for cpu in 0..CPU_SETSIZE {
118        // safety: cpu < CPU_SETSIZE by construction
119        if unsafe { libc::CPU_ISSET(cpu, &cpu_set) } {
120            cpus.push(CpuId::new(cpu)?);
121        }
122    }
123
124    Ok(cpus)
125}
126
127#[cfg(test)]
128mod tests {
129    use super::*;
130
131    #[test]
132    fn test_cpu_id_validation() {
133        let result = CpuId::new(CPU_SETSIZE);
134        assert_eq!(result.unwrap_err().raw_os_error(), Some(libc::EINVAL));
135    }
136
137    #[test]
138    fn test_cpu_affinity_returns_sorted() {
139        let cpus = cpu_affinity(None).expect("failed to query current CPU affinity");
140        assert!(
141            cpus.windows(2).all(|window| *window[0] <= *window[1]),
142            "cpu_affinity should return sorted CPU list"
143        );
144    }
145}
agave_cpu_utils/affinity.rs

agave_cpu_utils/
affinity.rs
