1use crate::process::Pid;
2use crate::{backend, io};
3
4/// `CpuSet` represents a bit-mask of CPUs.
5///
6/// `CpuSet`s are used by [`sched_setaffinity`] and [`sched_getaffinity`], for
7/// example.
8///
9/// # References
10///  - [Linux]
11///
12/// [Linux]: https://man7.org/linux/man-pages/man3/CPU_SET.3.html
13/// [`sched_setaffinity`]: crate::process::sched_setaffinity
14/// [`sched_getaffinity`]: crate::process::sched_getaffinity
15#[repr(C)]
16#[derive(Clone, Copy, Debug, Eq, Hash, PartialEq)]
17pub struct CpuSet {
18    cpu_set: backend::process::types::RawCpuSet,
19}
20
21impl CpuSet {
22    /// The maximum number of CPU in `CpuSet`.
23    pub const MAX_CPU: usize = backend::process::types::CPU_SETSIZE;
24
25    /// Create a new and empty `CpuSet`.
26    #[inline]
27    pub fn new() -> Self {
28        Self {
29            cpu_set: backend::process::types::raw_cpu_set_new(),
30        }
31    }
32
33    /// Test to see if a CPU is in the `CpuSet`.
34    ///
35    /// `field` is the CPU id to test.
36    #[inline]
37    pub fn is_set(&self, field: usize) -> bool {
38        backend::process::cpu_set::CPU_ISSET(field, &self.cpu_set)
39    }
40
41    /// Add a CPU to `CpuSet`.
42    ///
43    /// `field` is the CPU id to add.
44    #[inline]
45    pub fn set(&mut self, field: usize) {
46        backend::process::cpu_set::CPU_SET(field, &mut self.cpu_set)
47    }
48
49    /// Remove a CPU from `CpuSet`.
50    ///
51    /// `field` is the CPU id to remove.
52    #[inline]
53    pub fn unset(&mut self, field: usize) {
54        backend::process::cpu_set::CPU_CLR(field, &mut self.cpu_set)
55    }
56
57    /// Count the number of CPUs set in the `CpuSet`.
58    #[cfg(any(target_os = "android", target_os = "linux"))]
59    #[inline]
60    pub fn count(&self) -> u32 {
61        backend::process::cpu_set::CPU_COUNT(&self.cpu_set)
62    }
63
64    /// Zeroes the `CpuSet`.
65    #[inline]
66    pub fn clear(&mut self) {
67        backend::process::cpu_set::CPU_ZERO(&mut self.cpu_set)
68    }
69}
70
71impl Default for CpuSet {
72    #[inline]
73    fn default() -> Self {
74        Self::new()
75    }
76}
77
78/// `sched_setaffinity(pid, cpuset)`—Set a thread's CPU affinity mask.
79///
80/// `pid` is the thread ID to update. If pid is `None`, then the current thread
81/// is updated.
82///
83/// The `CpuSet` argument specifies the set of CPUs on which the thread will
84/// be eligible to run.
85///
86/// # References
87///  - [Linux]
88///
89/// [Linux]: https://man7.org/linux/man-pages/man2/sched_setaffinity.2.html
90#[inline]
91pub fn sched_setaffinity(pid: Option<Pid>, cpuset: &CpuSet) -> io::Result<()> {
92    backend::process::syscalls::sched_setaffinity(pid, &cpuset.cpu_set)
93}
94
95/// `sched_getaffinity(pid)`—Get a thread's CPU affinity mask.
96///
97/// `pid` is the thread ID to check. If pid is `None`, then the current thread
98/// is checked.
99///
100/// Returns the set of CPUs on which the thread is eligible to run.
101///
102/// # References
103///  - [Linux]
104///
105/// [Linux]: https://man7.org/linux/man-pages/man2/sched_getaffinity.2.html
106#[inline]
107pub fn sched_getaffinity(pid: Option<Pid>) -> io::Result<CpuSet> {
108    let mut cpuset = CpuSet::new();
109    backend::process::syscalls::sched_getaffinity(pid, &mut cpuset.cpu_set).and(Ok(cpuset))
110}
111