Profile by pid

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,
     );

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,

src/lib.rs

@@ -26,6 +26,11 @@ use inferno::{
 
 use signal_hook;
 
+pub enum Workload {
+    Command(Vec<String>),
+    Pid(u32),
+}
+
 #[cfg(target_os = "linux")]
 mod arch {
     use super::*;
@@ -37,7 +42,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 {
@@ -54,7 +59,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
     }
@@ -79,7 +92,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 {
@@ -102,8 +115,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
     }
@@ -146,10 +167,10 @@ fn terminated_by_error(status: ExitStatus) -> bool {
     !exit_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,
 ) {