feat(jstzd): print banner and progress bar while launching

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" }

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 JSTZ_BANNER: &str = r#"
+           __________
+           \  jstz  /
+            )______(
+            |""""""|_.-._,.---------.,_.-._
+            |      | | |               | | ''-.
+            |      |_| |_             _| |_..-'
+            |______| '-' `'---------'` '-'
+            )""""""(
+           /________\
+           `'------'`
+         .------------.
+        /______________\
+"#;
 
 /// The `main` function for running jstzd
 pub async fn main(config_path: &Option<String>) {
@@ -28,9 +44,22 @@ pub async fn main(config_path: &Option<String>) {
     }
 }
 
+// requiring a writer here so that we can test this function
+fn print_banner(writer: &mut impl Write) {
+    let _ = writeln!(writer, "{}", style(JSTZ_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,16 @@ async fn run(port: u16, config: JstzdConfig) {
     println!("Shutting down");
     server.stop().await.unwrap();
 }
+
+#[cfg(test)]
+mod lib_test {
+    #[test]
+    fn print_banner() {
+        let mut buf = vec![];
+        super::print_banner(&mut buf);
+        let s = String::from_utf8(buf).unwrap();
+        assert!(s.contains(super::JSTZ_BANNER));
+        assert!(s.contains(env!("CARGO_PKG_VERSION")));
+        assert!(s.contains(env!("CARGO_PKG_REPOSITORY")));
+    }
+}

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<bool> {
+        self.health_check_inner().await.0
+    }
+}
+
+impl Jstzd {
+    async fn health_check_inner(&self) -> (Result<bool>, Vec<Result<bool>>) {
         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,18 +390,78 @@ impl JstzdServer {
         }
     }
 
-    async fn spawn_jstzd(jstzd_config: JstzdConfig) -> Result<Jstzd> {
+    async fn spawn_jstzd(jstzd_config: JstzdConfig, print_info: bool) -> Result<Jstzd> {
         let mut jstzd = Jstzd::spawn(jstzd_config).await?;
 
-        let jstzd_healthy = retry(60, 500, || async { jstzd.health_check().await }).await;
+        let progress_bar = create_progress_bar(print_info);
+        let mut jstzd_healthy = false;
+        // 60 seconds
+        for _ in 0..120 {
+            let (overall_result, individual_results) = jstzd.health_check_inner().await;
+            jstzd_healthy = overall_result.unwrap_or_default();
+            let latest_progress = collect_progress(individual_results);
+            update_progress_bar(progress_bar.as_ref(), latest_progress);
+
+            if jstzd_healthy {
+                clear_progress_bar(progress_bar.as_ref());
+                break;
+            }
+
+            sleep(Duration::from_millis(500)).await;
+        }
+
         if !jstzd_healthy {
             let _ = jstzd.kill().await;
+            abandon_progress_bar(progress_bar.as_ref());
             bail!("jstzd never turns healthy");
         }
         Ok(jstzd)
     }
 }
 
+fn collect_progress(individual_results: Vec<Result<bool>>) -> u64 {
+    individual_results
+        .into_iter()
+        .fold(0, |acc, v| acc + v.unwrap_or_default() as u64)
+}
+
+fn create_progress_bar(print_info: bool) -> Option<ProgressBar> {
+    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}",
+                    )
+                    .unwrap(),
+            );
+            Some(v)
+        }
+        false => None,
+    }
+}
+
+fn clear_progress_bar(progress_bar: Option<&ProgressBar>) {
+    if let Some(bar) = progress_bar {
+        bar.finish_and_clear();
+    }
+}
+
+fn update_progress_bar(progress_bar: Option<&ProgressBar>, progress: u64) {
+    if let Some(bar) = progress_bar {
+        if progress > bar.position() {
+            bar.set_position(progress);
+        }
+    }
+}
+
+fn abandon_progress_bar(progress_bar: Option<&ProgressBar>) {
+    if let Some(b) = progress_bar {
+        b.abandon();
+    }
+}
+
 async fn health_check(state: &ServerState) -> bool {
     if let Some(v) = &state.server_handle {
         if !v.is_finished() {
@@ -459,3 +529,51 @@ async fn config_handler(
         None => http::StatusCode::NOT_FOUND.into_response(),
     }
 }
+
+#[cfg(test)]
+mod tests {
+    use indicatif::ProgressBar;
+
+    #[test]
+    fn collect_progress() {
+        assert_eq!(super::collect_progress(vec![Ok(true), Ok(false)]), 1);
+        assert_eq!(
+            super::collect_progress(vec![Ok(true), Err(anyhow::anyhow!(""))]),
+            1
+        );
+        assert_eq!(super::collect_progress(vec![Ok(true), Ok(true)]), 2);
+        assert_eq!(super::collect_progress(vec![Ok(false), Ok(false)]), 0);
+    }
+
+    #[test]
+    fn clear_progress_bar() {
+        let bar = ProgressBar::new(3);
+        super::clear_progress_bar(Some(&bar));
+        assert!(bar.is_finished());
+    }
+
+    #[test]
+    fn update_progress_bar() {
+        let bar = ProgressBar::new(3);
+        super::update_progress_bar(Some(&bar), 2);
+        assert_eq!(bar.position(), 2);
+
+        super::update_progress_bar(Some(&bar), 1);
+        assert_eq!(bar.position(), 2);
+    }
+
+    #[test]
+    fn create_progress_bar() {
+        assert!(super::create_progress_bar(false).is_none());
+
+        let bar = super::create_progress_bar(true).unwrap();
+        assert_eq!(bar.length().unwrap(), 4);
+    }
+
+    #[test]
+    fn abandon_progress_bar() {
+        let bar = ProgressBar::new(3);
+        super::abandon_progress_bar(Some(&bar));
+        assert!(bar.is_finished());
+    }
+}

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"
     );
 }