Merge pull request flamegraph-rs#53 from ASankaran/profile-by-pid

Profile by pid
leod · Dec 2, 2019 · d38adf9 · d38adf9
2 parents 1c34b2b + 5410cce
commit d38adf9
Show file tree

Hide file tree

Showing 3 changed files with 60 additions and 16 deletions.
diff --git a/src/bin/cargo-flamegraph.rs b/src/bin/cargo-flamegraph.rs
@@ -3,6 +3,8 @@ use std::path::{Path, PathBuf};
 
 use structopt::StructOpt;
 
+use flamegraph::Workload;
+
 #[derive(Debug, StructOpt)]
 #[structopt(raw(
     setting = "structopt::clap::AppSettings::TrailingVarArg"
@@ -285,8 +287,8 @@ fn main() {
         .take()
         .unwrap_or("flamegraph.svg".into());
 
-    flamegraph::generate_flamegraph_by_running_command(
-        workload,
+    flamegraph::generate_flamegraph_for_workload(
+        Workload::Command(workload),
         &flamegraph_filename,
         opt.root,
     );

diff --git a/src/bin/flamegraph.rs b/src/bin/flamegraph.rs
@@ -2,6 +2,8 @@ use std::path::PathBuf;
 
 use structopt::StructOpt;
 
+use flamegraph::Workload;
+
 #[derive(Debug, StructOpt)]
 #[structopt(raw(
     setting = "structopt::clap::AppSettings::TrailingVarArg"
@@ -19,16 +21,35 @@ struct Opt {
     #[structopt(long = "root")]
     root: bool,
 
+    /// Profile a running process by pid
+    #[structopt(
+        short = "p",
+        long = "pid"
+    )]
+    pid: Option<u32>,
+
     trailing_arguments: Vec<String>,
 }
 
-fn workload(opt: &Opt) -> Vec<String> {
-    if opt.trailing_arguments.is_empty() {
-        eprintln!("no workload given to generate a flamegraph for!");
-        std::process::exit(1);
-    }
+fn workload(opt: &Opt) -> Workload {
+    match opt.pid {
+        Some(p) => {
+            if !opt.trailing_arguments.is_empty() {
+                eprintln!("only a pid or command can be specified!");
+                std::process::exit(1);
+            }
+
+            Workload::Pid(p)
+        },
+        None => {
+            if opt.trailing_arguments.is_empty() {
+                eprintln!("no workload given to generate a flamegraph for!");
+                std::process::exit(1);
+            }
 
-    opt.trailing_arguments.clone()
+            Workload::Command(opt.trailing_arguments.clone())
+        }
+    }
 }
 
 fn main() {
@@ -41,7 +62,7 @@ fn main() {
         .take()
         .unwrap_or("flamegraph.svg".into());
 
-    flamegraph::generate_flamegraph_by_running_command(
+    flamegraph::generate_flamegraph_for_workload(
         workload,
         flamegraph_filename,
         opt.root,

diff --git a/src/lib.rs b/src/lib.rs
@@ -27,6 +27,11 @@ use inferno::{
 #[cfg(unix)]
 use signal_hook;
 
+pub enum Workload {
+    Command(Vec<String>),
+    Pid(u32),
+}
+
 #[cfg(target_os = "linux")]
 mod arch {
     use super::*;
@@ -38,7 +43,7 @@ mod arch {
          child command to exit";
 
     pub(crate) fn initial_command(
-        workload: Vec<String>,
+        workload: Workload,
         sudo: bool,
     ) -> Command {
         let mut command = if sudo {
@@ -55,7 +60,15 @@ mod arch {
             command.arg(arg);
         }
 
-        command.args(&workload);
+        match workload {
+            Workload::Command(c) => {
+                command.args(&c);
+            },
+            Workload::Pid(p) => {
+                command.arg("-p");
+                command.arg(p.to_string());
+            },
+        }
 
         command
     }
@@ -80,7 +93,7 @@ mod arch {
          child command to exit";
 
     pub(crate) fn initial_command(
-        workload: Vec<String>,
+        workload: Workload,
         sudo: bool,
     ) -> Command {
         let mut command = if sudo {
@@ -103,8 +116,16 @@ mod arch {
         command.arg("-o");
         command.arg("cargo-flamegraph.stacks");
 
-        command.arg("-c");
-        command.args(&workload);
+        match workload {
+            Workload::Command(c) => {
+                command.arg("-c");
+                command.args(&c);
+            },
+            Workload::Pid(p) => {
+                command.arg("-p");
+                command.arg(p.to_string());
+            },
+        }
 
         command
     }
@@ -147,10 +168,10 @@ fn terminated_by_error(status: ExitStatus) -> bool {
     !status.success()
 }
 
-pub fn generate_flamegraph_by_running_command<
+pub fn generate_flamegraph_for_workload<
     P: AsRef<std::path::Path>,
 >(
-    workload: Vec<String>,
+    workload: Workload,
     flamegraph_filename: P,
     sudo: bool,
 ) {