From 433439e4eb5db9edd07c46f4cdceff64ea21edcd 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 133ad3ea420b8..cff4ba47357c7 100644 --- a/library/std/src/sys/unix/thread.rs +++ b/library/std/src/sys/unix/thread.rs @@ -275,6 +275,14 @@ pub fn available_concurrency() -> 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")),