subms_mpsc_queue/features/
affinity.rs

1//! Thread CPU-affinity helpers.
2//!
3//! [`set_affinity`] pins the calling thread to the given set of
4//! logical CPU cores. Implementations:
5//!
6//! - **Linux:** `sched_setaffinity(0, ...)` via libc syscall.
7//! - **Windows:** `SetThreadAffinityMask(GetCurrentThread(), mask)`.
8//!   Mask is a 64-bit bitfield; cores >= 64 are rejected with
9//!   [`AffinityError::Unsupported`].
10//! - **Other (macOS, FreeBSD, *BSDs without sched_setaffinity):**
11//!   documented no-op returning [`AffinityError::Unsupported`].
12//!   macOS exposes only thread-policy affinity hints; we don't
13//!   pretend they're equivalent.
14//!
15//! The call only affects the current thread. Spawn-then-pin pattern:
16//!
17//! ```no_run
18//! # #[cfg(feature = "affinity")]
19//! use subms_mpsc_queue::set_affinity;
20//! # #[cfg(feature = "affinity")]
21//! std::thread::spawn(|| {
22//!     let _ = set_affinity(&[2]);
23//!     // ... hot loop pinned to core 2 ...
24//! });
25//! ```
26
27use std::fmt;
28
29/// Failure mode for [`set_affinity`].
30#[derive(Debug)]
31pub enum AffinityError {
32    /// The current platform does not support affinity pinning.
33    Unsupported,
34    /// A core index was out of range or invalid.
35    InvalidCore(usize),
36    /// The OS syscall returned an error code.
37    OsError(i32),
38}
39
40impl fmt::Display for AffinityError {
41    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
42        match self {
43            AffinityError::Unsupported => {
44                write!(f, "affinity pinning not supported on this platform")
45            }
46            AffinityError::InvalidCore(c) => write!(f, "invalid core index: {c}"),
47            AffinityError::OsError(e) => write!(f, "os error: {e}"),
48        }
49    }
50}
51
52impl std::error::Error for AffinityError {}
53
54/// Pin the calling thread to the given set of logical cores.
55///
56/// `cores` is the inclusion set; the OS chooses among them. Empty
57/// `cores` is rejected with `InvalidCore(0)` to avoid the
58/// platform-specific "no cores" trap (Linux interprets an empty
59/// affinity mask as "no schedulable CPU" and stalls the thread).
60pub fn set_affinity(cores: &[usize]) -> Result<(), AffinityError> {
61    if cores.is_empty() {
62        return Err(AffinityError::InvalidCore(0));
63    }
64    #[cfg(target_os = "linux")]
65    {
66        linux::set_affinity_linux(cores)
67    }
68    #[cfg(target_os = "windows")]
69    {
70        windows::set_affinity_windows(cores)
71    }
72    #[cfg(not(any(target_os = "linux", target_os = "windows")))]
73    {
74        // Reference the param so non-Linux/Windows builds don't warn.
75        let _ = cores;
76        Err(AffinityError::Unsupported)
77    }
78}
79
80#[cfg(target_os = "linux")]
81mod linux {
82    use super::AffinityError;
83    use std::mem;
84
85    // Manual libc shim - the recipe is zero-dep, so we declare just
86    // the syscall surface we need. cpu_set_t is a fixed-size 1024-bit
87    // bitmap on glibc; CPU_SET / CPU_ZERO are macros we emulate via
88    // direct bitfield writes.
89    #[repr(C)]
90    struct CpuSetT {
91        bits: [u64; 16], // 1024 bits = 128 bytes = 16 * u64
92    }
93
94    unsafe extern "C" {
95        fn sched_setaffinity(pid: u32, cpusetsize: usize, mask: *const CpuSetT) -> i32;
96    }
97
98    pub(super) fn set_affinity_linux(cores: &[usize]) -> Result<(), AffinityError> {
99        let mut set = CpuSetT { bits: [0u64; 16] };
100        for &c in cores {
101            if c >= 1024 {
102                return Err(AffinityError::InvalidCore(c));
103            }
104            set.bits[c / 64] |= 1u64 << (c % 64);
105        }
106        let rc = unsafe { sched_setaffinity(0, mem::size_of::<CpuSetT>(), &set as *const _) };
107        if rc == 0 {
108            Ok(())
109        } else {
110            Err(AffinityError::OsError(rc))
111        }
112    }
113}
114
115#[cfg(target_os = "windows")]
116mod windows {
117    use super::AffinityError;
118
119    type Handle = *mut core::ffi::c_void;
120    type DWordPtr = usize;
121
122    unsafe extern "system" {
123        fn GetCurrentThread() -> Handle;
124        fn SetThreadAffinityMask(thread: Handle, mask: DWordPtr) -> DWordPtr;
125    }
126
127    pub(super) fn set_affinity_windows(cores: &[usize]) -> Result<(), AffinityError> {
128        let mut mask: usize = 0;
129        let max_bits = usize::BITS as usize;
130        for &c in cores {
131            if c >= max_bits {
132                return Err(AffinityError::InvalidCore(c));
133            }
134            mask |= 1usize << c;
135        }
136        // SAFETY: GetCurrentThread returns a pseudo-handle; SetThreadAffinityMask is FFI-safe.
137        let prev = unsafe { SetThreadAffinityMask(GetCurrentThread(), mask) };
138        if prev == 0 {
139            Err(AffinityError::OsError(0))
140        } else {
141            Ok(())
142        }
143    }
144}
145
146#[cfg(test)]
147#[path = "affinity_tests.rs"]
148mod tests;
subms_mpsc_queue/features/affinity.rs

subms_mpsc_queue/features/
affinity.rs
