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::bail;
60use anyhow::Context;
61use anyhow::Result;
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    pub fn new() -> Cpumask {
112        Cpumask {
113            mask: bitvec![u64, Lsb0; 0; mask_width()],
114        }
115    }
116
117    /// Build a Cpumask object from a hexadecimal string.
118    pub fn from_str(cpumask: &str) -> Result<Cpumask> {
119        match cpumask {
120            "none" => {
121                let mask = bitvec![u64, Lsb0; 0; mask_width()];
122                return Ok(Self { mask });
123            }
124            "all" => {
125                let mask = bitvec![u64, Lsb0; 1; mask_width()];
126                return Ok(Self { mask });
127            }
128            _ => {}
129        }
130        let hex_str = {
131            let mut tmp_str = cpumask
132                .strip_prefix("0x")
133                .unwrap_or(cpumask)
134                .replace('_', "");
135            if tmp_str.len() % 2 != 0 {
136                tmp_str = "0".to_string() + &tmp_str;
137            }
138            tmp_str
139        };
140        let byte_vec =
141            hex::decode(&hex_str).with_context(|| format!("Failed to parse cpumask: {cpumask}"))?;
142
143        let mut mask = bitvec![u64, Lsb0; 0; mask_width()];
144        for (index, &val) in byte_vec.iter().rev().enumerate() {
145            let mut v = val;
146            while v != 0 {
147                let lsb = v.trailing_zeros() as usize;
148                v &= !(1 << lsb);
149                let cpu = index * 8 + lsb;
150                if cpu >= mask_width() {
151                    bail!(
152                        concat!(
153                            "Found cpu ({}) in cpumask ({}) which is larger",
154                            " than the number of cpus on the machine ({})"
155                        ),
156                        cpu,
157                        cpumask,
158                        mask_width()
159                    );
160                }
161                mask.set(cpu, true);
162            }
163        }
164
165        Ok(Self { mask })
166    }
167
168    pub fn from_cpulist(cpulist: &str) -> Result<Cpumask> {
169        let mut mask = Cpumask::new();
170        for cpu_id in read_cpulist(cpulist)? {
171            let _ = mask.set_cpu(cpu_id);
172        }
173
174        Ok(mask)
175    }
176
177    pub fn from_vec(vec: Vec<u64>) -> Self {
178        Self {
179            mask: BitVec::from_vec(vec),
180        }
181    }
182
183    pub fn from_bitvec(bitvec: BitVec<u64, Lsb0>) -> Self {
184        Self { mask: bitvec }
185    }
186
187    /// Return a slice of u64's whose bits reflect the Cpumask.
188    pub fn as_raw_slice(&self) -> &[u64] {
189        self.mask.as_raw_slice()
190    }
191
192    /// Return the mutable raw BitVec object backing the Cpumask.
193    pub fn as_raw_bitvec_mut(&mut self) -> &mut BitVec<u64, Lsb0> {
194        &mut self.mask
195    }
196
197    /// Return the raw BitVec object backing the Cpumask.
198    pub fn as_raw_bitvec(&self) -> &BitVec<u64, Lsb0> {
199        &self.mask
200    }
201
202    /// Set all bits in the Cpumask to 1
203    pub fn set_all(&mut self) {
204        self.mask.fill(true);
205    }
206
207    /// Set all bits in the Cpumask to 0
208    pub fn clear_all(&mut self) {
209        self.mask.fill(false);
210    }
211
212    /// Set a bit in the Cpumask. Returns an error if the specified CPU exceeds
213    /// the size of the Cpumask.
214    pub fn set_cpu(&mut self, cpu: usize) -> Result<()> {
215        self.check_cpu(cpu)?;
216        self.mask.set(cpu, true);
217        Ok(())
218    }
219
220    /// Clear a bit from the Cpumask. Returns an error if the specified CPU
221    /// exceeds the size of the Cpumask.
222    pub fn clear_cpu(&mut self, cpu: usize) -> Result<()> {
223        self.check_cpu(cpu)?;
224        self.mask.set(cpu, false);
225        Ok(())
226    }
227
228    /// Test whether the specified CPU bit is set in the Cpumask. If the CPU
229    /// exceeds the number of possible CPUs on the host, false is returned.
230    pub fn test_cpu(&self, cpu: usize) -> bool {
231        match self.mask.get(cpu) {
232            Some(bit) => *bit,
233            None => false,
234        }
235    }
236
237    /// Count the number of bits set in the Cpumask.
238    pub fn weight(&self) -> usize {
239        self.mask.count_ones()
240    }
241
242    /// Return true if the Cpumask has no bit set, false otherwise.
243    pub fn is_empty(&self) -> bool {
244        self.mask.count_ones() == 0
245    }
246
247    /// Return true if the Cpumask has all bits set, false otherwise.
248    pub fn is_full(&self) -> bool {
249        self.mask.count_ones() == mask_width()
250    }
251
252    /// The total size of the cpumask.
253    pub fn len(&self) -> usize {
254        mask_width()
255    }
256
257    /// Create a Cpumask that is the negation of the current Cpumask.
258    pub fn not(&self) -> Cpumask {
259        let mut new = self.clone();
260        new.mask = !new.mask;
261        new
262    }
263
264    /// Create a Cpumask that is the AND of the current Cpumask and another.
265    pub fn and(&self, other: &Cpumask) -> Cpumask {
266        let mut new = self.clone();
267        new.mask &= other.mask.clone();
268        new
269    }
270
271    /// Create a Cpumask that is the OR of the current Cpumask and another.
272    pub fn or(&self, other: &Cpumask) -> Cpumask {
273        let mut new = self.clone();
274        new.mask |= other.mask.clone();
275        new
276    }
277
278    /// Create a Cpumask that is the XOR of the current Cpumask and another.
279    pub fn xor(&self, other: &Cpumask) -> Cpumask {
280        let mut new = self.clone();
281        new.mask ^= other.mask.clone();
282        new
283    }
284
285    /// Iterate over each element of a Cpumask, and return the indices with bits
286    /// set.
287    ///
288    /// # Examples
289    ///
290    /// ```rust
291    /// use log::info;
292    /// use scx_utils::Cpumask;
293    /// let str = String::from("all");
294    /// let mask = Cpumask::from_str(&str).unwrap();
295    /// for cpu in mask.iter() {
296    ///     info!("cpu {} was set", cpu);
297    /// }
298    /// ```
299    pub fn iter(&self) -> CpumaskIterator<'_> {
300        CpumaskIterator {
301            mask: self,
302            index: 0,
303        }
304    }
305
306    /// Write out a CPU mask to a raw memory pointer. We normally use this as part of updating
307    /// the CPU masks on the BPF side.
308    ///
309    /// # Safety
310    ///
311    /// The caller must enforce that the bpfptr is valid and points to a valid region of
312    /// writable BPF memory that represents a CPU mask. This function is not thread safe.
313    pub unsafe fn write_to_ptr(&self, bpfptr: *mut u64, len: usize) -> Result<()> {
314        let cpumask_slice = self.as_raw_slice();
315        if len != cpumask_slice.len() {
316            bail!(
317                "BPF CPU mask has length {} u64s, Cpumask size is {}",
318                len,
319                cpumask_slice.len()
320            );
321        }
322
323        let ptr = bpfptr as *mut [u64; 64];
324        let bpfmask: &mut [u64; 64] = unsafe { &mut *ptr };
325        let (left, _) = bpfmask.split_at_mut(cpumask_slice.len());
326        left.clone_from_slice(cpumask_slice);
327
328        Ok(())
329    }
330
331    fn fmt_with(&self, f: &mut fmt::Formatter<'_>, case: char) -> fmt::Result {
332        let mut masks: Vec<u32> = self
333            .as_raw_slice()
334            .iter()
335            .flat_map(|x| [*x as u32, (x >> 32) as u32])
336            .collect();
337
338        // Throw out possible stray from u64 -> u32.
339        masks.truncate((mask_width()).div_ceil(32));
340
341        // Print the highest 32bit. Trim digits beyond mask_width().
342        let width = match (mask_width()).div_ceil(4) % 8 {
343            0 => 8,
344            v => v,
345        };
346        match case {
347            'x' => write!(f, "{:0width$x}", masks.pop().unwrap(), width = width)?,
348            'X' => write!(f, "{:0width$X}", masks.pop().unwrap(), width = width)?,
349            _ => unreachable!(),
350        }
351
352        // The rest in descending order.
353        for submask in masks.iter().rev() {
354            match case {
355                'x' => write!(f, ",{submask:08x}")?,
356                'X' => write!(f, ",{submask:08X}")?,
357                _ => unreachable!(),
358            }
359        }
360        Ok(())
361    }
362}
363
364pub fn read_cpulist(cpulist: &str) -> Result<Vec<usize>> {
365    let cpulist = cpulist.trim_end_matches('\0');
366    let cpu_groups: Vec<&str> = cpulist.split(',').collect();
367    let mut cpu_ids = vec![];
368    for group in cpu_groups.iter() {
369        let (min, max) = match sscanf!(group.trim(), "{usize}-{usize}") {
370            Ok((x, y)) => (x, y),
371            Err(_) => match sscanf!(group.trim(), "{usize}") {
372                Ok(x) => (x, x),
373                Err(_) => {
374                    bail!("Failed to parse cpulist {}", group.trim());
375                }
376            },
377        };
378        for i in min..(max + 1) {
379            cpu_ids.push(i);
380        }
381    }
382
383    Ok(cpu_ids)
384}
385
386pub struct CpumaskIterator<'a> {
387    mask: &'a Cpumask,
388    index: usize,
389}
390
391impl Iterator for CpumaskIterator<'_> {
392    type Item = usize;
393
394    fn next(&mut self) -> Option<Self::Item> {
395        while self.index < mask_width() {
396            let index = self.index;
397            self.index += 1;
398            let bit_val = self.mask.test_cpu(index);
399            if bit_val {
400                return Some(index);
401            }
402        }
403
404        None
405    }
406}
407
408impl fmt::Display for Cpumask {
409    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
410        self.fmt_with(f, 'x')
411    }
412}
413
414impl fmt::LowerHex for Cpumask {
415    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
416        self.fmt_with(f, 'x')
417    }
418}
419
420impl fmt::UpperHex for Cpumask {
421    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
422        self.fmt_with(f, 'X')
423    }
424}
425
426impl BitAndAssign<&Self> for Cpumask {
427    fn bitand_assign(&mut self, rhs: &Self) {
428        self.mask &= &rhs.mask;
429    }
430}
431
432impl BitOrAssign<&Self> for Cpumask {
433    fn bitor_assign(&mut self, rhs: &Self) {
434        self.mask |= &rhs.mask;
435    }
436}
437
438impl BitXorAssign<&Self> for Cpumask {
439    fn bitxor_assign(&mut self, rhs: &Self) {
440        self.mask ^= &rhs.mask;
441    }
442}
scx_utils/cpumask.rs

scx_utils/
cpumask.rs
