1 // Copyright 2019 The Chromium OS Authors. All rights reserved.
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, mem};
8
9 use libc::{
10 cpu_set_t, prctl, sched_getaffinity, sched_setaffinity, CPU_ISSET, CPU_SET, CPU_SETSIZE,
11 CPU_ZERO, EINVAL,
12 };
13
14 use super::{errno_result, Error, Result};
15
16 // This is needed because otherwise the compiler will complain that the
17 // impl doesn't reference any types from inside this crate.
18 struct CpuSet(cpu_set_t);
19
20 impl CpuSet {
new() -> CpuSet21 pub fn new() -> CpuSet {
22 // cpu_set_t is a C struct and can be safely initialized with zeroed memory.
23 let mut cpuset: cpu_set_t = unsafe { mem::MaybeUninit::zeroed().assume_init() };
24 // Safe because we pass a valid cpuset pointer.
25 unsafe { CPU_ZERO(&mut cpuset) };
26 CpuSet(cpuset)
27 }
28
to_cpus(&self) -> Vec<usize>29 pub fn to_cpus(&self) -> Vec<usize> {
30 let mut cpus = Vec::new();
31 for i in 0..(CPU_SETSIZE as usize) {
32 if unsafe { CPU_ISSET(i, &self.0) } {
33 cpus.push(i);
34 }
35 }
36 cpus
37 }
38 }
39
40 impl FromIterator<usize> for CpuSet {
from_iter<I: IntoIterator<Item = usize>>(cpus: I) -> Self41 fn from_iter<I: IntoIterator<Item = usize>>(cpus: I) -> Self {
42 // cpu_set_t is a C struct and can be safely initialized with zeroed memory.
43 let mut cpuset: cpu_set_t = unsafe { mem::zeroed() };
44 // Safe because we pass a valid cpuset pointer.
45 unsafe { CPU_ZERO(&mut cpuset) };
46 for cpu in cpus {
47 // Safe because we pass a valid cpuset pointer and cpu index.
48 unsafe { CPU_SET(cpu, &mut cpuset) };
49 }
50 CpuSet(cpuset)
51 }
52 }
53
54 /// Set the CPU affinity of the current thread to a given set of CPUs.
55 ///
56 /// # Examples
57 ///
58 /// Set the calling thread's CPU affinity so it will run on only CPUs
59 /// 0, 1, 5, and 6.
60 ///
61 /// ```
62 /// # use sys_util::set_cpu_affinity;
63 /// set_cpu_affinity(vec![0, 1, 5, 6]).unwrap();
64 /// ```
set_cpu_affinity<I: IntoIterator<Item = usize>>(cpus: I) -> Result<()>65 pub fn set_cpu_affinity<I: IntoIterator<Item = usize>>(cpus: I) -> Result<()> {
66 let CpuSet(cpuset) = cpus
67 .into_iter()
68 .map(|cpu| {
69 if cpu < CPU_SETSIZE as usize {
70 Ok(cpu)
71 } else {
72 Err(Error::new(EINVAL))
73 }
74 })
75 .collect::<Result<CpuSet>>()?;
76
77 // Safe because we pass 0 for the current thread, and cpuset is a valid pointer and only
78 // used for the duration of this call.
79 let res = unsafe { sched_setaffinity(0, mem::size_of_val(&cpuset), &cpuset) };
80
81 if res != 0 {
82 errno_result()
83 } else {
84 Ok(())
85 }
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 super::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 = match 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 #[cfg(feature = "chromeos")]
128 -1 => {
129 // Chrome OS has an pre-upstream version of this functionality which might work.
130 const PR_SET_CORE_SCHED: i32 = 0x200;
131 unsafe { prctl(PR_SET_CORE_SCHED, 1) }
132 }
133 ret => ret,
134 };
135 if ret == -1 {
136 let error = Error::last();
137 // prctl returns EINVAL for unknown functions, which we will ignore for now.
138 if error.errno() != libc::EINVAL {
139 return Err(error);
140 }
141 }
142 Ok(())
143 }
144