Add --duration and --sample_freq options

- Stop tracking flake.lock
- Stop tracking Cargo.lock

.gitignore

@@ -1,3 +1,4 @@
 /target
 flame.svg
-src/bpf/*_skel.rs
+src/bpf/*_skel.rs
+*.lock

Cargo.toml

@@ -10,7 +10,7 @@ memmap2 = "*"
 lazy_static = "*"
 anyhow = "*"
 thiserror = "*"
-libbpf-rs = {version = "*", features=["static"]}
+libbpf-rs = { version = "*", features = ["static"] }
 num_cpus = "*"
 perf-event-open-sys = "*"
 libc = "*"
@@ -20,12 +20,13 @@ procfs = "*"
 ring = "*"
 data-encoding = "*"
 page_size = "*"
-clap = { version = "*", features = ["derive"] }
+clap = { version = "4", features = ["derive", "string"] }
 blazesym = "0.2.0-alpha.10"
 tracing = "*"
 tracing-subscriber = "*"
 chrono = "*"
 inferno = "*"
+primal = "0.3"
 
 
 [build-dependencies]

src/main.rs

@@ -4,6 +4,7 @@ use clap::Parser;
 use inferno::flamegraph;
 use std::fmt::Write;
 use std::fs::File;
+use std::ops::RangeInclusive;
 use tracing::Level;
 use tracing_subscriber::fmt::format::FmtSpan;
 use tracing_subscriber::FmtSubscriber;
@@ -12,27 +13,69 @@ use lightswitch::object::build_id;
 use lightswitch::profiler::Collector;
 use lightswitch::profiler::Profiler;
 use lightswitch::unwind_info::{compact_printing_callback, UnwindInfoBuilder};
+use primal::is_prime;
 use std::error::Error;
 use std::path::PathBuf;
 
 use std::time::Duration;
 
+fn parse_duration(arg: &str) -> Result<Duration, std::num::ParseIntError> {
+    let seconds = arg.parse()?;
+    Ok(Duration::from_secs(seconds))
+}
+
+const SAMPLE_FREQ_RANGE: RangeInclusive<usize> = 1..=1001;
+
+fn sample_freq_in_range(s: &str) -> Result<u16, String> {
+    let sample_freq: usize = s
+        .parse()
+        .map_err(|_| format!("`{s}' isn't a valid frequency"))?;
+    if !(1..=1001).contains(&sample_freq) {
+        return Err(format!(
+            "sample frequency not in allowed range {}-{}",
+            SAMPLE_FREQ_RANGE.start(),
+            SAMPLE_FREQ_RANGE.end()
+        ));
+    }
+    if !is_prime(sample_freq.try_into().unwrap()) {
+        return Err(format!("sample frequency is not prime"));
+    }
+    Ok(sample_freq as u16)
+}
+
+// Used to prevent the default value for duration from being dropped before
+// we parse the command line args - probably a better way to do this
+fn max_duration_as_str() -> &'static str {
+    let max_duration: &'static str =
+        Box::leak(Duration::MAX.as_secs().to_string().into_boxed_str());
+    max_duration
+}
+
 #[derive(Parser, Debug)]
-struct Args {
+struct Cli {
+    /// Specific PIDs to profile
     #[arg(long)]
     pids: Vec<i32>,
     #[arg(long)]
     show_unwind_info: Option<String>,
     #[arg(long)]
     show_info: Option<String>,
-    #[arg(long)]
-    continuous: bool,
+    /// How long this agent will run in seconds
+    #[arg(short='D', long, default_value = max_duration_as_str(), value_parser = parse_duration)]
+    duration: Duration,
     #[arg(long, action=ArgAction::SetFalse)]
     filter_logs: bool,
+    // Verification for this option guarantees the only possible selections
+    // are prime numbers up to and including 1001
+    /// Per-CPU Sampling Frequency in Hz
+    #[arg(long, default_value_t = 19,
+      value_parser = sample_freq_in_range,
+    )]
+    sample_freq: u16,
 }
 
 fn main() -> Result<(), Box<dyn Error>> {
-    let args = Args::parse();
+    let args = Cli::parse();
 
     let subscriber = FmtSubscriber::builder()
         .with_max_level(if args.filter_logs {
@@ -59,16 +102,11 @@ fn main() -> Result<(), Box<dyn Error>> {
         return Ok(());
     }
 
-    let mut duration = Duration::MAX;
-    if !args.continuous {
-        duration = Duration::from_secs(5);
-    }
-
     let collector = Collector::new();
 
-    let mut p: Profiler<'_> = Profiler::new(false);
+    let mut p: Profiler<'_> = Profiler::new(false, args.duration, args.sample_freq);
     p.profile_pids(args.pids);
-    p.run(duration, collector.clone());
+    p.run(collector.clone());
 
     let profiles = collector.lock().unwrap().finish();
 
@@ -96,3 +134,9 @@ fn main() -> Result<(), Box<dyn Error>> {
 
     Ok(())
 }
+
+#[test]
+fn verify_cli() {
+    use clap::CommandFactory;
+    Cli::command().debug_assert()
+}

src/profiler.rs

@@ -332,6 +332,10 @@ pub struct Profiler<'bpf> {
     // Profile channel
     profile_send: Arc<Mutex<mpsc::Sender<RawAggregatedProfile>>>,
     profile_receive: Arc<Mutex<mpsc::Receiver<RawAggregatedProfile>>>,
+    // Duration of this profile
+    duration: Duration,
+    // Per-CPU Sampling Frequency of this profile in Hz
+    sample_freq: u16,
 }
 
 pub struct Collector {
@@ -390,7 +394,6 @@ impl Collector {
 }
 
 // Static config
-const SAMPLE_PERIOD_HZ: u64 = 200;
 const MAX_UNWIND_INFO_SHARDS: u64 = 50;
 const SHARD_CAPACITY: usize = MAX_UNWIND_TABLE_SIZE as usize;
 const PERF_BUFFER_PAGES: usize = 512 * 1024;
@@ -417,12 +420,12 @@ pub type SymbolizedAggregatedProfile = Vec<SymbolizedAggregatedSample>;
 
 impl Default for Profiler<'_> {
     fn default() -> Self {
-        Self::new(false)
+        Self::new(false, Duration::MAX, 19)
     }
 }
 
 impl Profiler<'_> {
-    pub fn new(bpf_debug: bool) -> Self {
+    pub fn new(bpf_debug: bool, duration: Duration, sample_freq: u16) -> Self {
         let mut skel_builder: ProfilerSkelBuilder = ProfilerSkelBuilder::default();
         skel_builder.obj_builder.debug(bpf_debug);
         let open_skel = skel_builder.open().expect("open skel");
@@ -468,6 +471,8 @@ impl Profiler<'_> {
             filter_pids,
             profile_send,
             profile_receive,
+            duration,
+            sample_freq,
         }
     }
 
@@ -485,7 +490,7 @@ impl Profiler<'_> {
             .expect("handle send");
     }
 
-    pub fn run(mut self, duration: Duration, collector: Arc<Mutex<Collector>>) {
+    pub fn run(mut self, collector: Arc<Mutex<Collector>>) {
         self.setup_perf_events();
         self.set_bpf_map_info();
 
@@ -527,7 +532,7 @@ impl Profiler<'_> {
         let mut time_since_last_scheduled_collection: Instant = Instant::now();
 
         loop {
-            if start_time.elapsed() >= duration {
+            if start_time.elapsed() >= self.duration {
                 debug!("done after running for {:?}", start_time.elapsed());
                 let profile = self.collect_profile();
                 self.send_profile(profile);
@@ -558,7 +563,7 @@ impl Profiler<'_> {
                             }
                         } */
                     } else {
-                        error!("unknow event {}", event.type_);
+                        error!("unknown event {}", event.type_);
                     }
                 }
                 Err(_) => {
@@ -1195,8 +1200,10 @@ impl Profiler<'_> {
     pub fn setup_perf_events(&mut self) {
         let mut prog_fds = Vec::new();
         for i in 0..num_cpus::get() {
-            let perf_fd = unsafe { setup_perf_event(i.try_into().unwrap(), SAMPLE_PERIOD_HZ) }
-                .expect("setup perf event");
+            let perf_fd = unsafe {
+                setup_perf_event(i.try_into().unwrap(), self.sample_freq.try_into().unwrap())
+            }
+            .expect("setup perf event");
             prog_fds.push(perf_fd);
         }