1 use crate::process::Pid;
2 use 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)]
17 pub struct CpuSet {
18     cpu_set: backend::process::types::RawCpuSet,
19 }
20 
21 impl 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]
new() -> Self27     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]
is_set(&self, field: usize) -> bool37     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]
set(&mut self, field: usize)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]
unset(&mut self, field: usize)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(linux_kernel)]
59     #[inline]
count(&self) -> u3260     pub fn count(&self) -> u32 {
61         backend::process::cpu_set::CPU_COUNT(&self.cpu_set)
62     }
63 
64     /// Zeroes the `CpuSet`.
65     #[inline]
clear(&mut self)66     pub fn clear(&mut self) {
67         backend::process::cpu_set::CPU_ZERO(&mut self.cpu_set)
68     }
69 }
70 
71 impl Default for CpuSet {
72     #[inline]
default() -> Self73     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 be
84 /// eligible to run.
85 ///
86 /// # References
87 ///  - [Linux]
88 ///
89 /// [Linux]: https://man7.org/linux/man-pages/man2/sched_setaffinity.2.html
90 #[inline]
sched_setaffinity(pid: Option<Pid>, cpuset: &CpuSet) -> io::Result<()>91 pub 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]
sched_getaffinity(pid: Option<Pid>) -> io::Result<CpuSet>107 pub 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 
112 /// `sched_getcpu()`—Get the CPU that the current thread is currently on.
113 ///
114 /// # References
115 ///  - [Linux]
116 ///  - [DragonFly BSD]
117 ///
118 /// [Linux]: https://man7.org/linux/man-pages/man2/sched_getcpu.2.html
119 /// [DragonFly BSD]: https://man.dragonflybsd.org/?command=sched_getcpu&section=2
120 // FreeBSD added `sched_getcpu` in 13.0.
121 #[cfg(any(linux_kernel, target_os = "dragonfly"))]
122 #[inline]
sched_getcpu() -> usize123 pub fn sched_getcpu() -> usize {
124     backend::process::syscalls::sched_getcpu()
125 }
126