orengine_utils/
numa.rs

1//! Provides utilities for working with NUMA (Non-Uniform Memory Access) nodes.
2//!
3//! This module offers functionality to:
4//! - Manage data per NUMA node with [`DataPerNUMANodeManager`]
5//! - Get information about available NUMA nodes
6//! - Control thread affinity to specific NUMA nodes
7//!
8//! # Example
9//!
10//! ```
11//! use orengine_utils::numa::{DataPerNUMANodeManager, get_current_thread_numa_node};
12//!
13//! // Create a manager with data for each NUMA node
14//! static MANAGER: DataPerNUMANodeManager<usize> = DataPerNUMANodeManager::from_arr([0; 64]);
15//!
16//! let numa_node = get_current_thread_numa_node();
17//! println!("Memory by {numa_node} NUMA node contains {}", MANAGER.get_ref_by_node(numa_node));
18//! ```
19
20use crate::hints::unwrap_or_bug_message_hint;
21use core::iter::Iterator;
22
23#[cfg(not(feature = "more_numa_nodes"))]
24pub const MAX_NUMA_NODES_SUPPORTED_: usize = 64;
25#[cfg(feature = "more_numa_nodes")]
26pub const MAX_NUMA_NODES_SUPPORTED_: usize = 1024;
27
28/// The maximum number of NUMA nodes supported by the library.
29///
30/// If a machine supports more NUMA nodes than this, panics may occur with `debug_assertions`
31/// or UB otherwise.
32pub const MAX_NUMA_NODES_SUPPORTED: usize = MAX_NUMA_NODES_SUPPORTED_;
33
34const NUMA_NODE_TOO_LARGE: &str = "this hardware supports more NUMA-nodes than expected, use the `more_numa_nodes` feature to increase the limit";
35
36/// Manages data per NUMA node.
37/// It allows storing data for each NUMA node and accessing it by the NUMA node ID.
38///
39/// # Example
40///
41/// ```rust
42/// use orengine_utils::numa::{DataPerNUMANodeManager, get_current_thread_numa_node};
43///
44/// // Create a manager with data for each NUMA node
45/// static MANAGER: DataPerNUMANodeManager<usize> = DataPerNUMANodeManager::from_arr([0; 64]);
46///
47/// let numa_node = get_current_thread_numa_node();
48/// println!("Memory by {numa_node} NUMA node contains {}", MANAGER.get_ref_by_node(numa_node));
49/// ```
50pub struct DataPerNUMANodeManager<T>([T; MAX_NUMA_NODES_SUPPORTED]);
51
52impl<T> DataPerNUMANodeManager<T> {
53    /// Creates a new manager from an array of data for each NUMA node.
54    pub const fn from_arr(inner: [T; MAX_NUMA_NODES_SUPPORTED]) -> Self {
55        Self(inner)
56    }
57
58    /// Gets a reference to the data for the specified NUMA node.
59    ///
60    /// # Panics
61    ///
62    /// Panics if the NUMA node ID is out of bounds.
63    pub fn get_ref_by_node(&self, numa_node: usize) -> &T {
64        unwrap_or_bug_message_hint(self.0.get(numa_node), NUMA_NODE_TOO_LARGE)
65    }
66
67    /// Gets a mutable reference to the data for the specified NUMA node.
68    ///
69    /// # Panics
70    ///
71    /// Panics if the NUMA node ID is out of bounds.
72    pub fn get_mut_by_node(&mut self, numa_node: usize) -> &mut T {
73        unwrap_or_bug_message_hint(self.0.get_mut(numa_node), NUMA_NODE_TOO_LARGE)
74    }
75
76    /// Returns an iterator over references to the data for all NUMA nodes.
77    pub fn iter(&self) -> impl Iterator<Item = &T> {
78        self.0.iter()
79    }
80
81    /// Returns an iterator over mutable references to the data for all NUMA nodes.
82    pub fn iter_mut(&mut self) -> impl Iterator<Item = &mut T> {
83        self.0.iter_mut()
84    }
85
86    /// Returns a pointer to the inner array.
87    pub fn as_ptr(&self) -> *const [T; MAX_NUMA_NODES_SUPPORTED] {
88        self.0.as_ptr().cast()
89    }
90}
91
92impl<T: Default> Default for DataPerNUMANodeManager<T> {
93    fn default() -> Self {
94        Self(core::array::from_fn(|_| T::default()))
95    }
96}
97
98/// Gets the NUMA node ID for the current thread.
99///
100/// Returns the NUMA node that the current thread is running on.
101/// If NUMA is not supported, returns 0.
102///
103/// # Examples
104///
105/// ```
106/// use orengine_utils::numa::get_current_thread_numa_node;
107///
108/// let node_id = get_current_thread_numa_node();
109/// println!("Current thread is on NUMA node {}", node_id);
110/// ```
111pub fn get_current_thread_numa_node() -> usize {
112    #[cfg(all(target_os = "linux", not(miri)))]
113    {
114        use core::mem::MaybeUninit;
115
116        let mut numa_node: MaybeUninit<u32> = MaybeUninit::uninit();
117
118        unsafe {
119            libc::syscall(
120                libc::SYS_getcpu,
121                core::ptr::null::<libc::c_void>(),
122                numa_node.as_mut_ptr(),
123                core::ptr::null::<libc::c_void>(),
124            );
125        }
126
127        unsafe { numa_node.assume_init() as usize }
128    }
129
130    #[cfg(any(not(target_os = "linux"), miri))]
131    {
132        0
133    }
134}
135
136#[cfg(all(test, not(miri)))]
137mod tests {
138    use super::*;
139    use alloc::vec::Vec;
140
141    #[test]
142    fn test_data_per_numa_node_manager_iterators() {
143        let mut arr = [1i32; MAX_NUMA_NODES_SUPPORTED];
144        for (i, item) in arr.iter_mut().enumerate().take(8) {
145            *item = i32::try_from(i + 1).unwrap();
146        }
147        let mut manager = DataPerNUMANodeManager::from_arr(arr);
148
149        // Test iter()
150        let values: Vec<i32> = manager.iter().copied().collect();
151        assert_eq!(values[0], 1);
152        assert_eq!(values[7], 8);
153        assert_eq!(values[8], 1); // rest should be default (1)
154
155        // Test iter_mut()
156        for val in manager.iter_mut().take(4) {
157            *val *= 2;
158        }
159        assert_eq!(*manager.get_ref_by_node(0), 2);
160        assert_eq!(*manager.get_ref_by_node(3), 8);
161        assert_eq!(*manager.get_ref_by_node(4), 5);
162
163        // Test iter_enumerated()
164        let enumerated: Vec<(usize, &i32)> = manager.iter().enumerate().collect();
165        assert_eq!(enumerated[0], (0, &2));
166        assert_eq!(enumerated[3], (3, &8));
167        assert_eq!(enumerated[4], (4, &5));
168
169        // Test iter_enumerated_mut()
170        for (node_id, val) in manager.iter_mut().enumerate() {
171            if node_id % 2 == 0 {
172                *val += 10;
173            }
174        }
175        assert_eq!(*manager.get_ref_by_node(0), 12);
176        assert_eq!(*manager.get_ref_by_node(1), 4);
177        assert_eq!(*manager.get_ref_by_node(2), 16);
178    }
179
180    #[test]
181    fn test_get_current_thread_numa_node() {
182        let node_id = get_current_thread_numa_node();
183        assert!(node_id < 1024, "node: {node_id}");
184    }
185
186    #[test]
187    fn test_data_per_numa_node_manager_bounds() {
188        let manager = DataPerNUMANodeManager::from_arr([0u8; MAX_NUMA_NODES_SUPPORTED]);
189
190        // Should work for valid indices
191        for i in 0..MAX_NUMA_NODES_SUPPORTED {
192            let _ref = manager.get_ref_by_node(i);
193        }
194    }
195
196    #[test]
197    fn test_common_case() {
198        let numa_node = get_current_thread_numa_node();
199        let manager = DataPerNUMANodeManager::from_arr([0u8; MAX_NUMA_NODES_SUPPORTED]);
200
201        assert_eq!(*manager.get_ref_by_node(numa_node), 0);
202    }
203}
orengine_utils/numa.rs

orengine_utils/
numa.rs
