From 858ed73a41bfc27d573495e4e8e067350da6f94f Mon Sep 17 00:00:00 2001 From: Francisco Gouveia Date: Sat, 26 Jul 2025 18:52:49 +0100 Subject: [PATCH 1/2] feat(process): create a `ProgressDrawTarget` (for `indicatif`) inside the `Process` --- src/cli/rustup_mode.rs | 12 ++---------- src/process.rs | 16 ++++++++++++++++ 2 files changed, 18 insertions(+), 10 deletions(-) diff --git a/src/cli/rustup_mode.rs b/src/cli/rustup_mode.rs index ebaba20cd4..dc613b8f8f 100644 --- a/src/cli/rustup_mode.rs +++ b/src/cli/rustup_mode.rs @@ -14,7 +14,7 @@ use clap::{Args, CommandFactory, Parser, Subcommand, ValueEnum, builder::Possibl use clap_complete::Shell; use console::style; use futures_util::stream::StreamExt; -use indicatif::{MultiProgress, ProgressBar, ProgressDrawTarget, ProgressStyle}; +use indicatif::{MultiProgress, ProgressBar, ProgressStyle}; use itertools::Itertools; use tracing::{info, trace, warn}; use tracing_subscriber::{EnvFilter, Registry, reload::Handle}; @@ -797,7 +797,6 @@ async fn default_( async fn check_updates(cfg: &Cfg<'_>, opts: CheckOpts) -> Result { let t = cfg.process.stdout().terminal(cfg.process); - let is_a_tty = t.is_a_tty(); let use_colors = matches!(t.color_choice(), ColorChoice::Auto | ColorChoice::Always); let mut update_available = false; let channels = cfg.list_channels()?; @@ -806,14 +805,7 @@ async fn check_updates(cfg: &Cfg<'_>, opts: CheckOpts) -> Result 0 { let multi_progress_bars = - MultiProgress::with_draw_target(match cfg.process.var("RUSTUP_TERM_PROGRESS_WHEN") { - Ok(s) if s.eq_ignore_ascii_case("always") => { - ProgressDrawTarget::term_like(Box::new(t)) - } - Ok(s) if s.eq_ignore_ascii_case("never") => ProgressDrawTarget::hidden(), - _ if is_a_tty => ProgressDrawTarget::term_like(Box::new(t)), - _ => ProgressDrawTarget::hidden(), - }); + MultiProgress::with_draw_target(cfg.process.progress_draw_target()); let channels = tokio_stream::iter(channels.into_iter()).map(|(name, distributable)| { let pb = multi_progress_bars.add(ProgressBar::new(1)); pb.set_style( diff --git a/src/process.rs b/src/process.rs index 9f4370a48c..acc32f1a80 100644 --- a/src/process.rs +++ b/src/process.rs @@ -14,6 +14,7 @@ use std::{ use std::{env, thread}; use anyhow::{Context, Result, bail}; +use indicatif::ProgressDrawTarget; #[cfg(feature = "test")] use tracing::subscriber::DefaultGuard; #[cfg(feature = "test")] @@ -151,6 +152,21 @@ impl Process { Process::TestProcess(p) => Ok(p.cwd.clone()), } } + + pub fn progress_draw_target(&self) -> ProgressDrawTarget { + match self { + Process::OsProcess(_) => (), + #[cfg(feature = "test")] + Process::TestProcess(_) => return ProgressDrawTarget::hidden(), + } + let t = self.stdout().terminal(self); + match self.var("RUSTUP_TERM_PROGRESS_WHEN") { + Ok(s) if s.eq_ignore_ascii_case("always") => ProgressDrawTarget::term_like(Box::new(t)), + Ok(s) if s.eq_ignore_ascii_case("never") => ProgressDrawTarget::hidden(), + _ if t.is_a_tty() => ProgressDrawTarget::term_like(Box::new(t)), + _ => ProgressDrawTarget::hidden(), + } + } } impl home::env::Env for Process { From 63e867b0473c6b4b15c25834824cac1728f967a4 Mon Sep 17 00:00:00 2001 From: Francisco Gouveia Date: Sat, 26 Jul 2025 17:39:01 +0100 Subject: [PATCH 2/2] feat(download_tracker): refactor in favor of `indicatif` Co-authored-by: 0xPoe --- src/cli/download_tracker.rs | 277 +++++------------------------------- src/utils/notifications.rs | 2 +- src/utils/units.rs | 92 +++--------- 3 files changed, 53 insertions(+), 318 deletions(-) diff --git a/src/cli/download_tracker.rs b/src/cli/download_tracker.rs index be9b95ed81..2ac4cca7e7 100644 --- a/src/cli/download_tracker.rs +++ b/src/cli/download_tracker.rs @@ -1,68 +1,34 @@ -use std::collections::VecDeque; -use std::fmt; -use std::io::Write; -use std::time::{Duration, Instant}; +use indicatif::{MultiProgress, ProgressBar, ProgressDrawTarget, ProgressStyle}; +use std::time::Duration; use crate::dist::Notification as In; use crate::notifications::Notification; -use crate::process::{Process, terminalsource}; +use crate::process::Process; use crate::utils::Notification as Un; -use crate::utils::units::{Size, Unit, UnitMode}; - -/// Keep track of this many past download amounts -const DOWNLOAD_TRACK_COUNT: usize = 5; /// Tracks download progress and displays information about it to a terminal. /// /// *not* safe for tracking concurrent downloads yet - it is basically undefined /// what will happen. pub(crate) struct DownloadTracker { - /// Content-Length of the to-be downloaded object. - content_len: Option, - /// Total data downloaded in bytes. - total_downloaded: usize, - /// Data downloaded this second. - downloaded_this_sec: usize, - /// Keeps track of amount of data downloaded every last few secs. - /// Used for averaging the download speed. NB: This does not necessarily - /// represent adjacent seconds; thus it may not show the average at all. - downloaded_last_few_secs: VecDeque, - /// Time stamp of the last second - last_sec: Option, - /// Time stamp of the start of the download - start_sec: Option, - term: terminalsource::ColorableTerminal, - /// Whether we displayed progress for the download or not. - /// - /// If the download is quick enough, we don't have time to - /// display the progress info. - /// In that case, we do not want to do some cleanup stuff we normally do. - /// - /// If we have displayed progress, this is the number of characters we - /// rendered, so we can erase it cleanly. - displayed_charcount: Option, - /// What units to show progress in - units: Vec, - /// Whether we display progress - display_progress: bool, - stdout_is_a_tty: bool, + /// MultiProgress bar for the downloads. + multi_progress_bars: MultiProgress, + /// ProgressBar for the current download. + progress_bar: ProgressBar, } impl DownloadTracker { /// Creates a new DownloadTracker. pub(crate) fn new_with_display_progress(display_progress: bool, process: &Process) -> Self { + let multi_progress_bars = MultiProgress::with_draw_target(if display_progress { + process.progress_draw_target() + } else { + ProgressDrawTarget::hidden() + }); + Self { - content_len: None, - total_downloaded: 0, - downloaded_this_sec: 0, - downloaded_last_few_secs: VecDeque::with_capacity(DOWNLOAD_TRACK_COUNT), - start_sec: None, - last_sec: None, - term: process.stdout().terminal(process), - displayed_charcount: None, - units: vec![Unit::B], - display_progress, - stdout_is_a_tty: process.stdout().is_a_tty(process), + multi_progress_bars, + progress_bar: ProgressBar::hidden(), } } @@ -70,222 +36,47 @@ impl DownloadTracker { match *n { Notification::Install(In::Utils(Un::DownloadContentLengthReceived(content_len))) => { self.content_length_received(content_len); - true } Notification::Install(In::Utils(Un::DownloadDataReceived(data))) => { - if self.stdout_is_a_tty { - self.data_received(data.len()); - } + self.data_received(data.len()); true } Notification::Install(In::Utils(Un::DownloadFinished)) => { self.download_finished(); true } - Notification::Install(In::Utils(Un::DownloadPushUnit(unit))) => { - self.push_unit(unit); - true - } - Notification::Install(In::Utils(Un::DownloadPopUnit)) => { - self.pop_unit(); - true - } + Notification::Install(In::Utils(Un::DownloadPushUnit(_))) => true, + Notification::Install(In::Utils(Un::DownloadPopUnit)) => true, _ => false, } } - /// Notifies self that Content-Length information has been received. + /// Sets the length for a new ProgressBar and gives it a style. pub(crate) fn content_length_received(&mut self, content_len: u64) { - self.content_len = Some(content_len as usize); + self.progress_bar.set_length(content_len); + self.progress_bar.set_style( + ProgressStyle::with_template( + "[{bar:40}] {bytes}/{total_bytes} ({bytes_per_sec}, ETA: {eta})", + ) + .unwrap() + .progress_chars("## "), + ); } /// Notifies self that data of size `len` has been received. pub(crate) fn data_received(&mut self, len: usize) { - self.total_downloaded += len; - self.downloaded_this_sec += len; - - let current_time = Instant::now(); - - match self.last_sec { - None => self.last_sec = Some(current_time), - Some(prev) => { - let elapsed = current_time.saturating_duration_since(prev); - if elapsed >= Duration::from_secs(1) { - if self.display_progress { - self.display(); - } - self.last_sec = Some(current_time); - if self.downloaded_last_few_secs.len() == DOWNLOAD_TRACK_COUNT { - self.downloaded_last_few_secs.pop_back(); - } - self.downloaded_last_few_secs - .push_front(self.downloaded_this_sec); - self.downloaded_this_sec = 0; - } - } + if self.progress_bar.is_hidden() && self.progress_bar.elapsed() >= Duration::from_secs(1) { + self.multi_progress_bars.add(self.progress_bar.clone()); } + self.progress_bar.inc(len as u64); } + /// Notifies self that the download has finished. pub(crate) fn download_finished(&mut self) { - if self.displayed_charcount.is_some() { - // Display the finished state - self.display(); - let _ = writeln!(self.term.lock()); - } - self.prepare_for_new_download(); - } - /// Resets the state to be ready for a new download. - fn prepare_for_new_download(&mut self) { - self.content_len = None; - self.total_downloaded = 0; - self.downloaded_this_sec = 0; - self.downloaded_last_few_secs.clear(); - self.start_sec = Some(Instant::now()); - self.last_sec = None; - self.displayed_charcount = None; - } - /// Display the tracked download information to the terminal. - fn display(&mut self) { - match self.start_sec { - // Maybe forgot to call `prepare_for_new_download` first - None => {} - Some(start_sec) => { - // Panic if someone pops the default bytes unit... - let unit = *self.units.last().unwrap(); - let total_h = Size::new(self.total_downloaded, unit, UnitMode::Norm); - let sum: usize = self.downloaded_last_few_secs.iter().sum(); - let len = self.downloaded_last_few_secs.len(); - let speed = if len > 0 { sum / len } else { 0 }; - let speed_h = Size::new(speed, unit, UnitMode::Rate); - let elapsed_h = Instant::now().saturating_duration_since(start_sec); - - // First, move to the start of the current line and clear it. - let _ = self.term.carriage_return(); - // We'd prefer to use delete_line() but on Windows it seems to - // sometimes do unusual things - // let _ = self.term.as_mut().unwrap().delete_line(); - // So instead we do: - if let Some(n) = self.displayed_charcount { - // This is not ideal as very narrow terminals might mess up, - // but it is more likely to succeed until term's windows console - // fixes whatever's up with delete_line(). - let _ = write!(self.term.lock(), "{}", " ".repeat(n)); - let _ = self.term.lock().flush(); - let _ = self.term.carriage_return(); - } - - let output = match self.content_len { - Some(content_len) => { - let content_len_h = Size::new(content_len, unit, UnitMode::Norm); - let percent = (self.total_downloaded as f64 / content_len as f64) * 100.; - let remaining = content_len - self.total_downloaded; - let eta_h = Duration::from_secs(if speed == 0 { - u64::MAX - } else { - (remaining / speed) as u64 - }); - format!( - "{} / {} ({:3.0} %) {} in {}{}", - total_h, - content_len_h, - percent, - speed_h, - elapsed_h.display(), - Eta(eta_h), - ) - } - None => format!( - "Total: {} Speed: {} Elapsed: {}", - total_h, - speed_h, - elapsed_h.display() - ), - }; - - let _ = write!(self.term.lock(), "{output}"); - // Since stdout is typically line-buffered and we don't print a newline, we manually flush. - let _ = self.term.lock().flush(); - self.displayed_charcount = Some(output.chars().count()); - } - } - } - - pub(crate) fn push_unit(&mut self, new_unit: Unit) { - self.units.push(new_unit); - } - - pub(crate) fn pop_unit(&mut self) { - self.units.pop(); - } -} - -struct Eta(Duration); - -impl fmt::Display for Eta { - fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { - match self.0 { - Duration::ZERO => Ok(()), - _ => write!(f, " ETA: {}", self.0.display()), - } - } -} - -trait DurationDisplay { - fn display(self) -> Display; -} - -impl DurationDisplay for Duration { - fn display(self) -> Display { - Display(self) - } -} - -/// Human readable representation of a `Duration`. -struct Display(Duration); - -impl fmt::Display for Display { - #[allow(clippy::many_single_char_names)] - fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { - const SECS_PER_YEAR: u64 = 60 * 60 * 24 * 365; - let secs = self.0.as_secs(); - if secs > SECS_PER_YEAR { - return f.write_str("Unknown"); - } - match format_dhms(secs) { - (0, 0, 0, s) => write!(f, "{s:2.0}s"), - (0, 0, m, s) => write!(f, "{m:2.0}m {s:2.0}s"), - (0, h, m, s) => write!(f, "{h:2.0}h {m:2.0}m {s:2.0}s"), - (d, h, m, s) => write!(f, "{d:3.0}d {h:2.0}h {m:2.0}m {s:2.0}s"), - } - } -} - -// we're doing modular arithmetic, treat as integer -fn format_dhms(sec: u64) -> (u64, u8, u8, u8) { - let (mins, sec) = (sec / 60, (sec % 60) as u8); - let (hours, mins) = (mins / 60, (mins % 60) as u8); - let (days, hours) = (hours / 24, (hours % 24) as u8); - (days, hours, mins, sec) -} - -#[cfg(test)] -mod tests { - use super::format_dhms; - - #[test] - fn download_tracker_format_dhms_test() { - assert_eq!(format_dhms(2), (0, 0, 0, 2)); - - assert_eq!(format_dhms(60), (0, 0, 1, 0)); - - assert_eq!(format_dhms(3_600), (0, 1, 0, 0)); - - assert_eq!(format_dhms(3_600 * 24), (1, 0, 0, 0)); - - assert_eq!(format_dhms(52_292), (0, 14, 31, 32)); - - assert_eq!(format_dhms(222_292), (2, 13, 44, 52)); + self.progress_bar.finish_and_clear(); + self.multi_progress_bars.remove(&self.progress_bar); + self.progress_bar = ProgressBar::hidden(); } } diff --git a/src/utils/notifications.rs b/src/utils/notifications.rs index 007bd39834..9ea4917277 100644 --- a/src/utils/notifications.rs +++ b/src/utils/notifications.rs @@ -89,7 +89,7 @@ impl Display for Notification<'_> { SetDefaultBufferSize(size) => write!( f, "using up to {} of RAM to unpack components", - units::Size::new(*size, units::Unit::B, units::UnitMode::Norm) + units::Size::new(*size, units::Unit::B) ), DownloadingFile(url, _) => write!(f, "downloading file from: '{url}'"), DownloadContentLengthReceived(len) => write!(f, "download size is: '{len}'"), diff --git a/src/utils/units.rs b/src/utils/units.rs index 365c99d7c5..e797bc48f0 100644 --- a/src/utils/units.rs +++ b/src/utils/units.rs @@ -6,22 +6,17 @@ pub enum Unit { IO, } -pub(crate) enum UnitMode { - Norm, - Rate, -} - /// Human readable size (some units) pub(crate) enum Size { - B(usize, UnitMode), - IO(usize, UnitMode), + B(usize), + IO(usize), } impl Size { - pub(crate) fn new(size: usize, unit: Unit, unitmode: UnitMode) -> Self { + pub(crate) fn new(size: usize, unit: Unit) -> Self { match unit { - Unit::B => Self::B(size, unitmode), - Unit::IO => Self::IO(size, unitmode), + Unit::B => Self::B(size), + Unit::IO => Self::IO(size), } } } @@ -29,16 +24,13 @@ impl Size { impl Display for Size { fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { match self { - Size::B(size, unitmode) => { + Size::B(size) => { const KI: f64 = 1024.0; const MI: f64 = KI * KI; const GI: f64 = KI * KI * KI; let size = *size as f64; - let suffix: String = match unitmode { - UnitMode::Norm => "".into(), - UnitMode::Rate => "/s".into(), - }; + let suffix: String = "".into(); if size >= GI { write!(f, "{:5.1} GiB{}", size / GI, suffix) @@ -50,16 +42,13 @@ impl Display for Size { write!(f, "{size:3.0} B{suffix}") } } - Size::IO(size, unitmode) => { + Size::IO(size) => { const K: f64 = 1000.0; const M: f64 = K * K; const G: f64 = K * K * K; let size = *size as f64; - let suffix: String = match unitmode { - UnitMode::Norm => "IO-ops".into(), - UnitMode::Rate => "IOPS".into(), - }; + let suffix: String = "IO-ops".into(); if size >= G { write!(f, "{:5.1} giga-{}", size / G, suffix) @@ -79,78 +68,33 @@ impl Display for Size { mod tests { #[test] fn unit_formatter_test() { - use crate::utils::units::{Size, Unit, UnitMode}; + use crate::utils::units::{Size, Unit}; // Test Bytes + assert_eq!(format!("{}", Size::new(1, Unit::B)), " 1 B"); + assert_eq!(format!("{}", Size::new(1024, Unit::B)), " 1.0 KiB"); assert_eq!( - format!("{}", Size::new(1, Unit::B, UnitMode::Norm)), - " 1 B" - ); - assert_eq!( - format!("{}", Size::new(1024, Unit::B, UnitMode::Norm)), - " 1.0 KiB" - ); - assert_eq!( - format!("{}", Size::new(1024usize.pow(2), Unit::B, UnitMode::Norm)), + format!("{}", Size::new(1024usize.pow(2), Unit::B)), " 1.0 MiB" ); assert_eq!( - format!("{}", Size::new(1024usize.pow(3), Unit::B, UnitMode::Norm)), + format!("{}", Size::new(1024usize.pow(3), Unit::B)), " 1.0 GiB" ); - // Test Bytes at given rate - assert_eq!( - format!("{}", Size::new(1, Unit::B, UnitMode::Rate)), - " 1 B/s" - ); - assert_eq!( - format!("{}", Size::new(1024, Unit::B, UnitMode::Rate)), - " 1.0 KiB/s" - ); - assert_eq!( - format!("{}", Size::new(1024usize.pow(2), Unit::B, UnitMode::Rate)), - " 1.0 MiB/s" - ); - assert_eq!( - format!("{}", Size::new(1024usize.pow(3), Unit::B, UnitMode::Rate)), - " 1.0 GiB/s" - ); - //Test I/O Operations + assert_eq!(format!("{}", Size::new(1, Unit::IO)), " 1 IO-ops"); assert_eq!( - format!("{}", Size::new(1, Unit::IO, UnitMode::Norm)), - " 1 IO-ops" - ); - assert_eq!( - format!("{}", Size::new(1000, Unit::IO, UnitMode::Norm)), + format!("{}", Size::new(1000, Unit::IO)), " 1.0 kilo-IO-ops" ); assert_eq!( - format!("{}", Size::new(1000usize.pow(2), Unit::IO, UnitMode::Norm)), + format!("{}", Size::new(1000usize.pow(2), Unit::IO)), " 1.0 mega-IO-ops" ); assert_eq!( - format!("{}", Size::new(1000usize.pow(3), Unit::IO, UnitMode::Norm)), + format!("{}", Size::new(1000usize.pow(3), Unit::IO)), " 1.0 giga-IO-ops" ); - - //Test I/O Operations at given rate - assert_eq!( - format!("{}", Size::new(1, Unit::IO, UnitMode::Rate)), - " 1 IOPS" - ); - assert_eq!( - format!("{}", Size::new(1000, Unit::IO, UnitMode::Rate)), - " 1.0 kilo-IOPS" - ); - assert_eq!( - format!("{}", Size::new(1000usize.pow(2), Unit::IO, UnitMode::Rate)), - " 1.0 mega-IOPS" - ); - assert_eq!( - format!("{}", Size::new(1000usize.pow(3), Unit::IO, UnitMode::Rate)), - " 1.0 giga-IOPS" - ); } }