imp(cli): add log filtering

Author: omarabid

pyroscope_cli/Cargo.lock

@@ -25,6 +25,7 @@ utils = { path = "utils" }
 cli = { path = "cli"}
 core = { path = "core" }
 human-panic = "1.0.3"
+pretty_env_logger = "0.4.0"
 better-panic = "0.3.0"
 log = "=0.4.11"
 

pyroscope_cli/Cargo.toml

@@ -8,7 +8,7 @@ edition = "2021"
 [dependencies]
 thiserror = "1.0.30"
 color-backtrace = "0.5.1"
-config = "0.11.0"
+config = "=0.11.0"
 lazy_static = "1.4.0"
 slog = "2.7.0"
 slog-term = "2.8.0"

pyroscope_cli/utils/Cargo.toml

@@ -20,20 +20,22 @@ pub struct AppConfig {
     pub pid: Option<i32>,
     pub spy_name: Spy,
     pub application_name: Option<String>,
-    // TODO: placeholder for future implementation
-    //pub auth_token: Option<String>,
     pub detect_subprocesses: Option<bool>,
     pub no_logging: Option<bool>,
     pub log_level: LogLevel,
     pub no_root_drop: Option<bool>,
-    pub pyspy_blocking: Option<bool>,
     pub rbspy_blocking: Option<bool>,
+    pub pyspy_blocking: Option<bool>,
+    pub pyspy_idle: Option<bool>,
+    pub pyspy_gil: Option<bool>,
+    pub pyspy_native: Option<bool>,
     pub sample_rate: Option<u32>,
     pub server_address: Option<String>,
     pub tag: Option<String>,
     // TODO: placeholder for future implementation
     //pub upstream_request_timeout: Option<String>,
     //pub upstream_threads: Option<u32>,
+    //pub auth_token: Option<String>,
     pub user_name: Option<u32>,
     pub group_name: Option<u32>,
     pub command: Option<String>,

pyroscope_cli/utils/src/app_config.rs

@@ -1,7 +1,9 @@
 use slog::o;
 use slog::Drain;
 
+use super::app_config::AppConfig;
 use super::error::Result;
+use super::types::LogLevel;
 
 pub fn setup_logging() -> Result<slog_scope::GlobalLoggerGuard> {
     // Setup Logging
@@ -15,9 +17,11 @@ pub fn default_root_logger() -> Result<slog::Logger> {
     // Create drains
     let drain = slog::Duplicate(default_discard()?, default_discard()?).fuse();
 
-    // TODO: Link logging to config
     // Merge drains
-    //let drain = slog::Duplicate(default_term_drain().unwrap_or(default_discard()?), drain).fuse();
+    let drain = slog::Duplicate(default_term_drain().unwrap_or(default_discard()?), drain).fuse();
+
+    // Filter drain
+    let drain = LevelFilter { drain }.fuse();
 
     // Create Logger
     let logger = slog::Logger::root(drain, o!("who" => "pyroscope-cli"));
@@ -34,10 +38,52 @@ fn default_discard() -> Result<slog_async::Async> {
 
 // term drain: Log to Terminal
 fn default_term_drain() -> Result<slog_async::Async> {
-    let plain = slog_term::PlainSyncDecorator::new(std::io::stdout());
-    let term = slog_term::FullFormat::new(plain);
+    let decorator = slog_term::TermDecorator::new().build();
+    let term = slog_term::FullFormat::new(decorator);
 
     let drain = slog_async::Async::default(term.build().fuse());
 
     Ok(drain)
 }
+
+struct LevelFilter<D> {
+    drain: D,
+}
+
+impl<D> Drain for LevelFilter<D>
+where
+    D: Drain,
+{
+    type Ok = Option<D::Ok>;
+    type Err = Option<D::Err>;
+
+    fn log(
+        &self, record: &slog::Record, values: &slog::OwnedKVList,
+    ) -> std::result::Result<Self::Ok, Self::Err> {
+        // Discard logs if no_logging flag is set
+        let logging = AppConfig::get::<bool>("no_logging").unwrap_or(false);
+        if logging {
+            return Ok(None);
+        }
+
+        // TODO: This is probably expensive (and should be cached)
+        let log_level = AppConfig::get::<LogLevel>("log_level").unwrap_or(LogLevel::Info);
+
+        // convert log level to slog level
+        let current_level = match log_level {
+            LogLevel::Trace => slog::Level::Trace,
+            LogLevel::Debug => slog::Level::Debug,
+            LogLevel::Info => slog::Level::Info,
+            LogLevel::Warn => slog::Level::Warning,
+            LogLevel::Error => slog::Level::Error,
+            LogLevel::Critical => slog::Level::Critical,
+        };
+
+        // check if log level is above current level
+        if record.level().is_at_least(current_level) {
+            self.drain.log(record, values).map(Some).map_err(Some)
+        } else {
+            Ok(None)
+        }
+    }
+}

pyroscope_cli/utils/src/logger.rs

@@ -6,6 +6,8 @@ use std::str::FromStr;
 
 #[derive(Debug, Serialize, Deserialize, Copy, Clone, ArgEnum)]
 pub enum LogLevel {
+    #[serde(rename = "trace")]
+    Trace,
     #[serde(rename = "debug")]
     Debug,
     #[serde(rename = "info")]
@@ -14,6 +16,8 @@ pub enum LogLevel {
     Warn,
     #[serde(rename = "error")]
     Error,
+    #[serde(rename = "critical")]
+    Critical,
 }
 
 impl FromStr for LogLevel {