1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144

// Copyright 2019 The ChromiumOS Authors
// 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;
use std::mem;

use libc::cpu_set_t;
use libc::prctl;
use libc::sched_getaffinity;
use libc::sched_setaffinity;
use libc::CPU_ISSET;
use libc::CPU_SET;
use libc::CPU_SETSIZE;
use libc::CPU_ZERO;
use libc::EINVAL;

use super::Error;
use super::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 {
        // SAFETY:
        // 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() };
        // SAFETY:
        // Safe because we pass a valid cpuset pointer.
        unsafe { CPU_ZERO(&mut cpuset) };
        CpuSet(cpuset)
    }

    #[allow(clippy::unnecessary_cast)]
    pub fn to_cpus(&self) -> Vec<usize> {
        let mut cpus = Vec::new();
        for i in 0..(CPU_SETSIZE as usize) {
            // SAFETY: Safe because `i` and `self.0` are valid.
            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 {
        let mut cpuset = CpuSet::new();
        for cpu in cpus {
            // SAFETY:
            // Safe because we pass a valid cpu index and cpuset.0 is a valid pointer.
            unsafe { CPU_SET(cpu, &mut cpuset.0) };
        }
        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 base::linux::set_cpu_affinity;
///   set_cpu_affinity(vec![0, 1, 5, 6]).unwrap();
/// ```
#[allow(clippy::unnecessary_cast)]
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>>()?;

    // SAFETY:
    // Safe because we pass 0 for the current thread, and cpuset is a valid pointer and only
    // used for the duration of this call.
    crate::syscall!(unsafe { sched_setaffinity(0, mem::size_of_val(&cpuset), &cpuset) })?;

    Ok(())
}

pub fn get_cpu_affinity() -> Result<Vec<usize>> {
    let mut cpu_set = CpuSet::new();

    // SAFETY:
    // 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.
    crate::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,
    }

    // SAFETY: Safe because we check the return value to prctl.
    let ret = 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
        )
    };
    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(())
}