affinity 0.1.2

Provides a consistent way to set core affinity for currently running threads and processes
Documentation 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65

use crate::Result;
use ::libc::*;
use std::mem::{size_of, zeroed};

pub fn set_thread_affinity(core_ids: &[usize]) -> Result<()> {
    let mut set: cpu_set_t = unsafe { zeroed() };
    unsafe {
        for core_id in core_ids {
            CPU_SET(*core_id, &mut set);
        }
    }

    if let Err(e) = _sched_setaffinity(0, size_of::<cpu_set_t>(), &set) {
        return Err(From::from(format!(
            "sched_setaffinity failed with errno {}",
            e
        )));
    }
    Ok(())
}

pub fn get_thread_affinity() -> Result<Vec<usize>> {
    let mut affinity = Vec::new();
    let mut set: cpu_set_t = unsafe { zeroed() };

    if let Err(e) = _sched_getaffinity(0, size_of::<cpu_set_t>(), &mut set) {
        return Err(From::from(format!(
            "sched_getaffinity failed with errno {}",
            e
        )));
    }

    for i in 0..CPU_SETSIZE as usize {
        if unsafe { CPU_ISSET(i, &set) } {
            affinity.push(i);
        }
    }

    Ok(affinity)
}

/* Wrappers around unsafe OS calls */
fn _sched_setaffinity(
    pid: pid_t,
    cpusetsize: usize,
    mask: &cpu_set_t,
) -> std::result::Result<(), i32> {
    let res = unsafe { sched_setaffinity(pid, cpusetsize, mask) };
    if res != 0 {
        return Err(errno::errno().into());
    }
    Ok(())
}

fn _sched_getaffinity(
    pid: pid_t,
    cpusetsize: usize,
    mask: &mut cpu_set_t,
) -> std::result::Result<(), i32> {
    let res = unsafe { sched_getaffinity(pid, cpusetsize, mask) };
    if res != 0 {
        return Err(errno::errno().into());
    }
    Ok(())
}