From 7c9611d124cc0bf127a0d83d75617445f70c422a Mon Sep 17 00:00:00 2001 From: Josh Triplett Date: Mon, 27 Sep 2021 13:05:44 -0700 Subject: [PATCH] Make std::thread::available_concurrency support process-limited number of CPUs Use libc::sched_getaffinity and count the number of CPUs in the returned mask. This handles cases where the process doesn't have access to all CPUs, such as when limited via taskset or similar. --- library/std/src/sys/unix/thread.rs | 8 ++++++++ 1 file changed, 8 insertions(+) diff --git a/library/std/src/sys/unix/thread.rs b/library/std/src/sys/unix/thread.rs index 6f4863057aba..b99eb2e553f0 100644 --- a/library/std/src/sys/unix/thread.rs +++ b/library/std/src/sys/unix/thread.rs @@ -275,6 +275,14 @@ pub fn available_parallelism() -> io::Result { target_os = "solaris", target_os = "illumos", ))] { + #[cfg(any(target_os = "android", target_os = "linux"))] + { + let mut set: libc::cpu_set_t = unsafe { mem::zeroed() }; + if unsafe { libc::sched_getaffinity(0, mem::size_of::(), &mut set) } == 0 { + let count = unsafe { libc::CPU_COUNT(&set) }; + return Ok(unsafe { NonZeroUsize::new_unchecked(count as usize) }); + } + } match unsafe { libc::sysconf(libc::_SC_NPROCESSORS_ONLN) } { -1 => Err(io::Error::last_os_error()), 0 => Err(io::Error::new_const(io::ErrorKind::NotFound, &"The number of hardware threads is not known for the target platform")),