From acdc9f4e0d5e93a9517b460eac54fc3b0bdf802c Mon Sep 17 00:00:00 2001
From: Roman Gershman <romange@gmail.com>
Date: Fri, 15 Oct 2021 14:57:31 +0300
Subject: [PATCH] Allow reading perf.data directly. Fixes #52

---
 src/bin/flamegraph.rs | 20 +++++++++++++++-----
 src/lib.rs            | 43 +++++++++++++++++++++++++------------------
 2 files changed, 40 insertions(+), 23 deletions(-)

diff --git a/src/bin/flamegraph.rs b/src/bin/flamegraph.rs
index 77731e4..d460d29 100644
--- a/src/bin/flamegraph.rs
+++ b/src/bin/flamegraph.rs
@@ -1,3 +1,5 @@
+use std::path::PathBuf;
+
 use anyhow::anyhow;
 use structopt::StructOpt;
 
@@ -19,6 +21,9 @@ struct Opt {
     #[structopt(flatten)]
     graph: flamegraph::Options,
 
+    #[structopt(parse(from_os_str), long = "perfdata", conflicts_with = "pid")]
+    perf_file: Option<PathBuf>,
+
     trailing_arguments: Vec<String>,
 }
 
@@ -41,11 +46,16 @@ fn main() -> anyhow::Result<()> {
 
     opt.graph.check()?;
 
-    let workload = 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")),
+    let workload = if let Some(perf_file) = opt.perf_file {
+        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")),
+        }
     };
     flamegraph::generate_flamegraph_for_workload(workload, opt.graph)
 }
diff --git a/src/lib.rs b/src/lib.rs
index bedda65..0c7b523 100644
--- a/src/lib.rs
+++ b/src/lib.rs
@@ -28,6 +28,7 @@ use inferno::{
 pub enum Workload {
     Command(Vec<String>),
     Pid(u32),
+    ReadPerf(String),
 }
 
 #[cfg(target_os = "linux")]
@@ -82,6 +83,7 @@ mod arch {
                 command.arg("-p");
                 command.arg(p.to_string());
             }
+            Workload::ReadPerf(_) => (),
         }
 
         (command, perf_output)
@@ -163,6 +165,7 @@ mod arch {
                 command.arg("-p");
                 command.arg(p.to_string());
             }
+            Workload::ReadPerf(_) => (),
         }
 
         (command, None)
@@ -231,27 +234,31 @@ pub fn generate_flamegraph_for_workload(
         signal_hook::low_level::register(SIGINT, || {}).expect("cannot register signal handler")
     };
 
-    let (mut command, perf_output) =
-        arch::initial_command(workload, opts.root, opts.frequency, opts.custom_cmd);
-    if opts.verbose {
-        println!("command {:?}", command);
-    }
+    let perf_output = if let Workload::ReadPerf(perf_file) = workload {
+        Some(perf_file)
+    } else {
+        let (mut command, perf_output) =
+            arch::initial_command(workload, opts.root, opts.frequency, opts.custom_cmd);
+        if opts.verbose {
+            println!("command {:?}", command);
+        }
 
-    let mut recorder = command.spawn().expect(arch::SPAWN_ERROR);
+        let mut recorder = command.spawn().expect(arch::SPAWN_ERROR);
 
-    let exit_status = recorder.wait().expect(arch::WAIT_ERROR);
+        let exit_status = recorder.wait().expect(arch::WAIT_ERROR);
 
-    #[cfg(unix)]
-    signal_hook::low_level::unregister(handler);
-
-    // only stop if perf exited unsuccessfully, but
-    // was not killed by a signal (assuming that the
-    // latter case usually means the user interrupted
-    // it in some way)
-    if terminated_by_error(exit_status) {
-        eprintln!("failed to sample program");
-        std::process::exit(1);
-    }
+        #[cfg(unix)]
+        signal_hook::low_level::unregister(handler);
+        // only stop if perf exited unsuccessfully, but
+        // was not killed by a signal (assuming that the
+        // latter case usually means the user interrupted
+        // it in some way)
+        if terminated_by_error(exit_status) {
+            eprintln!("failed to sample program");
+            std::process::exit(1);
+        }
+        perf_output
+    };
 
     let output = arch::output(perf_output, opts.script_no_inline)?;