scx_utils/
cpumask.rs

1// Copyright (c) Meta Platforms, Inc. and affiliates.
2
3// This software may be used and distributed according to the terms of the
4// GNU General Public License version 2.
5
6//! # SCX Cpumask
7//!
8//! A crate that allows creating, reading, and manipulating cpumasks.
9//!
10//! Cpumask
11//! -------
12//!
13//! A Cpumask object is simply a BitVec of u64's, along with a series of helper
14//! functions for creating, manipulating, and reading these BitVec objects.
15//!
16//! Empty Cpumasks can be created directly, or they can be created from a
17//! hexadecimal string:
18//!
19//!```no_run
20//!     use scx_utils::Cpumask;
21//!     let all_zeroes = Cpumask::new();
22//!     let from_str_mask = Cpumask::from_str("0xf0");
23//!```
24//!
25//! The hexadecimal string also supports the special values "none" and "all",
26//! respectively to specify no CPU (empty mask) or all CPUs (full mask):
27//!
28//!```
29//!     use scx_utils::Cpumask;
30//!     let str = String::from("none");
31//!     let all_zeroes = Cpumask::from_str(&str);
32//!
33//!     let str = String::from("all");
34//!     let all_ones = Cpumask::from_str(&str);
35//!```
36//!
37//! A Cpumask can be queried and updated using its helper functions:
38//!
39//!```rust
40//!     use log::info;
41//!     use scx_utils::Cpumask;
42//!     let str = String::from("none");
43//!     let mut mask = Cpumask::from_str(&str).unwrap();
44//!     info!("{:#?}", mask); // 32:<11111111000000001111111100000000>
45//!     assert!(!mask.test_cpu(0));
46//!     mask.set_cpu(0);
47//!     assert!(mask.test_cpu(0));
48//!
49//!     mask.clear_all();
50//!     info!("{:#?}", mask); // 32:<00000000000000000000000000000000>
51//!     assert!(!mask.test_cpu(0));
52//!
53//!     mask.set_all();
54//!     info!("{:#?}", mask); // 32:<11111111111111111111111111111111>
55//!     assert!(mask.test_cpu(0));
56//!```
57
58use crate::NR_CPU_IDS;
59use anyhow::Context;
60use anyhow::Result;
61use anyhow::bail;
62use bitvec::prelude::*;
63use sscanf::sscanf;
64use std::fmt;
65use std::ops::BitAndAssign;
66use std::ops::BitOrAssign;
67use std::ops::BitXorAssign;
68
69#[cfg(any(test, feature = "testutils"))]
70thread_local! {
71    /// Per-thread override for Cpumask width. 0 means use *NR_CPU_IDS.
72    /// Thread-local so parallel test threads don't interfere.
73    static MASK_WIDTH_OVERRIDE: std::cell::Cell<usize> = const { std::cell::Cell::new(0) };
74}
75
76/// Return the effective Cpumask width: the test override if set, else *NR_CPU_IDS.
77fn mask_width() -> usize {
78    #[cfg(any(test, feature = "testutils"))]
79    {
80        let ovr = MASK_WIDTH_OVERRIDE.with(|c| c.get());
81        if ovr > 0 {
82            return ovr;
83        }
84    }
85    *NR_CPU_IDS
86}
87
88/// Override the Cpumask width for the current thread. All subsequent
89/// Cpumask::new(), from_str(), and related calls on this thread will use
90/// this width instead of NR_CPU_IDS. Set to 0 to restore the default.
91#[cfg(any(test, feature = "testutils"))]
92pub fn set_cpumask_test_width(width: usize) {
93    MASK_WIDTH_OVERRIDE.with(|c| c.set(width));
94}
95
96#[derive(Debug, Eq, Clone, Hash, Ord, PartialEq, PartialOrd)]
97pub struct Cpumask {
98    mask: BitVec<u64, Lsb0>,
99}
100
101impl Cpumask {
102    fn check_cpu(&self, cpu: usize) -> Result<()> {
103        if cpu >= mask_width() {
104            bail!("Invalid CPU {} passed, max {}", cpu, mask_width());
105        }
106
107        Ok(())
108    }
109
110    /// Build a new empty Cpumask object.
111    // Default is intentionally not implemented; new() is the established API.
112    #[allow(clippy::new_without_default)]
113    pub fn new() -> Cpumask {
114        Cpumask {
115            mask: bitvec![u64, Lsb0; 0; mask_width()],
116        }
117    }
118
119    /// Build a Cpumask object from a hexadecimal string.
120    // Not std::str::FromStr: callers use Cpumask::from_str without the trait
121    // in scope, and the error type is anyhow::Error.
122    #[allow(clippy::should_implement_trait)]
123    pub fn from_str(cpumask: &str) -> Result<Cpumask> {
124        match cpumask {
125            "none" => {
126                let mask = bitvec![u64, Lsb0; 0; mask_width()];
127                return Ok(Self { mask });
128            }
129            "all" => {
130                let mask = bitvec![u64, Lsb0; 1; mask_width()];
131                return Ok(Self { mask });
132            }
133            _ => {}
134        }
135        let hex_str = {
136            let mut tmp_str = cpumask
137                .strip_prefix("0x")
138                .unwrap_or(cpumask)
139                .replace('_', "");
140            if !tmp_str.len().is_multiple_of(2) {
141                tmp_str = "0".to_string() + &tmp_str;
142            }
143            tmp_str
144        };
145        let byte_vec =
146            hex::decode(&hex_str).with_context(|| format!("Failed to parse cpumask: {cpumask}"))?;
147
148        let mut mask = bitvec![u64, Lsb0; 0; mask_width()];
149        for (index, &val) in byte_vec.iter().rev().enumerate() {
150            let mut v = val;
151            while v != 0 {
152                let lsb = v.trailing_zeros() as usize;
153                v &= !(1 << lsb);
154                let cpu = index * 8 + lsb;
155                if cpu >= mask_width() {
156                    bail!(
157                        concat!(
158                            "Found cpu ({}) in cpumask ({}) which is larger",
159                            " than the number of cpus on the machine ({})"
160                        ),
161                        cpu,
162                        cpumask,
163                        mask_width()
164                    );
165                }
166                mask.set(cpu, true);
167            }
168        }
169
170        Ok(Self { mask })
171    }
172
173    pub fn from_cpulist(cpulist: &str) -> Result<Cpumask> {
174        let mut mask = Cpumask::new();
175        for cpu_id in read_cpulist(cpulist)? {
176            let _ = mask.set_cpu(cpu_id);
177        }
178
179        Ok(mask)
180    }
181
182    /// Format the Cpumask as a compact CPU list string like "0-7,16-23".
183    /// Returns "none" if no CPUs are set.
184    pub fn to_cpulist(&self) -> String {
185        let cpus: Vec<usize> = self.iter().collect();
186        if cpus.is_empty() {
187            return String::from("none");
188        }
189
190        let mut ranges = Vec::new();
191        let mut start = cpus[0];
192        let mut end = cpus[0];
193
194        for &cpu in &cpus[1..] {
195            if cpu == end + 1 {
196                end = cpu;
197            } else {
198                ranges.push(if start == end {
199                    format!("{}", start)
200                } else {
201                    format!("{}-{}", start, end)
202                });
203                start = cpu;
204                end = cpu;
205            }
206        }
207
208        ranges.push(if start == end {
209            format!("{}", start)
210        } else {
211            format!("{}-{}", start, end)
212        });
213
214        ranges.join(",")
215    }
216
217    pub fn from_vec(vec: Vec<u64>) -> Self {
218        Self {
219            mask: BitVec::from_vec(vec),
220        }
221    }
222
223    pub fn from_bitvec(bitvec: BitVec<u64, Lsb0>) -> Self {
224        Self { mask: bitvec }
225    }
226
227    /// Return a slice of u64's whose bits reflect the Cpumask.
228    pub fn as_raw_slice(&self) -> &[u64] {
229        self.mask.as_raw_slice()
230    }
231
232    /// Return the mutable raw BitVec object backing the Cpumask.
233    pub fn as_raw_bitvec_mut(&mut self) -> &mut BitVec<u64, Lsb0> {
234        &mut self.mask
235    }
236
237    /// Return the raw BitVec object backing the Cpumask.
238    pub fn as_raw_bitvec(&self) -> &BitVec<u64, Lsb0> {
239        &self.mask
240    }
241
242    /// Set all bits in the Cpumask to 1
243    pub fn set_all(&mut self) {
244        self.mask.fill(true);
245    }
246
247    /// Set all bits in the Cpumask to 0
248    pub fn clear_all(&mut self) {
249        self.mask.fill(false);
250    }
251
252    /// Set a bit in the Cpumask. Returns an error if the specified CPU exceeds
253    /// the size of the Cpumask.
254    pub fn set_cpu(&mut self, cpu: usize) -> Result<()> {
255        self.check_cpu(cpu)?;
256        self.mask.set(cpu, true);
257        Ok(())
258    }
259
260    /// Clear a bit from the Cpumask. Returns an error if the specified CPU
261    /// exceeds the size of the Cpumask.
262    pub fn clear_cpu(&mut self, cpu: usize) -> Result<()> {
263        self.check_cpu(cpu)?;
264        self.mask.set(cpu, false);
265        Ok(())
266    }
267
268    /// Test whether the specified CPU bit is set in the Cpumask. If the CPU
269    /// exceeds the number of possible CPUs on the host, false is returned.
270    pub fn test_cpu(&self, cpu: usize) -> bool {
271        match self.mask.get(cpu) {
272            Some(bit) => *bit,
273            None => false,
274        }
275    }
276
277    /// Count the number of bits set in the Cpumask.
278    pub fn weight(&self) -> usize {
279        self.mask.count_ones()
280    }
281
282    /// Return true if the Cpumask has no bit set, false otherwise.
283    pub fn is_empty(&self) -> bool {
284        self.mask.count_ones() == 0
285    }
286
287    /// Return true if the Cpumask has all bits set, false otherwise.
288    pub fn is_full(&self) -> bool {
289        self.mask.count_ones() == mask_width()
290    }
291
292    /// The total size of the cpumask.
293    pub fn len(&self) -> usize {
294        mask_width()
295    }
296
297    /// Create a Cpumask that is the negation of the current Cpumask.
298    pub fn not(&self) -> Cpumask {
299        let mut new = self.clone();
300        new.mask = !new.mask;
301        new
302    }
303
304    /// Create a Cpumask that is the AND of the current Cpumask and another.
305    pub fn and(&self, other: &Cpumask) -> Cpumask {
306        let mut new = self.clone();
307        new.mask &= other.mask.clone();
308        new
309    }
310
311    /// Create a Cpumask that is the OR of the current Cpumask and another.
312    pub fn or(&self, other: &Cpumask) -> Cpumask {
313        let mut new = self.clone();
314        new.mask |= other.mask.clone();
315        new
316    }
317
318    /// Create a Cpumask that is the XOR of the current Cpumask and another.
319    pub fn xor(&self, other: &Cpumask) -> Cpumask {
320        let mut new = self.clone();
321        new.mask ^= other.mask.clone();
322        new
323    }
324
325    /// Iterate over each element of a Cpumask, and return the indices with bits
326    /// set.
327    ///
328    /// # Examples
329    ///
330    /// ```rust
331    /// use log::info;
332    /// use scx_utils::Cpumask;
333    /// let str = String::from("all");
334    /// let mask = Cpumask::from_str(&str).unwrap();
335    /// for cpu in mask.iter() {
336    ///     info!("cpu {} was set", cpu);
337    /// }
338    /// ```
339    pub fn iter(&self) -> CpumaskIterator<'_> {
340        CpumaskIterator {
341            mask: self,
342            index: 0,
343        }
344    }
345
346    /// Write out a CPU mask to a raw memory pointer. We normally use this as part of updating
347    /// the CPU masks on the BPF side.
348    ///
349    /// # Safety
350    ///
351    /// The caller must enforce that the bpfptr is valid and points to a valid region of
352    /// writable BPF memory that represents a CPU mask. This function is not thread safe.
353    pub unsafe fn write_to_ptr(&self, bpfptr: *mut u64, len: usize) -> Result<()> {
354        let cpumask_slice = self.as_raw_slice();
355        if len != cpumask_slice.len() {
356            bail!(
357                "BPF CPU mask has length {} u64s, Cpumask size is {}",
358                len,
359                cpumask_slice.len()
360            );
361        }
362
363        let ptr = bpfptr as *mut [u64; 64];
364        let bpfmask: &mut [u64; 64] = unsafe { &mut *ptr };
365        let (left, _) = bpfmask.split_at_mut(cpumask_slice.len());
366        left.clone_from_slice(cpumask_slice);
367
368        Ok(())
369    }
370
371    fn fmt_with(&self, f: &mut fmt::Formatter<'_>, case: char) -> fmt::Result {
372        let mut masks: Vec<u32> = self
373            .as_raw_slice()
374            .iter()
375            .flat_map(|x| [*x as u32, (x >> 32) as u32])
376            .collect();
377
378        // Throw out possible stray from u64 -> u32.
379        masks.truncate((mask_width()).div_ceil(32));
380
381        // Print the highest 32bit. Trim digits beyond mask_width().
382        let width = match (mask_width()).div_ceil(4) % 8 {
383            0 => 8,
384            v => v,
385        };
386        match case {
387            'x' => write!(f, "{:0width$x}", masks.pop().unwrap(), width = width)?,
388            'X' => write!(f, "{:0width$X}", masks.pop().unwrap(), width = width)?,
389            _ => unreachable!(),
390        }
391
392        // The rest in descending order.
393        for submask in masks.iter().rev() {
394            match case {
395                'x' => write!(f, ",{submask:08x}")?,
396                'X' => write!(f, ",{submask:08X}")?,
397                _ => unreachable!(),
398            }
399        }
400        Ok(())
401    }
402}
403
404pub fn read_cpulist(cpulist: &str) -> Result<Vec<usize>> {
405    let cpulist = cpulist.trim_end_matches('\0');
406    let cpu_groups: Vec<&str> = cpulist.split(',').collect();
407    let mut cpu_ids = vec![];
408    for group in cpu_groups.iter() {
409        let (min, max) = match sscanf!(group.trim(), "{usize}-{usize}") {
410            Some((x, y)) => (x, y),
411            None => match sscanf!(group.trim(), "{usize}") {
412                Some(x) => (x, x),
413                None => {
414                    bail!("Failed to parse cpulist {}", group.trim());
415                }
416            },
417        };
418        for i in min..(max + 1) {
419            cpu_ids.push(i);
420        }
421    }
422
423    Ok(cpu_ids)
424}
425
426pub struct CpumaskIterator<'a> {
427    mask: &'a Cpumask,
428    index: usize,
429}
430
431impl Iterator for CpumaskIterator<'_> {
432    type Item = usize;
433
434    fn next(&mut self) -> Option<Self::Item> {
435        while self.index < mask_width() {
436            let index = self.index;
437            self.index += 1;
438            let bit_val = self.mask.test_cpu(index);
439            if bit_val {
440                return Some(index);
441            }
442        }
443
444        None
445    }
446}
447
448impl fmt::Display for Cpumask {
449    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
450        self.fmt_with(f, 'x')
451    }
452}
453
454impl fmt::LowerHex for Cpumask {
455    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
456        self.fmt_with(f, 'x')
457    }
458}
459
460impl fmt::UpperHex for Cpumask {
461    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
462        self.fmt_with(f, 'X')
463    }
464}
465
466impl BitAndAssign<&Self> for Cpumask {
467    fn bitand_assign(&mut self, rhs: &Self) {
468        self.mask &= &rhs.mask;
469    }
470}
471
472impl BitOrAssign<&Self> for Cpumask {
473    fn bitor_assign(&mut self, rhs: &Self) {
474        self.mask |= &rhs.mask;
475    }
476}
477
478impl BitXorAssign<&Self> for Cpumask {
479    fn bitxor_assign(&mut self, rhs: &Self) {
480        self.mask ^= &rhs.mask;
481    }
482}
483
484#[cfg(test)]
485mod tests {
486    use super::*;
487
488    #[test]
489    fn test_to_cpulist_empty() {
490        let mask = Cpumask::new();
491        assert_eq!(mask.to_cpulist(), "none");
492    }
493
494    #[test]
495    fn test_to_cpulist_single_cpu() {
496        set_cpumask_test_width(16);
497        let mut mask = Cpumask::new();
498        mask.set_cpu(5).unwrap();
499        assert_eq!(mask.to_cpulist(), "5");
500    }
501
502    #[test]
503    fn test_to_cpulist_contiguous_range() {
504        set_cpumask_test_width(16);
505        let mut mask = Cpumask::new();
506        for cpu in 0..8 {
507            mask.set_cpu(cpu).unwrap();
508        }
509        assert_eq!(mask.to_cpulist(), "0-7");
510    }
511
512    #[test]
513    fn test_to_cpulist_multiple_ranges() {
514        set_cpumask_test_width(16);
515        let mut mask = Cpumask::new();
516        for cpu in 0..4 {
517            mask.set_cpu(cpu).unwrap();
518        }
519        for cpu in 8..12 {
520            mask.set_cpu(cpu).unwrap();
521        }
522        assert_eq!(mask.to_cpulist(), "0-3,8-11");
523    }
524
525    #[test]
526    fn test_to_cpulist_scattered() {
527        set_cpumask_test_width(16);
528        let mut mask = Cpumask::new();
529        mask.set_cpu(1).unwrap();
530        mask.set_cpu(3).unwrap();
531        mask.set_cpu(5).unwrap();
532        assert_eq!(mask.to_cpulist(), "1,3,5");
533    }
534
535    #[test]
536    fn test_to_cpulist_mixed() {
537        set_cpumask_test_width(16);
538        let mut mask = Cpumask::new();
539        mask.set_cpu(0).unwrap();
540        mask.set_cpu(1).unwrap();
541        mask.set_cpu(2).unwrap();
542        mask.set_cpu(5).unwrap();
543        mask.set_cpu(10).unwrap();
544        mask.set_cpu(11).unwrap();
545        assert_eq!(mask.to_cpulist(), "0-2,5,10-11");
546    }
547
548    #[test]
549    fn test_to_cpulist_roundtrip() {
550        set_cpumask_test_width(32);
551        let original = "0-3,8-11,16";
552        let mask = Cpumask::from_cpulist(original).unwrap();
553        assert_eq!(mask.to_cpulist(), original);
554    }
555}
scx_utils/cpumask.rs

scx_utils/
cpumask.rs
