diff --git a/src/bin/flamegraph.rs b/src/bin/flamegraph.rs index d551ee0..7877a47 100644 --- a/src/bin/flamegraph.rs +++ b/src/bin/flamegraph.rs @@ -20,6 +20,7 @@ struct Opt { #[clap(flatten)] graph: flamegraph::Options, + /// Read perf data from the given file. #[clap(long = "perfdata", conflicts_with = "pid")] perf_file: Option, diff --git a/src/lib.rs b/src/lib.rs index 58b0d03..45f72f3 100644 --- a/src/lib.rs +++ b/src/lib.rs @@ -26,9 +26,13 @@ use clap::{ }; use inferno::{collapse::Collapse, flamegraph::color::Palette, flamegraph::from_reader}; +/// Mode of operation. pub enum Workload { + /// Execute an executable with the given command and arguments. Command(Vec), + /// Profile a running process with the given PID. Pid(u32), + /// Read profiling data from a file. ReadPerf(String), } @@ -168,7 +172,7 @@ mod arch { } pub(crate) fn initial_command( - workload: Workload, + workload: &Workload, sudo: Option>, freq: u32, custom_cmd: Option, @@ -363,14 +367,24 @@ pub fn generate_flamegraph_for_workload(workload: Workload, opts: Options) -> an let perf_output = if let Workload::ReadPerf(perf_file) = workload { Some(perf_file) } else { - arch::initial_command( - workload, - sudo, - opts.frequency(), - opts.custom_cmd, - opts.verbose, - opts.ignore_status, - ) + let out = (0..opts.iterations.unwrap_or(1)).fold(String::new(), |mut out, _i| { + if let Some(iter_out) = arch::initial_command( + &workload, + sudo, + opts.frequency(), + opts.custom_cmd.clone(), + opts.verbose, + opts.ignore_status, + ) { + out.push_str(&iter_out) + } + out + }); + if out.is_empty() { + None + } else { + Some(out) + } }; #[cfg(unix)] @@ -487,6 +501,10 @@ pub struct Options { #[clap(short = 'F', long = "freq")] frequency: Option, + /// Number of runs for target binary. Defaults to 1. + #[clap(long)] + iterations: Option, + /// Custom command for invoking perf/dtrace #[clap(short, long = "cmd")] custom_cmd: Option,