| //! This crate manages CPU affinities. |
| //! |
| //! ## Example |
| //! |
| //! This example shows how to create a thread for each available processor and pin each thread to its corresponding processor. |
| //! |
| //! ``` |
| //! extern crate core_affinity; |
| //! |
| //! use std::thread; |
| //! |
| //! // Retrieve the IDs of all active CPU cores. |
| //! let core_ids = core_affinity::get_core_ids().unwrap(); |
| //! |
| //! // Create a thread for each active CPU core. |
| //! let handles = core_ids.into_iter().map(|id| { |
| //! thread::spawn(move || { |
| //! // Pin this thread to a single CPU core. |
| //! let res = core_affinity::set_for_current(id); |
| //! if (res) { |
| //! // Do more work after this. |
| //! } |
| //! }) |
| //! }).collect::<Vec<_>>(); |
| //! |
| //! for handle in handles.into_iter() { |
| //! handle.join().unwrap(); |
| //! } |
| //! ``` |
| |
| #[cfg(any( |
| target_os = "android", |
| target_os = "linux", |
| target_os = "macos", |
| target_os = "freebsd" |
| ))] |
| extern crate libc; |
| |
| #[cfg_attr(all(not(test), not(target_os = "macos")), allow(unused_extern_crates))] |
| extern crate num_cpus; |
| |
| /// This function tries to retrieve information |
| /// on all the "cores" on which the current thread |
| /// is allowed to run. |
| pub fn get_core_ids() -> Option<Vec<CoreId>> { |
| get_core_ids_helper() |
| } |
| |
| /// This function tries to pin the current |
| /// thread to the specified core. |
| /// |
| /// # Arguments |
| /// |
| /// * core_id - ID of the core to pin |
| pub fn set_for_current(core_id: CoreId) -> bool { |
| set_for_current_helper(core_id) |
| } |
| |
| /// This represents a CPU core. |
| #[derive(Copy, Clone, Debug, PartialEq, Eq)] |
| pub struct CoreId { |
| pub id: usize, |
| } |
| |
| // Linux Section |
| |
| #[cfg(any(target_os = "android", target_os = "linux"))] |
| #[inline] |
| fn get_core_ids_helper() -> Option<Vec<CoreId>> { |
| linux::get_core_ids() |
| } |
| |
| #[cfg(any(target_os = "android", target_os = "linux"))] |
| #[inline] |
| fn set_for_current_helper(core_id: CoreId) -> bool { |
| linux::set_for_current(core_id) |
| } |
| |
| #[cfg(any(target_os = "android", target_os = "linux"))] |
| mod linux { |
| use std::mem; |
| |
| use libc::{CPU_ISSET, CPU_SET, CPU_SETSIZE, cpu_set_t, sched_getaffinity, sched_setaffinity}; |
| |
| use super::CoreId; |
| |
| pub fn get_core_ids() -> Option<Vec<CoreId>> { |
| if let Some(full_set) = get_affinity_mask() { |
| let mut core_ids: Vec<CoreId> = Vec::new(); |
| |
| for i in 0..CPU_SETSIZE as usize { |
| if unsafe { CPU_ISSET(i, &full_set) } { |
| core_ids.push(CoreId{ id: i }); |
| } |
| } |
| |
| Some(core_ids) |
| } |
| else { |
| None |
| } |
| } |
| |
| pub fn set_for_current(core_id: CoreId) -> bool { |
| // Turn `core_id` into a `libc::cpu_set_t` with only |
| // one core active. |
| let mut set = new_cpu_set(); |
| |
| unsafe { CPU_SET(core_id.id, &mut set) }; |
| |
| // Set the current thread's core affinity. |
| let res = unsafe { |
| sched_setaffinity(0, // Defaults to current thread |
| mem::size_of::<cpu_set_t>(), |
| &set) |
| }; |
| res == 0 |
| } |
| |
| fn get_affinity_mask() -> Option<cpu_set_t> { |
| let mut set = new_cpu_set(); |
| |
| // Try to get current core affinity mask. |
| let result = unsafe { |
| sched_getaffinity(0, // Defaults to current thread |
| mem::size_of::<cpu_set_t>(), |
| &mut set) |
| }; |
| |
| if result == 0 { |
| Some(set) |
| } |
| else { |
| None |
| } |
| } |
| |
| fn new_cpu_set() -> cpu_set_t { |
| unsafe { mem::zeroed::<cpu_set_t>() } |
| } |
| |
| #[cfg(test)] |
| mod tests { |
| use num_cpus; |
| |
| use super::*; |
| |
| #[test] |
| fn test_linux_get_affinity_mask() { |
| match get_affinity_mask() { |
| Some(_) => {}, |
| None => { assert!(false); }, |
| } |
| } |
| |
| #[test] |
| fn test_linux_get_core_ids() { |
| match get_core_ids() { |
| Some(set) => { |
| assert_eq!(set.len(), num_cpus::get()); |
| }, |
| None => { assert!(false); }, |
| } |
| } |
| |
| #[test] |
| fn test_linux_set_for_current() { |
| let ids = get_core_ids().unwrap(); |
| |
| assert!(ids.len() > 0); |
| |
| let res = set_for_current(ids[0]); |
| assert_eq!(res, true); |
| |
| // Ensure that the system pinned the current thread |
| // to the specified core. |
| let mut core_mask = new_cpu_set(); |
| unsafe { CPU_SET(ids[0].id, &mut core_mask) }; |
| |
| let new_mask = get_affinity_mask().unwrap(); |
| |
| let mut is_equal = true; |
| |
| for i in 0..CPU_SETSIZE as usize { |
| let is_set1 = unsafe { |
| CPU_ISSET(i, &core_mask) |
| }; |
| let is_set2 = unsafe { |
| CPU_ISSET(i, &new_mask) |
| }; |
| |
| if is_set1 != is_set2 { |
| is_equal = false; |
| } |
| } |
| |
| assert!(is_equal); |
| } |
| } |
| } |
| |
| // Windows Section |
| |
| #[cfg(target_os = "windows")] |
| #[inline] |
| fn get_core_ids_helper() -> Option<Vec<CoreId>> { |
| windows::get_core_ids() |
| } |
| |
| #[cfg(target_os = "windows")] |
| #[inline] |
| fn set_for_current_helper(core_id: CoreId) -> bool { |
| windows::set_for_current(core_id) |
| } |
| |
| #[cfg(target_os = "windows")] |
| extern crate winapi; |
| |
| #[cfg(target_os = "windows")] |
| mod windows { |
| use winapi::shared::basetsd::{DWORD_PTR, PDWORD_PTR}; |
| use winapi::um::processthreadsapi::{GetCurrentProcess, GetCurrentThread}; |
| use winapi::um::winbase::{GetProcessAffinityMask, SetThreadAffinityMask}; |
| |
| use super::CoreId; |
| |
| pub fn get_core_ids() -> Option<Vec<CoreId>> { |
| if let Some(mask) = get_affinity_mask() { |
| // Find all active cores in the bitmask. |
| let mut core_ids: Vec<CoreId> = Vec::new(); |
| |
| for i in 0..64 as u64 { |
| let test_mask = 1 << i; |
| |
| if (mask & test_mask) == test_mask { |
| core_ids.push(CoreId { id: i as usize }); |
| } |
| } |
| |
| Some(core_ids) |
| } |
| else { |
| None |
| } |
| } |
| |
| pub fn set_for_current(core_id: CoreId) -> bool { |
| // Convert `CoreId` back into mask. |
| let mask: u64 = 1 << core_id.id; |
| |
| // Set core affinity for current thread. |
| let res = unsafe { |
| SetThreadAffinityMask( |
| GetCurrentThread(), |
| mask as DWORD_PTR |
| ) |
| }; |
| res != 0 |
| } |
| |
| fn get_affinity_mask() -> Option<u64> { |
| let mut system_mask: usize = 0; |
| let mut process_mask: usize = 0; |
| |
| let res = unsafe { |
| GetProcessAffinityMask( |
| GetCurrentProcess(), |
| &mut process_mask as PDWORD_PTR, |
| &mut system_mask as PDWORD_PTR |
| ) |
| }; |
| |
| // Successfully retrieved affinity mask |
| if res != 0 { |
| Some(process_mask as u64) |
| } |
| // Failed to retrieve affinity mask |
| else { |
| None |
| } |
| } |
| |
| #[cfg(test)] |
| mod tests { |
| use num_cpus; |
| |
| use super::*; |
| |
| #[test] |
| fn test_windows_get_core_ids() { |
| match get_core_ids() { |
| Some(set) => { |
| assert_eq!(set.len(), num_cpus::get()); |
| }, |
| None => { assert!(false); }, |
| } |
| } |
| |
| #[test] |
| fn test_windows_set_for_current() { |
| let ids = get_core_ids().unwrap(); |
| |
| assert!(ids.len() > 0); |
| |
| assert_ne!(set_for_current(ids[0]), 0); |
| } |
| } |
| } |
| |
| // MacOS Section |
| |
| #[cfg(target_os = "macos")] |
| #[inline] |
| fn get_core_ids_helper() -> Option<Vec<CoreId>> { |
| macos::get_core_ids() |
| } |
| |
| #[cfg(target_os = "macos")] |
| #[inline] |
| fn set_for_current_helper(core_id: CoreId) -> bool { |
| macos::set_for_current(core_id) |
| } |
| |
| #[cfg(target_os = "macos")] |
| mod macos { |
| use std::mem; |
| |
| use libc::{c_int, c_uint, c_void, pthread_self}; |
| |
| use num_cpus; |
| |
| use super::CoreId; |
| |
| type kern_return_t = c_int; |
| type integer_t = c_int; |
| type natural_t = c_uint; |
| type thread_t = c_uint; |
| type thread_policy_flavor_t = natural_t; |
| type mach_msg_type_number_t = natural_t; |
| |
| #[repr(C)] |
| struct thread_affinity_policy_data_t { |
| affinity_tag: integer_t, |
| } |
| |
| type thread_policy_t = *mut thread_affinity_policy_data_t; |
| |
| const THREAD_AFFINITY_POLICY: thread_policy_flavor_t = 4; |
| |
| extern { |
| fn thread_policy_set( |
| thread: thread_t, |
| flavor: thread_policy_flavor_t, |
| policy_info: thread_policy_t, |
| count: mach_msg_type_number_t, |
| ) -> kern_return_t; |
| } |
| |
| pub fn get_core_ids() -> Option<Vec<CoreId>> { |
| Some((0..(num_cpus::get())).into_iter() |
| .map(|n| CoreId { id: n as usize }) |
| .collect::<Vec<_>>()) |
| } |
| |
| pub fn set_for_current(core_id: CoreId) -> bool { |
| let THREAD_AFFINITY_POLICY_COUNT: mach_msg_type_number_t = |
| mem::size_of::<thread_affinity_policy_data_t>() as mach_msg_type_number_t / |
| mem::size_of::<integer_t>() as mach_msg_type_number_t; |
| |
| let mut info = thread_affinity_policy_data_t { |
| affinity_tag: core_id.id as integer_t, |
| }; |
| |
| let res = unsafe { |
| thread_policy_set( |
| pthread_self() as thread_t, |
| THREAD_AFFINITY_POLICY, |
| &mut info as thread_policy_t, |
| THREAD_AFFINITY_POLICY_COUNT |
| ) |
| }; |
| res == 0 |
| } |
| |
| #[cfg(test)] |
| mod tests { |
| use num_cpus; |
| |
| use super::*; |
| |
| #[test] |
| fn test_macos_get_core_ids() { |
| match get_core_ids() { |
| Some(set) => { |
| assert_eq!(set.len(), num_cpus::get()); |
| }, |
| None => { assert!(false); }, |
| } |
| } |
| |
| #[test] |
| fn test_macos_set_for_current() { |
| let ids = get_core_ids().unwrap(); |
| |
| assert!(ids.len() > 0); |
| |
| set_for_current(ids[0]); |
| } |
| } |
| } |
| |
| |
| // FreeBSD Section |
| |
| #[cfg(target_os = "freebsd")] |
| #[inline] |
| fn get_core_ids_helper() -> Option<Vec<CoreId>> { |
| freebsd::get_core_ids() |
| } |
| |
| #[cfg(target_os = "freebsd")] |
| #[inline] |
| fn set_for_current_helper(core_id: CoreId) -> bool { |
| freebsd::set_for_current(core_id) |
| } |
| |
| #[cfg(target_os = "freebsd")] |
| mod freebsd { |
| use std::mem; |
| |
| use libc::{ |
| cpuset_getaffinity, cpuset_setaffinity, cpuset_t, CPU_ISSET, |
| CPU_LEVEL_WHICH, CPU_SET, CPU_SETSIZE, CPU_WHICH_TID, |
| }; |
| |
| use super::CoreId; |
| |
| pub fn get_core_ids() -> Option<Vec<CoreId>> { |
| if let Some(full_set) = get_affinity_mask() { |
| let mut core_ids: Vec<CoreId> = Vec::new(); |
| |
| for i in 0..CPU_SETSIZE as usize { |
| if unsafe { CPU_ISSET(i, &full_set) } { |
| core_ids.push(CoreId { id: i }); |
| } |
| } |
| |
| Some(core_ids) |
| } else { |
| None |
| } |
| } |
| |
| pub fn set_for_current(core_id: CoreId) -> bool { |
| // Turn `core_id` into a `libc::cpuset_t` with only |
| // one core active. |
| let mut set = new_cpu_set(); |
| |
| unsafe { CPU_SET(core_id.id, &mut set) }; |
| |
| // Set the current thread's core affinity. |
| let res = unsafe { |
| // FreeBSD's sched_setaffinity currently operates on process id, |
| // therefore using cpuset_setaffinity instead. |
| cpuset_setaffinity( |
| CPU_LEVEL_WHICH, |
| CPU_WHICH_TID, |
| -1, // -1 == current thread |
| mem::size_of::<cpuset_t>(), |
| &set, |
| ) |
| }; |
| res == 0 |
| } |
| |
| fn get_affinity_mask() -> Option<cpuset_t> { |
| let mut set = new_cpu_set(); |
| |
| // Try to get current core affinity mask. |
| let result = unsafe { |
| // FreeBSD's sched_getaffinity currently operates on process id, |
| // therefore using cpuset_getaffinity instead. |
| cpuset_getaffinity( |
| CPU_LEVEL_WHICH, |
| CPU_WHICH_TID, |
| -1, // -1 == current thread |
| mem::size_of::<cpuset_t>(), |
| &mut set, |
| ) |
| }; |
| |
| if result == 0 { |
| Some(set) |
| } else { |
| None |
| } |
| } |
| |
| fn new_cpu_set() -> cpuset_t { |
| unsafe { mem::zeroed::<cpuset_t>() } |
| } |
| |
| #[cfg(test)] |
| mod tests { |
| use num_cpus; |
| |
| use super::*; |
| |
| #[test] |
| fn test_freebsd_get_affinity_mask() { |
| match get_affinity_mask() { |
| Some(_) => {} |
| None => { |
| assert!(false); |
| } |
| } |
| } |
| |
| #[test] |
| fn test_freebsd_get_core_ids() { |
| match get_core_ids() { |
| Some(set) => { |
| assert_eq!(set.len(), num_cpus::get()); |
| } |
| None => { |
| assert!(false); |
| } |
| } |
| } |
| |
| #[test] |
| fn test_freebsd_set_for_current() { |
| let ids = get_core_ids().unwrap(); |
| |
| assert!(ids.len() > 0); |
| |
| let res = set_for_current(ids[0]); |
| assert_eq!(res, true); |
| |
| // Ensure that the system pinned the current thread |
| // to the specified core. |
| let mut core_mask = new_cpu_set(); |
| unsafe { CPU_SET(ids[0].id, &mut core_mask) }; |
| |
| let new_mask = get_affinity_mask().unwrap(); |
| |
| let mut is_equal = true; |
| |
| for i in 0..CPU_SETSIZE as usize { |
| let is_set1 = unsafe { CPU_ISSET(i, &core_mask) }; |
| let is_set2 = unsafe { CPU_ISSET(i, &new_mask) }; |
| |
| if is_set1 != is_set2 { |
| is_equal = false; |
| } |
| } |
| |
| assert!(is_equal); |
| } |
| } |
| } |
| |
| // Stub Section |
| |
| #[cfg(not(any( |
| target_os = "linux", |
| target_os = "android", |
| target_os = "windows", |
| target_os = "macos", |
| target_os = "freebsd" |
| )))] |
| #[inline] |
| fn get_core_ids_helper() -> Option<Vec<CoreId>> { |
| None |
| } |
| |
| #[cfg(not(any( |
| target_os = "linux", |
| target_os = "android", |
| target_os = "windows", |
| target_os = "macos", |
| target_os = "freebsd" |
| )))] |
| #[inline] |
| fn set_for_current_helper(_core_id: CoreId) -> bool { |
| false |
| } |
| |
| #[cfg(test)] |
| mod tests { |
| use num_cpus; |
| |
| use super::*; |
| |
| // #[test] |
| // fn test_num_cpus() { |
| // println!("Num CPUs: {}", num_cpus::get()); |
| // println!("Num Physical CPUs: {}", num_cpus::get_physical()); |
| // } |
| |
| #[test] |
| fn test_get_core_ids() { |
| match get_core_ids() { |
| Some(set) => { |
| assert_eq!(set.len(), num_cpus::get()); |
| }, |
| None => { assert!(false); }, |
| } |
| } |
| |
| #[test] |
| fn test_set_for_current() { |
| let ids = get_core_ids().unwrap(); |
| |
| assert!(ids.len() > 0); |
| |
| set_for_current(ids[0]); |
| } |
| } |