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
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167

use std::ffi::c_void;

use crate::Result;

extern "system" {
    // Present in all windows processes (kernel32.dll)
    fn GetLastError() -> u32;
    fn GetCurrentProcess() -> *mut c_void;
    fn GetCurrentThread() -> *mut c_void;

    fn SetProcessAffinityMask(process_handle: *mut c_void, mask: usize) -> u32;
    fn GetProcessAffinityMask(
        process_handle: *mut c_void,
        cur_mask: *mut usize,
        system_mask: *mut usize,
    ) -> u32;

    fn SetThreadAffinityMask(thread_handle: *mut c_void, mask: usize) -> usize;
}
const ERROR_INVALID_PARAMETER: u32 = 0x57;

/// Binds the current __process__ to the specified core(s)
///
/// Note : This has a side effect of binding new child processes to the same cores
pub fn set_process_affinity<B: AsRef<[usize]>>(core_ids: B) -> Result<()> {
    let cur_proc = unsafe { GetCurrentProcess() };
    let mut wanted_mask = 0usize;

    // Create the bitmask
    for core_id in core_ids.as_ref() {
        wanted_mask |= 1usize << core_id;
    }

    //println!("Binding process to cores {:b}", wanted_mask);
    if let Err(last_error) = set_process_affinity_mask(cur_proc, wanted_mask) {
        return Err(From::from(format!(
            "SetProcessAffinityMask failed with error 0x{:x}",
            last_error
        )));
    }

    Ok(())
}
/// Returns a list of cores that the current __process__ is bound to
pub fn get_process_affinity() -> Result<Vec<usize>> {
    let mut affinity = Vec::new();
    let mut cur_core = 0usize;
    let mut cur_mask = 1usize;
    let cur_proc = unsafe { GetCurrentProcess() };

    match get_process_affinity_mask(cur_proc) {
        Ok(mask) => {
            while cur_mask != 0 {
                if cur_mask & mask != 0 {
                    affinity.push(cur_core);
                }
                cur_core += 1;
                cur_mask <<= 1;
            }
        }
        Err(last_error) => {
            return Err(From::from(format!(
                "GetProcessAffinityMask failed with error 0x{:x}",
                last_error
            )));
        }
    };

    Ok(affinity)
}

pub fn set_thread_affinity(core_ids: &[usize]) -> Result<()> {
    let cur_thread = unsafe { GetCurrentThread() };
    let mut wanted_mask = 0usize;

    // Create the bitmask
    for core_id in core_ids {
        wanted_mask |= 1usize << core_id;
    }

    //println!("Binding thread to cores {:b}", wanted_mask);
    if let Err(last_error) = set_thread_affinity_mask(cur_thread, wanted_mask) {
        return Err(From::from(format!(
            "SetThreadAffinityMask failed with error 0x{:x}",
            last_error
        )));
    }

    Ok(())
}
pub fn get_thread_affinity() -> Result<Vec<usize>> {
    let mut cur_mask = 1usize;
    let cur_thread = unsafe { GetCurrentThread() };
    let mut cur_core = 1usize;
    let mut affinity = Vec::new();

    // There is not GetThreadAffinityMask on Windows so we must use
    // the return value of SetThreadAffinityMask()
    while cur_mask != 0 {
        cur_core += 1;
        match set_thread_affinity_mask(cur_thread, cur_mask) {
            Ok(mask) => {
                // Restore the old affinity
                if let Err(e) = set_thread_affinity_mask(cur_thread, mask) {
                    return Err(From::from(format!(
                        "SetThreadAffinityMask failed with error 0x{:x}",
                        e
                    )));
                }

                // Count the set bits in the mask
                cur_core = 0;
                cur_mask = 1;
                while cur_mask != 0 {
                    if cur_mask & mask != 0 {
                        affinity.push(cur_core);
                    }
                    cur_core += 1;
                    cur_mask <<= 1;
                }

                break;
            }
            Err(last_error) => {
                // ERROR_INVALID_PARAMETER can happen if we try to set the thread affinity
                // to a core that our global process disabled (Through SetProcessAffinityMask)
                if last_error != ERROR_INVALID_PARAMETER {
                    return Err(From::from(format!(
                        "SetThreadAffinityMask failed with error 0x{:x}",
                        last_error
                    )));
                }

                // Try the next core
                cur_mask <<= 1;
                continue;
            }
        };
    }

    Ok(affinity)
}

/* Wrappers around unsafe OS calls */
fn set_thread_affinity_mask(thread: *mut c_void, mask: usize) -> std::result::Result<usize, u32> {
    let res = unsafe { SetThreadAffinityMask(thread, mask) };
    if res == 0 {
        return Err(unsafe { GetLastError() });
    }
    Ok(res)
}
fn set_process_affinity_mask(process: *mut c_void, mask: usize) -> std::result::Result<(), u32> {
    let res = unsafe { SetProcessAffinityMask(process, mask) };
    if res == 0 {
        return Err(unsafe { GetLastError() });
    }
    Ok(())
}
fn get_process_affinity_mask(process: *mut c_void) -> std::result::Result<usize, u32> {
    let mut cur_mask = 0usize;
    let mut system_mask = 0usize;
    let res = unsafe { GetProcessAffinityMask(process, &mut cur_mask, &mut system_mask) };
    if res == 0 {
        return Err(unsafe { GetLastError() });
    }
    Ok(cur_mask)
}