tuwunel_core/utils/sys/
compute.rs

1//! CPU topology, affinity, and parallelism utilities.
2//!
3//! The helpers inspect logical-core availability and derive sibling sets for
4//! simultaneous multithreading and hardware nodes. Platform-specific
5//! implementations provide available parallelism and current-CPU data.
6
7use std::{cell::Cell, fmt::Debug, path::PathBuf, sync::LazyLock};
8
9use crate::{Result, is_equal_to};
10
11type Id = usize;
12
13type Mask = u128;
14type Masks = [Mask; MASK_BITS];
15
16const MASK_BITS: usize = CORES_MAX;
17
18/// Maximum number of cores we support; for now limited to bits of our mask
19/// integral.
20pub const CORES_MAX: usize = 128;
21
22/// The mask of logical cores available to the process (at startup).
23static CORES_AVAILABLE: LazyLock<Mask> = LazyLock::new(|| into_mask(query_cores_available()));
24
25/// Stores the mask of logical-cores with thread/HT/SMT association. Each group
26/// here makes up a physical-core.
27static SMT_TOPOLOGY: LazyLock<Masks> = LazyLock::new(init_smt_topology);
28
29/// Stores the mask of logical-core associations on a node/socket. Bits are set
30/// for all logical cores within all physical cores of the node.
31static NODE_TOPOLOGY: LazyLock<Masks> = LazyLock::new(init_node_topology);
32
33thread_local! {
34	/// Tracks the affinity for this thread. This is updated when affinities
35	/// are set via our set_affinity() interface.
36	static CORE_AFFINITY: Cell<Mask> = const { Cell::new(0) };
37}
38
39/// Set the core affinity for this thread. The ID should be listed in
40/// CORES_AVAILABLE. Empty input is a no-op; prior affinity unchanged.
41#[tracing::instrument(
42	level = "debug",
43	skip_all,
44	fields(
45		id = ?std::thread::current().id(),
46		name = %std::thread::current().name().unwrap_or("None"),
47		set = ?ids.clone().collect::<Vec<_>>(),
48		CURRENT = %format!("[b{:b}]", CORE_AFFINITY.get()),
49		AVAILABLE = %format!("[b{:b}]", *CORES_AVAILABLE),
50	),
51)]
52pub fn set_affinity<I>(mut ids: I)
53where
54	I: Iterator<Item = Id> + Clone + Debug,
55{
56	use core_affinity::{CoreId, set_each_for_current, set_for_current};
57
58	let n = ids.clone().count();
59	let mask: Mask = ids.clone().fold(0, |mask, id| {
60		debug_assert!(is_core_available(id), "setting affinity to unavailable core");
61		mask | (1 << id)
62	});
63
64	if n > 1 {
65		set_each_for_current(ids.map(|id| CoreId { id }));
66	} else if n > 0 {
67		set_for_current(CoreId { id: ids.next().expect("n > 0") });
68	}
69
70	if mask.count_ones() > 0 {
71		CORE_AFFINITY.replace(mask);
72	}
73}
74
75/// Get the core affinity for this thread.
76pub fn get_affinity() -> impl Iterator<Item = Id> {
77	CORE_AFFINITY
78		.get()
79		.ne(&0)
80		.then_some(from_mask(CORE_AFFINITY.get()))
81		.or_else(|| Some(from_mask(*CORES_AVAILABLE)))
82		.into_iter()
83		.flatten()
84}
85
86/// List the cores sharing SMT-tier resources
87pub fn smt_siblings() -> impl Iterator<Item = Id> {
88	from_mask(get_affinity().fold(0_u128, |mask, id| {
89		mask | SMT_TOPOLOGY
90			.get(id)
91			.expect("ID must not exceed max cpus")
92	}))
93}
94
95/// List the cores sharing Node-tier resources relative to this threads current
96/// affinity.
97pub fn node_siblings() -> impl Iterator<Item = Id> {
98	from_mask(get_affinity().fold(0_u128, |mask, id| {
99		mask | NODE_TOPOLOGY
100			.get(id)
101			.expect("Id must not exceed max cpus")
102	}))
103}
104
105/// Get the cores sharing SMT resources relative to id.
106#[inline]
107pub fn smt_affinity(id: Id) -> impl Iterator<Item = Id> {
108	from_mask(
109		*SMT_TOPOLOGY
110			.get(id)
111			.expect("ID must not exceed max cpus"),
112	)
113}
114
115/// Get the cores sharing Node resources relative to id.
116#[inline]
117pub fn node_affinity(id: Id) -> impl Iterator<Item = Id> {
118	from_mask(
119		*NODE_TOPOLOGY
120			.get(id)
121			.expect("ID must not exceed max cpus"),
122	)
123}
124
125/// Get the number of threads which could execute in parallel based on hardware
126/// constraints of this system.
127#[cfg(not(target_os = "openbsd"))]
128#[inline]
129#[must_use]
130pub fn available_parallelism() -> usize { cores_available().count() }
131
132/// Reports the number of CPUs OpenBSD currently has online.
133///
134/// The count is read from the system directly rather than from the core mask
135/// recorded at startup, and is not narrowed by process affinity, which OpenBSD
136/// does not expose.
137#[cfg(target_os = "openbsd")]
138#[inline]
139#[must_use]
140pub fn available_parallelism() -> usize { num_cpus::get() }
141
142/// Gets the ID of the nth core available. This bijects our sequence of cores to
143/// actual ID's which may have gaps for cores which are not available.
144#[inline]
145#[must_use]
146pub fn nth_core_available(i: usize) -> Option<Id> { cores_available().nth(i) }
147
148/// Determine if core (by id) is available to the process.
149#[inline]
150#[must_use]
151pub fn is_core_available(id: Id) -> bool { cores_available().any(is_equal_to!(id)) }
152
153/// Get the list of cores available. The values were recorded at program start.
154#[inline]
155pub fn cores_available() -> impl Iterator<Item = Id> { from_mask(*CORES_AVAILABLE) }
156
157/// Returns the logical CPU currently executing the calling thread.
158///
159/// The value is a point-in-time scheduler observation and may become stale on
160/// the next instruction boundary. Linux obtains it through `sched_getcpu()`.
161#[cfg(target_os = "linux")]
162#[inline]
163pub fn getcpu() -> Result<usize> {
164	use crate::{Error, utils::math};
165
166	// SAFETY: This is part of an interface with many low-level calls taking many
167	// raw params, but it's unclear why this specific call is unsafe. Nevertheless
168	// the value obtained here is semantically unsafe because it can change on the
169	// instruction boundary trailing its own acquisition and also any other time.
170	let ret: i32 = unsafe { libc::sched_getcpu() };
171
172	#[cfg(target_os = "linux")]
173	// SAFETY: On modern linux systems with a vdso if we can optimize away the branch checking
174	// for error (see getcpu(2)) then this system call becomes a memory access.
175	unsafe {
176		std::hint::assert_unchecked(ret >= 0);
177	};
178
179	if ret == -1 {
180		return Err(Error::from_errno());
181	}
182
183	math::try_into(ret)
184}
185
186/// Reports that current-CPU queries are unsupported on this platform.
187///
188/// No scheduler observation is attempted. The result contains an I/O error with
189/// the unsupported error kind.
190#[cfg(not(target_os = "linux"))]
191#[inline]
192pub fn getcpu() -> Result<usize> { Err(crate::Error::Io(std::io::ErrorKind::Unsupported.into())) }
193
194#[cfg(not(target_os = "openbsd"))]
195fn query_cores_available() -> impl Iterator<Item = Id> {
196	core_affinity::get_core_ids()
197		.unwrap_or_default()
198		.into_iter()
199		.map(|core_id| core_id.id)
200}
201
202#[cfg(target_os = "openbsd")]
203fn query_cores_available() -> impl Iterator<Item = Id> { 0..num_cpus::get() }
204
205fn init_smt_topology() -> [Mask; MASK_BITS] { [Mask::default(); MASK_BITS] }
206
207fn init_node_topology() -> [Mask; MASK_BITS] { [Mask::default(); MASK_BITS] }
208
209fn into_mask<I>(ids: I) -> Mask
210where
211	I: Iterator<Item = Id>,
212{
213	ids.inspect(|&id| {
214		debug_assert!(id < MASK_BITS, "Core ID must be < Mask::BITS at least for now");
215	})
216	.fold(Mask::default(), |mask, id| mask | (1 << id))
217}
218
219fn from_mask(v: Mask) -> impl Iterator<Item = Id> {
220	(0..MASK_BITS).filter(move |&i| (v & (1 << i)) != 0)
221}
222
223fn _sys_path(id: usize, suffix: &str) -> PathBuf {
224	format!("/sys/devices/system/cpu/cpu{id}/{suffix}").into()
225}
tuwunel_core/utils/sys/compute.rs

tuwunel_core/utils/sys/
compute.rs
