bigoish 0.1.1

Test the computational complexity (big-O) of Rust algorithms
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
66
67
68
69
70
71
72
73
74
75

//! APIs for pinning a thread to a CPU core, currently Linux only.

use std::io::{Error as IoError, Result as IoResult};
use std::mem::{size_of, zeroed};

use libc::{CPU_SET, cpu_set_t, gettid, pid_t, sched_getaffinity, sched_getcpu, sched_setaffinity};

/// On drop, revert to the old thread scheduler affinity.
pub(crate) struct CoreAffinityGuard {
    thread: pid_t,
    old_affinity: cpu_set_t,
}

impl Drop for CoreAffinityGuard {
    fn drop(&mut self) {
        unsafe { sched_setaffinity(self.thread, size_of::<cpu_set_t>(), &self.old_affinity) };
    }
}

/// Get the CPU affinity of the current thread.
fn get_affinity() -> IoResult<cpu_set_t> {
    let thread = unsafe { gettid() };
    unsafe {
        let mut set: cpu_set_t = zeroed();
        let code = sched_getaffinity(thread, size_of::<cpu_set_t>(), &mut set);
        if code != 0 {
            return Err(std::io::Error::from_raw_os_error(code));
        }
        Ok(set)
    }
}

/// Pin the thread to the current core, until the returned guard is dropped.
#[must_use]
pub(crate) fn fix_core_affinity() -> IoResult<CoreAffinityGuard> {
    let thread = unsafe { gettid() };
    let old_affinity = get_affinity()?;
    unsafe {
        let mut new_set: cpu_set_t = zeroed();
        let cpu = sched_getcpu();
        CPU_SET(cpu as usize, &mut new_set);
        let code = sched_setaffinity(thread, size_of::<cpu_set_t>(), &new_set);
        if code != 0 {
            return Err(IoError::from_raw_os_error(code));
        }
    };
    Ok(CoreAffinityGuard {
        thread,
        old_affinity,
    })
}

#[cfg(test)]
mod tests {
    use std::{collections::HashSet, time::Duration};

    use super::*;

    #[test]
    fn core_pinning() {
        let affinity = get_affinity();
        let _guard = fix_core_affinity();
        let mut set = HashSet::new();
        for _ in 0..20 {
            std::thread::sleep(Duration::from_millis(10));
            set.insert(unsafe { libc::sched_getcpu() });
        }
        assert_eq!(set.len(), 1);
        drop(_guard);
        assert_eq!(
            format!("{:?}", affinity.unwrap()),
            format!("{:?}", get_affinity().unwrap())
        );
    }
}