kernelkit/
affinity.rs

1//! CPU affinity helpers with Linux fast paths and safe fallbacks elsewhere.
2
3use crate::{Error, Result};
4
5/// Pin the calling thread to a specific CPU core.
6/// # Errors
7/// Returns an error if pinning fails.
8pub fn pin_to_core(core_id: usize) -> Result<()> {
9    #[cfg(target_os = "linux")]
10    {
11        let cpu_set_words = usize::try_from(libc::CPU_SETSIZE)
12            .unwrap_or(1024)
13            .div_ceil(usize::BITS as usize);
14        let max_core = cpu_set_words * usize::BITS as usize;
15        if core_id >= max_core {
16            return Err(Error::System {
17                operation: "sched_setaffinity",
18                source: std::io::Error::from(std::io::ErrorKind::InvalidInput),
19            });
20        }
21
22        let mut mask = vec![0usize; cpu_set_words];
23        mask[core_id / usize::BITS as usize] |= 1usize << (core_id % usize::BITS as usize);
24
25        // SAFETY: affinity mask lives for the duration of the syscall.
26        let result = unsafe {
27            libc::sched_setaffinity(
28                0,
29                std::mem::size_of_val(mask.as_slice()),
30                mask.as_ptr().cast::<libc::cpu_set_t>(),
31            )
32        };
33        if result != 0 {
34            return Err(Error::System {
35                operation: "sched_setaffinity",
36                source: std::io::Error::last_os_error(),
37            });
38        }
39    }
40
41    #[cfg(not(target_os = "linux"))]
42    {
43        let _ = core_id;
44    }
45
46    Ok(())
47}
48
49/// Read the CPU affinity mask for a specific IRQ number.
50///
51/// Reads `/proc/irq/<irq>/smp_affinity` to determine which cores
52/// handle a given interrupt. Useful for co-locating scan threads
53/// with NVMe IRQ handlers on the same NUMA socket.
54///
55/// # Errors
56///
57/// Returns an error if the sysfs file cannot be read or parsed.
58pub fn read_irq_affinity(irq: u32) -> Result<Vec<u32>> {
59    #[cfg(target_os = "linux")]
60    {
61        let path = format!("/proc/irq/{irq}/smp_affinity");
62        let content = std::fs::read_to_string(&path).map_err(|source| Error::System {
63            operation: "read_irq_affinity",
64            source,
65        })?;
66        // smp_affinity is a comma-separated hex mask
67        let mut cores = Vec::new();
68        for (group_idx, chunk) in content.trim().split(',').rev().enumerate() {
69            let mask = u64::from_str_radix(chunk.trim(), 16).map_err(|_| Error::System {
70                operation: "read_irq_affinity",
71                source: std::io::Error::new(
72                    std::io::ErrorKind::InvalidData,
73                    format!("malformed affinity mask chunk: {chunk}"),
74                ),
75            })?;
76            for bit in 0..64 {
77                if mask & (1 << bit) != 0 {
78                    let core = u32::try_from(group_idx * 64 + bit).map_err(|_| Error::System {
79                        operation: "read_irq_affinity",
80                        source: std::io::Error::new(
81                            std::io::ErrorKind::InvalidData,
82                            "affinity mask core index overflow",
83                        ),
84                    })?;
85                    cores.push(core);
86                }
87            }
88        }
89        Ok(cores)
90    }
91    #[cfg(not(target_os = "linux"))]
92    {
93        let _ = irq;
94        Ok(Vec::new())
95    }
96}
97
98/// Pin the calling thread to a specific NUMA node's cores.
99/// # Errors
100/// Returns an error if pinning fails.
101pub fn pin_to_numa_node(node: u32) -> Result<()> {
102    #[cfg(target_os = "linux")]
103    {
104        crate::numa::pin_to_node(node)
105    }
106
107    #[cfg(not(target_os = "linux"))]
108    {
109        let _ = node;
110        Ok(())
111    }
112}
113
114#[cfg(test)]
115mod tests {
116    use super::{pin_to_core, pin_to_numa_node};
117
118    #[test]
119    fn affinity_helpers_are_non_fatal() {
120        let _ = pin_to_core(0);
121        let _ = pin_to_numa_node(0);
122    }
123}
kernelkit/affinity.rs

kernelkit/
affinity.rs
