affinity_linux/
lib.rs

1#[cfg(target_os = "linux")]
2use libc::{cpu_set_t, sched_getaffinity, sched_setaffinity, CPU_ISSET, CPU_SET, CPU_SETSIZE};
3use std::{
4    collections::HashSet,
5    io,
6    mem::{size_of, zeroed},
7};
8
9type Result<T> = std::result::Result<T, Box<dyn std::error::Error>>;
10
11#[cfg(target_os = "linux")]
12pub fn get_thread_affinity() -> Result<Option<HashSet<usize>>> {
13    let mut affinity = HashSet::new();
14    let mut set: cpu_set_t = unsafe { zeroed() };
15
16    let res = unsafe { sched_getaffinity(0, size_of::<cpu_set_t>(), &mut set) };
17    if res != 0 {
18        return Err(From::from(format!(
19            "sched_getaffinity failed with: {}",
20            io::Error::from(errno::errno())
21        )));
22    }
23
24    for i in 0..CPU_SETSIZE as usize {
25        if unsafe { CPU_ISSET(i, &set) } {
26            affinity.insert(i);
27        }
28    }
29
30    Ok(Some(affinity))
31}
32
33#[cfg(not(target_os = "linux"))]
34pub fn get_thread_affinity() -> Result<Option<HashSet<usize>>> {
35    None
36}
37
38#[cfg(target_os = "linux")]
39pub fn set_thread_affinity(core_ids: impl Iterator<Item = usize>) -> Result<()> {
40    let mut set: cpu_set_t = unsafe { zeroed() };
41    unsafe {
42        for core_id in core_ids {
43            CPU_SET(core_id, &mut set);
44        }
45    }
46
47    let res = unsafe { sched_setaffinity(0, size_of::<cpu_set_t>(), &set) };
48    if res != 0 {
49        return Err(From::from(format!(
50            "sched_setaffinity failed with: {}",
51            io::Error::from(errno::errno())
52        )));
53    }
54
55    Ok(())
56}
57
58#[cfg(not(target_os = "linux"))]
59pub fn set_thread_affinity(core_ids: impl Iterator<Item = usize>) -> Result<()> {
60    None
61}
62
63#[cfg(all(test, target_os = "linux"))]
64mod tests {
65    use super::{get_thread_affinity, set_thread_affinity};
66
67    #[test]
68    fn get() {
69        let cpus = get_thread_affinity();
70        assert!(matches!(cpus, Ok(Some(_))), "produce set");
71        let set = cpus.unwrap().unwrap();
72        assert!(!set.is_empty(), "not empty");
73    }
74
75    #[test]
76    fn set() {
77        assert!(set_thread_affinity([0].into_iter()).is_ok());
78    }
79}
affinity_linux/lib.rs

affinity_linux/
lib.rs
