tenferro_cpu/
affinity.rs

1use std::num::NonZeroUsize;
2
3/// Return a best-effort CPU count available to the current process.
4///
5/// This first tries an OS-standard process-affinity query when supported, then
6/// falls back to `std::thread::available_parallelism()`, and finally to `1`.
7///
8/// # Examples
9///
10/// ```
11/// let available = tenferro_cpu::available_parallelism();
12/// assert!(available >= 1);
13/// ```
14pub fn available_parallelism() -> usize {
15    process_cpu_affinity_count()
16        .or_else(standard_available_parallelism)
17        .unwrap_or(1)
18}
19
20/// Return the current process affinity mask size when the platform exposes a
21/// standard affinity API.
22///
23/// Platforms without an affinity query return `None`.
24///
25/// # Examples
26///
27/// ```
28/// let count = tenferro_cpu::process_cpu_affinity_count();
29/// if let Some(count) = count {
30///     assert!(count >= 1);
31/// }
32/// ```
33pub fn process_cpu_affinity_count() -> Option<usize> {
34    platform_process_cpu_affinity_count()
35}
36
37pub(crate) fn standard_available_parallelism() -> Option<usize> {
38    std::thread::available_parallelism()
39        .ok()
40        .map(NonZeroUsize::get)
41}
42
43#[cfg(any(target_os = "linux", target_os = "android", test))]
44fn count_affinity_mask_bits(mask: &[u8]) -> Option<usize> {
45    let count = mask.iter().map(|byte| byte.count_ones() as usize).sum();
46    (count > 0).then_some(count)
47}
48
49#[cfg(any(target_os = "linux", target_os = "android"))]
50const LINUX_EINVAL: i32 = 22;
51
52#[cfg(any(target_os = "linux", target_os = "android"))]
53fn linux_next_affinity_mask_bytes(mask_bytes: usize, errno: Option<i32>) -> Option<usize> {
54    (errno == Some(LINUX_EINVAL))
55        .then(|| mask_bytes.checked_mul(2))
56        .flatten()
57}
58
59#[cfg(any(target_os = "linux", target_os = "android"))]
60fn platform_process_cpu_affinity_count() -> Option<usize> {
61    unsafe extern "C" {
62        fn sched_getaffinity(pid: i32, cpusetsize: usize, mask: *mut core::ffi::c_void) -> i32;
63    }
64
65    const INITIAL_MASK_BYTES: usize = 128;
66
67    let mut mask_bytes = INITIAL_MASK_BYTES;
68    loop {
69        let mut mask = vec![0u8; mask_bytes];
70        let rc = unsafe {
71            sched_getaffinity(0, mask_bytes, mask.as_mut_ptr().cast::<core::ffi::c_void>())
72        };
73        if rc == 0 {
74            return count_affinity_mask_bits(&mask);
75        }
76
77        mask_bytes = linux_next_affinity_mask_bytes(
78            mask_bytes,
79            std::io::Error::last_os_error().raw_os_error(),
80        )?;
81    }
82}
83
84#[cfg(target_os = "windows")]
85fn platform_process_cpu_affinity_count() -> Option<usize> {
86    type Handle = *mut core::ffi::c_void;
87    type DwordPtr = usize;
88    type Word = u16;
89
90    unsafe extern "system" {
91        fn GetCurrentProcess() -> Handle;
92        fn GetProcessAffinityMask(
93            process: Handle,
94            process_affinity_mask: *mut DwordPtr,
95            system_affinity_mask: *mut DwordPtr,
96        ) -> i32;
97        fn GetActiveProcessorGroupCount() -> Word;
98        fn GetActiveProcessorCount(group_number: Word) -> u32;
99        fn GetProcessGroupAffinity(
100            process: Handle,
101            group_count: *mut Word,
102            group_array: *mut Word,
103        ) -> i32;
104    }
105
106    let process = unsafe { GetCurrentProcess() };
107    let system_group_count = unsafe { GetActiveProcessorGroupCount() };
108
109    if system_group_count <= 1 {
110        let mut process_mask = 0usize;
111        let mut system_mask = 0usize;
112        let ok = unsafe {
113            GetProcessAffinityMask(
114                process,
115                std::ptr::addr_of_mut!(process_mask),
116                std::ptr::addr_of_mut!(system_mask),
117            )
118        };
119        if ok != 0 {
120            let count = process_mask.count_ones() as usize;
121            return (count > 0).then_some(count);
122        }
123        let count = unsafe { GetActiveProcessorCount(0) } as usize;
124        return (count > 0).then_some(count);
125    }
126
127    let mut group_count: Word = 0;
128    let ok = unsafe {
129        GetProcessGroupAffinity(
130            process,
131            std::ptr::addr_of_mut!(group_count),
132            std::ptr::null_mut(),
133        )
134    };
135    if ok != 0 || group_count == 0 {
136        let count = unsafe { GetActiveProcessorCount(u16::MAX) } as usize;
137        return (count > 0).then_some(count);
138    }
139
140    let mut groups = vec![0u16; group_count as usize];
141    let ok = unsafe {
142        GetProcessGroupAffinity(
143            process,
144            std::ptr::addr_of_mut!(group_count),
145            groups.as_mut_ptr(),
146        )
147    };
148    if ok == 0 || group_count == 0 {
149        let count = unsafe { GetActiveProcessorCount(u16::MAX) } as usize;
150        return (count > 0).then_some(count);
151    }
152
153    if group_count == 1 {
154        let mut process_mask = 0usize;
155        let mut system_mask = 0usize;
156        let ok = unsafe {
157            GetProcessAffinityMask(
158                process,
159                std::ptr::addr_of_mut!(process_mask),
160                std::ptr::addr_of_mut!(system_mask),
161            )
162        };
163        if ok != 0 {
164            let count = process_mask.count_ones() as usize;
165            return (count > 0).then_some(count);
166        }
167    }
168
169    let count = groups
170        .into_iter()
171        .map(|group| unsafe { GetActiveProcessorCount(group) } as usize)
172        .sum();
173    (count > 0).then_some(count)
174}
175
176#[cfg(not(any(target_os = "linux", target_os = "android", target_os = "windows")))]
177fn platform_process_cpu_affinity_count() -> Option<usize> {
178    None
179}
180
181#[cfg(test)]
182mod tests;
tenferro_cpu/affinity.rs

tenferro_cpu/
affinity.rs
