From da1009923911020dfe7c6e7b430c0eb705cc0d66 Mon Sep 17 00:00:00 2001 From: Huan-Cheng Chang Date: Fri, 20 Dec 2024 13:03:40 +0000 Subject: [PATCH] feat(jstzd): print banner and progress bar while launching --- Cargo.lock | 2 + crates/jstzd/Cargo.toml | 2 + crates/jstzd/src/lib.rs | 62 ++++++++++++++++++++++++++++++- crates/jstzd/src/task/jstzd.rs | 64 +++++++++++++++++++++++++++----- crates/jstzd/tests/jstzd_test.rs | 4 +- 5 files changed, 122 insertions(+), 12 deletions(-) diff --git a/Cargo.lock b/Cargo.lock index 489628196..39ab394f6 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -2992,10 +2992,12 @@ dependencies = [ "bincode", "bollard", "clap 4.5.20", + "console", "futures", "futures-util", "hex", "http 1.1.0", + "indicatif", "jstz_crypto", "jstz_kernel", "jstz_node", diff --git a/crates/jstzd/Cargo.toml b/crates/jstzd/Cargo.toml index cc2f32e2c..046ef8b53 100644 --- a/crates/jstzd/Cargo.toml +++ b/crates/jstzd/Cargo.toml @@ -13,9 +13,11 @@ async-trait.workspace = true axum.workspace = true bollard.workspace = true clap.workspace = true +console.workspace = true futures.workspace = true futures-util.workspace = true http.workspace = true +indicatif.workspace = true jstz_crypto = { path = "../jstz_crypto" } jstz_node = {path = "../jstz_node"} octez = { path = "../octez" } diff --git a/crates/jstzd/src/lib.rs b/crates/jstzd/src/lib.rs index 0a2c4507d..a0e46a785 100644 --- a/crates/jstzd/src/lib.rs +++ b/crates/jstzd/src/lib.rs @@ -7,12 +7,28 @@ pub use config::BOOTSTRAP_CONTRACT_NAMES; pub mod jstz_rollup_path { include!(concat!(env!("OUT_DIR"), "/jstz_rollup_path.rs")); } +use console::style; +use std::io::{stdout, Write}; use std::process::exit; use tokio::signal::unix::{signal, SignalKind}; include!("../build_config.rs"); pub const JSTZ_ROLLUP_ADDRESS: &str = "sr1PuFMgaRUN12rKQ3J2ae5psNtwCxPNmGNK"; pub const JSTZ_NATIVE_BRIDGE_ADDRESS: &str = "KT1GFiPkkTjd14oHe6MrBPiRh5djzRkVWcni"; +const SANDBOX_BANNER: &str = r#" + __________ + \ jstz / + )______( + |""""""|_.-._,.---------.,_.-._ + | | | | | | ''-. + | |_| |_ _| |_..-' + |______| '-' `'---------'` '-' + )""""""( + /________\ + `'------'` + .------------. + /______________\ +"#; /// The `main` function for running jstzd pub async fn main(config_path: &Option) { @@ -28,9 +44,22 @@ pub async fn main(config_path: &Option) { } } +// requiring a writer here so that we can test this function +fn print_banner(writer: &mut impl Write) { + let _ = writeln!(writer, "{}", style(SANDBOX_BANNER).bold()); + let _ = writeln!( + writer, + " {} {}", + env!("CARGO_PKG_VERSION"), + style(env!("CARGO_PKG_REPOSITORY")).blue().bold() + ); + let _ = writeln!(writer); +} + async fn run(port: u16, config: JstzdConfig) { let mut server = JstzdServer::new(config, port); - if let Err(e) = server.run().await { + print_banner(&mut stdout()); + if let Err(e) = server.run(true).await { eprintln!("failed to run jstzd server: {:?}", e); let _ = server.stop().await; exit(1); @@ -47,3 +76,34 @@ async fn run(port: u16, config: JstzdConfig) { println!("Shutting down"); server.stop().await.unwrap(); } + +mod lib_test { + #[test] + fn print_banner() { + let mut s = vec![]; + super::print_banner(&mut s); + let esc = "\u{1b}"; + assert_eq!( + String::from_utf8(s).unwrap(), + format!( + r#"{esc}[1m + __________ + \ jstz / + )______( + |""""""|_.-._,.---------.,_.-._ + | | | | | | ''-. + | |_| |_ _| |_..-' + |______| '-' `'---------'` '-' + )""""""( + /________\ + `'------'` + .------------. + /______________\ +{esc}[0m + 0.1.0-alpha.0 {esc}[34m{esc}[1mhttps://github.com/jstz-dev/jstz{esc}[0m + +"# + ) + ); + } +} diff --git a/crates/jstzd/src/task/jstzd.rs b/crates/jstzd/src/task/jstzd.rs index e96a50b74..a5d05e47e 100644 --- a/crates/jstzd/src/task/jstzd.rs +++ b/crates/jstzd/src/task/jstzd.rs @@ -21,6 +21,7 @@ use axum::{ routing::{get, put}, Router, }; +use indicatif::{ProgressBar, ProgressStyle}; use jstz_node::config::JstzNodeConfig; use octez::r#async::{ baker::OctezBakerConfig, @@ -36,6 +37,7 @@ use tokio::{ net::TcpListener, sync::{oneshot, RwLock}, task::JoinHandle, + time::{sleep, Duration}, }; trait IntoShared { @@ -163,6 +165,12 @@ impl Task for Jstzd { } async fn health_check(&self) -> Result { + self.health_check_inner().await.0 + } +} + +impl Jstzd { + async fn health_check_inner(&self) -> (Result, Vec>) { let check_results = futures::future::join_all([ self.octez_node.read().await.health_check(), self.baker.read().await.health_check(), @@ -173,22 +181,23 @@ impl Task for Jstzd { let mut healthy = true; let mut err = vec![]; - for result in check_results { + for result in &check_results { match result { Err(e) => err.push(e), - Ok(v) => healthy = healthy && v, + Ok(v) => healthy = healthy && *v, } } if !err.is_empty() { - bail!("failed to perform health check: {:?}", err) + ( + Err(anyhow::anyhow!("failed to perform health check: {:?}", err)), + check_results, + ) } else { - Ok(healthy) + (Ok(healthy), check_results) } } -} -impl Jstzd { async fn import_activator(octez_client: &OctezClient) -> Result<()> { octez_client .import_secret_key(ACTIVATOR_ACCOUNT_ALIAS, ACTIVATOR_ACCOUNT_SK) @@ -316,7 +325,7 @@ impl JstzdServer { shutdown(&mut lock).await } - pub async fn run(&mut self) -> Result<()> { + pub async fn run(&mut self, print_info: bool) -> Result<()> { let jstzd = Self::spawn_jstzd( self.inner .state @@ -330,6 +339,7 @@ impl JstzdServer { "cannot run jstzd server without jstzd config" ))? .clone(), + print_info, ) .await?; self.inner.state.write().await.jstzd.replace(jstzd); @@ -380,12 +390,48 @@ impl JstzdServer { } } - async fn spawn_jstzd(jstzd_config: JstzdConfig) -> Result { + async fn spawn_jstzd(jstzd_config: JstzdConfig, print_info: bool) -> Result { let mut jstzd = Jstzd::spawn(jstzd_config).await?; - let jstzd_healthy = retry(60, 500, || async { jstzd.health_check().await }).await; + let progress_bar = match print_info { + true => { + let v = ProgressBar::new(4); + v.set_style(ProgressStyle::default_bar().template( + "{spinner:.green} [{bar:40.cyan/blue}] {pos:>7}/{len:7} {msg}", + )?); + Some(v) + } + false => None, + }; + + let mut jstzd_healthy = false; + let mut curr_progress = 0; + for _ in 0..120 { + let (overall_result, individual_results) = jstzd.health_check_inner().await; + jstzd_healthy = overall_result.unwrap_or_default(); + if let Some(ref bar) = progress_bar { + let count = individual_results + .into_iter() + .fold(0, |acc, v| acc + v.unwrap_or_default() as u64); + if count > curr_progress { + curr_progress = count; + bar.set_position(curr_progress); + } + + if jstzd_healthy { + bar.finish_and_clear(); + break; + } + } + + sleep(Duration::from_millis(500)).await; + } + if !jstzd_healthy { let _ = jstzd.kill().await; + if let Some(b) = progress_bar { + b.abandon(); + } bail!("jstzd never turns healthy"); } Ok(jstzd) diff --git a/crates/jstzd/tests/jstzd_test.rs b/crates/jstzd/tests/jstzd_test.rs index 9030e0641..380b99e9f 100644 --- a/crates/jstzd/tests/jstzd_test.rs +++ b/crates/jstzd/tests/jstzd_test.rs @@ -51,7 +51,7 @@ async fn jstzd_test() { ) .await; - jstzd.run().await.unwrap(); + jstzd.run(false).await.unwrap(); ensure_jstzd_components_are_up(&jstzd, &octez_node_rpc_endpoint, jstzd_port).await; ensure_rollup_is_logging_to(kernel_debug_file).await; @@ -78,7 +78,7 @@ async fn jstzd_test() { // calling `run` after calling `stop` should fail because all states should have been cleared assert_eq!( - jstzd.run().await.unwrap_err().to_string(), + jstzd.run(false).await.unwrap_err().to_string(), "cannot run jstzd server without jstzd config" ); }