fix: improve usage of indicatif for tracking progress (#101)

The multiprogress wasn't used properly and leading to conflicting progress bars (within our own progress bars, as well as the log wrapper).

This PR introduce a layer on top of `indicatif::MultiProgress` to properly handle sections of progress bars, where we can dynamically add/update/remove progress bars from any sections.

We can see in the demo that new sections + progress bars are added on the fly and that extra logs (e.g. info logs) are appended on top of the progress bars.

Progress are also grouped together based on their parent score.

Co-authored-by: Ian Letourneau <letourneau.ian@gmail.com>
Co-authored-by: johnride <jg@nationtech.io>
Reviewed-on: NationTech/harmony#101

harmony_cli/src/cli_logger.rs

@@ -2,18 +2,15 @@ use harmony::{
     instrumentation::{self, HarmonyEvent},
     topology::TopologyStatus,
 };
-use indicatif::{MultiProgress, ProgressBar};
+use indicatif::MultiProgress;
 use indicatif_log_bridge::LogWrapper;
-use std::{
-    collections::HashMap,
-    sync::{Arc, Mutex},
-};
+use std::sync::{Arc, Mutex};
 
-use crate::progress;
+use crate::progress::{IndicatifProgressTracker, ProgressTracker};
 
 pub fn init() -> tokio::task::JoinHandle<()> {
-    configure_logger();
-    let handle = tokio::spawn(handle_events());
+    let base_progress = configure_logger();
+    let handle = tokio::spawn(handle_events(base_progress));
 
     loop {
         if instrumentation::instrument(HarmonyEvent::HarmonyStarted).is_ok() {
@@ -24,32 +21,45 @@ pub fn init() -> tokio::task::JoinHandle<()> {
     handle
 }
 
-fn configure_logger() {
+fn configure_logger() -> MultiProgress {
     let logger =
         env_logger::Builder::from_env(env_logger::Env::default().default_filter_or("info")).build();
     let level = logger.filter();
-    let multi = MultiProgress::new();
-    LogWrapper::new(multi.clone(), logger).try_init().unwrap();
+    let progress = MultiProgress::new();
+
+    LogWrapper::new(progress.clone(), logger)
+        .try_init()
+        .unwrap();
     log::set_max_level(level);
+
+    progress
 }
 
-async fn handle_events() {
-    instrumentation::subscribe("Harmony CLI Logger", {
-        let sections: Arc<Mutex<HashMap<String, MultiProgress>>> =
-            Arc::new(Mutex::new(HashMap::new()));
-        let progress_bars: Arc<Mutex<HashMap<String, ProgressBar>>> =
-            Arc::new(Mutex::new(HashMap::new()));
+async fn handle_events(base_progress: MultiProgress) {
+    let progress_tracker = Arc::new(IndicatifProgressTracker::new(base_progress.clone()));
+    let preparing_topology = Arc::new(Mutex::new(false));
+    let current_score: Arc<Mutex<Option<String>>> = Arc::new(Mutex::new(None));
 
+    instrumentation::subscribe("Harmony CLI Logger", {
         move |event| {
-            let sections_clone = Arc::clone(&sections);
-            let progress_bars_clone = Arc::clone(&progress_bars);
+            let progress_tracker = Arc::clone(&progress_tracker);
+            let preparing_topology = Arc::clone(&preparing_topology);
+            let current_score = Arc::clone(&current_score);
 
             async move {
-                let mut sections = sections_clone.lock().unwrap();
-                let mut progress_bars = progress_bars_clone.lock().unwrap();
+                let mut preparing_topology = preparing_topology.lock().unwrap();
+                let mut current_score = current_score.lock().unwrap();
 
                 match event {
                     HarmonyEvent::HarmonyStarted => {}
+                    HarmonyEvent::HarmonyFinished => {
+                        progress_tracker.add_section(
+                            "harmony-summary",
+                            &format!("\n{} Harmony completed\n\n", crate::theme::EMOJI_HARMONY),
+                        );
+                        progress_tracker.add_section("harmony-finished", "\n\n");
+                        return false;
+                    }
                     HarmonyEvent::TopologyStateChanged {
                         topology,
                         status,
@@ -60,111 +70,95 @@ async fn handle_events() {
                         match status {
                             TopologyStatus::Queued => {}
                             TopologyStatus::Preparing => {
-                                let section = progress::new_section(format!(
-                                    "{} Preparing environment: {topology}...",
-                                    crate::theme::EMOJI_TOPOLOGY,
-                                ));
-                                (*sections).insert(section_key, section);
+                                progress_tracker.add_section(
+                                    &section_key,
+                                    &format!(
+                                        "\n{} Preparing environment: {topology}...",
+                                        crate::theme::EMOJI_TOPOLOGY
+                                    ),
+                                );
+                                (*preparing_topology) = true;
                             }
                             TopologyStatus::Success => {
-                                let section = (*sections).get(&section_key).unwrap();
-                                let progress = progress::add_spinner(section, "".into());
-
-                                progress::success(
-                                    section,
-                                    Some(progress),
-                                    message.unwrap_or("".into()),
-                                );
-
-                                (*sections).remove(&section_key);
+                                (*preparing_topology) = false;
+                                progress_tracker.add_task(&section_key, "topology-success", "");
+                                progress_tracker
+                                    .finish_task("topology-success", &message.unwrap_or("".into()));
                             }
                             TopologyStatus::Noop => {
-                                let section = (*sections).get(&section_key).unwrap();
-                                let progress = progress::add_spinner(section, "".into());
-
-                                progress::skip(
-                                    section,
-                                    Some(progress),
-                                    message.unwrap_or("".into()),
-                                );
-
-                                (*sections).remove(&section_key);
+                                (*preparing_topology) = false;
+                                progress_tracker.add_task(&section_key, "topology-skip", "");
+                                progress_tracker
+                                    .skip_task("topology-skip", &message.unwrap_or("".into()));
                             }
                             TopologyStatus::Error => {
-                                let section = (*sections).get(&section_key).unwrap();
-                                let progress = progress::add_spinner(section, "".into());
-
-                                progress::error(
-                                    section,
-                                    Some(progress),
-                                    message.unwrap_or("".into()),
-                                );
-
-                                (*sections).remove(&section_key);
+                                progress_tracker.add_task(&section_key, "topology-error", "");
+                                (*preparing_topology) = false;
+                                progress_tracker
+                                    .fail_task("topology-error", &message.unwrap_or("".into()));
                             }
                         }
                     }
                     HarmonyEvent::InterpretExecutionStarted {
+                        execution_id: task_key,
                         topology,
-                        interpret,
+                        interpret: _,
                         score,
                         message,
                     } => {
-                        let section_key = if (*sections).contains_key(&topology_key(&topology)) {
+                        let is_key_topology = (*preparing_topology)
+                            && progress_tracker.contains_section(&topology_key(&topology));
+                        let is_key_current_score = current_score.is_some()
+                            && progress_tracker
+                                .contains_section(&score_key(&current_score.clone().unwrap()));
+                        let is_key_score = progress_tracker.contains_section(&score_key(&score));
+
+                        let section_key = if is_key_topology {
                             topology_key(&topology)
-                        } else if (*sections).contains_key(&score_key(&score)) {
-                            score_key(&interpret)
+                        } else if is_key_current_score {
+                            score_key(&current_score.clone().unwrap())
+                        } else if is_key_score {
+                            score_key(&score)
                         } else {
+                            (*current_score) = Some(score.clone());
                             let key = score_key(&score);
-                            let section = progress::new_section(format!(
-                                "\n{} Interpreting score: {score}...",
-                                crate::theme::EMOJI_SCORE,
-                            ));
-                            (*sections).insert(key.clone(), section);
+                            progress_tracker.add_section(
+                                &key,
+                                &format!(
+                                    "{} Interpreting score: {score}...",
+                                    crate::theme::EMOJI_SCORE
+                                ),
+                            );
                             key
                         };
-                        let section = (*sections).get(&section_key).unwrap();
-                        let progress_bar = progress::add_spinner(section, message);
 
-                        (*progress_bars).insert(interpret_key(&interpret), progress_bar);
+                        progress_tracker.add_task(&section_key, &task_key, &message);
                     }
                     HarmonyEvent::InterpretExecutionFinished {
-                        topology,
-                        interpret,
+                        execution_id: task_key,
+                        topology: _,
+                        interpret: _,
                         score,
                         outcome,
                     } => {
-                        let has_topology = (*sections).contains_key(&topology_key(&topology));
-                        let section_key = if has_topology {
-                            topology_key(&topology)
-                        } else {
-                            score_key(&score)
-                        };
-
-                        let section = (*sections).get(&section_key).unwrap();
-                        let progress_bar =
-                            (*progress_bars).get(&interpret_key(&interpret)).cloned();
-
-                        let _ = section.clear();
+                        if current_score.is_some() && current_score.clone().unwrap() == score {
+                            (*current_score) = None;
+                        }
 
                         match outcome {
                             Ok(outcome) => match outcome.status {
                                 harmony::interpret::InterpretStatus::SUCCESS => {
-                                    progress::success(section, progress_bar, outcome.message)
+                                    progress_tracker.finish_task(&task_key, &outcome.message);
                                 }
                                 harmony::interpret::InterpretStatus::NOOP => {
-                                    progress::skip(section, progress_bar, outcome.message)
+                                    progress_tracker.skip_task(&task_key, &outcome.message);
                                 }
-                                _ => progress::error(section, progress_bar, outcome.message),
+                                _ => progress_tracker.fail_task(&task_key, &outcome.message),
                             },
                             Err(err) => {
-                                progress::error(section, progress_bar, err.to_string());
+                                progress_tracker.fail_task(&task_key, &err.to_string());
                             }
                         }
-
-                        if !has_topology {
-                            (*progress_bars).remove(&section_key);
-                        }
                     }
                 }
                 true
@@ -181,7 +175,3 @@ fn topology_key(topology: &str) -> String {
 fn score_key(score: &str) -> String {
     format!("score-{score}")
 }
-
-fn interpret_key(interpret: &str) -> String {
-    format!("interpret-{interpret}")
-}

harmony_cli/src/lib.rs

@@ -1,5 +1,6 @@
 use clap::Parser;
 use clap::builder::ArgPredicate;
+use harmony::instrumentation;
 use harmony::inventory::Inventory;
 use harmony::maestro::Maestro;
 use harmony::{score::Score, topology::Topology};
@@ -97,6 +98,7 @@ pub async fn run<T: Topology + Send + Sync + 'static>(
 
     let result = init(maestro, args_struct).await;
 
+    instrumentation::instrument(instrumentation::HarmonyEvent::HarmonyFinished).unwrap();
     let _ = tokio::try_join!(cli_logger_handle);
     result
 }

harmony_cli/src/progress.rs

@@ -1,50 +1,163 @@
+use indicatif::{MultiProgress, ProgressBar};
+use std::collections::HashMap;
+use std::sync::{Arc, Mutex};
 use std::time::Duration;
 
-use indicatif::{MultiProgress, ProgressBar};
-
-pub fn new_section(title: String) -> MultiProgress {
-    let multi_progress = MultiProgress::new();
-    let _ = multi_progress.println(title);
-
-    multi_progress
+pub trait ProgressTracker: Send + Sync {
+    fn contains_section(&self, id: &str) -> bool;
+    fn add_section(&self, id: &str, message: &str);
+    fn add_task(&self, section_id: &str, task_id: &str, message: &str);
+    fn finish_task(&self, id: &str, message: &str);
+    fn fail_task(&self, id: &str, message: &str);
+    fn skip_task(&self, id: &str, message: &str);
+    fn clear(&self);
 }
 
-pub fn add_spinner(multi_progress: &MultiProgress, message: String) -> ProgressBar {
-    let progress = multi_progress.add(ProgressBar::new_spinner());
-
-    progress.set_style(crate::theme::SPINNER_STYLE.clone());
-    progress.set_message(message);
-    progress.enable_steady_tick(Duration::from_millis(100));
-
-    progress
+struct Section {
+    header_index: usize,
+    task_count: usize,
+    pb: ProgressBar,
 }
 
-pub fn success(multi_progress: &MultiProgress, progress: Option<ProgressBar>, message: String) {
-    if let Some(progress) = progress {
-        multi_progress.remove(&progress)
+struct IndicatifProgressTrackerState {
+    sections: HashMap<String, Section>,
+    tasks: HashMap<String, ProgressBar>,
+    pb_count: usize,
+}
+
+#[derive(Clone)]
+pub struct IndicatifProgressTracker {
+    mp: MultiProgress,
+    state: Arc<Mutex<IndicatifProgressTrackerState>>,
+}
+
+impl IndicatifProgressTracker {
+    pub fn new(base: MultiProgress) -> Self {
+        // The indicatif log bridge will insert a progress bar at the top.
+        // To prevent our first section from being erased, we need to create
+        // a dummy progress bar as our first progress bar.
+        let _ = base.clear();
+        let log_pb = base.add(ProgressBar::new(1));
+
+        let mut sections = HashMap::new();
+        sections.insert(
+            "__log__".into(),
+            Section {
+                header_index: 0,
+                task_count: 0,
+                pb: log_pb.clone(),
+            },
+        );
+
+        let mut tasks = HashMap::new();
+        tasks.insert("__log__".into(), log_pb);
+
+        let state = Arc::new(Mutex::new(IndicatifProgressTrackerState {
+            sections,
+            tasks,
+            pb_count: 1,
+        }));
+
+        Self { mp: base, state }
+    }
+}
+
+impl ProgressTracker for IndicatifProgressTracker {
+    fn add_section(&self, id: &str, message: &str) {
+        let mut state = self.state.lock().unwrap();
+
+        let header_pb = self
+            .mp
+            .add(ProgressBar::new(1).with_style(crate::theme::SECTION_STYLE.clone()));
+        header_pb.finish_with_message(message.to_string());
+
+        let header_index = state.pb_count;
+        state.pb_count += 1;
+
+        state.sections.insert(
+            id.to_string(),
+            Section {
+                header_index,
+                task_count: 0,
+                pb: header_pb,
+            },
+        );
     }
 
-    let progress = multi_progress.add(ProgressBar::new_spinner());
-    progress.set_style(crate::theme::SUCCESS_SPINNER_STYLE.clone());
-    progress.finish_with_message(message);
-}
+    fn add_task(&self, section_id: &str, task_id: &str, message: &str) {
+        let mut state = self.state.lock().unwrap();
 
-pub fn error(multi_progress: &MultiProgress, progress: Option<ProgressBar>, message: String) {
-    if let Some(progress) = progress {
-        multi_progress.remove(&progress)
+        let insertion_index = {
+            let current_section = state
+                .sections
+                .get(section_id)
+                .expect("Section ID not found");
+            current_section.header_index + current_section.task_count + 1 // +1 to insert after header
+        };
+
+        let pb = self.mp.insert(insertion_index, ProgressBar::new_spinner());
+        pb.set_style(crate::theme::SPINNER_STYLE.clone());
+        pb.set_prefix("  ");
+        pb.set_message(message.to_string());
+        pb.enable_steady_tick(Duration::from_millis(80));
+
+        state.pb_count += 1;
+
+        let section = state
+            .sections
+            .get_mut(section_id)
+            .expect("Section ID not found");
+        section.task_count += 1;
+
+        // We inserted a new progress bar, so we must update the header_index
+        // for all subsequent sections.
+        for (id, s) in state.sections.iter_mut() {
+            if id != section_id && s.header_index >= insertion_index {
+                s.header_index += 1;
+            }
+        }
+
+        state.tasks.insert(task_id.to_string(), pb);
     }
 
-    let progress = multi_progress.add(ProgressBar::new_spinner());
-    progress.set_style(crate::theme::ERROR_SPINNER_STYLE.clone());
-    progress.finish_with_message(message);
-}
-
-pub fn skip(multi_progress: &MultiProgress, progress: Option<ProgressBar>, message: String) {
-    if let Some(progress) = progress {
-        multi_progress.remove(&progress)
+    fn finish_task(&self, id: &str, message: &str) {
+        let state = self.state.lock().unwrap();
+        if let Some(pb) = state.tasks.get(id) {
+            pb.set_style(crate::theme::SUCCESS_SPINNER_STYLE.clone());
+            pb.finish_with_message(message.to_string());
+        }
     }
 
-    let progress = multi_progress.add(ProgressBar::new_spinner());
-    progress.set_style(crate::theme::SKIP_SPINNER_STYLE.clone());
-    progress.finish_with_message(message);
+    fn fail_task(&self, id: &str, message: &str) {
+        let state = self.state.lock().unwrap();
+        if let Some(pb) = state.tasks.get(id) {
+            pb.set_style(crate::theme::ERROR_SPINNER_STYLE.clone());
+            pb.finish_with_message(message.to_string());
+        }
+    }
+
+    fn skip_task(&self, id: &str, message: &str) {
+        let state = self.state.lock().unwrap();
+        if let Some(pb) = state.tasks.get(id) {
+            pb.set_style(crate::theme::SKIP_SPINNER_STYLE.clone());
+            pb.finish_with_message(message.to_string());
+        }
+    }
+
+    fn contains_section(&self, id: &str) -> bool {
+        let state = self.state.lock().unwrap();
+        state.sections.contains_key(id)
+    }
+
+    fn clear(&self) {
+        let mut state = self.state.lock().unwrap();
+
+        state.tasks.values().for_each(|p| self.mp.remove(p));
+        state.tasks.clear();
+        state.sections.values().for_each(|s| self.mp.remove(&s.pb));
+        state.sections.clear();
+        state.pb_count = 0;
+
+        let _ = self.mp.clear();
+    }
 }

harmony_cli/src/theme.rs

@@ -11,8 +11,11 @@ pub static EMOJI_TOPOLOGY: Emoji<'_, '_> = Emoji("📦", "");
 pub static EMOJI_SCORE: Emoji<'_, '_> = Emoji("🎶", "");
 
 lazy_static! {
+    pub static ref SECTION_STYLE: ProgressStyle = ProgressStyle::default_spinner()
+        .template("{wide_msg:.bold}")
+        .unwrap();
     pub static ref SPINNER_STYLE: ProgressStyle = ProgressStyle::default_spinner()
-        .template("    {spinner:.green} {msg}")
+        .template("    {spinner:.green} {wide_msg}")
         .unwrap()
         .tick_strings(&["⠋", "⠙", "⠹", "⠸", "⠼", "⠴", "⠦", "⠧", "⠇", "⠏"]);
     pub static ref SUCCESS_SPINNER_STYLE: ProgressStyle = SPINNER_STYLE