gdt_cpus/affinity_mask.rs
1//! CPU affinity mask for specifying sets of logical processors.
2//!
3//! This module provides the [`AffinityMask`] type, which represents a set of
4//! logical processor IDs that a thread may run on. Unlike pinning to a single
5//! core, affinity masks allow threads to migrate between a specified set of
6//! cores, reducing scheduling latency while still constraining execution.
7
8/// A cross-platform CPU affinity mask representing a set of logical processors.
9///
10/// An `AffinityMask` specifies which logical processors (CPU threads) a thread
11/// is allowed to execute on. This is useful for:
12///
13/// - Restricting latency-sensitive threads to performance cores (P-cores)
14/// - Restricting background threads to efficiency cores (E-cores)
15/// - Allowing thread migration within a subset of cores to reduce scheduling latency
16///
17/// # Example
18///
19/// ```
20/// use gdt_cpus::AffinityMask;
21///
22/// // Create a mask for cores 0, 1, and 2
23/// let mask = AffinityMask::from_cores(&[0, 1, 2]);
24/// assert_eq!(mask.count(), 3);
25/// assert!(mask.contains(0));
26/// assert!(mask.contains(1));
27/// assert!(mask.contains(2));
28/// assert!(!mask.contains(3));
29/// ```
30///
31/// # Platform Behavior
32///
33/// When used with [`set_thread_affinity`](crate::set_thread_affinity):
34///
35/// - **Linux/Windows**: The mask is applied directly to constrain thread execution
36/// - **macOS**: Returns `Error::Unsupported`; use QoS classes instead
37#[derive(Clone, PartialEq, Eq, Default)]
38pub struct AffinityMask {
39 /// Bitset stored as multiple u64 words to support >64 cores.
40 /// bits[0] contains cores 0-63, bits[1] contains cores 64-127, etc.
41 bits: Vec<u64>,
42}
43
44impl AffinityMask {
45 /// Creates an empty affinity mask with no cores set.
46 ///
47 /// # Example
48 ///
49 /// ```
50 /// use gdt_cpus::AffinityMask;
51 ///
52 /// let mask = AffinityMask::empty();
53 /// assert!(mask.is_empty());
54 /// assert_eq!(mask.count(), 0);
55 /// ```
56 pub fn empty() -> Self {
57 Self::default()
58 }
59
60 /// Creates an affinity mask with a single core set.
61 ///
62 /// # Arguments
63 ///
64 /// * `logical_core_id` - The logical processor ID to include in the mask
65 ///
66 /// # Example
67 ///
68 /// ```
69 /// use gdt_cpus::AffinityMask;
70 ///
71 /// let mask = AffinityMask::single(5);
72 /// assert_eq!(mask.count(), 1);
73 /// assert!(mask.contains(5));
74 /// ```
75 pub fn single(logical_core_id: usize) -> Self {
76 let mut mask = Self::empty();
77
78 mask.add(logical_core_id);
79
80 mask
81 }
82
83 /// Creates an affinity mask from a slice of core IDs.
84 ///
85 /// # Arguments
86 ///
87 /// * `core_ids` - Slice of logical processor IDs to include
88 ///
89 /// # Example
90 ///
91 /// ```
92 /// use gdt_cpus::AffinityMask;
93 ///
94 /// let mask = AffinityMask::from_cores(&[0, 2, 4, 6]);
95 /// assert_eq!(mask.count(), 4);
96 /// assert!(mask.contains(0));
97 /// assert!(!mask.contains(1));
98 /// assert!(mask.contains(2));
99 /// ```
100 pub fn from_cores(core_ids: &[usize]) -> Self {
101 let mut mask = Self::empty();
102
103 for &id in core_ids {
104 mask.add(id);
105 }
106
107 mask
108 }
109
110 /// Adds a logical core to the mask.
111 ///
112 /// If the core is already in the mask, this is a no-op.
113 ///
114 /// # Arguments
115 ///
116 /// * `logical_core_id` - The logical processor ID to add
117 ///
118 /// # Example
119 ///
120 /// ```
121 /// use gdt_cpus::AffinityMask;
122 ///
123 /// let mut mask = AffinityMask::empty();
124 /// mask.add(0);
125 /// mask.add(1);
126 /// assert_eq!(mask.count(), 2);
127 /// ```
128 pub fn add(&mut self, logical_core_id: usize) {
129 let word_idx = logical_core_id / 64;
130 let bit_idx = logical_core_id % 64;
131
132 if word_idx >= self.bits.len() {
133 self.bits.resize(word_idx + 1, 0);
134 }
135
136 self.bits[word_idx] |= 1u64 << bit_idx;
137 }
138
139 /// Removes a logical core from the mask.
140 ///
141 /// If the core is not in the mask, this is a no-op.
142 ///
143 /// # Arguments
144 ///
145 /// * `logical_core_id` - The logical processor ID to remove
146 ///
147 /// # Example
148 ///
149 /// ```
150 /// use gdt_cpus::AffinityMask;
151 ///
152 /// let mut mask = AffinityMask::from_cores(&[0, 1, 2]);
153 /// mask.remove(1);
154 /// assert_eq!(mask.count(), 2);
155 /// assert!(!mask.contains(1));
156 /// ```
157 pub fn remove(&mut self, logical_core_id: usize) {
158 let word_idx = logical_core_id / 64;
159 let bit_idx = logical_core_id % 64;
160
161 if word_idx < self.bits.len() {
162 self.bits[word_idx] &= !(1u64 << bit_idx);
163 }
164 }
165
166 /// Checks if a logical core is in the mask.
167 ///
168 /// # Arguments
169 ///
170 /// * `logical_core_id` - The logical processor ID to check
171 ///
172 /// # Returns
173 ///
174 /// `true` if the core is in the mask, `false` otherwise.
175 ///
176 /// # Example
177 ///
178 /// ```
179 /// use gdt_cpus::AffinityMask;
180 ///
181 /// let mask = AffinityMask::from_cores(&[0, 2, 4]);
182 /// assert!(mask.contains(0));
183 /// assert!(!mask.contains(1));
184 /// assert!(mask.contains(2));
185 /// ```
186 pub fn contains(&self, logical_core_id: usize) -> bool {
187 let word_idx = logical_core_id / 64;
188 let bit_idx = logical_core_id % 64;
189
190 if word_idx >= self.bits.len() {
191 return false;
192 }
193
194 (self.bits[word_idx] & (1u64 << bit_idx)) != 0
195 }
196
197 /// Returns the number of cores in the mask.
198 ///
199 /// # Example
200 ///
201 /// ```
202 /// use gdt_cpus::AffinityMask;
203 ///
204 /// let mask = AffinityMask::from_cores(&[0, 1, 2, 3]);
205 /// assert_eq!(mask.count(), 4);
206 /// ```
207 pub fn count(&self) -> usize {
208 self.bits.iter().map(|w| w.count_ones() as usize).sum()
209 }
210
211 /// Returns `true` if the mask is empty (no cores set).
212 ///
213 /// # Example
214 ///
215 /// ```
216 /// use gdt_cpus::AffinityMask;
217 ///
218 /// let empty = AffinityMask::empty();
219 /// assert!(empty.is_empty());
220 ///
221 /// let non_empty = AffinityMask::single(0);
222 /// assert!(!non_empty.is_empty());
223 /// ```
224 pub fn is_empty(&self) -> bool {
225 self.bits.iter().all(|&w| w == 0)
226 }
227
228 /// Returns an iterator over the core IDs in the mask.
229 ///
230 /// # Example
231 ///
232 /// ```
233 /// use gdt_cpus::AffinityMask;
234 ///
235 /// let mask = AffinityMask::from_cores(&[1, 3, 5]);
236 /// let cores: Vec<usize> = mask.iter().collect();
237 /// assert_eq!(cores, vec![1, 3, 5]);
238 /// ```
239 pub fn iter(&self) -> impl Iterator<Item = usize> + '_ {
240 self.bits.iter().enumerate().flat_map(|(word_idx, &word)| {
241 (0..64).filter_map(move |bit_idx| {
242 if (word & (1u64 << bit_idx)) != 0 {
243 Some(word_idx * 64 + bit_idx)
244 } else {
245 None
246 }
247 })
248 })
249 }
250
251 /// Returns the union of this mask with another.
252 ///
253 /// The resulting mask contains all cores that are in either mask.
254 ///
255 /// # Example
256 ///
257 /// ```
258 /// use gdt_cpus::AffinityMask;
259 ///
260 /// let a = AffinityMask::from_cores(&[0, 1]);
261 /// let b = AffinityMask::from_cores(&[1, 2]);
262 /// let union = a.union(&b);
263 /// assert_eq!(union.count(), 3);
264 /// assert!(union.contains(0));
265 /// assert!(union.contains(1));
266 /// assert!(union.contains(2));
267 /// ```
268 pub fn union(&self, other: &AffinityMask) -> AffinityMask {
269 let max_len = self.bits.len().max(other.bits.len());
270 let mut bits = vec![0u64; max_len];
271
272 for (i, &word) in self.bits.iter().enumerate() {
273 bits[i] |= word;
274 }
275 for (i, &word) in other.bits.iter().enumerate() {
276 bits[i] |= word;
277 }
278
279 AffinityMask { bits }
280 }
281
282 /// Returns the intersection of this mask with another.
283 ///
284 /// The resulting mask contains only cores that are in both masks.
285 ///
286 /// # Example
287 ///
288 /// ```
289 /// use gdt_cpus::AffinityMask;
290 ///
291 /// let a = AffinityMask::from_cores(&[0, 1, 2]);
292 /// let b = AffinityMask::from_cores(&[1, 2, 3]);
293 /// let intersection = a.intersection(&b);
294 /// assert_eq!(intersection.count(), 2);
295 /// assert!(intersection.contains(1));
296 /// assert!(intersection.contains(2));
297 /// ```
298 pub fn intersection(&self, other: &AffinityMask) -> AffinityMask {
299 let min_len = self.bits.len().min(other.bits.len());
300 let bits: Vec<u64> = self.bits[..min_len]
301 .iter()
302 .zip(other.bits[..min_len].iter())
303 .map(|(&a, &b)| a & b)
304 .collect();
305
306 AffinityMask { bits }
307 }
308
309 /// Returns the first 64 cores as a raw `u64` bitmask.
310 ///
311 /// This is useful for platform APIs that only support 64 cores.
312 /// Cores beyond index 63 are not included.
313 ///
314 /// # Example
315 ///
316 /// ```
317 /// use gdt_cpus::AffinityMask;
318 ///
319 /// let mask = AffinityMask::from_cores(&[0, 1, 63]);
320 /// let raw = mask.as_raw_u64();
321 /// assert_eq!(raw, 0x8000_0000_0000_0003);
322 /// ```
323 pub fn as_raw_u64(&self) -> u64 {
324 self.bits.first().copied().unwrap_or(0)
325 }
326
327 /// Returns the raw bits as a slice.
328 ///
329 /// Each element represents 64 cores: `bits[0]` = cores 0-63,
330 /// `bits[1]` = cores 64-127, etc.
331 pub fn as_raw_bits(&self) -> &[u64] {
332 &self.bits
333 }
334}
335
336impl std::fmt::Debug for AffinityMask {
337 fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
338 let cores: Vec<usize> = self.iter().collect();
339
340 f.debug_struct("AffinityMask")
341 .field("cores", &cores)
342 .field("count", &cores.len())
343 .finish()
344 }
345}
346
347impl std::fmt::Display for AffinityMask {
348 fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
349 let cores: Vec<usize> = self.iter().collect();
350
351 if cores.is_empty() {
352 write!(f, "AffinityMask(empty)")
353 } else {
354 write!(f, "AffinityMask({:?})", cores)
355 }
356 }
357}
358
359impl FromIterator<usize> for AffinityMask {
360 fn from_iter<T: IntoIterator<Item = usize>>(iter: T) -> Self {
361 let mut mask = AffinityMask::empty();
362
363 for id in iter {
364 mask.add(id);
365 }
366
367 mask
368 }
369}
370
371impl IntoIterator for &AffinityMask {
372 type Item = usize;
373 type IntoIter = std::vec::IntoIter<usize>;
374
375 fn into_iter(self) -> Self::IntoIter {
376 self.iter().collect::<Vec<_>>().into_iter()
377 }
378}
379
380#[cfg(test)]
381mod tests {
382 use super::*;
383
384 #[test]
385 fn test_empty_mask() {
386 let mask = AffinityMask::empty();
387 assert!(mask.is_empty());
388 assert_eq!(mask.count(), 0);
389 assert!(!mask.contains(0));
390 }
391
392 #[test]
393 fn test_single_core() {
394 let mask = AffinityMask::single(5);
395 assert!(!mask.is_empty());
396 assert_eq!(mask.count(), 1);
397 assert!(mask.contains(5));
398 assert!(!mask.contains(4));
399 }
400
401 #[test]
402 fn test_from_cores() {
403 let mask = AffinityMask::from_cores(&[0, 2, 4, 6]);
404 assert_eq!(mask.count(), 4);
405 assert!(mask.contains(0));
406 assert!(!mask.contains(1));
407 assert!(mask.contains(2));
408 }
409
410 #[test]
411 fn test_add_remove() {
412 let mut mask = AffinityMask::empty();
413 mask.add(0);
414 mask.add(1);
415 assert_eq!(mask.count(), 2);
416
417 mask.remove(0);
418 assert_eq!(mask.count(), 1);
419 assert!(!mask.contains(0));
420 assert!(mask.contains(1));
421 }
422
423 #[test]
424 fn test_high_core_ids() {
425 let mut mask = AffinityMask::empty();
426 mask.add(0);
427 mask.add(64);
428 mask.add(128);
429
430 assert_eq!(mask.count(), 3);
431 assert!(mask.contains(0));
432 assert!(mask.contains(64));
433 assert!(mask.contains(128));
434 assert!(!mask.contains(63));
435 assert!(!mask.contains(65));
436 }
437
438 #[test]
439 fn test_iter() {
440 let mask = AffinityMask::from_cores(&[1, 3, 5, 64, 65]);
441 let cores: Vec<usize> = mask.iter().collect();
442 assert_eq!(cores, vec![1, 3, 5, 64, 65]);
443 }
444
445 #[test]
446 fn test_as_raw_u64() {
447 let mask = AffinityMask::from_cores(&[0, 1, 63]);
448 assert_eq!(mask.as_raw_u64(), 0x8000_0000_0000_0003);
449 }
450
451 #[test]
452 fn test_from_iterator() {
453 let mask: AffinityMask = vec![0, 2, 4].into_iter().collect();
454 assert_eq!(mask.count(), 3);
455 assert!(mask.contains(0));
456 assert!(mask.contains(2));
457 assert!(mask.contains(4));
458 }
459}