From 3717646366a91ee9ca0b574973ec090670f9a4f8 Mon Sep 17 00:00:00 2001 From: Yoshua Wuyts Date: Sat, 18 Jul 2020 17:11:46 +0200 Subject: [PATCH] Add std::thread::available_concurrency --- .../std/src/thread/available_concurrency.rs | 157 ++++++++++++++++++ library/std/src/thread/mod.rs | 6 + library/test/src/helpers/concurrency.rs | 103 +----------- library/test/src/lib.rs | 1 + 4 files changed, 166 insertions(+), 101 deletions(-) create mode 100644 library/std/src/thread/available_concurrency.rs diff --git a/library/std/src/thread/available_concurrency.rs b/library/std/src/thread/available_concurrency.rs new file mode 100644 index 00000000000..4e805e4f599 --- /dev/null +++ b/library/std/src/thread/available_concurrency.rs @@ -0,0 +1,157 @@ +use crate::io; +use crate::num::NonZeroUsize; + +/// Returns the number of hardware threads available to the program. +/// +/// This value should be considered only a hint. +/// +/// # Platform-specific behavior +/// +/// If interpreted as the number of actual hardware threads, it may undercount on +/// Windows systems with more than 64 hardware threads. If interpreted as the +/// available concurrency for that process, it may overcount on Windows systems +/// when limited by a process wide affinity mask or job object limitations, and +/// it may overcount on Linux systems when limited by a process wide affinity +/// mask or affected by cgroups limits. +/// +/// # Errors +/// +/// This function will return an error in the following situations, but is not +/// limited to just these cases: +/// +/// - If the number of hardware threads is not known for the target platform. +/// - The process lacks permissions to view the number of hardware threads +/// available. +/// +/// # Examples +/// +/// ``` +/// # #![allow(dead_code)] +/// #![feature(available_concurrency)] +/// use std::thread; +/// +/// let count = thread::available_concurrency().map(|n| n.get()).unwrap_or(1); +/// ``` +#[unstable(feature = "available_concurrency", issue = "74479")] +pub fn available_concurrency() -> io::Result { + available_concurrency_internal() +} + +cfg_if::cfg_if! { + if #[cfg(windows)] { + #[allow(nonstandard_style)] + fn available_concurrency_internal() -> io::Result { + #[repr(C)] + struct SYSTEM_INFO { + wProcessorArchitecture: u16, + wReserved: u16, + dwPageSize: u32, + lpMinimumApplicationAddress: *mut u8, + lpMaximumApplicationAddress: *mut u8, + dwActiveProcessorMask: *mut u8, + dwNumberOfProcessors: u32, + dwProcessorType: u32, + dwAllocationGranularity: u32, + wProcessorLevel: u16, + wProcessorRevision: u16, + } + extern "system" { + fn GetSystemInfo(info: *mut SYSTEM_INFO) -> i32; + } + let res = unsafe { + let mut sysinfo = crate::mem::zeroed(); + GetSystemInfo(&mut sysinfo); + sysinfo.dwNumberOfProcessors as usize + }; + match res { + 0 => Err(io::Error::new(io::ErrorKind::NotFound, "The number of hardware threads is not known for the target platform")), + cpus => Ok(unsafe { NonZeroUsize::new_unchecked(cpus) }), + } + } + } else if #[cfg(any( + target_os = "android", + target_os = "cloudabi", + target_os = "emscripten", + target_os = "fuchsia", + target_os = "ios", + target_os = "linux", + target_os = "macos", + target_os = "solaris", + target_os = "illumos", + ))] { + fn available_concurrency_internal() -> io::Result { + match unsafe { libc::sysconf(libc::_SC_NPROCESSORS_ONLN) } { + -1 => Err(io::Error::last_os_error()), + 0 => Err(io::Error::new(io::ErrorKind::NotFound, "The number of hardware threads is not known for the target platform")), + cpus => Ok(unsafe { NonZeroUsize::new_unchecked(cpus as usize) }), + } + } + } else if #[cfg(any(target_os = "freebsd", target_os = "dragonfly", target_os = "netbsd"))] { + fn available_concurrency_internal() -> io::Result { + use crate::ptr; + + let mut cpus: libc::c_uint = 0; + let mut cpus_size = crate::mem::size_of_val(&cpus); + + unsafe { + cpus = libc::sysconf(libc::_SC_NPROCESSORS_ONLN) as libc::c_uint; + } + + // Fallback approach in case of errors or no hardware threads. + if cpus < 1 { + let mut mib = [libc::CTL_HW, libc::HW_NCPU, 0, 0]; + let res = unsafe { + libc::sysctl( + mib.as_mut_ptr(), + 2, + &mut cpus as *mut _ as *mut _, + &mut cpus_size as *mut _ as *mut _, + ptr::null_mut(), + 0, + ) + }; + + // Handle errors if any. + if res == -1 { + return Err(io::Error::last_os_error()); + } else if cpus == 0 { + return Err(io::Error::new(io::ErrorKind::NotFound, "The number of hardware threads is not known for the target platform")); + } + } + Ok(unsafe { NonZeroUsize::new_unchecked(cpus as usize) }) + } + } else if #[cfg(target_os = "openbsd")] { + fn available_concurrency_internal() -> io::Result { + use crate::ptr; + + let mut cpus: libc::c_uint = 0; + let mut cpus_size = crate::mem::size_of_val(&cpus); + let mut mib = [libc::CTL_HW, libc::HW_NCPU, 0, 0]; + + let res = unsafe { + libc::sysctl( + mib.as_mut_ptr(), + 2, + &mut cpus as *mut _ as *mut _, + &mut cpus_size as *mut _ as *mut _, + ptr::null_mut(), + 0, + ) + }; + + // Handle errors if any. + if res == -1 { + return Err(io::Error::last_os_error()); + } else if cpus == 0 { + return Err(io::Error::new(io::ErrorKind::NotFound, "The number of hardware threads is not known for the target platform")); + } + + Ok(unsafe { NonZeroUsize::new_unchecked(cpus as usize) }) + } + } else { + // FIXME: implement on vxWorks, Redox, HermitCore, Haiku, l4re + fn available_concurrency_internal() -> io::Result { + Err(io::Error::new(io::ErrorKind::NotFound, "The number of hardware threads is not known for the target platform")) + } + } +} diff --git a/library/std/src/thread/mod.rs b/library/std/src/thread/mod.rs index 087175bb92a..45c10266ba2 100644 --- a/library/std/src/thread/mod.rs +++ b/library/std/src/thread/mod.rs @@ -175,9 +175,15 @@ #[macro_use] mod local; +#[unstable(feature = "available_concurrency", issue = "74479")] +mod available_concurrency; + #[stable(feature = "rust1", since = "1.0.0")] pub use self::local::{AccessError, LocalKey}; +#[unstable(feature = "available_concurrency", issue = "74479")] +pub use available_concurrency::available_concurrency; + // The types used by the thread_local! macro to access TLS keys. Note that there // are two types, the "OS" type and the "fast" type. The OS thread local key // type is accessed via platform-specific API calls and is slow, while the fast diff --git a/library/test/src/helpers/concurrency.rs b/library/test/src/helpers/concurrency.rs index 7ca27bf0dc1..7e9bd50f556 100644 --- a/library/test/src/helpers/concurrency.rs +++ b/library/test/src/helpers/concurrency.rs @@ -1,6 +1,7 @@ //! Helper module which helps to determine amount of threads to be used //! during tests execution. use std::env; +use std::thread; #[allow(deprecated)] pub fn get_concurrency() -> usize { @@ -12,106 +13,6 @@ pub fn get_concurrency() -> usize { _ => panic!("RUST_TEST_THREADS is `{}`, should be a positive integer.", s), } } - Err(..) => num_cpus(), - } -} - -cfg_if::cfg_if! { - if #[cfg(windows)] { - #[allow(nonstandard_style)] - fn num_cpus() -> usize { - #[repr(C)] - struct SYSTEM_INFO { - wProcessorArchitecture: u16, - wReserved: u16, - dwPageSize: u32, - lpMinimumApplicationAddress: *mut u8, - lpMaximumApplicationAddress: *mut u8, - dwActiveProcessorMask: *mut u8, - dwNumberOfProcessors: u32, - dwProcessorType: u32, - dwAllocationGranularity: u32, - wProcessorLevel: u16, - wProcessorRevision: u16, - } - extern "system" { - fn GetSystemInfo(info: *mut SYSTEM_INFO) -> i32; - } - unsafe { - let mut sysinfo = std::mem::zeroed(); - GetSystemInfo(&mut sysinfo); - sysinfo.dwNumberOfProcessors as usize - } - } - } else if #[cfg(any( - target_os = "android", - target_os = "cloudabi", - target_os = "emscripten", - target_os = "fuchsia", - target_os = "ios", - target_os = "linux", - target_os = "macos", - target_os = "solaris", - target_os = "illumos", - ))] { - fn num_cpus() -> usize { - unsafe { libc::sysconf(libc::_SC_NPROCESSORS_ONLN) as usize } - } - } else if #[cfg(any(target_os = "freebsd", target_os = "dragonfly", target_os = "netbsd"))] { - fn num_cpus() -> usize { - use std::ptr; - - let mut cpus: libc::c_uint = 0; - let mut cpus_size = std::mem::size_of_val(&cpus); - - unsafe { - cpus = libc::sysconf(libc::_SC_NPROCESSORS_ONLN) as libc::c_uint; - } - if cpus < 1 { - let mut mib = [libc::CTL_HW, libc::HW_NCPU, 0, 0]; - unsafe { - libc::sysctl( - mib.as_mut_ptr(), - 2, - &mut cpus as *mut _ as *mut _, - &mut cpus_size as *mut _ as *mut _, - ptr::null_mut(), - 0, - ); - } - if cpus < 1 { - cpus = 1; - } - } - cpus as usize - } - } else if #[cfg(target_os = "openbsd")] { - fn num_cpus() -> usize { - use std::ptr; - - let mut cpus: libc::c_uint = 0; - let mut cpus_size = std::mem::size_of_val(&cpus); - let mut mib = [libc::CTL_HW, libc::HW_NCPU, 0, 0]; - - unsafe { - libc::sysctl( - mib.as_mut_ptr(), - 2, - &mut cpus as *mut _ as *mut _, - &mut cpus_size as *mut _ as *mut _, - ptr::null_mut(), - 0, - ); - } - if cpus < 1 { - cpus = 1; - } - cpus as usize - } - } else { - // FIXME: implement on vxWorks, Redox, HermitCore, Haiku, l4re - fn num_cpus() -> usize { - 1 - } + Err(..) => thread::available_concurrency().map(|n| n.get()).unwrap_or(1), } } diff --git a/library/test/src/lib.rs b/library/test/src/lib.rs index b0b81f85fe0..9c5bb8957b5 100644 --- a/library/test/src/lib.rs +++ b/library/test/src/lib.rs @@ -24,6 +24,7 @@ #![feature(rustc_private)] #![feature(nll)] #![feature(bool_to_option)] +#![feature(available_concurrency)] #![feature(set_stdio)] #![feature(panic_unwind)] #![feature(staged_api)]