gammalooprs/utils/
progress.rs

1use std::{fmt, time::Duration};
2
3use indicatif::{FormattedDuration, ProgressState, ProgressStyle};
4
5const LONG_RUNNING_PROGRESS_TEMPLATE: &str = "[{elapsed_precise} | ETA: {eta_precise}] {bar:40.cyan/blue} {pos:>7}/{len:7} ({percent}%) {msg}";
6
7fn write_eta(eta: Duration, writer: &mut dyn fmt::Write) -> fmt::Result {
8    // Indicatif saturates a floating-point ETA overflow to `u64::MAX` seconds.
9    // Rendering that sentinel literally produces a meaningless 213503982334601-day ETA.
10    if eta.as_secs() == u64::MAX {
11        writer.write_str("N/A")
12    } else {
13        write!(writer, "{}", FormattedDuration(eta))
14    }
15}
16
17fn estimate_average_eta(
18    elapsed: Duration,
19    completed_steps: u64,
20    total_steps: Option<u64>,
21) -> Option<Duration> {
22    let total_steps = total_steps?;
23    if completed_steps == 0 {
24        return None;
25    }
26    if completed_steps >= total_steps {
27        return Some(Duration::ZERO);
28    }
29
30    let remaining_steps = total_steps - completed_steps;
31    let eta_nanos = elapsed
32        .as_nanos()
33        .checked_mul(u128::from(remaining_steps))?
34        / u128::from(completed_steps);
35    let eta_seconds = eta_nanos / 1_000_000_000;
36    if eta_seconds > u128::from(u64::MAX) {
37        return None;
38    }
39    Some(Duration::from_secs(eta_seconds as u64))
40}
41
42fn write_eta_after_warmup(
43    eta: Option<Duration>,
44    completed_steps: u64,
45    minimum_completed_steps: u64,
46    is_finished: bool,
47    writer: &mut dyn fmt::Write,
48) -> fmt::Result {
49    if !is_finished && completed_steps < minimum_completed_steps {
50        writer.write_str("N/A")
51    } else if let Some(eta) = eta {
52        write_eta(eta, writer)
53    } else {
54        writer.write_str("N/A")
55    }
56}
57
58fn write_safe_eta(state: &ProgressState, writer: &mut dyn fmt::Write) {
59    let is_finished =
60        state.is_finished() || state.len().is_some_and(|length| state.pos() >= length);
61    let _ = write_eta_after_warmup(
62        estimate_average_eta(state.elapsed(), state.pos(), state.len()),
63        state.pos(),
64        1,
65        is_finished,
66        writer,
67    );
68}
69
70/// Standard style for bounded, long-running progress bars.
71///
72/// The custom ETA renderer uses average completed-work throughput and keeps
73/// unavailable or overflowing estimates out of terminal output.
74pub fn long_running_progress_style() -> ProgressStyle {
75    ProgressStyle::with_template(LONG_RUNNING_PROGRESS_TEMPLATE)
76        .expect("long-running progress bar template should be valid")
77        .with_key("eta_precise", write_safe_eta)
78}
79
80/// Long-running progress style that withholds ETA until enough work units have
81/// completed for the rate estimate to be meaningful.
82pub fn long_running_progress_style_with_eta_warmup(minimum_completed_steps: u64) -> ProgressStyle {
83    ProgressStyle::with_template(LONG_RUNNING_PROGRESS_TEMPLATE)
84        .expect("long-running progress bar template should be valid")
85        .with_key(
86            "eta_precise",
87            move |state: &ProgressState, writer: &mut dyn fmt::Write| {
88                let is_finished =
89                    state.is_finished() || state.len().is_some_and(|length| state.pos() >= length);
90                let _ = write_eta_after_warmup(
91                    estimate_average_eta(state.elapsed(), state.pos(), state.len()),
92                    state.pos(),
93                    minimum_completed_steps,
94                    is_finished,
95                    writer,
96                );
97            },
98        )
99}
100
101#[cfg(test)]
102mod tests {
103    use super::*;
104
105    #[test]
106    fn saturated_eta_is_rendered_as_unavailable() {
107        let mut rendered = String::new();
108        write_eta(Duration::new(u64::MAX, 0), &mut rendered).unwrap();
109        assert_eq!(rendered, "N/A");
110    }
111
112    #[test]
113    fn finite_eta_keeps_precise_duration_format() {
114        let mut rendered = String::new();
115        write_eta(Duration::from_secs(3_661), &mut rendered).unwrap();
116        assert_eq!(rendered, "01:01:01");
117    }
118
119    #[test]
120    fn eta_is_hidden_until_warmup_completes() {
121        let mut rendered = String::new();
122        write_eta_after_warmup(
123            Some(Duration::from_secs(10_533)),
124            1,
125            10,
126            false,
127            &mut rendered,
128        )
129        .unwrap();
130        assert_eq!(rendered, "N/A");
131
132        rendered.clear();
133        write_eta_after_warmup(Some(Duration::from_secs(600)), 10, 10, false, &mut rendered)
134            .unwrap();
135        assert_eq!(rendered, "00:10:00");
136    }
137
138    #[test]
139    fn average_eta_uses_completed_work_instead_of_idle_ticks() {
140        assert_eq!(
141            estimate_average_eta(Duration::from_secs(655), 21, Some(332)),
142            Some(Duration::from_secs(9_700)),
143        );
144        assert_eq!(
145            estimate_average_eta(Duration::from_secs(655), 0, Some(332)),
146            None
147        );
148    }
149}
gammalooprs/utils/progress.rs

gammalooprs/utils/
progress.rs
