This is an automated email from the ASF dual-hosted git repository.
zuston pushed a commit to branch master
in repository https://gitbox.apache.org/repos/asf/incubator-uniffle.git
The following commit(s) were added to refs/heads/master by this push:
new 541388b53 [#1237] feat(rust): support populating args by clap (#1236)
541388b53 is described below
commit 541388b5322488b2b449f3a20440db07313272f4
Author: Junfan Zhang <[email protected]>
AuthorDate: Tue Oct 17 09:34:08 2023 +0800
[#1237] feat(rust): support populating args by clap (#1236)
### What changes were proposed in this pull request?
This PR is to support populating args by clap.
### Why are the changes needed?
Before this, the config path is specified by the environment key, which is
not elegant.
### Does this PR introduce _any_ user-facing change?
Yes
### How was this patch tested?
Existing tests.
---
rust/experimental/server/Cargo.lock | 1 +
rust/experimental/server/Cargo.toml | 1 +
rust/experimental/server/README.md | 2 +-
rust/experimental/server/src/config.rs | 16 ++++++++++------
rust/experimental/server/src/main.rs | 18 +++++++++++++++++-
5 files changed, 30 insertions(+), 8 deletions(-)
diff --git a/rust/experimental/server/Cargo.lock
b/rust/experimental/server/Cargo.lock
index 73344c2bb..3d21a2437 100644
--- a/rust/experimental/server/Cargo.lock
+++ b/rust/experimental/server/Cargo.lock
@@ -3175,6 +3175,7 @@ dependencies = [
"async-trait",
"await-tree",
"bytes 1.5.0",
+ "clap",
"console-subscriber",
"crc32fast",
"croaring",
diff --git a/rust/experimental/server/Cargo.toml
b/rust/experimental/server/Cargo.toml
index 9f03bd468..ba8f20ee5 100644
--- a/rust/experimental/server/Cargo.toml
+++ b/rust/experimental/server/Cargo.toml
@@ -93,6 +93,7 @@ tokio-console = "0.1.8"
console-subscriber = "0.1.9"
pin-project-lite = "0.2.8"
signal-hook = "0.3.17"
+clap = "3.0.14"
[dependencies.hdrs]
version = "0.3.0"
diff --git a/rust/experimental/server/README.md
b/rust/experimental/server/README.md
index ad5515d42..52ef6a3a5 100644
--- a/rust/experimental/server/README.md
+++ b/rust/experimental/server/README.md
@@ -141,7 +141,7 @@ cargo build --features hdfs --release
```
2. worker run with tokio-console. the log level of `trace` must be enabled
```shell
- WORKER_IP={ip} RUST_LOG=trace WORKER_CONFIG_PATH=./config.toml
./uniffle-worker
+ WORKER_IP={ip} RUST_LOG=trace ./uniffle-worker -c ./config.toml
```
3. tokio-console client side connect
```shell
diff --git a/rust/experimental/server/src/config.rs
b/rust/experimental/server/src/config.rs
index ed2864693..85e37d0a7 100644
--- a/rust/experimental/server/src/config.rs
+++ b/rust/experimental/server/src/config.rs
@@ -232,6 +232,15 @@ impl StorageType {
const CONFIG_FILE_PATH_KEY: &str = "WORKER_CONFIG_PATH";
impl Config {
+ pub fn from(cfg_path: &str) -> Self {
+ let path = Path::new(cfg_path);
+
+ // Read the file content as a string
+ let file_content = fs::read_to_string(path).expect("Failed to read
file");
+
+ toml::from_str(&file_content).unwrap()
+ }
+
pub fn create_from_env() -> Config {
let path = match std::env::var(CONFIG_FILE_PATH_KEY) {
Ok(val) => val,
@@ -241,12 +250,7 @@ impl Config {
),
};
- let path = Path::new(&path);
-
- // Read the file content as a string
- let file_content = fs::read_to_string(path).expect("Failed to read
file");
-
- toml::from_str(&file_content).unwrap()
+ Config::from(&path)
}
}
diff --git a/rust/experimental/server/src/main.rs
b/rust/experimental/server/src/main.rs
index 5191e4084..b92c453ac 100644
--- a/rust/experimental/server/src/main.rs
+++ b/rust/experimental/server/src/main.rs
@@ -32,6 +32,7 @@ use crate::signal::details::wait_for_signal;
use crate::util::{gen_worker_uid, get_local_ip};
use anyhow::Result;
+use clap::{App, Arg};
use log::info;
use std::net::{IpAddr, Ipv4Addr, SocketAddr};
use std::time::Duration;
@@ -166,7 +167,22 @@ fn init_log(log: &LogConfig) -> WorkerGuard {
}
fn main() -> Result<()> {
- let config = Config::create_from_env();
+ let args_match = App::new("Uniffle Worker")
+ .version("0.9.0-SNAPSHOT")
+ .about("Rust based shuffle server for Apache Uniffle")
+ .arg(
+ Arg::with_name("config")
+ .short('c')
+ .long("config")
+ .value_name("FILE")
+ .default_value("./config.toml")
+ .help("Sets a custom config file")
+ .takes_value(true),
+ )
+ .get_matches();
+
+ let config_path = args_match.value_of("config").unwrap_or("./config.toml");
+ let config = Config::from(config_path);
let runtime_manager = RuntimeManager::from(config.runtime_config.clone());
