implement libc::sched_getaffinity and libc::sched_setaffinity

Author: folkertdev

src/tools/miri/src/bin/miri.rs

@@ -592,6 +592,9 @@ fn main() {
             let num_cpus = param
                 .parse::<u32>()
                 .unwrap_or_else(|err| show_error!("-Zmiri-num-cpus requires a `u32`: {}", err));
+            if !(1..=miri::MAX_CPUS).contains(&usize::try_from(num_cpus).unwrap()) {
+                show_error!("-Zmiri-num-cpus must be in the range 1..={}", miri::MAX_CPUS);
+            }
             miri_config.num_cpus = num_cpus;
         } else if let Some(param) = arg.strip_prefix("-Zmiri-force-page-size=") {
             let page_size = param.parse::<u64>().unwrap_or_else(|err| {

src/tools/miri/src/concurrency/cpu_affinity.rs

@@ -0,0 +1,95 @@
+use crate::bug;
+use rustc_target::abi::Endian;
+
+/// The maximum number of CPUs supported by miri.
+///
+/// This value is compatible with the libc `CPU_SETSIZE` constant and corresponds to the number
+/// of CPUs that a `cpu_set_t` can contain.
+///
+/// Real machines can have more CPUs than this number, and there exist APIs to set their affinity,
+/// but this is not currently supported by miri.
+pub const MAX_CPUS: usize = 1024;
+
+/// A thread's CPU affinity mask determines the set of CPUs on which it is eligible to run.
+// the actual representation depends on the target's endianness and pointer width.
+// See CpuAffinityMask::set for details
+#[derive(Clone)]
+pub(crate) struct CpuAffinityMask([u8; Self::CPU_MASK_BYTES]);
+
+impl CpuAffinityMask {
+    pub(crate) const CPU_MASK_BYTES: usize = MAX_CPUS / 8;
+
+    pub fn new(target: &rustc_target::spec::Target, cpu_count: u32) -> Self {
+        let mut this = Self([0; Self::CPU_MASK_BYTES]);
+
+        // the default affinity mask includes only the available CPUs
+        for i in 0..cpu_count as usize {
+            this.set(target, i);
+        }
+
+        this
+    }
+
+    pub fn chunk_size(target: &rustc_target::spec::Target) -> u64 {
+        // The actual representation of the CpuAffinityMask is [c_ulong; _], in practice either
+        //
+        // - [u32; 32] on 32-bit platforms
+        // - [u64; 16] everywhere else
+
+        // FIXME: this should be `size_of::<core::ffi::c_ulong>()`
+        u64::from(target.pointer_width / 8)
+    }
+
+    fn set(&mut self, target: &rustc_target::spec::Target, cpu: usize) {
+        // we silently ignore CPUs that are out of bounds. This matches the behavior of
+        // `sched_setaffinity` with a mask that specifies more than `CPU_SETSIZE` CPUs.
+        if cpu >= MAX_CPUS {
+            return;
+        }
+
+        // The actual representation of the CpuAffinityMask is [c_ulong; _], in practice either
+        //
+        // - [u32; 32] on 32-bit platforms
+        // - [u64; 16] everywhere else
+        //
+        // Within the array elements, we need to use the endianness of the target.
+        match Self::chunk_size(target) {
+            4 => {
+                let start = cpu / 32 * 4; // first byte of the correct u32
+                let chunk = self.0[start..].first_chunk_mut::<4>().unwrap();
+                let offset = cpu % 32;
+                *chunk = match target.options.endian {
+                    Endian::Little => (u32::from_le_bytes(*chunk) | 1 << offset).to_le_bytes(),
+                    Endian::Big => (u32::from_be_bytes(*chunk) | 1 << offset).to_be_bytes(),
+                };
+            }
+            8 => {
+                let start = cpu / 64 * 8; // first byte of the correct u64
+                let chunk = self.0[start..].first_chunk_mut::<8>().unwrap();
+                let offset = cpu % 64;
+                *chunk = match target.options.endian {
+                    Endian::Little => (u64::from_le_bytes(*chunk) | 1 << offset).to_le_bytes(),
+                    Endian::Big => (u64::from_be_bytes(*chunk) | 1 << offset).to_be_bytes(),
+                };
+            }
+            other => bug!("other chunk sizes are not supported: {other}"),
+        };
+    }
+
+    pub fn as_slice(&self) -> &[u8] {
+        self.0.as_slice()
+    }
+
+    pub fn from_array(
+        target: &rustc_target::spec::Target,
+        cpu_count: u32,
+        bytes: [u8; Self::CPU_MASK_BYTES],
+    ) -> Option<Self> {
+        // mask by what CPUs are actually available
+        let default = Self::new(target, cpu_count);
+        let masked = std::array::from_fn(|i| bytes[i] & default.0[i]);
+
+        // at least one thread must be set for the input to be valid
+        masked.iter().any(|b| *b != 0).then_some(Self(masked))
+    }
+}

src/tools/miri/src/concurrency/mod.rs

@@ -1,3 +1,4 @@
+pub mod cpu_affinity;
 pub mod data_race;
 pub mod init_once;
 mod range_object_map;

src/tools/miri/src/concurrency/thread.rs

@@ -936,6 +936,11 @@ pub trait EvalContextExt<'tcx>: crate::MiriInterpCxExt<'tcx> {
         // After this all accesses will be treated as occurring in the new thread.
         let old_thread_id = this.machine.threads.set_active_thread_id(new_thread_id);
 
+        // The child inherits its parent's cpu affinity.
+        if let Some(cpuset) = this.machine.thread_cpu_affinity.get(&old_thread_id).cloned() {
+            this.machine.thread_cpu_affinity.insert(new_thread_id, cpuset);
+        }
+
         // Perform the function pointer load in the new thread frame.
         let instance = this.get_ptr_fn(start_routine)?.as_instance()?;
 

src/tools/miri/src/lib.rs

@@ -129,6 +129,7 @@ pub use crate::borrow_tracker::{
 };
 pub use crate::clock::{Clock, Instant};
 pub use crate::concurrency::{
+    cpu_affinity::MAX_CPUS,
     data_race::{AtomicFenceOrd, AtomicReadOrd, AtomicRwOrd, AtomicWriteOrd, EvalContextExt as _},
     init_once::{EvalContextExt as _, InitOnceId},
     sync::{CondvarId, EvalContextExt as _, MutexId, RwLockId, SynchronizationObjects},

src/tools/miri/src/machine.rs

@@ -30,6 +30,7 @@ use rustc_target::spec::abi::Abi;
 
 use crate::{
     concurrency::{
+        cpu_affinity::{self, CpuAffinityMask},
         data_race::{self, NaReadType, NaWriteType},
         weak_memory,
     },
@@ -471,6 +472,12 @@ pub struct MiriMachine<'tcx> {
 
     /// The set of threads.
     pub(crate) threads: ThreadManager<'tcx>,
+
+    /// Stores which thread is eligible to run on which CPUs.
+    /// This has no effect at all, it is just tracked to produce the correct result
+    /// in `sched_getaffinity`
+    pub(crate) thread_cpu_affinity: FxHashMap<ThreadId, CpuAffinityMask>,
+
     /// The state of the primitive synchronization objects.
     pub(crate) sync: SynchronizationObjects,
 
@@ -627,6 +634,20 @@ impl<'tcx> MiriMachine<'tcx> {
         let stack_addr = if tcx.pointer_size().bits() < 32 { page_size } else { page_size * 32 };
         let stack_size =
             if tcx.pointer_size().bits() < 32 { page_size * 4 } else { page_size * 16 };
+        assert!(
+            usize::try_from(config.num_cpus).unwrap() <= cpu_affinity::MAX_CPUS,
+            "miri only supports up to {} CPUs, but {} were configured",
+            cpu_affinity::MAX_CPUS,
+            config.num_cpus
+        );
+        let threads = ThreadManager::default();
+        let mut thread_cpu_affinity = FxHashMap::default();
+        if matches!(&*tcx.sess.target.os, "linux" | "freebsd" | "android") {
+            thread_cpu_affinity.insert(
+                threads.active_thread(),
+                CpuAffinityMask::new(&tcx.sess.target, config.num_cpus),
+            );
+        }
         MiriMachine {
             tcx,
             borrow_tracker,
@@ -644,7 +665,8 @@ impl<'tcx> MiriMachine<'tcx> {
             fds: shims::FdTable::new(config.mute_stdout_stderr),
             dirs: Default::default(),
             layouts,
-            threads: ThreadManager::default(),
+            threads,
+            thread_cpu_affinity,
             sync: SynchronizationObjects::default(),
             static_roots: Vec::new(),
             profiler,
@@ -765,6 +787,7 @@ impl VisitProvenance for MiriMachine<'_> {
         #[rustfmt::skip]
         let MiriMachine {
             threads,
+            thread_cpu_affinity: _,
             sync: _,
             tls,
             env_vars,

src/tools/miri/src/shims/unix/foreign_items.rs

@@ -3,8 +3,10 @@ use std::str;
 
 use rustc_middle::ty::layout::LayoutOf;
 use rustc_span::Symbol;
+use rustc_target::abi::Size;
 use rustc_target::spec::abi::Abi;
 
+use crate::concurrency::cpu_affinity::CpuAffinityMask;
 use crate::shims::alloc::EvalContextExt as _;
 use crate::shims::unix::*;
 use crate::*;
@@ -571,6 +573,101 @@ pub trait EvalContextExt<'tcx>: crate::MiriInterpCxExt<'tcx> {
                 let result = this.nanosleep(req, rem)?;
                 this.write_scalar(Scalar::from_i32(result), dest)?;
             }
+            "sched_getaffinity" => {
+                // Currently this function does not exist on all Unixes, e.g. on macOS.
+                if !matches!(&*this.tcx.sess.target.os, "linux" | "freebsd" | "android") {
+                    throw_unsup_format!(
+                        "`sched_getaffinity` is not supported on {}",
+                        this.tcx.sess.target.os
+                    );
+                }
+
+                let [pid, cpusetsize, mask] =
+                    this.check_shim(abi, Abi::C { unwind: false }, link_name, args)?;
+                let pid = this.read_scalar(pid)?.to_u32()?;
+                let cpusetsize = this.read_target_usize(cpusetsize)?;
+                let mask = this.read_pointer(mask)?;
+
+                // TODO: when https://github.com/rust-lang/miri/issues/3730 is fixed this should use its notion of tid/pid
+                let thread_id = match pid {
+                    0 => this.active_thread(),
+                    _ => throw_unsup_format!("`sched_getaffinity` is only supported with a pid of 0 (indicating the current thread)"),
+                };
+
+                // The actual representation of the CpuAffinityMask is [c_ulong; _], in practice either
+                //
+                // - [u32; 32] on 32-bit platforms
+                // - [u64; 16] everywhere else
+                let chunk_size = CpuAffinityMask::chunk_size(&this.tcx.sess.target);
+
+                if this.ptr_is_null(mask)? {
+                    let einval = this.eval_libc("EFAULT");
+                    this.set_last_error(einval)?;
+                    this.write_scalar(Scalar::from_i32(-1), dest)?;
+                } else if cpusetsize == 0 || cpusetsize.checked_rem(chunk_size).unwrap() != 0 {
+                    // we only copy whole chunks of size_of::<c_ulong>()
+                    let einval = this.eval_libc("EINVAL");
+                    this.set_last_error(einval)?;
+                    this.write_scalar(Scalar::from_i32(-1), dest)?;
+                } else if let Some(cpuset) = this.machine.thread_cpu_affinity.get(&thread_id) {
+                    let cpuset = cpuset.clone();
+                    // we only copy whole chunks of size_of::<c_ulong>()
+                    let byte_count = Ord::min(cpuset.as_slice().len(), cpusetsize.try_into().unwrap());
+                    this.write_bytes_ptr(mask, cpuset.as_slice()[..byte_count].iter().copied())?;
+                    this.write_scalar(Scalar::from_i32(0), dest)?;
+                } else {
+                    // The thread whose ID is pid could not be found
+                    let einval = this.eval_libc("ESRCH");
+                    this.set_last_error(einval)?;
+                    this.write_scalar(Scalar::from_i32(-1), dest)?;
+                }
+            }
+            "sched_setaffinity" => {
+                // Currently this function does not exist on all Unixes, e.g. on macOS.
+                if !matches!(&*this.tcx.sess.target.os, "linux" | "freebsd" | "android") {
+                    throw_unsup_format!(
+                        "`sched_setaffinity` is not supported on {}",
+                        this.tcx.sess.target.os
+                    );
+                }
+
+                let [pid, cpusetsize, mask] =
+                    this.check_shim(abi, Abi::C { unwind: false }, link_name, args)?;
+                let pid = this.read_scalar(pid)?.to_u32()?;
+                let cpusetsize = this.read_target_usize(cpusetsize)?;
+                let mask = this.read_pointer(mask)?;
+
+                // TODO: when https://github.com/rust-lang/miri/issues/3730 is fixed this should use its notion of tid/pid
+                let thread_id = match pid {
+                    0 => this.active_thread(),
+                    _ => throw_unsup_format!("`sched_setaffinity` is only supported with a pid of 0 (indicating the current thread)"),
+                };
+
+                #[allow(clippy::map_entry)]
+                if this.ptr_is_null(mask)? {
+                    let einval = this.eval_libc("EFAULT");
+                    this.set_last_error(einval)?;
+                    this.write_scalar(Scalar::from_i32(-1), dest)?;
+                } else {
+                    // NOTE: cpusetsize might be smaller than `CpuAffinityMask::CPU_MASK_BYTES`
+                    let bits_slice = this.read_bytes_ptr_strip_provenance(mask, Size::from_bytes(cpusetsize))?;
+                    // This ignores the bytes beyond `CpuAffinityMask::CPU_MASK_BYTES`
+                    let bits_array: [u8;CpuAffinityMask::CPU_MASK_BYTES] =
+                        std::array::from_fn(|i| bits_slice.get(i).copied().unwrap_or(0));
+                    match CpuAffinityMask::from_array(&this.tcx.sess.target, this.machine.num_cpus, bits_array) {
+                        Some(cpuset) => {
+                            this.machine.thread_cpu_affinity.insert(thread_id, cpuset);
+                            this.write_scalar(Scalar::from_i32(0), dest)?;
+                        }
+                        None => {
+                            // The intersection between the mask and the available CPUs was empty.
+                            let einval = this.eval_libc("EINVAL");
+                            this.set_last_error(einval)?;
+                            this.write_scalar(Scalar::from_i32(-1), dest)?;
+                        }
+                    }
+                }
+            }
 
             // Miscellaneous
             "isatty" => {

src/tools/miri/src/shims/unix/linux/foreign_items.rs

@@ -178,19 +178,6 @@ pub trait EvalContextExt<'tcx>: crate::MiriInterpCxExt<'tcx> {
 
                 this.write_scalar(Scalar::from_i32(SIGRTMAX), dest)?;
             }
-            "sched_getaffinity" => {
-                // This shim isn't useful, aside from the fact that it makes `num_cpus`
-                // fall back to `sysconf` where it will successfully determine the number of CPUs.
-                let [pid, cpusetsize, mask] =
-                    this.check_shim(abi, Abi::C { unwind: false }, link_name, args)?;
-                this.read_scalar(pid)?.to_i32()?;
-                this.read_target_usize(cpusetsize)?;
-                this.deref_pointer_as(mask, this.libc_ty_layout("cpu_set_t"))?;
-                // FIXME: we just return an error.
-                let einval = this.eval_libc("EINVAL");
-                this.set_last_error(einval)?;
-                this.write_scalar(Scalar::from_i32(-1), dest)?;
-            }
 
             // Incomplete shims that we "stub out" just to get pre-main initialization code to work.
             // These shims are enabled only when the caller is in the standard library.

src/tools/miri/tests/fail-dep/libc/affinity.rs

@@ -0,0 +1,17 @@
+//@ignore-target-windows: only very limited libc on Windows
+//@ignore-target-apple: `sched_setaffinity` is not supported on macOS
+//@compile-flags: -Zmiri-disable-isolation -Zmiri-num-cpus=4
+
+fn main() {
+    use libc::{cpu_set_t, sched_setaffinity};
+
+    use std::mem::size_of;
+
+    // If pid is zero, then the calling thread is used.
+    const PID: i32 = 0;
+
+    let cpuset: cpu_set_t = unsafe { core::mem::MaybeUninit::zeroed().assume_init() };
+
+    let err = unsafe { sched_setaffinity(PID, size_of::<cpu_set_t>() + 1, &cpuset) }; //~ ERROR: memory access failed
+    assert_eq!(err, 0);
+}

src/tools/miri/tests/fail-dep/libc/affinity.stderr

@@ -0,0 +1,20 @@
+error: Undefined Behavior: memory access failed: ALLOC has size 128, so pointer to 129 bytes starting at offset 0 is out-of-bounds
+  --> $DIR/affinity.rs:LL:CC
+   |
+LL |     let err = unsafe { sched_setaffinity(PID, size_of::<cpu_set_t>() + 1, &cpuset) };
+   |                        ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ memory access failed: ALLOC has size 128, so pointer to 129 bytes starting at offset 0 is out-of-bounds
+   |
+   = help: this indicates a bug in the program: it performed an invalid operation, and caused Undefined Behavior
+   = help: see https://doc.rust-lang.org/nightly/reference/behavior-considered-undefined.html for further information
+help: ALLOC was allocated here:
+  --> $DIR/affinity.rs:LL:CC
+   |
+LL |     let cpuset: cpu_set_t = unsafe { core::mem::MaybeUninit::zeroed().assume_init() };
+   |         ^^^^^^
+   = note: BACKTRACE (of the first span):
+   = note: inside `main` at $DIR/affinity.rs:LL:CC
+
+note: some details are omitted, run with `MIRIFLAGS=-Zmiri-backtrace=full` for a verbose backtrace
+
+error: aborting due to 1 previous error
+