diff --git a/appmgr/Cargo.lock b/appmgr/Cargo.lock index 66e52ce10..084d50f46 100644 --- a/appmgr/Cargo.lock +++ b/appmgr/Cargo.lock @@ -804,6 +804,7 @@ dependencies = [ "rpc-toolkit", "rust-argon2", "scopeguard", + "sequence_trie", "serde", "serde_json", "serde_yaml", @@ -2580,6 +2581,12 @@ dependencies = [ "pest", ] +[[package]] +name = "sequence_trie" +version = "0.3.6" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "1ee22067b7ccd072eeb64454b9c6e1b33b61cd0d49e895fd48676a184580e0c3" + [[package]] name = "serde" version = "1.0.130" diff --git a/appmgr/Cargo.toml b/appmgr/Cargo.toml index 7e616c848..89e128052 100644 --- a/appmgr/Cargo.toml +++ b/appmgr/Cargo.toml @@ -89,6 +89,7 @@ rpassword = "5.0.1" rpc-toolkit = { version = "*", path = "../rpc-toolkit/rpc-toolkit" } rust-argon2 = "0.8.3" scopeguard = "1.1" # because avahi-sys fucks your shit up +sequence_trie = "0.3.6" serde = { version = "1.0.130", features = ["derive", "rc"] } serde_cbor = { package = "ciborium", version = "0.1.0" } serde_json = "1.0.68" diff --git a/appmgr/src/bin/embassyd.rs b/appmgr/src/bin/embassyd.rs index c0f7b0bd5..26f25600d 100644 --- a/appmgr/src/bin/embassyd.rs +++ b/appmgr/src/bin/embassyd.rs @@ -1,3 +1,4 @@ +use std::collections::HashMap; use std::time::Duration; use anyhow::anyhow; @@ -36,8 +37,9 @@ fn err_to_500(e: Error) -> Response { async fn inner_main( cfg_path: Option<&str>, log_level: LevelFilter, + module_logging: HashMap, ) -> Result, Error> { - let rpc_ctx = RpcContext::init(cfg_path, log_level).await?; + let rpc_ctx = RpcContext::init(cfg_path, log_level, module_logging).await?; let mut shutdown_recv = rpc_ctx.shutdown.subscribe(); let sig_handler_ctx = rpc_ctx.clone(); @@ -252,6 +254,12 @@ fn main() { .multiple(true) .takes_value(false), ) + .arg( + clap::Arg::with_name("log-module") + .long("log-module") + .multiple(true) + .takes_value(true), + ) .get_matches(); // initializes the bootstrap logger, this will be replaced with the EmbassyLogger later @@ -262,6 +270,19 @@ fn main() { 3 => log::LevelFilter::Debug, _ => log::LevelFilter::Trace, }; + let module_logging = matches + .values_of("log-module") + .into_iter() + .flatten() + .filter_map(|s| s.split_once("=")) + .map(|(m, l)| { + ( + m.to_owned(), + l.parse() + .expect(&format!("Invalid 'log-module' argument: {}", l)), + ) + }) + .collect::>(); let cfg_path = matches.value_of("config"); let res = { @@ -270,7 +291,7 @@ fn main() { .build() .expect("failed to initialize runtime"); rt.block_on(async { - match inner_main(cfg_path, filter).await { + match inner_main(cfg_path, filter, module_logging).await { Ok(a) => Ok(a), Err(e) => { (|| async { diff --git a/appmgr/src/context/rpc.rs b/appmgr/src/context/rpc.rs index ed6f38fea..e1c4207d7 100644 --- a/appmgr/src/context/rpc.rs +++ b/appmgr/src/context/rpc.rs @@ -1,5 +1,5 @@ use std::borrow::Cow; -use std::collections::VecDeque; +use std::collections::{HashMap, VecDeque}; use std::net::{IpAddr, SocketAddr}; use std::ops::Deref; use std::path::{Path, PathBuf}; @@ -130,11 +130,17 @@ impl RpcContext { pub async fn init>( cfg_path: Option

, log_level: LevelFilter, + module_logging: HashMap, ) -> Result { let base = RpcContextConfig::load(cfg_path).await?; let log_epoch = Arc::new(AtomicU64::new(rand::random())); - let logger = - EmbassyLogger::init(log_level, log_epoch.clone(), base.log_server.clone(), false); + let logger = EmbassyLogger::init( + log_level, + log_epoch.clone(), + base.log_server.clone(), + false, + module_logging, + ); let (shutdown, _) = tokio::sync::broadcast::channel(1); let secret_store = base.secret_store().await?; let db = base.db(&secret_store).await?; diff --git a/appmgr/src/util/logger.rs b/appmgr/src/util/logger.rs index 58e230568..66d15ba57 100644 --- a/appmgr/src/util/logger.rs +++ b/appmgr/src/util/logger.rs @@ -4,8 +4,84 @@ use std::sync::Arc; use log::{set_boxed_logger, set_max_level, LevelFilter, Metadata, Record}; use reqwest::{Client, Url}; +use sequence_trie::SequenceTrie; use stderrlog::{StdErrLog, Timestamp}; +#[derive(Clone, Debug)] +pub struct ModuleMap { + trie: SequenceTrie, +} +impl ModuleMap { + fn new(vals: HashMap) -> Self { + let mut module_trie = SequenceTrie::new(); + for (k, v) in vals { + let mut include_submodules = false; + let module_key = k + .split("::") + .take_while(|&s| { + if s == "*" { + include_submodules = true; + false + } else { + true + } + }) + .collect::>(); + match module_trie.get_node(module_key.clone()) { + None => match module_trie.get_ancestor_node(module_key.clone()) { + None => { + module_trie.insert(module_key, (v, include_submodules)); + } + Some(ancestor) => match ancestor.value() { + None => { + module_trie.insert(module_key, (v, include_submodules)); + } + Some((_, sub)) => { + if !sub { + module_trie.insert(module_key, (v, include_submodules)); + } + } + }, + }, + Some(n) => match n.value() { + None => { + module_trie.insert(module_key.clone(), (v, include_submodules)); + if include_submodules { + let new_node = module_trie.get_node_mut(module_key).unwrap(); // we just inserted it + let child_keys = new_node + .children_with_keys() + .into_iter() + .map(|x| x.0.clone()) + .collect::>(); + for c in child_keys { + new_node.remove(std::iter::once(&c)); + } + } + } + Some(_) => unreachable!("Trie build failed on 'impossible' duplicate"), + }, + } + } + ModuleMap { trie: module_trie } + } + fn level_for<'a>(&'a self, k: &str) -> &'a LevelFilter { + let module_key = k.split("::"); + match self.trie.get(module_key.clone()) { + None => match self.trie.get_ancestor(module_key) { + None => &LevelFilter::Off, + Some((level_filter, include_submodules)) => { + if *include_submodules { + level_filter + } else { + &LevelFilter::Off + } + } + }, + Some((level_filter, _)) => level_filter, + } + } +} + #[derive(Clone)] pub struct EmbassyLogger { log_level: log::LevelFilter, @@ -13,6 +89,7 @@ pub struct EmbassyLogger { logger: StdErrLog, sharing: Arc, share_dest: Url, + module_map: ModuleMap, } impl EmbassyLogger { pub fn init( @@ -20,6 +97,7 @@ impl EmbassyLogger { log_epoch: Arc, share_dest: Option, share_errors: bool, + module_map: HashMap, ) -> Self { let share_dest = match share_dest { None => Url::parse("https://beta-registry-0-3.start9labs.com/error-logs").unwrap(), // TODO @@ -41,6 +119,7 @@ impl EmbassyLogger { logger, sharing: Arc::new(AtomicBool::new(share_errors)), share_dest: share_dest, + module_map: ModuleMap::new(module_map), }; set_boxed_logger(Box::new(embassy_logger.clone())).unwrap(); set_max_level(log_level); @@ -53,10 +132,17 @@ impl EmbassyLogger { impl log::Log for EmbassyLogger { fn enabled(&self, metadata: &Metadata) -> bool { - self.logger.enabled(metadata) + let top = metadata.target().split("::").next().unwrap(); + if vec!["embassy", "embassyd", "embassy-cli", "embassy-sdk"].contains(&top) { + metadata.level() <= self.log_level + } else { + &metadata.level() <= self.module_map.level_for(metadata.target()) + } } fn log(&self, record: &Record) { - self.logger.log(record); + if self.enabled(record.metadata()) { + self.logger.log(record); + } if self.sharing.load(Ordering::SeqCst) { if record.level() <= log::Level::Warn { let mut body = HashMap::new(); @@ -89,3 +175,53 @@ pub async fn order_level() { pub fn module() { println!("{}", module_path!()) } + +proptest::proptest! { + #[test] + fn submodules_handled_by_parent(s0 in "[a-z][a-z0-9_]+", s1 in "[a-z][a-z0-9_]+", level in filter_strategy()) { + proptest::prop_assume!(level > LevelFilter::Off); + let mut hm = HashMap::new(); + hm.insert(format!("{}::*", s0.clone()), level); + let mod_map = ModuleMap::new(hm); + proptest::prop_assert_eq!(mod_map.level_for(&format!("{}::{}", s0, s1)), &level) + } + #[test] + fn submodules_ignored_by_parent(s0 in "[a-z][a-z0-9_]+", s1 in "[a-z][a-z0-9_]+", level in filter_strategy()) { + proptest::prop_assume!(level > LevelFilter::Off); + let mut hm = HashMap::new(); + hm.insert(s0.clone(), level); + let mod_map = ModuleMap::new(hm); + proptest::prop_assert_eq!(mod_map.level_for(&format!("{}::{}", s0, s1)), &LevelFilter::Off) + } + #[test] + fn duplicate_insertion_ignored(s0 in "[a-z][a-z0-9_]+", s1 in "[a-z][a-z0-9_]+", level in filter_strategy()) { + proptest::prop_assume!(level > LevelFilter::Off); + let mut hm = HashMap::new(); + hm.insert(format!("{}::*", s0.clone()), level); + let sub = format!("{}::{}", s0, s1); + hm.insert(sub.clone(), level); + let mod_map = ModuleMap::new(hm); + proptest::prop_assert_eq!(mod_map.trie.get(sub.split("::")), None) + } + #[test] + fn parent_child_simul(s0 in "[a-z][a-z0-9_]+", s1 in "[a-z][a-z0-9_]+", level0 in filter_strategy(), level1 in filter_strategy()) { + let mut hm = HashMap::new(); + hm.insert(s0.clone(), level0); + let sub = format!("{}::{}", s0, s1); + hm.insert(sub.clone(), level1); + let mod_map = ModuleMap::new(hm); + proptest::prop_assert_eq!(mod_map.level_for(&s0), &level0); + proptest::prop_assert_eq!(mod_map.level_for(&sub), &level1); + } +} +fn filter_strategy() -> impl proptest::strategy::Strategy { + use proptest::strategy::Just; + proptest::prop_oneof![ + Just(LevelFilter::Off), + Just(LevelFilter::Error), + Just(LevelFilter::Warn), + Just(LevelFilter::Info), + Just(LevelFilter::Debug), + Just(LevelFilter::Trace), + ] +}