imp(cli): add gil/idle/native pyspy arguments

Author: omarabid

pyroscope_cli/cli/src/lib.rs

@@ -94,6 +94,13 @@ enum Commands {
             parse(try_from_str)
         )]
         pid: i32,
+        #[clap(
+            name = "rbspy_blocking",
+            long = "rbspy-blocking",
+            value_name = "RBSPY_BLOCKING",
+            help = "enable blocking mode for rbspy"
+        )]
+        rbspy_blocking: bool,
         #[clap(
             name = "pyspy_blocking",
             long = "pyspy-blocking",
@@ -102,12 +109,26 @@ enum Commands {
         )]
         pyspy_blocking: bool,
         #[clap(
-            name = "rbspy_blocking",
-            long = "rbspy-blocking",
-            value_name = "RBSPY_BLOCKING",
-            help = "enable blocking mode for rbspy"
+            name = "pyspy_idle",
+            long = "pyspy-idle",
+            value_name = "PYSPY_IDLE",
+            help = "include idle threads for pyspy"
         )]
-        rbspy_blocking: bool,
+        pyspy_idle: bool,
+        #[clap(
+            name = "pyspy_gil",
+            long = "pyspy-gil",
+            value_name = "PYSPY_GIL",
+            help = "enable GIL mode for pyspy"
+        )]
+        pyspy_gil: bool,
+        #[clap(
+            name = "pyspy_native",
+            long = "pyspy-native",
+            value_name = "PYSPY_NATIVE",
+            help = "enable native extensions profiling for pyspy"
+        )]
+        pyspy_native: bool,
         #[clap(
             name = "sample_rate",
             long = "sample-rate",
@@ -211,6 +232,13 @@ enum Commands {
         //help = "disable permissions drop when ran under root. use this one if you want to run your command as root"
         //)]
         //no_root_drop: bool,
+        #[clap(
+            name = "rbspy_blocking",
+            long = "rbspy-blocking",
+            value_name = "RBSPY_BLOCKING",
+            help = "enable blocking mode for rbspy"
+        )]
+        rbspy_blocking: bool,
         #[clap(
             name = "pyspy_blocking",
             long = "pyspy-blocking",
@@ -219,12 +247,26 @@ enum Commands {
         )]
         pyspy_blocking: bool,
         #[clap(
-            name = "rbspy_blocking",
-            long = "rbspy-blocking",
-            value_name = "RBSPY_BLOCKING",
-            help = "enable blocking mode for rbspy"
+            name = "pyspy_idle",
+            long = "pyspy-idle",
+            value_name = "PYSPY_IDLE",
+            help = "include idle threads for pyspy"
         )]
-        rbspy_blocking: bool,
+        pyspy_idle: bool,
+        #[clap(
+            name = "pyspy_gil",
+            long = "pyspy-gil",
+            value_name = "PYSPY_GIL",
+            help = "enable GIL mode for pyspy"
+        )]
+        pyspy_gil: bool,
+        #[clap(
+            name = "pyspy_native",
+            long = "pyspy-native",
+            value_name = "PYSPY_NATIVE",
+            help = "enable native extensions profiling for pyspy"
+        )]
+        pyspy_native: bool,
         #[clap(
             name = "sample_rate",
             long = "sample-rate",

pyroscope_cli/core/src/commands.rs

@@ -1,10 +1,6 @@
 use utils::app_config::AppConfig;
 use utils::error::{Error, Result};
-use utils::types::{LogLevel, Spy};
-
-use pyroscope::PyroscopeAgent;
-use pyroscope_pyspy::{Pyspy, PyspyConfig};
-use pyroscope_rbspy::{Rbspy, RbspyConfig};
+use utils::types::Spy;
 
 use ctrlc;
 use std::sync::mpsc::channel;

pyroscope_cli/core/src/profiler.rs

@@ -21,8 +21,11 @@ impl Profiler {
         let sample_rate: u32 = AppConfig::get::<u32>("sample_rate")?;
 
         // TODO: CLI should probably unify this into a single argument
-        let pyspy_blocking: bool = AppConfig::get::<bool>("pyspy_blocking")?;
         let rbspy_blocking: bool = AppConfig::get::<bool>("rbspy_blocking")?;
+        let pyspy_blocking: bool = AppConfig::get::<bool>("pyspy_blocking")?;
+        let pyspy_idle: bool = AppConfig::get::<bool>("pyspy_idle")?;
+        let pyspy_gil: bool = AppConfig::get::<bool>("pyspy_gil")?;
+        let pyspy_native: bool = AppConfig::get::<bool>("pyspy_native")?;
 
         let detect_subprocesses: bool = AppConfig::get::<bool>("detect_subprocesses")?;
 
@@ -34,7 +37,10 @@ impl Profiler {
                 let config = PyspyConfig::new(pid)
                     .sample_rate(sample_rate)
                     .lock_process(pyspy_blocking)
-                    .with_subprocesses(detect_subprocesses);
+                    .with_subprocesses(detect_subprocesses)
+                    .include_idle(pyspy_idle)
+                    .gil_only(pyspy_gil)
+                    .native(pyspy_native);
                 let backend = Pyspy::new(config);
                 PyroscopeAgent::builder(server_address, app_name)
                     .backend(backend)

pyroscope_cli/src/resources/default_config.toml

@@ -3,6 +3,9 @@ log_level = "info"
 no_logging = false
 rbspy_blocking = false
 pyspy_blocking = false
+pyspy_idle = false
+pyspy_gil = false
+pyspy_native = false
 sample_rate = 100
 server_address = "http://localhost:4040"
 debug = false