From ae9dba8c8ca39564f5cbd65b39093c9c432b9341 Mon Sep 17 00:00:00 2001 From: Qihang Cai Date: Thu, 18 Jul 2024 17:51:17 +0800 Subject: [PATCH 1/4] fix: can't exit with `Ctrl+C` after integrating with `ztm` Signed-off-by: Qihang Cai --- mega/Cargo.toml | 1 + mega/src/cli.rs | 5 +++++ 2 files changed, 6 insertions(+) diff --git a/mega/Cargo.toml b/mega/Cargo.toml index 97a004acd..6cb33524f 100644 --- a/mega/Cargo.toml +++ b/mega/Cargo.toml @@ -29,6 +29,7 @@ rand = { workspace = true } smallvec = { workspace = true } config = { workspace = true } shadow-rs = { workspace = true } +ctrlc = "3.4.4" [dev-dependencies] reqwest = { workspace = true , features = ["stream", "json"] } diff --git a/mega/src/cli.rs b/mega/src/cli.rs index fd496a4fe..7c38c28f3 100644 --- a/mega/src/cli.rs +++ b/mega/src/cli.rs @@ -42,6 +42,11 @@ pub fn parse(args: Option>) -> MegaResult { init_log(&config.log); + ctrlc::set_handler(move || { + tracing::info!("Received Ctrl-C signal, exiting..."); + std::process::exit(0); + }).unwrap(); + let (cmd, subcommand_args) = match matches.subcommand() { Some((cmd, args)) => (cmd, args), _ => { From 7f831103e64c4318d8442ffecdd2119a61e1591a Mon Sep 17 00:00:00 2001 From: Qihang Cai Date: Fri, 19 Jul 2024 00:05:16 +0800 Subject: [PATCH 2/4] add `${}` syntax for `Config` & add `base_dir` & support env_var Signed-off-by: Qihang Cai --- common/Cargo.toml | 1 + common/src/config.rs | 75 ++++++++++++++++++++++++++++++++++++++++++-- mega/config.toml | 16 ++++++---- 3 files changed, 84 insertions(+), 8 deletions(-) diff --git a/common/Cargo.toml b/common/Cargo.toml index d8343cd1d..b72dcae6f 100644 --- a/common/Cargo.toml +++ b/common/Cargo.toml @@ -17,3 +17,4 @@ clap = { workspace = true, features = ["derive"] } idgenerator = { workspace = true } serde = { workspace = true } config = { workspace = true } +envsubst = "0.2.1" diff --git a/common/src/config.rs b/common/src/config.rs index 319245864..d494a7688 100644 --- a/common/src/config.rs +++ b/common/src/config.rs @@ -1,10 +1,16 @@ +use std::cell::RefCell; +use std::collections::HashMap; use c::{ConfigError, FileFormat}; use config as c; use serde::{Deserialize, Serialize}; use std::path::PathBuf; +use std::sync::Arc; +use config::builder::DefaultState; +use config::Source; #[derive(Serialize, Deserialize, Debug, Default, Clone)] pub struct Config { + pub base_dir: PathBuf, pub log: LogConfig, pub database: DbConfig, pub ssh: SshConfig, @@ -17,8 +23,11 @@ pub struct Config { impl Config { pub fn new(path: &str) -> Result { - let builder = c::Config::builder().add_source(c::File::new(path, FileFormat::Toml)); - let config = builder.build().unwrap(); + let builder = c::Config::builder() + .add_source(c::File::new(path, FileFormat::Toml)) + .add_source(c::Environment::with_prefix("mega")); // e.g. MEGA_BASE_DIR == base_dir + // support ${} variable substitution + let config = variable_placeholder_substitute(builder); Config::from_config(config) } @@ -29,6 +38,68 @@ impl Config { } } +/// supports braces-delimited variables (i.e. ${foo}). +/// - only top-level variables are supported. +/// - only support `String` type. +/// - do not support recursive substitution +fn variable_placeholder_substitute(mut builder: c::ConfigBuilder) -> c::Config { + // `Config::set` is deprecated, use `ConfigBuilder::set_override` instead + let config = builder.clone().build().unwrap(); + let mut vars = HashMap::new(); + // top-level variables + for (k, mut v) in config.collect().unwrap() { // a copy + if let c::ValueKind::String(s) = &v.kind { + if envsubst::is_templated(&s) { // do not support recursive substitution + let new_str = envsubst::substitute(s, &vars).unwrap(); + v.kind = c::ValueKind::String(new_str); + builder = builder.set_override(k, v).unwrap(); + } else { + println!("{}: {:?}", k, s); + vars.insert(k, s.clone()); + } + } + } + // second-level or nested variables + let map = Arc::new(RefCell::new(HashMap::new())); + for (k, v) in config.collect().unwrap() { + if let c::ValueKind::Table(_) = v.kind { + let map_c = map.clone(); + traverse_config(&k, &v, Arc::new(move |key, value| { + if let c::ValueKind::String(str) = &value.kind { + if envsubst::is_templated(&str) { + map_c.borrow_mut().insert(key.to_string(), value.clone()); + } + } + })); + } + } + + // do substitution: ${} -> value + for (k, mut v) in Arc::try_unwrap(map).unwrap().into_inner() { + let old_str = v.clone().into_string().unwrap(); + let new_str = envsubst::substitute(&old_str, &vars).unwrap(); + println!("{}: {} -> {}", k, old_str, &new_str); + v.kind = c::ValueKind::String(new_str); + builder = builder.set_override(&k, v).unwrap(); + } + + builder.build().unwrap() +} + +/// visitor pattern: traverse each config & execute the closure `f` +fn traverse_config(key: &str, value: &c::Value, f: Arc) { + match &value.kind { + c::ValueKind::Table(table) => { + for (k, v) in table.iter() { + // join keys by '.' + let new_key = if key.is_empty() { k.clone() } else { format!("{}.{}", key, k) }; + traverse_config(&new_key, v, f.clone()); + } + } + _ => f(key, value), + } +} + #[derive(Serialize, Deserialize, Debug, Clone)] pub struct LogConfig { pub log_path: PathBuf, diff --git a/mega/config.toml b/mega/config.toml index e30d25f86..e8417e79d 100644 --- a/mega/config.toml +++ b/mega/config.toml @@ -1,8 +1,12 @@ +# the directory where the data files is located, such as logs, database, etc. +# can be overrided by environment variable `MEGA_BASE_DIR` +base_dir = "/tmp/.mega" + # Filling the following environment variables with values you set ## Logging Configuration [log] # The path which log file is saved -log_path = "/tmp/.mega/logs" +log_path = "${base_dir}/logs" # log level level = "debug" @@ -17,7 +21,7 @@ print_std = true db_type = "sqlite" # used for sqlite -db_path = "/tmp/.mega/mega.db" +db_path = "${base_dir}/mega.db" # database connection url db_url = "postgres://mega:mega@localhost:5432/mega" @@ -33,7 +37,7 @@ sqlx_logging = false [ssh] -ssh_key_path = "/tmp/.mega/ssh" +ssh_key_path = "${base_dir}/ssh" [storage] # raw object stroage type, can be `local` or `remote` @@ -43,9 +47,9 @@ raw_obj_storage_type = "LOCAL" big_obj_threshold = 1024 # set the local path of the project storage -raw_obj_local_path = "/tmp/.mega/objects" +raw_obj_local_path = "${base_dir}/objects" -lfs_obj_local_path = "/tmp/.mega/lfs" +lfs_obj_local_path = "${base_dir}/lfs" obs_access_key = "" obs_secret_key = "" @@ -68,7 +72,7 @@ import_dir = "/third-part" pack_decode_mem_size = 4 # The location where the object stored when the memory used by decode exceeds the limit -pack_decode_cache_path = "/tmp/.mega/cache" +pack_decode_cache_path = "${base_dir}/cache" clean_cache_after_decode = true From 15bb9fc34b74295ff9dfd41c3ea8001e3c375233 Mon Sep 17 00:00:00 2001 From: Qihang Cai Date: Fri, 19 Jul 2024 14:52:26 +0800 Subject: [PATCH 3/4] support recursive substitute for `${}` in config Signed-off-by: Qihang Cai --- common/src/config.rs | 71 +++++++++++++++++++++++++------------------- 1 file changed, 40 insertions(+), 31 deletions(-) diff --git a/common/src/config.rs b/common/src/config.rs index d494a7688..223572ce6 100644 --- a/common/src/config.rs +++ b/common/src/config.rs @@ -4,9 +4,9 @@ use c::{ConfigError, FileFormat}; use config as c; use serde::{Deserialize, Serialize}; use std::path::PathBuf; -use std::sync::Arc; +use std::rc::Rc; use config::builder::DefaultState; -use config::Source; +use config::{Source, ValueKind}; #[derive(Serialize, Deserialize, Debug, Default, Clone)] pub struct Config { @@ -38,62 +38,71 @@ impl Config { } } -/// supports braces-delimited variables (i.e. ${foo}). -/// - only top-level variables are supported. +/// supports braces-delimited variables (i.e. ${foo}) in config. +/// ### Example: +/// ```toml +/// base_dir = "/tmp/.mega" +/// [log] +/// log_path = "${base_dir}/logs" +/// ``` +/// ### Limitations: /// - only support `String` type. -/// - do not support recursive substitution +/// - vars apply from up to down fn variable_placeholder_substitute(mut builder: c::ConfigBuilder) -> c::Config { // `Config::set` is deprecated, use `ConfigBuilder::set_override` instead - let config = builder.clone().build().unwrap(); + let config = builder.clone().build().unwrap(); // initial config let mut vars = HashMap::new(); // top-level variables for (k, mut v) in config.collect().unwrap() { // a copy - if let c::ValueKind::String(s) = &v.kind { - if envsubst::is_templated(&s) { // do not support recursive substitution - let new_str = envsubst::substitute(s, &vars).unwrap(); - v.kind = c::ValueKind::String(new_str); - builder = builder.set_override(k, v).unwrap(); + if let ValueKind::String(str) = &v.kind { + if envsubst::is_templated(str) { + let new_str = envsubst::substitute(str, &vars).unwrap(); + v.kind = ValueKind::String(new_str.clone()); + builder = builder.set_override(&k, v).unwrap(); + vars.insert(k, new_str); } else { - println!("{}: {:?}", k, s); - vars.insert(k, s.clone()); + vars.insert(k, str.clone()); } } } // second-level or nested variables - let map = Arc::new(RefCell::new(HashMap::new())); + // extract all config k-v + let map = Rc::new(RefCell::new(HashMap::new())); for (k, v) in config.collect().unwrap() { - if let c::ValueKind::Table(_) = v.kind { + if let ValueKind::Table(_) = v.kind { let map_c = map.clone(); - traverse_config(&k, &v, Arc::new(move |key, value| { - if let c::ValueKind::String(str) = &value.kind { - if envsubst::is_templated(&str) { - map_c.borrow_mut().insert(key.to_string(), value.clone()); - } + traverse_config(&k, &v, &move |key: &str, value: &c::Value| { + if let ValueKind::String(_) = value.kind { + map_c.borrow_mut().insert(key.to_string(), value.clone()); } - })); + }); } } - // do substitution: ${} -> value - for (k, mut v) in Arc::try_unwrap(map).unwrap().into_inner() { - let old_str = v.clone().into_string().unwrap(); - let new_str = envsubst::substitute(&old_str, &vars).unwrap(); - println!("{}: {} -> {}", k, old_str, &new_str); - v.kind = c::ValueKind::String(new_str); - builder = builder.set_override(&k, v).unwrap(); + // do substitution: ${} -> real value + for (k, mut v) in Rc::try_unwrap(map).unwrap().into_inner() { + let mut str = v.clone().into_string().unwrap(); + if envsubst::is_templated(&str) { + let new_str = envsubst::substitute(&str, &vars).unwrap(); + println!("{}: {} -> {}", k, str, &new_str); + v.kind = ValueKind::String(new_str.clone()); + builder = builder.set_override(&k, v).unwrap(); + str = new_str; + } + vars.insert(k, str); } builder.build().unwrap() } /// visitor pattern: traverse each config & execute the closure `f` -fn traverse_config(key: &str, value: &c::Value, f: Arc) { +fn traverse_config(key: &str, value: &c::Value, f: &impl Fn(&str, &c::Value)) { match &value.kind { - c::ValueKind::Table(table) => { + ValueKind::Table(table) => { for (k, v) in table.iter() { // join keys by '.' let new_key = if key.is_empty() { k.clone() } else { format!("{}.{}", key, k) }; - traverse_config(&new_key, v, f.clone()); + traverse_config(&new_key, v, f); } } _ => f(key, value), From a2ca472e81becccbbfc46cf5d27ad6d8988333ec Mon Sep 17 00:00:00 2001 From: Qihang Cai Date: Sat, 20 Jul 2024 16:33:16 +0800 Subject: [PATCH 4/4] fix: param `--config`, type `String` to `PathBuf`, to support quotation (`""`) Signed-off-by: Qihang Cai --- mega/src/cli.rs | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/mega/src/cli.rs b/mega/src/cli.rs index 7c38c28f3..0cffcc849 100644 --- a/mega/src/cli.rs +++ b/mega/src/cli.rs @@ -1,7 +1,7 @@ //! Cli module is responsible for parsing command line arguments and executing the appropriate. use std::env; - +use std::path::PathBuf; use clap::{Arg, ArgMatches, Command}; use tracing_subscriber::fmt::writer::MakeWriterExt; @@ -31,8 +31,8 @@ pub fn parse(args: Option>) -> MegaResult { // Get the path to the config file in the current directory let config_path = current_dir.join("config.toml"); - let config = if let Some(path) = matches.get_one::("config").cloned() { - Config::new(path.as_str()).unwrap() + let config = if let Some(path) = matches.get_one::("config").cloned() { + Config::new(path.to_str().unwrap()).unwrap() } else if config_path.exists() { Config::new(config_path.to_str().unwrap()).unwrap() } else {