144 lines
4.6 KiB
Rust
144 lines
4.6 KiB
Rust
// Copyright 2019 The Chromium OS Authors. All rights reserved.
|
|
// Use of this source code is governed by a BSD-style license that can be
|
|
// found in the LICENSE file.
|
|
|
|
//! Wrappers for CPU affinity functions.
|
|
|
|
use std::{iter::FromIterator, mem};
|
|
|
|
use libc::{
|
|
cpu_set_t, prctl, sched_getaffinity, sched_setaffinity, CPU_ISSET, CPU_SET, CPU_SETSIZE,
|
|
CPU_ZERO, EINVAL,
|
|
};
|
|
|
|
use super::{errno_result, Error, Result};
|
|
|
|
// This is needed because otherwise the compiler will complain that the
|
|
// impl doesn't reference any types from inside this crate.
|
|
struct CpuSet(cpu_set_t);
|
|
|
|
impl CpuSet {
|
|
pub fn new() -> CpuSet {
|
|
// cpu_set_t is a C struct and can be safely initialized with zeroed memory.
|
|
let mut cpuset: cpu_set_t = unsafe { mem::MaybeUninit::zeroed().assume_init() };
|
|
// Safe because we pass a valid cpuset pointer.
|
|
unsafe { CPU_ZERO(&mut cpuset) };
|
|
CpuSet(cpuset)
|
|
}
|
|
|
|
pub fn to_cpus(&self) -> Vec<usize> {
|
|
let mut cpus = Vec::new();
|
|
for i in 0..(CPU_SETSIZE as usize) {
|
|
if unsafe { CPU_ISSET(i, &self.0) } {
|
|
cpus.push(i);
|
|
}
|
|
}
|
|
cpus
|
|
}
|
|
}
|
|
|
|
impl FromIterator<usize> for CpuSet {
|
|
fn from_iter<I: IntoIterator<Item = usize>>(cpus: I) -> Self {
|
|
// cpu_set_t is a C struct and can be safely initialized with zeroed memory.
|
|
let mut cpuset: cpu_set_t = unsafe { mem::zeroed() };
|
|
// Safe because we pass a valid cpuset pointer.
|
|
unsafe { CPU_ZERO(&mut cpuset) };
|
|
for cpu in cpus {
|
|
// Safe because we pass a valid cpuset pointer and cpu index.
|
|
unsafe { CPU_SET(cpu, &mut cpuset) };
|
|
}
|
|
CpuSet(cpuset)
|
|
}
|
|
}
|
|
|
|
/// Set the CPU affinity of the current thread to a given set of CPUs.
|
|
///
|
|
/// # Examples
|
|
///
|
|
/// Set the calling thread's CPU affinity so it will run on only CPUs
|
|
/// 0, 1, 5, and 6.
|
|
///
|
|
/// ```
|
|
/// # use sys_util::set_cpu_affinity;
|
|
/// set_cpu_affinity(vec![0, 1, 5, 6]).unwrap();
|
|
/// ```
|
|
pub fn set_cpu_affinity<I: IntoIterator<Item = usize>>(cpus: I) -> Result<()> {
|
|
let CpuSet(cpuset) = cpus
|
|
.into_iter()
|
|
.map(|cpu| {
|
|
if cpu < CPU_SETSIZE as usize {
|
|
Ok(cpu)
|
|
} else {
|
|
Err(Error::new(EINVAL))
|
|
}
|
|
})
|
|
.collect::<Result<CpuSet>>()?;
|
|
|
|
// Safe because we pass 0 for the current thread, and cpuset is a valid pointer and only
|
|
// used for the duration of this call.
|
|
let res = unsafe { sched_setaffinity(0, mem::size_of_val(&cpuset), &cpuset) };
|
|
|
|
if res != 0 {
|
|
errno_result()
|
|
} else {
|
|
Ok(())
|
|
}
|
|
}
|
|
|
|
pub fn get_cpu_affinity() -> Result<Vec<usize>> {
|
|
let mut cpu_set = CpuSet::new();
|
|
|
|
// Safe because we pass 0 for the current thread, and cpu_set.0 is a valid pointer and only
|
|
// used for the duration of this call.
|
|
super::syscall!(unsafe { sched_getaffinity(0, mem::size_of_val(&cpu_set.0), &mut cpu_set.0) })?;
|
|
|
|
Ok(cpu_set.to_cpus())
|
|
}
|
|
|
|
/// Enable experimental core scheduling for the current thread.
|
|
///
|
|
/// If successful, the kernel should not schedule this thread with any other thread within the same
|
|
/// SMT core. Because this is experimental, this will return success on kernels which do not support
|
|
/// this function.
|
|
pub fn enable_core_scheduling() -> Result<()> {
|
|
const PR_SCHED_CORE: i32 = 62;
|
|
const PR_SCHED_CORE_CREATE: i32 = 1;
|
|
|
|
#[allow(clippy::upper_case_acronyms, non_camel_case_types, dead_code)]
|
|
/// Specifies the scope of the pid parameter of `PR_SCHED_CORE`.
|
|
enum pid_type {
|
|
/// `PID` refers to threads.
|
|
PIDTYPE_PID,
|
|
/// `TGPID` refers to a process.
|
|
PIDTYPE_TGID,
|
|
/// `TGPID` refers to a process group.
|
|
PIDTYPE_PGID,
|
|
}
|
|
|
|
let ret = match unsafe {
|
|
prctl(
|
|
PR_SCHED_CORE,
|
|
PR_SCHED_CORE_CREATE,
|
|
0, // id of target task, 0 indicates current task
|
|
pid_type::PIDTYPE_PID as i32, // PID scopes to this thread only
|
|
0, // ignored by PR_SCHED_CORE_CREATE command
|
|
)
|
|
} {
|
|
#[cfg(feature = "chromeos")]
|
|
-1 => {
|
|
// Chrome OS has an pre-upstream version of this functionality which might work.
|
|
const PR_SET_CORE_SCHED: i32 = 0x200;
|
|
unsafe { prctl(PR_SET_CORE_SCHED, 1) }
|
|
}
|
|
ret => ret,
|
|
};
|
|
if ret == -1 {
|
|
let error = Error::last();
|
|
// prctl returns EINVAL for unknown functions, which we will ignore for now.
|
|
if error.errno() != libc::EINVAL {
|
|
return Err(error);
|
|
}
|
|
}
|
|
Ok(())
|
|
}
|