more changes

flake.nix

@@ -37,6 +37,7 @@
             glibc
             glibc.static
             protobuf
+            openssl
           ];
           nativeBuildInputs = with pkgs; [
             pkg-config
@@ -74,7 +75,6 @@
               # Debugging
               strace
               gdb
-              openssl
               # Other tools
               skopeo
               cargo-edit

proto/src/profile.rs

@@ -5,13 +5,11 @@ pub mod pprof {
 }
 
 use anyhow::{anyhow, Result};
-use pprof::Label;
 use std::collections::hash_map::Entry;
 use std::collections::HashMap;
 
-pub struct Pprof {
+pub struct PprofBuilder {
     freq_in_hz: i64,
-    function_idx: u64,
 
     known_mappings: HashMap<u64, u64>,
     mappings: Vec<pprof::Mapping>,
@@ -30,14 +28,14 @@ pub struct Pprof {
 
 pub enum LabelStringOrNumber {
     String(String),
+    /// Value and unit.
     Number(i64, String),
 }
 
-impl Pprof {
-    pub fn new() -> Self {
+impl PprofBuilder {
+    pub fn with_frequency(freq_in_hz: i64) -> Self {
         Self {
-            freq_in_hz: 27,
-            function_idx: 1,
+            freq_in_hz,
 
             known_mappings: HashMap::new(),
             mappings: Vec::new(),
@@ -337,7 +335,7 @@ mod tests {
 
     #[test]
     fn test_string_table() {
-        let mut pprof = Pprof::new();
+        let mut pprof = PprofBuilder::with_freqency(27);
         assert_eq!(pprof.get_or_insert_string("hi"), 1);
         assert_eq!(pprof.get_or_insert_string("salut"), 2);
         assert_eq!(pprof.string_table, vec!["", "hi", "salut"]);
@@ -350,7 +348,7 @@ mod tests {
 
     #[test]
     fn test_mappings() {
-        let mut pprof = Pprof::new();
+        let mut pprof = PprofBuilder::with_freqency(27);
         assert_eq!(
             pprof.add_mapping(0, 0x100, 0x200, 0x0, "file.so", "sha256-abc"),
             1
@@ -368,7 +366,7 @@ mod tests {
 
     #[test]
     fn test_locations() {
-        let mut pprof = Pprof::new();
+        let mut pprof = PprofBuilder::with_freqency(27);
         let _ = pprof.add_line("hahahaha-first-line");
         let (line, function_id) = pprof.add_line("test-line");
 
@@ -395,7 +393,7 @@ mod tests {
 
     #[test]
     fn test_sample() {
-        let mut pprof = Pprof::new();
+        let mut pprof = PprofBuilder::with_freqency(27);
         let labels = vec![
             pprof.new_label("key", LabelStringOrNumber::String("value".into())),
             pprof.new_label("key", LabelStringOrNumber::Number(123, "pid".into())),
@@ -431,7 +429,7 @@ mod tests {
 
         let mut rng = rand::thread_rng();
 
-        let mut pprof = Pprof::new();
+        let mut pprof = PprofBuilder::with_freqency(27);
         // Let's say we have this profile
         let raw_samples = vec![
             (vec![123], 200),

src/collector.rs

@@ -25,19 +25,49 @@ pub trait Collector {
     );
 }
 
-pub type ThreadSafeCollector = Arc<Mutex<dyn Collector + Send>>;
+pub type ThreadSafeCollector = Arc<Mutex<Box<dyn Collector + Send>>>;
 
+#[derive(Default)]
+pub struct NullCollector {
+    procs: HashMap<i32, ProcessInfo>,
+    objs: HashMap<ExecutableId, ObjectFileInfo>,
+}
+
+impl NullCollector {
+    pub fn new() -> Self {
+        Self::default()
+    }
+}
+
+impl Collector for NullCollector {
+    fn collect(
+        &mut self,
+        _profile: RawAggregatedProfile,
+        _procs: &HashMap<i32, ProcessInfo>,
+        _objs: &HashMap<ExecutableId, ObjectFileInfo>,
+    ) {
+    }
+
+    fn finish(
+        &self,
+    ) -> (
+        RawAggregatedProfile,
+        &HashMap<i32, ProcessInfo>,
+        &HashMap<ExecutableId, ObjectFileInfo>,
+    ) {
+        (RawAggregatedProfile::new(), &self.procs, &self.objs)
+    }
+}
+
+#[derive(Default)]
 pub struct StreamingCollector {
     procs: HashMap<i32, ProcessInfo>,
     objs: HashMap<ExecutableId, ObjectFileInfo>,
 }
 
 impl StreamingCollector {
-    pub fn new() -> ThreadSafeCollector {
-        Arc::new(Mutex::new(Self {
-            procs: HashMap::new(),
-            objs: HashMap::new(),
-        }))
+    pub fn new() -> Self {
+        Self::default()
     }
 }
 
@@ -72,19 +102,16 @@ impl Collector for StreamingCollector {
     }
 }
 
+#[derive(Default)]
 pub struct AggregatorCollector {
     profiles: Vec<RawAggregatedProfile>,
     procs: HashMap<i32, ProcessInfo>,
     objs: HashMap<ExecutableId, ObjectFileInfo>,
 }
 
 impl AggregatorCollector {
-    pub fn new() -> ThreadSafeCollector {
-        Arc::new(Mutex::new(Self {
-            profiles: Vec::new(),
-            procs: HashMap::new(),
-            objs: HashMap::new(),
-        }))
+    pub fn new() -> Self {
+        Self::default()
     }
 }
 

src/main.rs

@@ -5,12 +5,12 @@ use std::io::Write;
 use std::ops::RangeInclusive;
 use std::panic;
 use std::path::PathBuf;
+use std::sync::{Arc, Mutex};
 use std::time::Duration;
 
 use clap::Parser;
 use inferno::flamegraph;
-use lightswitch::collector::AggregatorCollector;
-use lightswitch::collector::StreamingCollector;
+use lightswitch::collector::{AggregatorCollector, Collector, NullCollector, StreamingCollector};
 use nix::unistd::Uid;
 use primal::is_prime;
 use prost::Message;
@@ -91,12 +91,12 @@ enum ProfileFormat {
     #[default]
     FlameGraph,
     Pprof,
-    /// Do not produce a profile. Used for kernel tests.
-    None,
 }
 
 #[derive(PartialEq, clap::ValueEnum, Debug, Clone, Default)]
 enum ProfileSender {
+    /// Discard the profile. Used for kernel tests.
+    None,
     #[default]
     LocalDisk,
     Remote,
@@ -220,10 +220,13 @@ fn main() -> Result<(), Box<dyn Error>> {
     }
 
     // TODO: change collector based on symbolizer type and whether continuous or one shot
-    let collector = match args.sender {
-        ProfileSender::LocalDisk => AggregatorCollector::new(),
-        ProfileSender::Remote => StreamingCollector::new(),
-    };
+    let collector = Arc::new(Mutex::new(match args.sender {
+        ProfileSender::None => Box::new(NullCollector::new()) as Box<dyn Collector + Send>,
+        ProfileSender::LocalDisk => {
+            Box::new(AggregatorCollector::new()) as Box<dyn Collector + Send>
+        }
+        ProfileSender::Remote => Box::new(StreamingCollector::new()) as Box<dyn Collector + Send>,
+    }));
 
     let mut p: Profiler<'_> = Profiler::new(
         args.libbpf_logs,
@@ -237,7 +240,7 @@ fn main() -> Result<(), Box<dyn Error>> {
     let collector = collector.lock().unwrap();
     let (raw_profile, procs, objs) = collector.finish();
 
-    // If we need to send the profile to the backend we are done.
+    // If we need to send the profile to the backend there's nothing else to do.
     if args.sender == ProfileSender::Remote {
         return Ok(());
     }
@@ -277,9 +280,6 @@ fn main() -> Result<(), Box<dyn Error>> {
                 }
             }
         }
-        ProfileFormat::None => {
-            // Do nothing
-        }
     }
 
     Ok(())
@@ -307,7 +307,7 @@ mod tests {
         let actual = String::from_utf8(cmd.unwrap().stdout).unwrap();
         insta::assert_yaml_snapshot!(actual, @r###"
         ---
-        "Usage: lightswitch [OPTIONS]\n\nOptions:\n      --pids <PIDS>\n          Specific PIDs to profile\n\n      --tids <TIDS>\n          Specific TIDs to profile (these can be outside the PIDs selected above)\n\n      --show-unwind-info <PATH_TO_BINARY>\n          Show unwind info for given binary\n\n      --show-info <PATH_TO_BINARY>\n          Show build ID for given binary\n\n  -D, --duration <DURATION>\n          How long this agent will run in seconds\n          \n          [default: 18446744073709551615]\n\n      --libbpf-logs\n          Enable libbpf logs. This includes the BPF verifier output\n\n      --bpf-logging\n          Enable BPF programs logging\n\n      --logging <LOGGING>\n          Set lightswitch's logging level\n          \n          [default: info]\n          [possible values: trace, debug, info, warn, error]\n\n      --sample-freq <SAMPLE_FREQ_IN_HZ>\n          Per-CPU Sampling Frequency in Hz\n          \n          [default: 19]\n\n      --profile-format <PROFILE_FORMAT>\n          Output file for Flame Graph in SVG format\n          \n          [default: flame-graph]\n\n          Possible values:\n          - flame-graph\n          - pprof\n          - none:        Do not produce a profile. Used for kernel tests\n\n      --profile-name <PROFILE_NAME>\n          Name for the generated profile\n\n      --sender <SENDER>\n          Where to store the profile. If remote is chosen [...]\n          \n          [default: local-disk]\n          [possible values: local-disk, remote]\n\n  -h, --help\n          Print help (see a summary with '-h')\n"
+        "Usage: lightswitch [OPTIONS]\n\nOptions:\n      --pids <PIDS>\n          Specific PIDs to profile\n\n      --tids <TIDS>\n          Specific TIDs to profile (these can be outside the PIDs selected above)\n\n      --show-unwind-info <PATH_TO_BINARY>\n          Show unwind info for given binary\n\n      --show-info <PATH_TO_BINARY>\n          Show build ID for given binary\n\n  -D, --duration <DURATION>\n          How long this agent will run in seconds\n          \n          [default: 18446744073709551615]\n\n      --libbpf-logs\n          Enable libbpf logs. This includes the BPF verifier output\n\n      --bpf-logging\n          Enable BPF programs logging\n\n      --logging <LOGGING>\n          Set lightswitch's logging level\n          \n          [default: info]\n          [possible values: trace, debug, info, warn, error]\n\n      --sample-freq <SAMPLE_FREQ_IN_HZ>\n          Per-CPU Sampling Frequency in Hz\n          \n          [default: 19]\n\n      --profile-format <PROFILE_FORMAT>\n          Output file for Flame Graph in SVG format\n          \n          [default: flame-graph]\n          [possible values: flame-graph, pprof]\n\n      --profile-name <PROFILE_NAME>\n          Name for the generated profile\n\n      --sender <SENDER>\n          Where to store the profile. If remote is chosen [...]\n          \n          [default: local-disk]\n\n          Possible values:\n          - none:       Discard the profile. Used for kernel tests\n          - local-disk\n          - remote\n\n  -h, --help\n          Print help (see a summary with '-h')\n"
         "###);
     }
 

src/profile.rs

@@ -1,5 +1,5 @@
 use proto::profile::LabelStringOrNumber;
-use proto::profile::Pprof;
+use proto::profile::PprofBuilder;
 use std::collections::HashMap;
 use std::fmt::Write;
 use std::path::PathBuf;
@@ -21,8 +21,8 @@ pub fn to_proto(
     profile: SymbolizedAggregatedProfile,
     procs: &HashMap<i32, ProcessInfo>,
     objs: &HashMap<ExecutableId, ObjectFileInfo>,
-) -> Pprof {
-    let mut pprof = Pprof::new();
+) -> PprofBuilder {
+    let mut pprof = PprofBuilder::with_frequency(27);
 
     for sample in profile {
         let pid = sample.pid;
@@ -31,15 +31,15 @@ pub fn to_proto(
         let mut location_ids = Vec::new();
 
         for frame in kstack {
+            // TODO: Add real values, read kernel build ID, etc.
             let mapping_id: u64 = pprof.add_mapping(
-                0x1000000, // TODO
+                0x1000000,
                 0xFFFFFFFF,
                 0xFFFFFFFF,
                 0x0,
-                "[kernel]",
-                "fake_kernel_build_id", // TODO
+                "[kernel]", // Special value.
+                "fake_kernel_build_id",
             );
-            println!("{}", frame.name);
 
             let (line, _) = pprof.add_line(&frame.name);
             let location =
@@ -51,12 +51,12 @@ pub fn to_proto(
             let addr = frame.virtual_address;
 
             let Some(info) = procs.get(&pid) else {
-                //r.push("<could not find process>".to_string());
+                // r.push("<could not find process>".to_string());
                 continue;
             };
 
             let Some(mapping) = info.mappings.find_mapping(addr) else {
-                //r.push("<could not find mapping>".to_string());
+                // r.push("<could not find mapping>".to_string());
                 continue;
             };
 

tests/integration_test.rs

@@ -1,9 +1,10 @@
 use std::io;
 use std::io::Write;
 use std::process::{Child, Command, Stdio};
+use std::sync::{Arc, Mutex};
 use std::time::Duration;
 
-use lightswitch::collector::Collector;
+use lightswitch::collector::{AggregatorCollector, Collector};
 use lightswitch::profile::symbolize_profile;
 use lightswitch::profiler::Profiler;
 use lightswitch::profiler::SymbolizedAggregatedProfile;
@@ -93,7 +94,9 @@ fn test_integration() {
     build_test_binary("cpp-progs");
     let cpp_proc = TestProcess::new("main_cpp_clang_O1");
 
-    let collector = Collector::new();
+    let collector = Arc::new(Mutex::new(
+        Box::new(AggregatorCollector::new()) as Box<dyn Collector + Send>
+    ));
     let mut p = Profiler::new(bpf_test_debug, bpf_test_debug, Duration::from_secs(5), 999);
     p.profile_pids(vec![cpp_proc.pid()]);
     p.run(collector.clone());

vm.nix

@@ -87,32 +87,32 @@ let
       [[target]]
       name = "Fedora 5.15"
       kernel = "${kernel_5_15}/bzImage"
-      command = "${lightswitch}/bin/lightswitch --duration 0 --profile-format=none"
+      command = "${lightswitch}/bin/lightswitch --duration 0 --sender=none"
 
       [[target]]
       name = "Fedora 6.0"
       kernel = "${kernel_6_0}/bzImage"
-      command = "${lightswitch}/bin/lightswitch --duration 0 --profile-format=none"
+      command = "${lightswitch}/bin/lightswitch --duration 0 --sender=none"
 
       [[target]]
       name = "Fedora 6.2"
       kernel = "${kernel_6_2}/bzImage"
-      command = "${lightswitch}/bin/lightswitch --duration 0 --profile-format=none"
+      command = "${lightswitch}/bin/lightswitch --duration 0 --sender=none"
 
       [[target]]
       name = "Fedora 6.6"
       kernel = "${kernel_6_6}/bzImage"
-      command = "${lightswitch}/bin/lightswitch --duration 0 --profile-format=none"
+      command = "${lightswitch}/bin/lightswitch --duration 0 --sender=none"
 
       [[target]]
       name = "Upstream 6.8.7"
       kernel = "${kernel_6_8_7}/bzImage"
-      command = "${lightswitch}/bin/lightswitch --duration 0 --profile-format=none"
+      command = "${lightswitch}/bin/lightswitch --duration 0 --sender=none"
 
       [[target]]
       name = "Upstream v6.9-rc5"
       kernel = "${kernel_6_9_rc5}/bzImage"
-      command = "${lightswitch}/bin/lightswitch --duration 0 --profile-format=none"
+      command = "${lightswitch}/bin/lightswitch --duration 0 --sender=none"
     '';
     nativeBuildInputs = [ ];
     installPhase = ''