Skip to content

Commit

Permalink
feat:cli arg to config and test
Browse files Browse the repository at this point in the history
Signed-off-by: Chen Kai <[email protected]>
  • Loading branch information
GrapeBaBa committed Aug 27, 2024
1 parent 2db5234 commit e584c6f
Show file tree
Hide file tree
Showing 4 changed files with 284 additions and 17 deletions.
1 change: 1 addition & 0 deletions bin/archiver/src/api.rs
Original file line number Diff line number Diff line change
Expand Up @@ -166,6 +166,7 @@ mod tests {
origin_block: *blob_test_helper::ORIGIN_BLOCK,
beacon_config: Default::default(),
storage_config: Default::default(),
log_config: Default::default(),
};
let archiver = Archiver::new(beacon_client.clone(), storage, config, shutdown_rx);
(archiver, beacon_client)
Expand Down
5 changes: 5 additions & 0 deletions bin/archiver/src/archiver.rs
Original file line number Diff line number Diff line change
Expand Up @@ -13,6 +13,7 @@ use tokio::sync::Mutex;
use tokio::time::{interval, sleep};
use tracing::log::{debug, error, info, trace};

use crate::LogConfig;
use blob_archiver_beacon::beacon_client;
use blob_archiver_beacon::beacon_client::BeaconClient;
use blob_archiver_storage::storage;
Expand Down Expand Up @@ -56,6 +57,8 @@ pub struct Config {
pub beacon_config: beacon_client::Config,

pub storage_config: storage::Config,

pub log_config: LogConfig,
}

