bind_threads/
bind_threads.rs

1extern crate hwloc;
2extern crate libc;
3#[cfg(target_os = "windows")]
4extern crate kernel32;
5#[cfg(target_os = "windows")]
6extern crate winapi;
7
8use hwloc::{Topology, ObjectType, CPUBIND_THREAD, CpuSet};
9use std::thread;
10use std::sync::{Arc, Mutex};
11
12/// Example which spawns one thread per core and then assigns it to each.
13///
14/// Example Output with 2 cores (no HT) on linux:
15///
16/// ```
17/// Found 2 cores.
18/// Thread 0: Before Some(0-1), After Some(0)
19/// Thread 1: Before Some(0-1), After Some(1)
20/// ```
21fn main() {
22    let topo = Arc::new(Mutex::new(Topology::new()));
23
24    // Grab the number of cores in a block so that the lock is removed once
25    // the end of the block is reached.
26    let num_cores = {
27        let topo_rc = topo.clone();
28        let topo_locked = topo_rc.lock().unwrap();
29        (*topo_locked)
30            .objects_with_type(&ObjectType::Core)
31            .unwrap()
32            .len()
33    };
34    println!("Found {} cores.", num_cores);
35
36    // Spawn one thread for each and pass the topology down into scope.
37    let handles: Vec<_> = (0..num_cores)
38        .map(|i| {
39            let child_topo = topo.clone();
40            thread::spawn(move || {
41                // Get the current thread id and lock the topology to use.
42                let tid = get_thread_id();
43                let mut locked_topo = child_topo.lock().unwrap();
44
45                // Thread binding before explicit set.
46                let before = locked_topo.get_cpubind_for_thread(tid, CPUBIND_THREAD);
47
48                // load the cpuset for the given core index.
49                let mut bind_to = cpuset_for_core(&*locked_topo, i);
50
51                // Get only one logical processor (in case the core is SMT/hyper-threaded).
52                bind_to.singlify();
53
54                // Set the binding.
55                locked_topo
56                    .set_cpubind_for_thread(tid, bind_to, CPUBIND_THREAD)
57                    .unwrap();
58
59                // Thread binding after explicit set.
60                let after = locked_topo.get_cpubind_for_thread(tid, CPUBIND_THREAD);
61                println!("Thread {}: Before {:?}, After {:?}", i, before, after);
62            })
63        })
64        .collect();
65
66    // Wait for all threads to complete before ending the program.
67    for h in handles {
68        h.join().unwrap();
69    }
70}
71
72/// Load the `CpuSet` for the given core index.
73fn cpuset_for_core(topology: &Topology, idx: usize) -> CpuSet {
74    let cores = (*topology).objects_with_type(&ObjectType::Core).unwrap();
75    match cores.get(idx) {
76        Some(val) => val.cpuset().unwrap(),
77        None => panic!("No Core found with id {}", idx),
78    }
79}
80
81/// Helper method to get the thread id through libc, with current rust stable (1.5.0) its not
82/// possible otherwise I think.
83#[cfg(not(target_os = "windows"))]
84fn get_thread_id() -> libc::pthread_t {
85    unsafe { libc::pthread_self() }
86}
87
88#[cfg(target_os = "windows")]
89fn get_thread_id() -> winapi::winnt::HANDLE {
90    unsafe { kernel32::GetCurrentThread() }
91}