Skip to content

Commit

Permalink
Make std::thread::available_concurrency support process-limited numbe…
Browse files Browse the repository at this point in the history
…r 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.
  • Loading branch information
joshtriplett committed Oct 4, 2021
1 parent 2b6ed3b commit 433439e
Showing 1 changed file with 8 additions and 0 deletions.
8 changes: 8 additions & 0 deletions library/std/src/sys/unix/thread.rs
Original file line number Diff line number Diff line change
Expand Up @@ -275,6 +275,14 @@ pub fn available_concurrency() -> 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")),
Expand Down

0 comments on commit 433439e

Please sign in to comment.