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.

diff --git a/library/std/src/sys/unix/thread.rs b/library/std/src/sys/unix/thread.rs
index 6f4863057aba431b0635eb11a4671e335e573913..b99eb2e553f08908effdf7385961767f75b30c88 100644 (file)
--- 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<NonZeroUsize> {
             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::<libc::cpu_set_t>(), &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")),