nuclei/rt/runtime/
cores.rs

1//! The cores a runtime's threads run on: [`Builder::worker_cores`] and
2//! [`Builder::blocking_cores`], checked against the cores the process may
3//! run on, and the affinity masks made from them.
4
5use std::io;
6
7use super::{Builder, CpuMask};
8
9impl Builder {
10    /// The cores the workers run on, in the order the workers take them:
11    /// worker `k` runs on `cores[k % cores.len()]`, pinned to it, or, with
12    /// [`pin_workers(false)`](Self::pin_workers), on any core of the list.
13    /// The worker count defaults to the list's length. Every core must be
14    /// one the process may run on and the list must not be empty;
15    /// [`build`](Self::build) fails with `InvalidInput` otherwise. Pinning
16    /// takes effect on Linux and Windows; elsewhere the list sets the
17    /// worker count and is checked against the cores the machine has.
18    pub fn worker_cores(mut self, cores: impl IntoIterator<Item = usize>) -> Self {
19        self.worker_cores = Some(cores.into_iter().collect());
20        self
21    }
22
23    /// The cores the blocking pool's threads run on, any of them (default:
24    /// the cores of the thread that builds the runtime). Checked as
25    /// [`worker_cores`](Self::worker_cores) are, and in effect on Linux and
26    /// Windows.
27    pub fn blocking_cores(mut self, cores: impl IntoIterator<Item = usize>) -> Self {
28        self.blocking_cores = Some(cores.into_iter().collect());
29        self
30    }
31}
32
33/// The cores the process may run on: its affinity mask's (Linux, Windows),
34/// or every core of the machine where no mask is known.
35fn process_cores() -> Vec<usize> {
36    let cores = super::allowed_cores();
37    if !cores.is_empty() {
38        return cores;
39    }
40    (0..std::thread::available_parallelism().map_or(1, usize::from)).collect()
41}
42
43/// Refuses an empty list of cores, or one naming a core the process may not
44/// run on, as the configuration of `setting`.
45pub(super) fn validate(setting: &str, cores: Option<&[usize]>) -> io::Result<()> {
46    let Some(cores) = cores else {
47        return Ok(());
48    };
49    if cores.is_empty() {
50        return Err(crate::config::refuse(setting, "empty", "at least one core"));
51    }
52    let allowed = process_cores();
53    if let Some(core) = cores.iter().find(|core| !allowed.contains(core)) {
54        return Err(crate::config::refuse(
55            setting,
56            format_args!("a list naming core {core}"),
57            format_args!("a list of the cores this process may run on ({allowed:?})"),
58        ));
59    }
60    Ok(())
61}
62
63/// The affinity mask of `cores`; `None` where threads take no mask.
64pub(super) fn mask_of(cores: &[usize]) -> Option<CpuMask> {
65    #[cfg(target_os = "linux")]
66    {
67        // SAFETY: a zeroed set of our own, and plain writes into it.
68        let mut set: libc::cpu_set_t = unsafe { core::mem::zeroed() };
69        for &core in cores.iter().filter(|&&c| c < libc::CPU_SETSIZE as usize) {
70            // SAFETY: as above; the core is inside the set.
71            unsafe { libc::CPU_SET(core, &mut set) };
72        }
73        Some(set)
74    }
75    #[cfg(windows)]
76    {
77        Some(
78            cores
79                .iter()
80                .filter(|&&c| c < usize::BITS as usize)
81                .fold(0usize, |mask, &c| mask | (1 << c)),
82        )
83    }
84    #[cfg(not(any(target_os = "linux", windows)))]
85    {
86        let _ = cores;
87        None
88    }
89}
nuclei/rt/runtime/cores.rs

nuclei/rt/runtime/
cores.rs
