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;
8 use std::mem;
9
10 use libc::{cpu_set_t, sched_setaffinity, CPU_SET, CPU_SETSIZE, CPU_ZERO, EINVAL};
11
12 use crate::{errno_result, Error, Result};
13
14 // This is needed because otherwise the compiler will complain that the
15 // impl doesn't reference any types from inside this crate.
16 struct CpuSet(cpu_set_t);
17
18 impl FromIterator<usize> for CpuSet {
from_iter<I: IntoIterator<Item = usize>>(cpus: I) -> Self19 fn from_iter<I: IntoIterator<Item = usize>>(cpus: I) -> Self {
20 // cpu_set_t is a C struct and can be safely initialized with zeroed memory.
21 let mut cpuset: cpu_set_t = unsafe { mem::zeroed() };
22 // Safe because we pass a valid cpuset pointer.
23 unsafe { CPU_ZERO(&mut cpuset) };
24 for cpu in cpus {
25 // Safe because we pass a valid cpuset pointer and cpu index.
26 unsafe { CPU_SET(cpu, &mut cpuset) };
27 }
28 CpuSet(cpuset)
29 }
30 }
31
32 /// Set the CPU affinity of the current thread to a given set of CPUs.
33 ///
34 /// # Examples
35 ///
36 /// Set the calling thread's CPU affinity so it will run on only CPUs
37 /// 0, 1, 5, and 6.
38 ///
39 /// ```
40 /// # use sys_util::set_cpu_affinity;
41 /// set_cpu_affinity(vec![0, 1, 5, 6]).unwrap();
42 /// ```
set_cpu_affinity<I: IntoIterator<Item = usize>>(cpus: I) -> Result<()>43 pub fn set_cpu_affinity<I: IntoIterator<Item = usize>>(cpus: I) -> Result<()> {
44 let CpuSet(cpuset) = cpus
45 .into_iter()
46 .map(|cpu| {
47 if cpu < CPU_SETSIZE as usize {
48 Ok(cpu)
49 } else {
50 Err(Error::new(EINVAL))
51 }
52 })
53 .collect::<Result<CpuSet>>()?;
54
55 // Safe because we pass 0 for the current thread, and cpuset is a valid pointer and only
56 // used for the duration of this call.
57 let res = unsafe { sched_setaffinity(0, mem::size_of_val(&cpuset), &cpuset) };
58
59 if res != 0 {
60 errno_result()
61 } else {
62 Ok(())
63 }
64 }
65
66 /// Enable experimental core scheduling for the current thread.
67 ///
68 /// If succesful, the kernel should not schedule this thread with any other thread within the same
69 /// SMT core.
70 #[cfg(feature = "chromeos")]
enable_core_scheduling() -> Result<()>71 pub fn enable_core_scheduling() -> Result<()> {
72 use libc::prctl;
73 const PR_SET_CORE_SCHED: i32 = 0x200;
74 let ret = unsafe { prctl(PR_SET_CORE_SCHED, 1) };
75 if ret == -1 {
76 errno_result()
77 } else {
78 Ok(())
79 }
80 }
81