rustix/process/
sched.rs

1use crate::process::Pid;
2use crate::{backend, io};
3use core::{fmt, hash};
4
5/// `CpuSet` represents a bit-mask of CPUs.
6///
7/// `CpuSet`s are used by [`sched_setaffinity`] and [`sched_getaffinity`], for
8/// example.
9///
10/// # References
11///  - [Linux]
12///
13/// [Linux]: https://man7.org/linux/man-pages/man3/CPU_SET.3.html
14/// [`sched_setaffinity`]: crate::process::sched_setaffinity
15/// [`sched_getaffinity`]: crate::process::sched_getaffinity
16#[repr(C)]
17#[derive(Clone, Copy)]
18pub struct CpuSet {
19    cpu_set: backend::process::types::RawCpuSet,
20}
21
22impl CpuSet {
23    /// The maximum number of CPU in `CpuSet`.
24    pub const MAX_CPU: usize = backend::process::types::CPU_SETSIZE;
25
26    /// Create a new and empty `CpuSet`.
27    #[inline]
28    pub fn new() -> Self {
29        Self {
30            cpu_set: backend::process::types::raw_cpu_set_new(),
31        }
32    }
33
34    /// Test to see if a CPU is in the `CpuSet`.
35    ///
36    /// `field` is the CPU id to test.
37    #[inline]
38    pub fn is_set(&self, field: usize) -> bool {
39        backend::process::cpu_set::CPU_ISSET(field, &self.cpu_set)
40    }
41
42    /// Add a CPU to `CpuSet`.
43    ///
44    /// `field` is the CPU id to add.
45    #[inline]
46    pub fn set(&mut self, field: usize) {
47        backend::process::cpu_set::CPU_SET(field, &mut self.cpu_set)
48    }
49
50    /// Remove a CPU from `CpuSet`.
51    ///
52    /// `field` is the CPU id to remove.
53    #[inline]
54    pub fn unset(&mut self, field: usize) {
55        backend::process::cpu_set::CPU_CLR(field, &mut self.cpu_set)
56    }
57
58    /// Count the number of CPUs set in the `CpuSet`.
59    #[cfg(linux_kernel)]
60    #[inline]
61    pub fn count(&self) -> u32 {
62        backend::process::cpu_set::CPU_COUNT(&self.cpu_set)
63    }
64
65    /// Zeroes the `CpuSet`.
66    #[inline]
67    pub fn clear(&mut self) {
68        backend::process::cpu_set::CPU_ZERO(&mut self.cpu_set)
69    }
70}
71
72impl Default for CpuSet {
73    #[inline]
74    fn default() -> Self {
75        Self::new()
76    }
77}
78
79impl fmt::Debug for CpuSet {
80    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
81        write!(f, "CpuSet {{")?;
82        let mut first = true;
83        for i in 0..Self::MAX_CPU {
84            if self.is_set(i) {
85                if first {
86                    write!(f, " ")?;
87                    first = false;
88                } else {
89                    write!(f, ", ")?;
90                }
91                write!(f, "cpu{}", i)?;
92            }
93        }
94        write!(f, " }}")
95    }
96}
97
98impl hash::Hash for CpuSet {
99    fn hash<H: hash::Hasher>(&self, state: &mut H) {
100        for i in 0..Self::MAX_CPU {
101            self.is_set(i).hash(state);
102        }
103    }
104}
105
106impl Eq for CpuSet {}
107
108impl PartialEq for CpuSet {
109    fn eq(&self, other: &Self) -> bool {
110        backend::process::cpu_set::CPU_EQUAL(&self.cpu_set, &other.cpu_set)
111    }
112}
113
114/// `sched_setaffinity(pid, cpuset)`—Set a thread's CPU affinity mask.
115///
116/// `pid` is the thread ID to update. If pid is `None`, then the current thread
117/// is updated.
118///
119/// The `CpuSet` argument specifies the set of CPUs on which the thread will be
120/// eligible to run.
121///
122/// # References
123///  - [Linux]
124///
125/// [Linux]: https://man7.org/linux/man-pages/man2/sched_setaffinity.2.html
126#[inline]
127pub fn sched_setaffinity(pid: Option<Pid>, cpuset: &CpuSet) -> io::Result<()> {
128    backend::process::syscalls::sched_setaffinity(pid, &cpuset.cpu_set)
129}
130
131/// `sched_getaffinity(pid)`—Get a thread's CPU affinity mask.
132///
133/// `pid` is the thread ID to check. If pid is `None`, then the current thread
134/// is checked.
135///
136/// Returns the set of CPUs on which the thread is eligible to run.
137///
138/// # References
139///  - [Linux]
140///
141/// [Linux]: https://man7.org/linux/man-pages/man2/sched_getaffinity.2.html
142#[inline]
143pub fn sched_getaffinity(pid: Option<Pid>) -> io::Result<CpuSet> {
144    let mut cpuset = CpuSet::new();
145    backend::process::syscalls::sched_getaffinity(pid, &mut cpuset.cpu_set).and(Ok(cpuset))
146}
147
148/// `sched_getcpu()`—Get the CPU that the current thread is currently on.
149///
150/// # References
151///  - [Linux]
152///  - [DragonFly BSD]
153///
154/// [Linux]: https://man7.org/linux/man-pages/man3/sched_getcpu.3.html
155/// [DragonFly BSD]: https://man.dragonflybsd.org/?command=sched_getcpu&section=2
156// FreeBSD added `sched_getcpu` in 13.0.
157#[cfg(any(linux_kernel, target_os = "dragonfly"))]
158#[inline]
159pub fn sched_getcpu() -> usize {
160    backend::process::syscalls::sched_getcpu()
161}
rustix/process/sched.rs

rustix/process/
sched.rs
