#[cfg(unix)]
use std::os::unix::process::ExitStatusExt;
#[cfg(target_os = "macos")]
use std::{
borrow::Cow,
io::{BufRead, Error, ErrorKind},
};
use std::{
env,
fs::File,
io::{BufReader, BufWriter, Cursor, Read, Write},
path::PathBuf,
process::{exit, Command, ExitStatus, Stdio},
str::FromStr,
};
use anyhow::{anyhow, bail, Context};
use clap::{
builder::{PossibleValuesParser, TypedValueParser},
Args,
};
#[cfg(not(any(target_os = "linux", target_os = "macos")))]
use inferno::collapse::dtrace::{Folder, Options as CollapseOptions};
#[cfg(target_os = "linux")]
use inferno::collapse::perf::{Folder, Options as CollapseOptions};
#[cfg(target_os = "macos")]
use inferno::collapse::xctrace::Folder;
use inferno::{collapse::Collapse, flamegraph::color::Palette, flamegraph::from_reader};
#[cfg(target_os = "macos")]
use quick_xml::events::Event;
#[cfg(not(target_os = "macos"))]
use rustc_demangle::demangle_stream;
#[cfg(target_os = "macos")]
use rustc_demangle::try_demangle;
#[cfg(unix)]
use signal_hook::consts::{SIGINT, SIGTERM};
pub enum Workload {
Command(Vec<String>),
Pid(Vec<u32>),
ReadPerf(PathBuf),
}
#[cfg(target_os = "linux")]
mod arch {
use std::fmt::Write;
use std::time::Duration;
use indicatif::{ProgressBar, ProgressStyle};
use super::*;
pub const SPAWN_ERROR: &str = "could not spawn perf";
pub const WAIT_ERROR: &str = "unable to wait for perf child command to exit";
pub(crate) fn initial_command(
workload: Workload,
sudo: Option<Option<&str>>,
freq: u32,
compress_level: Option<u8>,
custom_cmd: Option<String>,
verbose: bool,
ignore_status: bool,
) -> anyhow::Result<Option<PathBuf>> {
let perf = if let Ok(path) = env::var("PERF") {
path
} else {
if Command::new("perf")
.arg("--help")
.stderr(Stdio::null())
.stdout(Stdio::null())
.status()
.is_err()
{
bail!("perf is not installed or not present in $PATH");
}
String::from("perf")
};
let mut command = sudo_command(&perf, sudo);
let args = custom_cmd.unwrap_or_else(|| {
let mut args = format!("record -F {freq} --call-graph dwarf,64000 -g");
if let Some(z) = compress_level {
_ = write!(args, " -z={z}");
}
args
});
let mut perf_output = None;
let mut args = args.split_whitespace();
while let Some(arg) = args.next() {
command.arg(arg);
if arg == "-o" {
let next_arg = args.next().context("missing '-o' argument")?;
command.arg(next_arg);
perf_output = Some(PathBuf::from(next_arg));
}
}
let perf_output = match perf_output {
Some(path) => path,
None => {
command.arg("-o");
command.arg("perf.data");
PathBuf::from("perf.data")
}
};
match workload {
Workload::Command(c) => {
command.args(&c);
}
Workload::Pid(p) => {
if let Some((first, pids)) = p.split_first() {
let mut arg = first.to_string();
for pid in pids {
arg.push(',');
arg.push_str(&pid.to_string());
}
command.arg("-p");
command.arg(arg);
}
}
Workload::ReadPerf(_) => (),
}
run(command, verbose, ignore_status);
Ok(Some(perf_output))
}
pub fn output(
perf_output: Option<PathBuf>,
script_no_inline: bool,
sudo: Option<Option<&str>>,
) -> anyhow::Result<Vec<u8>> {
let perf = env::var("PERF").unwrap_or_else(|_| "perf".to_string());
let mut command = sudo_command(&perf, sudo);
command.arg("script");
command.arg("--force");
if script_no_inline {
command.arg("--no-inline");
}
if let Some(perf_output) = perf_output {
command.arg("-i");
command.arg(perf_output);
}
let spinner = ProgressBar::new_spinner().with_prefix("Running perf script");
spinner.set_style(
ProgressStyle::with_template("{prefix} [{elapsed}]: {spinner:.green}").unwrap(),
);
spinner.enable_steady_tick(Duration::from_millis(500));
let result = command.output().context("unable to call perf script");
spinner.finish();
let output = result?;
if !output.status.success() {
bail!(
"unable to run 'perf script': ({}) {}",
output.status,
std::str::from_utf8(&output.stderr)?
);
}
Ok(output.stdout)
}
}
#[cfg(target_os = "macos")]
mod arch {
use super::*;
pub const SPAWN_ERROR: &str = "could not spawn xctrace";
pub const WAIT_ERROR: &str = "unable to wait for xctrace record child command to exit";
pub(crate) fn initial_command(
workload: Workload,
sudo: Option<Option<&str>>,
freq: u32,
_compress_level: Option<u8>,
custom_cmd: Option<String>,
verbose: bool,
ignore_status: bool,
) -> anyhow::Result<Option<PathBuf>> {
if freq != 997 {
bail!("xctrace doesn't support custom frequency");
}
if custom_cmd.is_some() {
bail!("xctrace doesn't support custom command");
}
let xctrace = env::var("XCTRACE").unwrap_or_else(|_| "xctrace".to_string());
let trace_file = PathBuf::from("cargo-flamegraph.trace");
let mut command = sudo_command(&xctrace, sudo);
command
.arg("record")
.arg("--template")
.arg("Time Profiler")
.arg("--output")
.arg(&trace_file);
match workload {
Workload::Command(args) => {
command
.arg("--target-stdout")
.arg("-")
.arg("--launch")
.arg("--")
.args(args);
}
Workload::Pid(pid) => {
match &*pid {
[pid] => {
command.arg("--attach").arg(pid.to_string());
}
_ => {
bail!("xctrace only supports profiling a single process at a time");
}
}
}
Workload::ReadPerf(_) => {}
}
run(command, verbose, ignore_status);
Ok(Some(trace_file))
}
pub fn output(
trace_file: Option<PathBuf>,
script_no_inline: bool,
_sudo: Option<Option<&str>>,
) -> anyhow::Result<Vec<u8>> {
if script_no_inline {
bail!("--no-inline is only supported on Linux");
}
let xctrace = env::var("XCTRACE").unwrap_or_else(|_| "xctrace".to_string());
let trace_file = trace_file.context("no trace file found.")?;
let output = Command::new(xctrace)
.arg("export")
.arg("--input")
.arg(&trace_file)
.arg("--xpath")
.arg(r#"/trace-toc/*/data/table[@schema="time-profile"]"#)
.output()
.context("run xctrace export failed.")?;
std::fs::remove_dir_all(&trace_file)
.with_context(|| anyhow!("remove trace({}) failed.", trace_file.to_string_lossy()))?;
if !output.status.success() {
bail!(
"unable to run 'xctrace export': ({}) {}",
output.status,
String::from_utf8_lossy(&output.stderr)
);
}
Ok(output.stdout)
}
}
#[cfg(not(any(target_os = "linux", target_os = "macos")))]
mod arch {
use super::*;
pub const SPAWN_ERROR: &str = "could not spawn dtrace";
pub const WAIT_ERROR: &str = "unable to wait for dtrace child command to exit";
#[cfg(target_os = "macos")]
fn base_dtrace_command(sudo: Option<Option<&str>>) -> Command {
let mut command = sudo_command("arch", sudo);
#[cfg(target_pointer_width = "64")]
command.arg("-64".to_string());
#[cfg(target_pointer_width = "32")]
command.arg("-32".to_string());
command.arg(env::var("DTRACE").unwrap_or_else(|_| "dtrace".to_string()));
command
}
#[cfg(not(target_os = "macos"))]
fn base_dtrace_command(sudo: Option<Option<&str>>) -> Command {
let dtrace = env::var("DTRACE").unwrap_or_else(|_| "dtrace".to_string());
sudo_command(&dtrace, sudo)
}
pub(crate) fn initial_command(
workload: Workload,
sudo: Option<Option<&str>>,
freq: u32,
_compress_level: Option<u8>,
custom_cmd: Option<String>,
verbose: bool,
ignore_status: bool,
) -> anyhow::Result<Option<PathBuf>> {
let mut command = base_dtrace_command(sudo);
let dtrace_script = custom_cmd.unwrap_or(format!(
"profile-{freq} /pid == $target/ \
{{ @[ustack(100)] = count(); }}",
));
command.arg("-x");
command.arg("ustackframes=100");
command.arg("-n");
command.arg(&dtrace_script);
command.arg("-o");
command.arg("cargo-flamegraph.stacks");
match workload {
Workload::Command(c) => {
let mut escaped = String::new();
for (i, arg) in c.iter().enumerate() {
if i > 0 {
escaped.push(' ');
}
escaped.push_str(&arg.replace(' ', "\\ "));
}
command.arg("-c");
command.arg(&escaped);
#[cfg(target_os = "windows")]
{
let mut help_test = crate::arch::base_dtrace_command(None);
let dtrace_found = help_test
.arg("--help")
.stderr(Stdio::null())
.stdout(Stdio::null())
.status()
.is_ok();
if !dtrace_found {
let mut command_builder = Command::new(&c[0]);
command_builder.args(&c[1..]);
print_command(&command_builder, verbose);
let trace = blondie::trace_command(command_builder, false)
.map_err(|err| anyhow!("could not find dtrace and could not profile using blondie: {err:?}"))?;
let f = std::fs::File::create("./cargo-flamegraph.stacks")
.context("unable to create temporary file 'cargo-flamegraph.stacks'")?;
let mut f = std::io::BufWriter::new(f);
trace.write_dtrace(&mut f).map_err(|err| {
anyhow!("unable to write dtrace output to 'cargo-flamegraph.stacks': {err:?}")
})?;
return Ok(None);
}
}
}
Workload::Pid(p) => {
for p in p {
command.arg("-p");
command.arg(p.to_string());
}
}
Workload::ReadPerf(_) => (),
}
run(command, verbose, ignore_status);
Ok(None)
}
pub fn output(
_: Option<PathBuf>,
script_no_inline: bool,
sudo: Option<Option<&str>>,
) -> anyhow::Result<Vec<u8>> {
if script_no_inline {
bail!("--no-inline is only supported on Linux");
}
if sudo.is_some() {
#[cfg(unix)]
if let Ok(user) = env::var("USER") {
Command::new("sudo")
.args(["chown", user.as_str(), "cargo-flamegraph.stacks"])
.spawn()
.expect(arch::SPAWN_ERROR)
.wait()
.expect(arch::WAIT_ERROR);
}
}
let mut buf = vec![];
let mut f = File::open("cargo-flamegraph.stacks")
.context("failed to open dtrace output file 'cargo-flamegraph.stacks'")?;
f.read_to_end(&mut buf)
.context("failed to read dtrace expected output file 'cargo-flamegraph.stacks'")?;
std::fs::remove_file("cargo-flamegraph.stacks")
.context("unable to remove temporary file 'cargo-flamegraph.stacks'")?;
let string = String::from_utf8_lossy(&buf);
let reencoded_buf = string.as_bytes().to_owned();
if reencoded_buf != buf {
println!("Lossily converted invalid utf-8 found in cargo-flamegraph.stacks");
}
Ok(reencoded_buf)
}
}
fn sudo_command(command: &str, sudo: Option<Option<&str>>) -> Command {
let sudo = match sudo {
Some(sudo) => sudo,
None => return Command::new(command),
};
let mut c = Command::new("sudo");
if let Some(sudo_args) = sudo {
c.arg(sudo_args);
}
c.arg(command);
c
}
fn run(mut command: Command, verbose: bool, ignore_status: bool) {
print_command(&command, verbose);
let mut recorder = command.spawn().expect(arch::SPAWN_ERROR);
let exit_status = recorder.wait().expect(arch::WAIT_ERROR);
if !ignore_status && terminated_by_error(exit_status) {
eprintln!(
"failed to sample program, exited with code: {:?}",
exit_status.code()
);
exit(1);
}
}
#[cfg(unix)]
fn terminated_by_error(status: ExitStatus) -> bool {
status
.signal() .is_none_or(|code| code != SIGINT && code != SIGTERM)
&& !status.success()
&& !(cfg!(target_os = "macos") && status.code() == Some(54))
}
#[cfg(not(unix))]
fn terminated_by_error(status: ExitStatus) -> bool {
!status.success()
}
fn print_command(cmd: &Command, verbose: bool) {
if verbose {
println!("command {:?}", cmd);
}
}
#[cfg(target_os = "macos")]
fn demangle_stream<R: BufRead, W: Write>(input: &mut R, output: &mut W) -> std::io::Result<()> {
let mut reader = quick_xml::Reader::from_reader(input);
let mut writer = quick_xml::Writer::new(output);
let mut buf = Vec::new();
loop {
buf.clear();
let start = match reader.read_event_into(&mut buf) {
Ok(Event::Eof) => break,
Ok(Event::Start(start)) => start,
Ok(Event::Empty(start)) => start,
Ok(el) => {
writer.write_event(el)?;
continue;
}
Err(err) => return Err(Error::new(ErrorKind::InvalidData, err)),
};
let mut new = start.clone();
new.clear_attributes();
for attr in start.attributes() {
let mut attr = attr.map_err(|err| Error::new(ErrorKind::InvalidData, err))?;
let mangled = String::from_utf8_lossy(attr.value.as_ref());
if let Ok(demangled) = try_demangle(&mangled) {
let demangled = format!("{demangled:#}");
attr.value = match quick_xml::escape::escape(demangled) {
Cow::Borrowed(s) => Cow::Borrowed(s.as_bytes()),
Cow::Owned(s) => Cow::Owned(s.into_bytes()),
};
}
new.push_attribute(attr);
}
writer.write_event(Event::Start(new))?;
}
Ok(())
}
pub fn generate_flamegraph_for_workload(workload: Workload, opts: Options) -> anyhow::Result<()> {
#[cfg(unix)]
let handler = unsafe {
signal_hook::low_level::register(SIGINT, || {}).expect("cannot register signal handler")
};
let sudo = opts.root.as_ref().map(|inner| inner.as_deref());
let perf_output = if let Workload::ReadPerf(perf_file) = workload {
Some(perf_file)
} else {
#[cfg(target_os = "linux")]
let compression = opts.compression_level;
#[cfg(not(target_os = "linux"))]
let compression = None;
arch::initial_command(
workload,
sudo,
opts.frequency(),
compression,
opts.custom_cmd,
opts.verbose,
opts.ignore_status,
)?
};
#[cfg(unix)]
signal_hook::low_level::unregister(handler);
let output = arch::output(perf_output, opts.script_no_inline, sudo)?;
let mut demangled_output = vec![];
#[cfg(not(target_os = "macos"))]
let demangle_result = demangle_stream(&mut Cursor::new(output), &mut demangled_output, false);
#[cfg(target_os = "macos")]
let demangle_result = demangle_stream(&mut Cursor::new(output), &mut demangled_output);
demangle_result.context("unable to demangle")?;
let perf_reader = BufReader::new(&*demangled_output);
let mut collapsed = vec![];
let collapsed_writer = BufWriter::new(&mut collapsed);
#[cfg(target_os = "linux")]
let mut folder = {
let mut collapse_options = CollapseOptions::default();
collapse_options.skip_after = opts.flamegraph_options.skip_after.clone();
Folder::from(collapse_options)
};
#[cfg(target_os = "macos")]
let mut folder = Folder::default();
#[cfg(not(any(target_os = "linux", target_os = "macos")))]
let mut folder = {
let collapse_options = CollapseOptions::default();
Folder::from(collapse_options)
};
folder
.collapse(perf_reader, collapsed_writer)
.context("unable to collapse generated profile data")?;
if let Some(command) = opts.post_process {
let command_vec = shlex::split(&command)
.ok_or_else(|| anyhow!("unable to parse post-process command"))?;
let mut child = Command::new(
command_vec
.first()
.ok_or_else(|| anyhow!("unable to parse post-process command"))?,
)
.args(command_vec.get(1..).unwrap_or(&[]))
.stdin(Stdio::piped())
.stdout(Stdio::piped())
.spawn()
.with_context(|| format!("unable to execute {:?}", command_vec))?;
let mut stdin = child
.stdin
.take()
.ok_or_else(|| anyhow::anyhow!("unable to capture post-process stdin"))?;
let mut stdout = child
.stdout
.take()
.ok_or_else(|| anyhow::anyhow!("unable to capture post-process stdout"))?;
let thread_handle = std::thread::spawn(move || -> anyhow::Result<_> {
let mut collapsed_processed = Vec::new();
stdout.read_to_end(&mut collapsed_processed).context(
"unable to read the processed stacks from the stdout of the post-process process",
)?;
Ok(collapsed_processed)
});
stdin
.write_all(&collapsed)
.context("unable to write the raw stacks to the stdin of the post-process process")?;
drop(stdin);
anyhow::ensure!(
child.wait()?.success(),
"post-process exited with a non zero exit code"
);
collapsed = thread_handle.join().unwrap()?;
}
let collapsed_reader = BufReader::new(&*collapsed);
let flamegraph_filename = opts.output;
println!("writing flamegraph to {:?}", flamegraph_filename);
let flamegraph_file = File::create(&flamegraph_filename)
.context("unable to create flamegraph.svg output file")?;
let flamegraph_writer = BufWriter::new(flamegraph_file);
let mut inferno_opts = opts.flamegraph_options.into_inferno();
from_reader(&mut inferno_opts, collapsed_reader, flamegraph_writer)
.context("unable to generate a flamegraph from the collapsed stack data")?;
if opts.open {
opener::open(&flamegraph_filename).context(format!(
"failed to open '{}'",
flamegraph_filename.display()
))?;
}
Ok(())
}
#[derive(Debug, Args)]
pub struct Options {
#[clap(short, long)]
pub verbose: bool,
#[clap(short, long, default_value = "flamegraph.svg")]
output: PathBuf,
#[clap(long)]
open: bool,
#[clap(long, value_name = "SUDO FLAGS")]
pub root: Option<Option<String>>,
#[clap(short = 'F', long = "freq")]
frequency: Option<u32>,
#[cfg(target_os = "linux")]
#[clap(short = 'z', long, num_args=0..=1, default_missing_value = "1")]
compression_level: Option<u8>,
#[clap(short, long = "cmd")]
custom_cmd: Option<String>,
#[clap(flatten)]
flamegraph_options: FlamegraphOptions,
#[clap(long)]
ignore_status: bool,
#[clap(long = "no-inline")]
script_no_inline: bool,
#[clap(long)]
post_process: Option<String>,
}
impl Options {
pub fn check(&self) -> anyhow::Result<()> {
match self.frequency.is_some() && self.custom_cmd.is_some() {
true => Err(anyhow!(
"Cannot pass both a custom command and a frequency."
)),
false => Ok(()),
}
}
pub fn frequency(&self) -> u32 {
self.frequency.unwrap_or(997)
}
}
#[derive(Debug, Args)]
pub struct FlamegraphOptions {
#[clap(long, value_name = "STRING")]
pub title: Option<String>,
#[clap(long, value_name = "STRING")]
pub subtitle: Option<String>,
#[clap(long)]
pub deterministic: bool,
#[clap(short, long)]
pub inverted: bool,
#[clap(long)]
pub reverse: bool,
#[clap(long, value_name = "STRING")]
pub notes: Option<String>,
#[clap(long, default_value = "0.01", value_name = "FLOAT")]
pub min_width: f64,
#[clap(long)]
pub image_width: Option<usize>,
#[clap(
long,
value_parser = PossibleValuesParser::new(Palette::VARIANTS).map(|s| Palette::from_str(&s).unwrap())
)]
pub palette: Option<Palette>,
#[cfg(target_os = "linux")]
#[clap(long, value_name = "FUNCTION")]
pub skip_after: Vec<String>,
#[clap(long = "flamechart", conflicts_with = "reverse")]
pub flame_chart: bool,
}
impl FlamegraphOptions {
pub fn into_inferno(self) -> inferno::flamegraph::Options<'static> {
let mut options = inferno::flamegraph::Options::default();
if let Some(title) = self.title {
options.title = title;
}
options.subtitle = self.subtitle;
options.deterministic = self.deterministic;
if self.inverted {
options.direction = inferno::flamegraph::Direction::Inverted;
}
options.reverse_stack_order = self.reverse;
options.notes = self.notes.unwrap_or_default();
options.min_width = self.min_width;
options.image_width = self.image_width;
if let Some(palette) = self.palette {
options.colors = palette;
}
options.flame_chart = self.flame_chart;
options
}
}