Auto merge of #89310 - joshtriplett:available-concurrency-affinity, r=m-ou-se

Make `std:🧵: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.

This also covers cgroup cpusets.
This commit is contained in:
bors 2021-11-07 11:53:25 +00:00
commit fecfc0e6cc
1 changed files with 8 additions and 0 deletions

View File

@ -275,6 +275,14 @@ pub fn available_parallelism() -> io::Result<NonZeroUsize> {
target_os = "solaris", target_os = "solaris",
target_os = "illumos", 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) } { match unsafe { libc::sysconf(libc::_SC_NPROCESSORS_ONLN) } {
-1 => Err(io::Error::last_os_error()), -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")), 0 => Err(io::Error::new_const(io::ErrorKind::NotFound, &"The number of hardware threads is not known for the target platform")),