From 557960d98a1d7f613cf8b81df3f01cbdec141793 Mon Sep 17 00:00:00 2001 From: marshmallow Date: Sun, 23 Nov 2025 10:43:17 +1100 Subject: [PATCH] add status bar --- Cargo.lock | 71 +++++++++++- wire/cli/src/apply.rs | 19 +++- wire/cli/src/main.rs | 5 +- wire/cli/src/tracing_setup.rs | 41 +++++-- wire/lib/Cargo.toml | 2 + wire/lib/src/commands/pty/mod.rs | 32 +++--- wire/lib/src/hive/node.rs | 10 ++ wire/lib/src/lib.rs | 23 +++- wire/lib/src/status.rs | 179 +++++++++++++++++++++++++++++++ 9 files changed, 355 insertions(+), 27 deletions(-) create mode 100644 wire/lib/src/status.rs diff --git a/Cargo.lock b/Cargo.lock index 82d8dcb..c2cb651 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -385,6 +385,15 @@ dependencies = [ "zeroize", ] +[[package]] +name = "deranged" +version = "0.5.5" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "ececcb659e7ba858fb4f10388c250a7252eb0a27373f1a72b8748afdd248e587" +dependencies = [ + "powerfmt", +] + [[package]] name = "derive_more" version = "2.0.1" @@ -515,7 +524,7 @@ source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "778e2ac28f6c47af28e4907f13ffd1e1ddbd400980a9abd7c8df189bf578a5ad" dependencies = [ "libc", - "windows-sys 0.52.0", + "windows-sys 0.60.2", ] [[package]] @@ -736,6 +745,15 @@ version = "0.5.2" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "fc0fef456e4baa96da950455cd02c081ca953b141298e41db3fc7e36b1da849c" +[[package]] +name = "humanize-duration" +version = "0.0.7" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "9d17650201754d4f79437bc2fa1f157273a7779de1391c818b275d6bb40998b9" +dependencies = [ + "time", +] + [[package]] name = "icu_collections" version = "2.0.0" @@ -945,6 +963,7 @@ dependencies = [ "futures", "gethostname", "gjson", + "humanize-duration", "im", "itertools", "key_agent", @@ -963,6 +982,7 @@ dependencies = [ "strip-ansi-escapes", "syn 2.0.110", "tempdir", + "termion", "thiserror 2.0.17", "tokio", "tokio-util", @@ -1169,6 +1189,12 @@ dependencies = [ "windows-sys 0.52.0", ] +[[package]] +name = "num-conv" +version = "0.1.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "51d515d32fb182ee37cda2ccdcb92950d6a3c2893aa280e540671c2cd0f3b1d9" + [[package]] name = "num-traits" version = "0.2.19" @@ -1200,6 +1226,12 @@ dependencies = [ "syn 2.0.110", ] +[[package]] +name = "numtoa" +version = "0.2.4" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "6aa2c4e539b869820a2b82e1aef6ff40aa85e65decdd5185e83fb4b1249cd00f" + [[package]] name = "object" version = "0.36.7" @@ -1322,6 +1354,12 @@ dependencies = [ "zerovec", ] +[[package]] +name = "powerfmt" +version = "0.2.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "439ee305def115ba05938db6eb1644ff94165c5ab5e9420d1c1bcedbba909391" + [[package]] name = "ppv-lite86" version = "0.2.21" @@ -1588,7 +1626,7 @@ dependencies = [ "errno", "libc", "linux-raw-sys", - "windows-sys 0.52.0", + "windows-sys 0.60.2", ] [[package]] @@ -1917,6 +1955,16 @@ dependencies = [ "windows-sys 0.59.0", ] +[[package]] +name = "termion" +version = "4.0.6" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "f44138a9ae08f0f502f24104d82517ef4da7330c35acd638f1f29d3cd5475ecb" +dependencies = [ + "libc", + "numtoa", +] + [[package]] name = "textwrap" version = "0.16.2" @@ -1982,6 +2030,25 @@ dependencies = [ "cfg-if", ] +[[package]] +name = "time" +version = "0.3.44" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "91e7d9e3bb61134e77bde20dd4825b97c010155709965fedf0f49bb138e52a9d" +dependencies = [ + "deranged", + "num-conv", + "powerfmt", + "serde", + "time-core", +] + +[[package]] +name = "time-core" +version = "0.1.6" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "40868e7c1d2f0b8d73e4a8c7f0ff63af4f6d19be117e90bd73eb1d62cf831c6b" + [[package]] name = "tinystr" version = "0.8.1" diff --git a/wire/cli/src/apply.rs b/wire/cli/src/apply.rs index 30d06ca..b0df6e3 100644 --- a/wire/cli/src/apply.rs +++ b/wire/cli/src/apply.rs @@ -5,10 +5,11 @@ use futures::{FutureExt, StreamExt}; use itertools::{Either, Itertools}; use lib::hive::node::{Context, GoalExecutor, Name, StepState, should_apply_locally}; use lib::hive::{Hive, HiveLocation}; +use lib::status::STATUS; use lib::{SubCommandModifiers, errors::HiveLibError}; use miette::{Diagnostic, IntoDiagnostic, Result}; use std::collections::HashSet; -use std::io::Read; +use std::io::{Read, stderr}; use std::sync::Arc; use thiserror::Error; use tracing::{Span, error, info}; @@ -85,7 +86,7 @@ pub async fn apply( }, ); - let mut set = hive + let selected_nodes: Vec<_> = hive .nodes .iter_mut() .filter(|(name, node)| { @@ -93,6 +94,17 @@ pub async fn apply( || names.contains(name) || node.tags.iter().any(|tag| tags.contains(tag)) }) + .collect(); + + STATUS.lock().add_many( + &selected_nodes + .iter() + .map(|(name, _)| *name) + .collect::>(), + ); + + let mut set = selected_nodes + .into_iter() .map(|(name, node)| { info!("Resolved {:?} to include {}", args.on, name); @@ -143,6 +155,9 @@ pub async fn apply( std::mem::drop(header_span); if !errors.is_empty() { + // clear the status bar if we are about to print error messages + STATUS.lock().clear(&mut stderr()); + return Err(NodeErrors( errors .into_iter() diff --git a/wire/cli/src/main.rs b/wire/cli/src/main.rs index 59e2456..4918dc5 100644 --- a/wire/cli/src/main.rs +++ b/wire/cli/src/main.rs @@ -1,6 +1,9 @@ // SPDX-License-Identifier: AGPL-3.0-or-later // Copyright 2024-2025 wire Contributors +#![feature(sync_nonpoison)] +#![feature(nonpoison_mutex)] + use std::process::Command; use crate::cli::Cli; @@ -35,7 +38,7 @@ async fn main() -> Result<()> { let args = Cli::parse(); let modifiers = args.to_subcommand_modifiers(); - setup_logging(&args.verbose); + setup_logging(&args.verbose, !&args.no_progress); #[cfg(debug_assertions)] if args.markdown_help { diff --git a/wire/cli/src/tracing_setup.rs b/wire/cli/src/tracing_setup.rs index abb4c9a..43b1085 100644 --- a/wire/cli/src/tracing_setup.rs +++ b/wire/cli/src/tracing_setup.rs @@ -3,11 +3,14 @@ use std::{ collections::VecDeque, - io::{self, Stderr, Write, stderr}, + io::{self, Stderr, Write, stderr}, time::Duration, }; use clap_verbosity_flag::{LogLevel, Verbosity}; -use lib::STDIN_CLOBBER_LOCK; +use lib::{ + STDIN_CLOBBER_LOCK, + status::{STATUS}, +}; use owo_colors::{OwoColorize, Stream, Style}; use tracing::{Level, Subscriber}; use tracing_log::AsTrace; @@ -42,13 +45,14 @@ impl NonClobberingWriter { } } + /// expects the caller to write the status line fn dump_previous(&mut self) -> Result<(), io::Error> { + STATUS.lock().clear(&mut self.stderr); + for buf in self.queue.iter().rev() { self.stderr.write(buf).map(|_| ())?; } - self.stderr.flush()?; - Ok(()) } } @@ -58,7 +62,7 @@ impl Write for NonClobberingWriter { if let 1.. = STDIN_CLOBBER_LOCK.available_permits() { self.dump_previous().map(|()| 0)?; - self.stderr.write(buf) + STATUS.lock().write_above_status(buf, &mut self.stderr) } else { self.queue.push_front(buf.to_vec()); @@ -229,12 +233,37 @@ where } } +async fn status_tick_worker() { + let mut interval = tokio::time::interval(Duration::from_secs(1)); + let mut stderr = stderr(); + + loop { + interval.tick().await; + + if STDIN_CLOBBER_LOCK.available_permits() < 1 { + continue; + } + + let mut status = STATUS.lock(); + + status.clear(&mut stderr); + status.write_status(&mut stderr); + } +} + /// Set up logging for the application /// Uses `WireFieldFormat` if -v was never passed -pub fn setup_logging(verbosity: &Verbosity) { +pub fn setup_logging(verbosity: &Verbosity, show_progress: bool) { let filter = verbosity.log_level_filter().as_trace(); let registry = tracing_subscriber::registry(); + STATUS.lock().show_progress(show_progress); + + // spawn worker to tick the status bar + if show_progress { + tokio::spawn(status_tick_worker()); + } + if verbosity.is_present() { let layer = tracing_subscriber::fmt::layer() .without_time() diff --git a/wire/lib/Cargo.toml b/wire/lib/Cargo.toml index f3f5e34..f762b45 100644 --- a/wire/lib/Cargo.toml +++ b/wire/lib/Cargo.toml @@ -37,6 +37,8 @@ aho-corasick = "1.1.4" num_enum = "0.7.5" gjson = "0.8.1" owo-colors = { workspace = true } +termion = "4.0.6" +humanize-duration = "0.0.7" [dev-dependencies] tempdir = "0.3" diff --git a/wire/lib/src/commands/pty/mod.rs b/wire/lib/src/commands/pty/mod.rs index 68bf3dc..8a399ed 100644 --- a/wire/lib/src/commands/pty/mod.rs +++ b/wire/lib/src/commands/pty/mod.rs @@ -2,6 +2,7 @@ // Copyright 2024-2025 wire Contributors use crate::commands::pty::output::{WatchStdoutArguments, handle_pty_stdout}; +use crate::status::STATUS; use aho_corasick::PatternID; use itertools::Itertools; use nix::sys::termios::{LocalFlags, SetArg, Termios, tcgetattr, tcsetattr}; @@ -10,6 +11,7 @@ use nix::unistd::write as posix_write; use portable_pty::{CommandBuilder, NativePtySystem, PtyPair, PtySize}; use rand::distr::Alphabetic; use std::collections::VecDeque; +use std::io::stderr; use std::sync::{LazyLock, Mutex}; use std::{ io::{Read, Write}, @@ -23,7 +25,7 @@ use tracing::{Span, debug, trace}; use crate::commands::CommandArguments; use crate::commands::pty::input::watch_stdin_from_user; use crate::errors::CommandError; -use crate::{STDIN_CLOBBER_LOCK, SubCommandModifiers}; +use crate::{SubCommandModifiers, aquire_stdin_lock}; use crate::{ commands::{ChildOutputMode, WireCommandChip}, errors::HiveLibError, @@ -153,7 +155,7 @@ pub(crate) async fn interactive_command_with_env>( command.env(key, value); } - let clobber_guard = STDIN_CLOBBER_LOCK.acquire().await.unwrap(); + let clobber_guard = aquire_stdin_lock().await; let _guard = StdinTermiosAttrGuard::new().map_err(HiveLibError::CommandError)?; let child = pty_pair .slave @@ -247,17 +249,21 @@ fn print_authenticate_warning>( return Ok(()); } - eprintln!( - "{} | Authenticate for \"sudo {}\":", - arguments - .target - .map_or(Ok("localhost (!)".to_string()), |target| Ok(format!( - "{}@{}:{}", - target.user, - target.get_preferred_host()?, - target.port - )))?, - arguments.command_string.as_ref() + let _ = STATUS.lock().write_above_status( + &format!( + "{} | Authenticate for \"sudo {}\":\n", + arguments + .target + .map_or(Ok("localhost (!)".to_string()), |target| Ok(format!( + "{}@{}:{}", + target.user, + target.get_preferred_host()?, + target.port + )))?, + arguments.command_string.as_ref() + ) + .into_bytes(), + &mut stderr(), ); Ok(()) diff --git a/wire/lib/src/hive/node.rs b/wire/lib/src/hive/node.rs index 45575da..73005e9 100644 --- a/wire/lib/src/hive/node.rs +++ b/wire/lib/src/hive/node.rs @@ -21,6 +21,7 @@ use crate::hive::steps::evaluate::Evaluate; use crate::hive::steps::keys::{Key, Keys, PushKeyAgent, UploadKeyAt}; use crate::hive::steps::ping::Ping; use crate::hive::steps::push::{PushBuildOutput, PushEvaluatedOutput}; +use crate::status::STATUS; use crate::{EvalGoal, StrictHostKeyChecking, SubCommandModifiers}; use super::HiveLibError; @@ -293,6 +294,7 @@ pub struct StepState { pub key_agent_directory: Option, } +#[allow(clippy::struct_excessive_bools)] pub struct Context<'a> { pub name: &'a Name, pub node: &'a mut Node, @@ -434,6 +436,10 @@ impl<'a> GoalExecutor<'a> { progress = format!("{}/{length}", position + 1) ); + STATUS + .lock() + .set_node_step(self.context.name, step.to_string()); + if let Err(err) = step.execute(&mut self.context).await.inspect_err(|_| { error!("Failed to execute `{step}`"); }) { @@ -446,10 +452,14 @@ impl<'a> GoalExecutor<'a> { return Ok(()); } + STATUS.lock().mark_node_failed(self.context.name); + return Err(err); } } + STATUS.lock().mark_node_succeeded(self.context.name); + Ok(()) } } diff --git a/wire/lib/src/lib.rs b/wire/lib/src/lib.rs index 9e2e6e0..a018501 100644 --- a/wire/lib/src/lib.rs +++ b/wire/lib/src/lib.rs @@ -3,15 +3,25 @@ #![feature(assert_matches)] #![feature(iter_intersperse)] +#![feature(sync_nonpoison)] +#![feature(nonpoison_mutex)] -use std::{io::IsTerminal, sync::LazyLock}; +use std::{ + io::{IsTerminal, stderr}, + sync::LazyLock, +}; -use tokio::sync::Semaphore; +use tokio::sync::{AcquireError, Semaphore, SemaphorePermit}; -use crate::{errors::HiveLibError, hive::node::Name}; +use crate::{ + errors::HiveLibError, + hive::node::Name, + status::{STATUS}, +}; pub mod commands; pub mod hive; +pub mod status; #[cfg(test)] mod test_macros; @@ -54,3 +64,10 @@ pub enum EvalGoal<'a> { } pub static STDIN_CLOBBER_LOCK: LazyLock = LazyLock::new(|| Semaphore::new(1)); + +pub async fn aquire_stdin_lock<'a>() -> Result, AcquireError> { + let result = STDIN_CLOBBER_LOCK.acquire().await?; + STATUS.lock().wipe_out(&mut stderr()); + + Ok(result) +} diff --git a/wire/lib/src/status.rs b/wire/lib/src/status.rs new file mode 100644 index 0000000..43f489b --- /dev/null +++ b/wire/lib/src/status.rs @@ -0,0 +1,179 @@ +// SPDX-License-Identifier: AGPL-3.0-or-later +// Copyright 2024-2025 wire Contributors + +use owo_colors::OwoColorize; +use std::{fmt::Write, time::Instant}; +use termion::{clear, cursor}; + +use crate::{STDIN_CLOBBER_LOCK, hive::node::Name}; + +use std::{ + collections::HashMap, + sync::{LazyLock, nonpoison::Mutex}, +}; + +#[derive(Default)] +pub enum NodeStatus { + #[default] + Pending, + Running(String), + Succeeded, + Failed, +} + +pub struct Status { + statuses: HashMap, + began: Instant, + show_progress: bool +} + +/// global status used for the progress bar in the cli crate +pub static STATUS: LazyLock> = LazyLock::new(|| Mutex::new(Status::new())); + +impl Status { + fn new() -> Self { + Self { + statuses: HashMap::default(), + began: Instant::now(), + show_progress: false + } + } + + pub const fn show_progress(&mut self, show_progress: bool) { + self.show_progress = show_progress; + } + + pub fn add_many(&mut self, names: &[&Name]) { + self.statuses.extend( + names + .iter() + .map(|name| (name.0.to_string(), NodeStatus::Pending)), + ); + } + + pub fn set_node_step(&mut self, node: &Name, step: String) { + self.statuses + .insert(node.0.to_string(), NodeStatus::Running(step)); + } + + pub fn mark_node_failed(&mut self, node: &Name) { + self.statuses.insert(node.0.to_string(), NodeStatus::Failed); + } + + pub fn mark_node_succeeded(&mut self, node: &Name) { + self.statuses + .insert(node.0.to_string(), NodeStatus::Succeeded); + } + + #[must_use] + fn num_finished(&self) -> usize { + self.statuses + .iter() + .filter(|(_, status)| matches!(status, NodeStatus::Succeeded | NodeStatus::Failed)) + .count() + } + + #[must_use] + fn num_running(&self) -> usize { + self.statuses + .iter() + .filter(|(_, status)| matches!(status, NodeStatus::Running(..))) + .count() + } + + #[must_use] + fn num_failed(&self) -> usize { + self.statuses + .iter() + .filter(|(_, status)| matches!(status, NodeStatus::Failed)) + .count() + } + + #[must_use] + pub fn get_msg(&self) -> String { + if self.statuses.is_empty() { + return String::new(); + } + + let mut msg = format!("[{} / {}", self.num_finished(), self.statuses.len(),); + + let num_failed = self.num_failed(); + let num_running = self.num_running(); + + let failed = if num_failed >= 1 { + Some(format!("{} Failed", num_failed.red())) + } else { + None + }; + + let running = if num_running >= 1 { + Some(format!("{} Deploying", num_running.blue())) + } else { + None + }; + + let _ = match (failed, running) { + (None, None) => write!(&mut msg, ""), + (Some(message), None) | (None, Some(message)) => write!(&mut msg, " ({message})"), + (Some(failed), Some(running)) => write!(&mut msg, " ({failed}, {running})"), + }; + + let _ = write!(&mut msg, "]"); + + let _ = write!( + &mut msg, + " {}s", + self.began + .elapsed() + .as_secs() + ); + + msg + } + + pub fn clear(&self, writer: &mut T) { + if !self.show_progress { + return; + } + + let _ = write!(writer, "{}", cursor::Save); + // let _ = write!(writer, "{}", cursor::Down(1)); + let _ = write!(writer, "{}", cursor::Left(999)); + let _ = write!(writer, "{}", clear::CurrentLine); + } + + /// used when there is an interactive prompt + pub fn wipe_out(&self, writer: &mut T) { + if !self.show_progress { + return; + } + + let _ = write!(writer, "{}", cursor::Save); + let _ = write!(writer, "{}", cursor::Left(999)); + let _ = write!(writer, "{}", clear::CurrentLine); + let _ = writer.flush(); + } + + pub fn write_status(&mut self, writer: &mut T) { + if self.show_progress { + let _ = write!(writer, "{}", self.get_msg()); + } + } + + pub fn write_above_status( + &mut self, + buf: &[u8], + writer: &mut T, + ) -> std::io::Result { + if STDIN_CLOBBER_LOCK.available_permits() != 1 { + // skip + return Ok(0); + } + + self.clear(writer); + let written = writer.write(buf)?; + self.write_status(writer); + + Ok(written) + } +} -- 2.51.2