pub struct Archiver {
Expand Down Expand Up @@ -500,6 +503,7 @@ mod tests {
origin_block: *blob_test_helper::ORIGIN_BLOCK,
beacon_config: Default::default(),
storage_config: Default::default(),
log_config: Default::default(),
};
let archiver = Archiver::new(beacon_client.clone(), storage, config, shutdown_rx);
(archiver, beacon_client)
Expand Down Expand Up @@ -1825,6 +1829,7 @@ mod tests {
origin_block: *blob_test_helper::ORIGIN_BLOCK,
beacon_config: Default::default(),
storage_config: Default::default(),
log_config: Default::default(),
};
let archiver = Archiver::new(
Arc::new(Mutex::new(beacon_client_eth2)),
Expand Down
279 changes: 264 additions & 15 deletions bin/archiver/src/main.rs
Original file line number Diff line number Diff line change
@@ -1,24 +1,30 @@
use crate::archiver::{Archiver, Config, STARTUP_FETCH_BLOB_MAXIMUM_RETRIES};
use again::RetryPolicy;
use blob_archiver_beacon::beacon_client::BeaconClientEth2;
use blob_archiver_beacon::blob_test_helper;
use blob_archiver_storage::fs::FSStorage;
use clap::Parser;
use eth2::types::BlockId;
use eth2::{BeaconNodeHttpClient, SensitiveUrl, Timeouts};
use serde::Serialize;
use std::fs;
use std::path::PathBuf;
use std::str::FromStr;
use std::sync::Arc;
use std::time::Duration;

use again::RetryPolicy;
use clap::Parser;
use eth2::types::{BlockId, Hash256};
use eth2::{BeaconNodeHttpClient, SensitiveUrl, Timeouts};
use serde::{Deserialize, Serialize};
use tokio::sync::Mutex;
use tracing::log::error;
use tracing_appender::rolling::{RollingFileAppender, Rotation};
use tracing_subscriber::layer::SubscriberExt;
use tracing_subscriber::util::SubscriberInitExt;
use tracing_subscriber::{fmt, EnvFilter};

use blob_archiver_beacon::beacon_client::BeaconClientEth2;
use blob_archiver_beacon::{beacon_client, blob_test_helper};
use blob_archiver_storage::fs::FSStorage;
use blob_archiver_storage::s3::S3Config;
use blob_archiver_storage::storage;
use blob_archiver_storage::storage::StorageType;

use crate::archiver::{Archiver, Config, STARTUP_FETCH_BLOB_MAXIMUM_RETRIES};

mod api;
mod archiver;

Expand All @@ -27,7 +33,9 @@ static INIT: std::sync::Once = std::sync::Once::new();
#[tokio::main]
async fn main() {
let args = CliArgs::parse();
println!("{}", args.verbose);

let config = args.to_config();
println!("{:#?}", config);
init_logging(0, None, None);
let beacon_client = BeaconNodeHttpClient::new(
SensitiveUrl::from_str("https://ethereum-beacon-api.publicnode.com").unwrap(),
Expand All @@ -42,6 +50,7 @@ async fn main() {
origin_block: *blob_test_helper::ORIGIN_BLOCK,
beacon_config: Default::default(),
storage_config: Default::default(),
log_config: Default::default(),
};
let archiver = Archiver::new(
Arc::new(Mutex::new(beacon_client_eth2)),
Expand Down Expand Up @@ -128,14 +137,14 @@ pub fn setup_tracing(

#[derive(Parser, Serialize)]
pub struct CliArgs {
#[clap(short='v', long, action = clap::ArgAction::Count, default_value = "3")]
#[clap(short = 'v', long, action = clap::ArgAction::Count, default_value = "3")]
verbose: u8,

#[clap(long, default_value = "logs")]
log_dir: String,
#[clap(long)]
log_dir: Option<String>,

#[clap(long, default_value = "DAILY")]
log_rotation: String,
#[clap(long, help = "Log rotation values: DAILY, HOURLY, MINUTELY, NEVER")]
log_rotation: Option<String>,

#[clap(long, required = true)]
beacon_endpoint: String,
Expand Down Expand Up @@ -169,3 +178,243 @@ pub struct CliArgs {
#[clap(long)]
fs_dir: Option<String>,
}

impl CliArgs {
fn to_config(&self) -> Config {
let log_config = LogConfig {
log_dir: self.log_dir.as_ref().map(PathBuf::from),
log_rotation: self.log_rotation.clone(),
verbose: self.verbose,
};

let beacon_config = beacon_client::Config {
beacon_endpoint: self.beacon_endpoint.clone(),
beacon_client_timeout: Duration::from_secs(self.beacon_client_timeout),
};

let storage_type = StorageType::from_str(self.storage_type.as_str()).unwrap();

let s3_config = if storage_type == StorageType::S3 {
Some(S3Config {
endpoint: self.s3_endpoint.clone().unwrap(),
bucket: self.s3_bucket.clone().unwrap(),
path: self.s3_path.clone().unwrap(),
compression: self.s3_compress.unwrap(),
})
} else {
None
};

let fs_dir = self.fs_dir.as_ref().map(PathBuf::from);

let storage_config = storage::Config {
storage_type,
s3_config,
fs_dir,
};

let mut padded_hex = self
.origin_block
.strip_prefix("0x")
.unwrap_or(&self.origin_block)
.to_string();
padded_hex = format!("{:0<64}", padded_hex);
let origin_block = Hash256::from_slice(&hex::decode(padded_hex).expect("Invalid hex"));

Config {
poll_interval: Duration::from_secs(self.poll_interval),
listen_addr: self.listen_addr.clone(),
origin_block,
beacon_config,
storage_config,
log_config,
}
}
}

#[allow(dead_code)]
fn to_rotation(s: &str) -> Rotation {
match s {
"DAILY" => Rotation::DAILY,
"HOURLY" => Rotation::HOURLY,
"MINUTELY" => Rotation::MINUTELY,
_ => Rotation::NEVER,
}
}

#[cfg(test)]
mod tests {
use std::path::PathBuf;
use std::time::Duration;

use eth2::types::Hash256;

use blob_archiver_storage::s3::S3Config;

use crate::storage::StorageType;
use crate::CliArgs;

#[test]
fn test_cli_args_to_config_s3_storage() {
let cli_args = CliArgs {
verbose: 3,
log_dir: Some("logs".to_string()),
log_rotation: Some("DAILY".to_string()),
beacon_endpoint: "http://localhost:5052".to_string(),
beacon_client_timeout: 10,
poll_interval: 6,
listen_addr: "0.0.0.0:8000".to_string(),
origin_block: "0x1234".to_string(),
storage_type: "s3".to_string(),
s3_endpoint: Some("https://s3.amazonaws.com".to_string()),
s3_bucket: Some("my-bucket".to_string()),
s3_path: Some("my-path".to_string()),
s3_compress: Some(true),
fs_dir: None,
};

let config = cli_args.to_config();

assert_eq!(config.poll_interval, Duration::from_secs(6));
assert_eq!(config.listen_addr, "0.0.0.0:8000");
assert_eq!(
config.origin_block,
Hash256::from_slice(
&hex::decode("1234000000000000000000000000000000000000000000000000000000000000")
.unwrap()
)
);
assert_eq!(
config.beacon_config.beacon_endpoint,
"http://localhost:5052"
);
assert_eq!(
config.beacon_config.beacon_client_timeout,
Duration::from_secs(10)
);
assert_eq!(config.storage_config.storage_type, StorageType::S3);
assert_eq!(
config.storage_config.s3_config,
Some(S3Config {
endpoint: "https://s3.amazonaws.com".to_string(),
bucket: "my-bucket".to_string(),
path: "my-path".to_string(),
compression: true,
})
);
assert_eq!(config.storage_config.fs_dir, None);
}

#[test]
fn test_cli_args_to_config_fs_storage() {
let cli_args = CliArgs {
verbose: 3,
log_dir: Some("logs".to_string()),
log_rotation: Some("DAILY".to_string()),
beacon_endpoint: "http://localhost:5052".to_string(),
beacon_client_timeout: 10,
poll_interval: 6,
listen_addr: "0.0.0.0:8000".to_string(),
origin_block: "0xabcd".to_string(),
storage_type: "fs".to_string(),
s3_endpoint: None,
s3_bucket: None,
s3_path: None,
s3_compress: None,
fs_dir: Some("/path/to/storage".to_string()),
};

let config = cli_args.to_config();

assert_eq!(config.poll_interval, Duration::from_secs(6));
assert_eq!(config.listen_addr, "0.0.0.0:8000");
assert_eq!(
config.origin_block,
Hash256::from_slice(
&hex::decode("abcd000000000000000000000000000000000000000000000000000000000000")
.unwrap()
)
);
assert_eq!(
config.beacon_config.beacon_endpoint,
"http://localhost:5052"
);
assert_eq!(
config.beacon_config.beacon_client_timeout,
Duration::from_secs(10)
);
assert_eq!(config.storage_config.storage_type, StorageType::FS);
assert_eq!(config.storage_config.s3_config, None);
assert_eq!(
config.storage_config.fs_dir,
Some(PathBuf::from("/path/to/storage"))
);
}

#[test]
fn test_cli_args_to_config_origin_block_padding() {
let cli_args = CliArgs {
verbose: 3,
log_dir: Some("logs".to_string()),
log_rotation: Some("DAILY".to_string()),
beacon_endpoint: "http://localhost:5052".to_string(),
beacon_client_timeout: 10,
poll_interval: 6,
listen_addr: "0.0.0.0:8000".to_string(),
origin_block: "0x1".to_string(),
storage_type: "fs".to_string(),
s3_endpoint: None,
s3_bucket: None,
s3_path: None,
s3_compress: None,
fs_dir: Some("/path/to/storage".to_string()),
};

let config = cli_args.to_config();

assert_eq!(
config.origin_block,
Hash256::from_slice(
&hex::decode("1000000000000000000000000000000000000000000000000000000000000000")
.unwrap()
)
);
}

#[test]
fn test_cli_args_to_config_origin_block_without_prefix() {
let cli_args = CliArgs {
verbose: 3,
log_dir: Some("logs".to_string()),
log_rotation: Some("DAILY".to_string()),
beacon_endpoint: "http://localhost:5052".to_string(),
beacon_client_timeout: 10,
poll_interval: 6,
listen_addr: "0.0.0.0:8000".to_string(),
origin_block: "ff".to_string(),
storage_type: "fs".to_string(),
s3_endpoint: None,
s3_bucket: None,
s3_path: None,
s3_compress: None,
fs_dir: Some("/path/to/storage".to_string()),
};

let config = cli_args.to_config();

assert_eq!(
config.origin_block,
Hash256::from_slice(
&hex::decode("ff00000000000000000000000000000000000000000000000000000000000000")
.unwrap()
)
);
}
}

#[derive(Debug, PartialEq, Eq, Clone, Default, Serialize, Deserialize)]
pub struct LogConfig {
log_dir: Option<PathBuf>,
log_rotation: Option<String>,
verbose: u8,
}
16 changes: 14 additions & 2 deletions crates/storage/src/storage.rs
Original file line number Diff line number Diff line change
Expand Up @@ -20,11 +20,23 @@ pub enum StorageType {
S3,
}

impl std::str::FromStr for StorageType {
type Err = String;

fn from_str(s: &str) -> Result<Self, Self::Err> {
match s.to_lowercase().as_str() {
"fs" => Ok(StorageType::FS),
"s3" => Ok(StorageType::S3),
_ => Err(format!("Invalid storage type: {}", s)),
}
}
}

#[derive(Debug, PartialEq, Eq, Clone, Default, Serialize, Deserialize)]
pub struct Config {
pub storage_type: StorageType,
pub s3_config: S3Config,
pub fs_dir: PathBuf,
pub s3_config: Option<S3Config>,
pub fs_dir: Option<PathBuf>,
}

#[derive(Debug, PartialEq, Clone, Serialize, Deserialize)]
Expand Down

0 comments on commit e584c6f

Please sign in to comment.