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

scx_utils/
cpumask.rs
