Add get_online_cpus() utiilty function (#39)

* Add get_online_cpus and tests for it

* Use get_online_cpus() in right context

* Replace libbpf_rs::num_possible_cpus with get_online_cpus everywhere

* Use possible CPUs where needed - better comments

* Return errors instead of panic'ing

* Correct error returns for anyhow context

Cargo.toml

@@ -26,12 +26,13 @@ tracing-subscriber = "0.3.18"
 chrono = "0.4.37"
 inferno = "0.11.19"
 primal = "0.3.2"
-nix = { version = "0.28.0", features = ["user"]}
+nix = { version = "0.28.0", features = ["user"] }
 
 [dev-dependencies]
 assert_cmd = { version = "2.0.14" }
 insta = { version = "1.38.0", features = ["yaml"] }
-rstest = "0.18.2"
+rstest = "0.19.0"
+tempdir = "0.3.7"
 
 [build-dependencies]
 bindgen = "0.69.4"

src/profiler.rs

@@ -22,7 +22,7 @@ use crate::collector::*;
 use crate::object::{BuildId, ObjectFile};
 use crate::perf_events::setup_perf_event;
 use crate::unwind_info::{in_memory_unwind_info, remove_redundant, remove_unnecesary_markers};
-use crate::util::summarize_address_range;
+use crate::util::{get_online_cpus, summarize_address_range};
 
 pub enum TracerEvent {
     ProcessExit(i32),
@@ -234,7 +234,9 @@ impl Profiler<'_> {
     }
 
     pub fn run(mut self, collector: Arc<Mutex<Collector>>) {
-        let num_cpus = num_possible_cpus().expect("get possible CPUs") as u64;
+        // In this case, we only want to calculate maximum sampling buffer sizes based on the
+        // number of online CPUs, NOT possible CPUs, when they differ - which is often.
+        let num_cpus = get_online_cpus().expect("get online CPUs").len() as u64;
         let max_samples_per_session =
             self.sample_freq as u64 * num_cpus * self.session_duration.as_secs();
         if max_samples_per_session >= MAX_AGGREGATED_STACKS_ENTRIES.into() {
@@ -465,6 +467,9 @@ impl Profiler<'_> {
         let value = unsafe { plain::as_bytes(&default) };
 
         let mut values: Vec<Vec<u8>> = Vec::new();
+        // This is a place where you need to know the POSSIBLE, not ONLINE CPUs, because eBPF's
+        // internals require setting up certain buffers for all possible CPUs, even if the CPUs
+        // don't all exist.
         let num_cpus = num_possible_cpus().expect("get possible CPUs") as u64;
         for _ in 0..num_cpus {
             values.push(value.to_vec());
@@ -1107,7 +1112,7 @@ impl Profiler<'_> {
 
     pub fn setup_perf_events(&mut self) {
         let mut prog_fds = Vec::new();
-        for i in 0..num_possible_cpus().expect("get possible CPUs") {
+        for i in get_online_cpus().expect("get online CPUs") {
             let perf_fd =
                 unsafe { setup_perf_event(i.try_into().unwrap(), self.sample_freq as u64) }
                     .expect("setup perf event");

src/util.rs

@@ -1,3 +1,6 @@
+use anyhow::{Context, Error};
+use std::{fs::File, io::Read};
+
 #[derive(Debug, PartialEq)]
 pub struct AddressBlockRange {
     pub addr: u64,
@@ -33,9 +36,48 @@ pub fn summarize_address_range(low: u64, high: u64) -> Vec<AddressBlockRange> {
     res
 }
 
+fn _read_cpu_range(path: &str) -> Result<Vec<u32>, Error> {
+    let mut cpus: Vec<_> = vec![];
+    let mut fh = File::open(path)?;
+    let mut cpu_range_str = String::new();
+    fh.read_to_string(&mut cpu_range_str)?;
+
+    for cpu_range in cpu_range_str.split(',') {
+        let rangeop_result = cpu_range.find('-');
+        match rangeop_result {
+            None => cpus.push(
+                cpu_range
+                    .trim_end()
+                    .parse::<u32>()
+                    .with_context(|| "Failed to parse lone CPU".to_string())?,
+            ),
+            Some(index) => {
+                let start = cpu_range[..index]
+                    .trim_end()
+                    .parse::<u32>()
+                    .with_context(|| "Failed to parse starting CPU".to_string())?;
+                let end = cpu_range[index + 1..]
+                    .trim_end()
+                    .parse::<u32>()
+                    .with_context(|| "Failed to parse ending CPU".to_string())?;
+                cpus.extend(start..end + 1);
+            }
+        }
+    }
+
+    Ok(cpus)
+}
+
+pub fn get_online_cpus() -> Result<Vec<u32>, Error> {
+    let cpus: Vec<u32> = _read_cpu_range("/sys/devices/system/cpu/online")?;
+
+    Ok(cpus)
+}
+
 #[cfg(test)]
 mod tests {
     use std::mem::size_of;
+    use tempdir::TempDir;
 
     use libbpf_rs::libbpf_sys;
     use libbpf_rs::MapFlags;
@@ -155,4 +197,50 @@ mod tests {
             assert_eq!(parsed.executable_id, mapping2.executable_id);
         }
     }
+
+    #[test]
+    fn cpu_ranges_to_list() {
+        use std::io::Seek;
+        use std::io::Write;
+
+        let tmp_dir = TempDir::new("cpu_devs").unwrap();
+        let file_path = tmp_dir.path().join("online");
+        let mut tmp_file = File::create(file_path.clone()).unwrap();
+        let file_str = file_path.to_str().unwrap();
+
+        writeln!(tmp_file, "0").unwrap();
+        let cpus = _read_cpu_range(file_str).unwrap();
+        assert_eq!(cpus, vec![0]);
+
+        tmp_file.rewind().unwrap();
+        writeln!(tmp_file, "0-7").unwrap();
+
+        let cpus = _read_cpu_range(file_str).unwrap();
+        assert_eq!(cpus, (0..=7).collect::<Vec<_>>());
+
+        tmp_file.rewind().unwrap();
+        writeln!(tmp_file, "0-7,16-23").unwrap();
+
+        let cpus = _read_cpu_range(file_str).unwrap();
+        let expected = (0..=7).chain(16..=23).collect::<Vec<_>>();
+
+        assert_eq!(cpus, expected);
+
+        tmp_file.rewind().unwrap();
+        writeln!(tmp_file, "0-1,3,7-9,48,49").unwrap();
+
+        let cpus = _read_cpu_range(file_str).unwrap();
+        assert_eq!(
+            cpus,
+            (0..=1)
+                .chain(3..=3)
+                .chain(7..=9)
+                .chain(48..=48)
+                .chain(49..=49)
+                .collect::<Vec<_>>()
+        );
+
+        drop(tmp_file);
+        tmp_dir.close().expect("tempdir should be destroyed");
+    }
 }