base/sys/linux/
sched.rs

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
7use std::iter::FromIterator;
8use std::mem;
9
10use libc::cpu_set_t;
11use libc::prctl;
12use libc::sched_getaffinity;
13use libc::sched_setaffinity;
14use libc::CPU_ISSET;
15use libc::CPU_SET;
16use libc::CPU_SETSIZE;
17use libc::CPU_ZERO;
18use libc::EINVAL;
19
20use super::Error;
21use 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.
25struct CpuSet(cpu_set_t);
26
27impl CpuSet {
28    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)]
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
51impl FromIterator<usize> for CpuSet {
52    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)]
75pub 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
95pub 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.
111pub 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}
base/sys/linux/sched.rs

base/sys/linux/
sched.rs
