diff --git a/Cargo.lock b/Cargo.lock index e437542..e2f60b8 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -390,11 +390,12 @@ checksum = "248e3bacc7dc6baa3b21e405ee045c3047101a49145e7e9eca583ab4c2ca5345" [[package]] name = "dashmap" -version = "5.5.3" +version = "6.0.1" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "978747c1d849a7d2ee5e8adc0159961c48fb7e5db2f06af6723b80123bb53856" +checksum = "804c8821570c3f8b70230c2ba75ffa5c0f9a4189b9a432b6656c536712acae28" dependencies = [ "cfg-if", + "crossbeam-utils", "hashbrown", "lock_api", "once_cell", @@ -490,6 +491,7 @@ dependencies = [ "cargo_metadata", "clap", "clap_complete", + "humantime", "inferno", "opener", "shlex", @@ -648,6 +650,12 @@ version = "1.0.3" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "df3b46402a9d5adb4c86a0cf463f42e19994e3ee891101b1841f30a545cb49a9" +[[package]] +name = "humantime" +version = "2.1.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "9a3a5bfb195931eeb336b2a7b4d761daec841b97f947d34394601737a7bba5e4" + [[package]] name = "hyper" version = "0.14.28" @@ -708,9 +716,9 @@ dependencies = [ [[package]] name = "inferno" -version = "0.11.19" +version = "0.11.21" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "321f0f839cd44a4686e9504b0a62b4d69a50b62072144c71c68f5873c167b8d9" +checksum = "232929e1d75fe899576a3d5c7416ad0d88dbfbb3c3d6aa00873a7408a50ddb88" dependencies = [ "ahash", "crossbeam-channel", diff --git a/Cargo.toml b/Cargo.toml index a44cb02..60d80e1 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -23,7 +23,8 @@ anyhow = "1.0.43" cargo_metadata = "0.18" clap = { version = "4.0.11", features = ["derive"] } clap_complete = "4.0.2" -inferno = { version = "0.11.0", default_features = false, features = ["multithreaded", "nameattr"] } +humantime = "2.1.0" +inferno = { version = "0.11.21", default_features = false, features = ["multithreaded", "nameattr"] } opener = "0.7.1" shlex = "1.1.0" diff --git a/src/bin/flamegraph.rs b/src/bin/flamegraph.rs index d551ee0..b6caf51 100644 --- a/src/bin/flamegraph.rs +++ b/src/bin/flamegraph.rs @@ -1,6 +1,6 @@ use std::path::PathBuf; -use anyhow::anyhow; +use anyhow::bail; use clap::{CommandFactory, Parser}; use clap_complete::Shell; @@ -46,11 +46,20 @@ fn main() -> anyhow::Result<()> { let path = perf_file.to_str().unwrap(); Workload::ReadPerf(path.to_string()) } else { - match (opt.pid, opt.trailing_arguments.is_empty()) { - (Some(p), true) => Workload::Pid(p), - (None, false) => Workload::Command(opt.trailing_arguments.clone()), - (Some(_), false) => return Err(anyhow!("cannot pass in command with --pid")), - (None, true) => return Err(anyhow!("no workload given to generate a flamegraph for")), + match ( + opt.pid, + opt.graph.global(), + opt.trailing_arguments.is_empty(), + ) { + (Some(_), _, false) => bail!("cannot pass in command with --pid"), + (Some(_), true, _) => bail!("cannot specify both --global and --pid"), + (_, true, false) => bail!("cannot pass in command with --global"), + + (Some(p), false, true) => Workload::Pid(p), + (None, false, false) => Workload::Command(opt.trailing_arguments.clone()), + (None, true, true) => Workload::Global, + + (None, false, true) => bail!("no workload given to generate a flamegraph for"), } }; flamegraph::generate_flamegraph_for_workload(workload, opt.graph) diff --git a/src/lib.rs b/src/lib.rs index 58b0d03..71bc3d2 100644 --- a/src/lib.rs +++ b/src/lib.rs @@ -19,7 +19,7 @@ use inferno::collapse::dtrace::{Folder, Options as CollapseOptions}; #[cfg(unix)] use signal_hook::consts::{SIGINT, SIGTERM}; -use anyhow::{anyhow, Context}; +use anyhow::{anyhow, bail, Context}; use clap::{ builder::{PossibleValuesParser, TypedValueParser}, Args, @@ -30,6 +30,7 @@ pub enum Workload { Command(Vec), Pid(u32), ReadPerf(String), + Global, } #[cfg(target_os = "linux")] @@ -42,11 +43,13 @@ mod arch { pub(crate) fn initial_command( workload: Workload, sudo: Option>, - freq: u32, - custom_cmd: Option, - verbose: bool, - ignore_status: bool, + opts: &Options, ) -> Option { + let freq = opts.frequency(); + let custom_cmd = opts.custom_cmd.clone(); + let verbose = opts.verbose; + let ignore_status = opts.ignore_status; + let perf = if let Ok(path) = env::var("PERF") { path } else { @@ -77,7 +80,16 @@ mod arch { } } + // This is checked earlier in Options::check + if opts.kernel { + unimplemented!("kernel sampling is not implemented using perf"); + } + match workload { + Workload::Global => { + // This is also checked earlier in Options::check + unimplemented!("global sampling is not implemented using perf") + } Workload::Command(c) => { command.args(&c); } @@ -118,7 +130,7 @@ mod arch { let output = command.output().context("unable to call perf script")?; if !output.status.success() { - anyhow::bail!(format!( + bail!(format!( "unable to run 'perf script': ({}) {}", output.status, std::str::from_utf8(&output.stderr)? @@ -170,21 +182,52 @@ mod arch { pub(crate) fn initial_command( workload: Workload, sudo: Option>, - freq: u32, - custom_cmd: Option, - verbose: bool, - ignore_status: bool, + opts: &Options, ) -> Option { + let freq = opts.frequency(); + let custom_cmd = opts.custom_cmd.clone(); + let global = opts.global; + let kernel = opts.kernel; + let verbose = opts.verbose; + let ignore_status = opts.ignore_status; + let mut command = base_dtrace_command(sudo); + let stack = if kernel { + "stack(100),ustack(100)" + } else { + "ustack(100)" + }; + let filter = if global { "" } else { "/pid == $target/" }; + + let timeout = if let Some(t) = opts.timeout { + format!("tick-{}ms {{ exit(0); }}", t.as_millis()) + } else { + "".to_owned() + }; + let dtrace_script = custom_cmd.unwrap_or(format!( - "profile-{freq} /pid == $target/ \ - {{ @[ustack(100)] = count(); }}", + "profile-{freq} {filter} \ + {{ @[{stack}] = count(); }} \ + {timeout}", )); command.arg("-x"); command.arg("ustackframes=100"); + // Adjust the max number of saved process handles, to avoid having to + // constantly reload symbol tables. + #[cfg(target_os = "illumos")] + if global { + command.arg("-x"); + command.arg("pgmax=1024"); + } + + if kernel { + command.arg("-x"); + command.arg("stackframes=100"); + } + command.arg("-n"); command.arg(&dtrace_script); @@ -192,6 +235,7 @@ mod arch { command.arg("cargo-flamegraph.stacks"); match workload { + Workload::Global => (), Workload::Command(c) => { let mut escaped = String::new(); for (i, arg) in c.iter().enumerate() { @@ -363,14 +407,7 @@ 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, - ) + arch::initial_command(workload, sudo, &opts) }; #[cfg(unix)] @@ -506,23 +543,54 @@ pub struct Options { /// stdout. #[clap(long)] post_process: Option, + + /// Capture a trace of the entire system + #[clap(long)] + global: bool, + + /// Include kernel stack frames + #[clap(long)] + kernel: bool, + + /// Amount of time to sample for + #[clap(long)] + timeout: Option, } impl Options { pub fn check(&self) -> anyhow::Result<()> { // Manually checking conflict because structopts `conflicts_with` leads // to a panic in completion generation for zsh at the moment (see #158) - match self.frequency.is_some() && self.custom_cmd.is_some() { - true => Err(anyhow!( - "Cannot pass both a custom command and a frequency." - )), - false => Ok(()), + if self.custom_cmd.is_some() { + if self.frequency.is_some() { + bail!("Cannot pass both a custom command and a frequency.") + } else if self.global { + bail!("Cannot specify global tracing when a custom command is also provided"); + } else if self.kernel { + bail!("Cannot specify kernel tracing when a custom command is also provided"); + } + } + + if cfg!(target_os = "linux") { + if self.global { + bail!("Global tracing is only supported using DTrace backend"); + } else if self.kernel { + bail!("Kernel tracing is only supported using DTrace backend"); + } else if self.timeout.is_some() { + bail!("Timeout is only supported using DTrace backend"); + } } + + Ok(()) } pub fn frequency(&self) -> u32 { self.frequency.unwrap_or(997) } + + pub fn global(&self) -> bool { + self.global + } } #[derive(Debug, Args)]