miri/shims/
cpu_affinity.rs

1use rustc_abi::{Endian, Size};
2use rustc_middle::ty::layout::LayoutOf;
3use rustc_target::spec::Os;
4
5use crate::*;
6
7/// The maximum number of CPUs supported by miri.
8///
9/// This value is compatible with the libc `CPU_SETSIZE` constant and corresponds to the number
10/// of CPUs that a `cpu_set_t` can contain.
11///
12/// Real machines can have more CPUs than this number, and there exist APIs to set their affinity,
13/// but this is not currently supported by miri.
14pub const MAX_CPUS: usize = 1024;
15
16/// A thread's CPU affinity mask determines the set of CPUs on which it is eligible to run.
17// the actual representation depends on the target's endianness and pointer width.
18// See CpuAffinityMask::set for details
19#[derive(Clone)]
20pub struct CpuAffinityMask([u8; Self::CPU_MASK_BYTES]);
21
22impl CpuAffinityMask {
23    pub(crate) const CPU_MASK_BYTES: usize = MAX_CPUS / 8;
24
25    pub fn new<'tcx>(cx: &impl LayoutOf<'tcx>, cpu_count: u32) -> Self {
26        let mut this = Self([0; Self::CPU_MASK_BYTES]);
27
28        // the default affinity mask includes only the available CPUs
29        for i in 0..cpu_count.to_usize() {
30            this.set(cx, i);
31        }
32
33        this
34    }
35
36    pub fn chunk_size<'tcx>(cx: &impl LayoutOf<'tcx>) -> u64 {
37        // The actual representation of the CpuAffinityMask is [c_ulong; _].
38        let ulong = helpers::path_ty_layout(cx, &["core", "ffi", "c_ulong"]);
39        ulong.size.bytes()
40    }
41
42    fn set<'tcx>(&mut self, cx: &impl LayoutOf<'tcx>, cpu: usize) {
43        // we silently ignore CPUs that are out of bounds. This matches the behavior of
44        // `sched_setaffinity` with a mask that specifies more than `CPU_SETSIZE` CPUs.
45        if cpu >= MAX_CPUS {
46            return;
47        }
48
49        // The actual representation of the CpuAffinityMask is [c_ulong; _].
50        // Within the array elements, we need to use the endianness of the target.
51        let target = &cx.tcx().sess.target;
52        #[expect(clippy::arithmetic_side_effects)] // we checked above that `cpu` is small enough
53        match Self::chunk_size(cx) {
54            4 => {
55                let start = cpu / 32 * 4; // first byte of the correct u32
56                let chunk = self.0[start..].first_chunk_mut::<4>().unwrap();
57                let offset = cpu % 32;
58                *chunk = match target.options.endian {
59                    Endian::Little => (u32::from_le_bytes(*chunk) | (1 << offset)).to_le_bytes(),
60                    Endian::Big => (u32::from_be_bytes(*chunk) | (1 << offset)).to_be_bytes(),
61                };
62            }
63            8 => {
64                let start = cpu / 64 * 8; // first byte of the correct u64
65                let chunk = self.0[start..].first_chunk_mut::<8>().unwrap();
66                let offset = cpu % 64;
67                *chunk = match target.options.endian {
68                    Endian::Little => (u64::from_le_bytes(*chunk) | (1 << offset)).to_le_bytes(),
69                    Endian::Big => (u64::from_be_bytes(*chunk) | (1 << offset)).to_be_bytes(),
70                };
71            }
72            other => bug!("chunk size not supported: {other}"),
73        };
74    }
75
76    pub fn as_slice(&self) -> &[u8] {
77        self.0.as_slice()
78    }
79
80    pub fn from_array<'tcx>(
81        cx: &impl LayoutOf<'tcx>,
82        cpu_count: u32,
83        bytes: [u8; Self::CPU_MASK_BYTES],
84    ) -> Option<Self> {
85        // mask by what CPUs are actually available
86        let default = Self::new(cx, cpu_count);
87        let masked = std::array::from_fn(|i| bytes[i] & default.0[i]);
88
89        // at least one thread must be set for the input to be valid
90        masked.iter().any(|b| *b != 0).then_some(Self(masked))
91    }
92}
93
94impl<'tcx> EvalContextExt<'tcx> for crate::MiriInterpCx<'tcx> {}
95pub trait EvalContextExt<'tcx>: crate::MiriInterpCxExt<'tcx> {
96    fn sched_getaffinity(
97        &mut self,
98        pid: &OpTy<'tcx>,
99        cpusetsize: &OpTy<'tcx>,
100        mask: &OpTy<'tcx>,
101        dest: &MPlaceTy<'tcx>,
102    ) -> InterpResult<'tcx> {
103        let this = self.eval_context_mut();
104
105        let pid = this.read_scalar(pid)?.to_u32()?;
106        let cpusetsize = this.read_target_usize(cpusetsize)?;
107        let mask = this.read_pointer(mask)?;
108
109        if this.machine.thread_cpu_affinity.is_none() {
110            throw_unsup_format!("`sched_getaffinity` is not supported on #![no_core] programs")
111        }
112
113        let thread_id = if pid == 0 {
114            this.active_thread()
115        } else if matches!(this.tcx.sess.target.os, Os::Linux | Os::Android) {
116            // On Linux/Android, pid can be a TID as returned by `gettid`.
117            let Some(thread_id) = this.get_thread_id_from_linux_tid(pid) else {
118                this.set_errno_and_return_neg1(LibcError("ESRCH"), dest)?;
119                return interp_ok(());
120            };
121            thread_id
122        } else {
123            throw_unsup_format!(
124                "`sched_getaffinity` is only supported with a pid of 0 (indicating the current thread) on non-Linux platforms"
125            )
126        };
127
128        // The mask is stored in chunks, and the size must be a whole number of chunks.
129        let chunk_size = CpuAffinityMask::chunk_size(this);
130
131        if this.ptr_is_null(mask)? {
132            this.set_errno_and_return_neg1(LibcError("EFAULT"), dest)?;
133        } else if cpusetsize == 0 || cpusetsize.checked_rem(chunk_size).unwrap() != 0 {
134            // we only copy whole chunks of size_of::<c_ulong>()
135            this.set_errno_and_return_neg1(LibcError("EINVAL"), dest)?;
136        } else if let Some(cpuset) =
137            this.machine.thread_cpu_affinity.as_ref().unwrap().get(&thread_id)
138        {
139            let cpuset = cpuset.clone();
140            // we only copy whole chunks of size_of::<c_ulong>()
141            let byte_count = Ord::min(cpuset.as_slice().len(), cpusetsize.try_into().unwrap());
142            this.write_bytes_ptr(mask, cpuset.as_slice()[..byte_count].iter().copied())?;
143            this.write_null(dest)?;
144        } else {
145            unreachable!("we validated the thread ID above");
146        }
147
148        interp_ok(())
149    }
150
151    fn sched_setaffinity(
152        &mut self,
153        pid: &OpTy<'tcx>,
154        cpusetsize: &OpTy<'tcx>,
155        mask: &OpTy<'tcx>,
156        dest: &MPlaceTy<'tcx>,
157    ) -> InterpResult<'tcx> {
158        let this = self.eval_context_mut();
159
160        let pid = this.read_scalar(pid)?.to_u32()?;
161        let cpusetsize = this.read_target_usize(cpusetsize)?;
162        let mask = this.read_pointer(mask)?;
163
164        if this.machine.thread_cpu_affinity.is_none() {
165            throw_unsup_format!("`sched_setaffinity` is not supported on #![no_core] programs")
166        }
167
168        let thread_id = if pid == 0 {
169            this.active_thread()
170        } else if matches!(this.tcx.sess.target.os, Os::Linux | Os::Android) {
171            // On Linux/Android, pid can be a TID as returned by `gettid`.
172            let Some(thread_id) = this.get_thread_id_from_linux_tid(pid) else {
173                this.set_errno_and_return_neg1(LibcError("ESRCH"), dest)?;
174                return interp_ok(());
175            };
176            thread_id
177        } else {
178            throw_unsup_format!(
179                "`sched_setaffinity` is only supported with a pid of 0 (indicating the current thread) on non-Linux platforms"
180            )
181        };
182
183        if this.ptr_is_null(mask)? {
184            this.set_errno_and_return_neg1(LibcError("EFAULT"), dest)?;
185        } else {
186            // NOTE: cpusetsize might be smaller than `CpuAffinityMask::CPU_MASK_BYTES`.
187            // Any unspecified bytes are treated as zero here (none of the CPUs are configured).
188            // This is not exactly documented, so we assume that this is the behavior in practice.
189            let bits_slice =
190                this.read_bytes_ptr_strip_provenance(mask, Size::from_bytes(cpusetsize))?;
191            // This ignores the bytes beyond `CpuAffinityMask::CPU_MASK_BYTES`
192            let bits_array: [u8; CpuAffinityMask::CPU_MASK_BYTES] =
193                std::array::from_fn(|i| bits_slice.get(i).copied().unwrap_or(0));
194            match CpuAffinityMask::from_array(this, this.machine.num_cpus, bits_array) {
195                Some(cpuset) => {
196                    this.machine.thread_cpu_affinity.as_mut().unwrap().insert(thread_id, cpuset);
197                    this.write_null(dest)?;
198                }
199                None => {
200                    // The intersection between the mask and the available CPUs was empty.
201                    this.set_errno_and_return_neg1(LibcError("EINVAL"), dest)?;
202                }
203            }
204        }
205
206        interp_ok(())
207    }
208}
miri/shims/cpu_affinity.rs

miri/shims/
cpu_affinity.rs
