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