1 // Copyright 2019 The ChromiumOS Authors
2 // Use of this source code is governed by a BSD-style license that can be
3 // found in the LICENSE file.
4
5 //! Wrappers for CPU affinity functions.
6
7 use std::iter::FromIterator;
8 use std::mem;
9
10 use libc::cpu_set_t;
11 use libc::prctl;
12 use libc::sched_getaffinity;
13 use libc::sched_setaffinity;
14 use libc::CPU_ISSET;
15 use libc::CPU_SET;
16 use libc::CPU_SETSIZE;
17 use libc::CPU_ZERO;
18 use libc::EINVAL;
19
20 use super::Error;
21 use super::Result;
22
23 // This is needed because otherwise the compiler will complain that the
24 // impl doesn't reference any types from inside this crate.
25 struct CpuSet(cpu_set_t);
26
27 impl CpuSet {
new() -> CpuSet28 pub fn new() -> CpuSet {
29 // SAFETY:
30 // cpu_set_t is a C struct and can be safely initialized with zeroed memory.
31 let mut cpuset: cpu_set_t = unsafe { mem::MaybeUninit::zeroed().assume_init() };
32 // SAFETY:
33 // Safe because we pass a valid cpuset pointer.
34 unsafe { CPU_ZERO(&mut cpuset) };
35 CpuSet(cpuset)
36 }
37
38 #[allow(clippy::unnecessary_cast)]
to_cpus(&self) -> Vec<usize>39 pub fn to_cpus(&self) -> Vec<usize> {
40 let mut cpus = Vec::new();
41 for i in 0..(CPU_SETSIZE as usize) {
42 // SAFETY: Safe because `i` and `self.0` are valid.
43 if unsafe { CPU_ISSET(i, &self.0) } {
44 cpus.push(i);
45 }
46 }
47 cpus
48 }
49 }
50
51 impl FromIterator<usize> for CpuSet {
from_iter<I: IntoIterator<Item = usize>>(cpus: I) -> Self52 fn from_iter<I: IntoIterator<Item = usize>>(cpus: I) -> Self {
53 let mut cpuset = CpuSet::new();
54 for cpu in cpus {
55 // SAFETY:
56 // Safe because we pass a valid cpu index and cpuset.0 is a valid pointer.
57 unsafe { CPU_SET(cpu, &mut cpuset.0) };
58 }
59 cpuset
60 }
61 }
62
63 /// Set the CPU affinity of the current thread to a given set of CPUs.
64 ///
65 /// # Examples
66 ///
67 /// Set the calling thread's CPU affinity so it will run on only CPUs
68 /// 0, 1, 5, and 6.
69 ///
70 /// ```
71 /// # use base::linux::set_cpu_affinity;
72 /// set_cpu_affinity(vec![0, 1, 5, 6]).unwrap();
73 /// ```
74 #[allow(clippy::unnecessary_cast)]
set_cpu_affinity<I: IntoIterator<Item = usize>>(cpus: I) -> Result<()>75 pub fn set_cpu_affinity<I: IntoIterator<Item = usize>>(cpus: I) -> Result<()> {
76 let CpuSet(cpuset) = cpus
77 .into_iter()
78 .map(|cpu| {
79 if cpu < CPU_SETSIZE as usize {
80 Ok(cpu)
81 } else {
82 Err(Error::new(EINVAL))
83 }
84 })
85 .collect::<Result<CpuSet>>()?;
86
87 // SAFETY:
88 // Safe because we pass 0 for the current thread, and cpuset is a valid pointer and only
89 // used for the duration of this call.
90 crate::syscall!(unsafe { sched_setaffinity(0, mem::size_of_val(&cpuset), &cpuset) })?;
91
92 Ok(())
93 }
94
get_cpu_affinity() -> Result<Vec<usize>>95 pub fn get_cpu_affinity() -> Result<Vec<usize>> {
96 let mut cpu_set = CpuSet::new();
97
98 // SAFETY:
99 // Safe because we pass 0 for the current thread, and cpu_set.0 is a valid pointer and only
100 // used for the duration of this call.
101 crate::syscall!(unsafe { sched_getaffinity(0, mem::size_of_val(&cpu_set.0), &mut cpu_set.0) })?;
102
103 Ok(cpu_set.to_cpus())
104 }
105
106 /// Enable experimental core scheduling for the current thread.
107 ///
108 /// If successful, the kernel should not schedule this thread with any other thread within the same
109 /// SMT core. Because this is experimental, this will return success on kernels which do not support
110 /// this function.
enable_core_scheduling() -> Result<()>111 pub fn enable_core_scheduling() -> Result<()> {
112 const PR_SCHED_CORE: i32 = 62;
113 const PR_SCHED_CORE_CREATE: i32 = 1;
114
115 #[allow(clippy::upper_case_acronyms, non_camel_case_types, dead_code)]
116 /// Specifies the scope of the pid parameter of `PR_SCHED_CORE`.
117 enum pid_type {
118 /// `PID` refers to threads.
119 PIDTYPE_PID,
120 /// `TGPID` refers to a process.
121 PIDTYPE_TGID,
122 /// `TGPID` refers to a process group.
123 PIDTYPE_PGID,
124 }
125
126 // SAFETY: Safe because we check the return value to prctl.
127 let ret = unsafe {
128 prctl(
129 PR_SCHED_CORE,
130 PR_SCHED_CORE_CREATE,
131 0, // id of target task, 0 indicates current task
132 pid_type::PIDTYPE_PID as i32, // PID scopes to this thread only
133 0, // ignored by PR_SCHED_CORE_CREATE command
134 )
135 };
136 if ret == -1 {
137 let error = Error::last();
138 // prctl returns EINVAL for unknown functions, which we will ignore for now.
139 if error.errno() != libc::EINVAL {
140 return Err(error);
141 }
142 }
143 Ok(())
144 }
145