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 // cpu_set_t is a C struct and can be safely initialized with zeroed memory.
30 let mut cpuset: cpu_set_t = unsafe { mem::MaybeUninit::zeroed().assume_init() };
31 // Safe because we pass a valid cpuset pointer.
32 unsafe { CPU_ZERO(&mut cpuset) };
33 CpuSet(cpuset)
34 }
35
to_cpus(&self) -> Vec<usize>36 pub fn to_cpus(&self) -> Vec<usize> {
37 let mut cpus = Vec::new();
38 for i in 0..(CPU_SETSIZE as usize) {
39 if unsafe { CPU_ISSET(i, &self.0) } {
40 cpus.push(i);
41 }
42 }
43 cpus
44 }
45 }
46
47 impl FromIterator<usize> for CpuSet {
from_iter<I: IntoIterator<Item = usize>>(cpus: I) -> Self48 fn from_iter<I: IntoIterator<Item = usize>>(cpus: I) -> Self {
49 let mut cpuset = CpuSet::new();
50 for cpu in cpus {
51 // Safe because we pass a valid cpu index and cpuset.0 is a valid pointer.
52 unsafe { CPU_SET(cpu, &mut cpuset.0) };
53 }
54 cpuset
55 }
56 }
57
58 /// Set the CPU affinity of the current thread to a given set of CPUs.
59 ///
60 /// # Examples
61 ///
62 /// Set the calling thread's CPU affinity so it will run on only CPUs
63 /// 0, 1, 5, and 6.
64 ///
65 /// ```
66 /// # use base::platform::set_cpu_affinity;
67 /// set_cpu_affinity(vec![0, 1, 5, 6]).unwrap();
68 /// ```
set_cpu_affinity<I: IntoIterator<Item = usize>>(cpus: I) -> Result<()>69 pub fn set_cpu_affinity<I: IntoIterator<Item = usize>>(cpus: I) -> Result<()> {
70 let CpuSet(cpuset) = cpus
71 .into_iter()
72 .map(|cpu| {
73 if cpu < CPU_SETSIZE as usize {
74 Ok(cpu)
75 } else {
76 Err(Error::new(EINVAL))
77 }
78 })
79 .collect::<Result<CpuSet>>()?;
80
81 // Safe because we pass 0 for the current thread, and cpuset is a valid pointer and only
82 // used for the duration of this call.
83 crate::syscall!(unsafe { sched_setaffinity(0, mem::size_of_val(&cpuset), &cpuset) })?;
84
85 Ok(())
86 }
87
get_cpu_affinity() -> Result<Vec<usize>>88 pub fn get_cpu_affinity() -> Result<Vec<usize>> {
89 let mut cpu_set = CpuSet::new();
90
91 // Safe because we pass 0 for the current thread, and cpu_set.0 is a valid pointer and only
92 // used for the duration of this call.
93 crate::syscall!(unsafe { sched_getaffinity(0, mem::size_of_val(&cpu_set.0), &mut cpu_set.0) })?;
94
95 Ok(cpu_set.to_cpus())
96 }
97
98 /// Enable experimental core scheduling for the current thread.
99 ///
100 /// If successful, the kernel should not schedule this thread with any other thread within the same
101 /// SMT core. Because this is experimental, this will return success on kernels which do not support
102 /// this function.
enable_core_scheduling() -> Result<()>103 pub fn enable_core_scheduling() -> Result<()> {
104 const PR_SCHED_CORE: i32 = 62;
105 const PR_SCHED_CORE_CREATE: i32 = 1;
106
107 #[allow(clippy::upper_case_acronyms, non_camel_case_types, dead_code)]
108 /// Specifies the scope of the pid parameter of `PR_SCHED_CORE`.
109 enum pid_type {
110 /// `PID` refers to threads.
111 PIDTYPE_PID,
112 /// `TGPID` refers to a process.
113 PIDTYPE_TGID,
114 /// `TGPID` refers to a process group.
115 PIDTYPE_PGID,
116 }
117
118 let ret = unsafe {
119 prctl(
120 PR_SCHED_CORE,
121 PR_SCHED_CORE_CREATE,
122 0, // id of target task, 0 indicates current task
123 pid_type::PIDTYPE_PID as i32, // PID scopes to this thread only
124 0, // ignored by PR_SCHED_CORE_CREATE command
125 )
126 };
127 if ret == -1 {
128 let error = Error::last();
129 // prctl returns EINVAL for unknown functions, which we will ignore for now.
130 if error.errno() != libc::EINVAL {
131 return Err(error);
132 }
133 }
134 Ok(())
135 }
136