Skip to main content

videre_core/
progress.rs

1use indicatif::{ProgressBar, ProgressStyle};
2use std::io::IsTerminal;
3use std::sync::atomic::{AtomicU64, Ordering};
4
5/// Reports progress for a batch of N items as an in-place bar (brew/docker/
6/// npm style) when stderr is a terminal, or periodic plain-text lines when
7/// it isn't (piped to a file, CI log), so a long run never looks hung in a
8/// log file, without per-item spam either way. `silent` suppresses the bar
9/// and periodic lines entirely, but NOT error output (see `println`) or the
10/// caller's own decision about whether to print a final summary.
11///
12/// Does not track elapsed time itself: callers that need it (e.g.
13/// `faces.rs`, whose summary spans both detection and clustering, not just
14/// the `Progress`-tracked detection phase) should use their own `Instant`
15/// spanning whatever the summary needs to cover.
16///
17/// Safe to share across threads: every method takes `&self`, so a single
18/// `Progress` value can be ticked concurrently from multiple `rayon`
19/// worker threads (e.g. from inside a `.par_iter()` closure) with no
20/// external `Arc`/`Mutex` wrapping needed at the call site.
21pub struct Progress {
22    total: u64,
23    done: AtomicU64,
24    mode: Mode,
25}
26
27enum Mode {
28    Bar(ProgressBar),
29    /// Non-TTY fallback: print one line every LOG_INTERVAL ticks.
30    Plain,
31    /// --silent: no bar, no periodic lines. Errors still print (see println).
32    Silent,
33}
34
35const LOG_INTERVAL: u64 = 25;
36
37impl Progress {
38    /// Creates a progress reporter for `total` items. When stderr is a TTY,
39    /// renders an in-place bar. When it isn't, falls back to one plain-text
40    /// line every `LOG_INTERVAL` items. `silent` suppresses both.
41    pub fn new(total: u64, silent: bool) -> Self {
42        let mode = if silent {
43            Mode::Silent
44        } else if std::io::stderr().is_terminal() {
45            let bar = ProgressBar::new(total);
46            bar.set_style(
47                ProgressStyle::with_template("{bar:40} {percent}%")
48                    .unwrap()
49                    .progress_chars("=> "),
50            );
51            Mode::Bar(bar)
52        } else {
53            Mode::Plain
54        };
55        Progress { total, done: AtomicU64::new(0), mode }
56    }
57
58    /// Advance by one item. Safe to call concurrently from multiple threads
59    /// (e.g. from inside a `rayon` `.par_iter()` closure) via a shared
60    /// `&Progress`, no external synchronization needed.
61    pub fn tick(&self) {
62        self.tick_by(1);
63    }
64
65    /// Advance by `n` items at once (for callers that complete work in
66    /// batches rather than one item at a time, e.g. `videre embed`'s
67    /// chunked pipeline). `n` must not exceed the number of items remaining
68    /// toward `total` (mirrors the same implicit contract `tick()` already
69    /// has: callers are responsible for not calling it more times, or with
70    /// a larger cumulative `n`, than `total` allows). Safe to call
71    /// concurrently from multiple threads, same as `tick()`.
72    pub fn tick_by(&self, n: u64) {
73        let before = self.done.fetch_add(n, Ordering::Relaxed);
74        let after = before + n;
75        match &self.mode {
76            Mode::Bar(bar) => bar.set_position(after),
77            Mode::Plain => {
78                if after / LOG_INTERVAL != before / LOG_INTERVAL || after == self.total {
79                    eprintln!("{}/{} images processed", after, self.total);
80                }
81            }
82            Mode::Silent => {}
83        }
84    }
85
86    /// Print a line that survives an active progress bar without corrupting
87    /// its rendering. Always prints, regardless of `silent`, matches the
88    /// existing unconditional behavior of per-image error messages
89    /// (`detect failed ...`, `embed_batch failed ...`, `write failed ...`),
90    /// which must stay visible even under --silent since they indicate data
91    /// loss, not routine progress.
92    pub fn println(&self, msg: &str) {
93        match &self.mode {
94            Mode::Bar(bar) => bar.println(msg),
95            Mode::Plain | Mode::Silent => eprintln!("{msg}"),
96        }
97    }
98
99    /// Clears the bar (if any) so the final summary prints cleanly below it
100    /// rather than being overwritten. Does not print anything itself, the
101    /// caller assembles and prints its own summary line(s).
102    pub fn finish(self) {
103        if let Mode::Bar(bar) = self.mode {
104            bar.finish_and_clear();
105        }
106    }
107}
108
109#[cfg(test)]
110mod tests {
111    use super::*;
112
113    #[test]
114    fn silent_mode_tick_does_not_panic() {
115        let p = Progress::new(10, true);
116        for _ in 0..10 {
117            p.tick();
118        }
119        p.finish();
120    }
121
122    #[test]
123    fn silent_mode_println_still_prints() {
124        // println() must not panic in silent mode; it always writes to
125        // stderr regardless of `silent` (verified by not panicking here -
126        // capturing stderr output itself is not practical in a unit test).
127        let p = Progress::new(5, true);
128        p.println("an error message");
129    }
130
131    #[test]
132    fn zero_total_does_not_panic() {
133        let p = Progress::new(0, true);
134        p.tick();
135        p.finish();
136    }
137
138    #[test]
139    fn silent_mode_tick_by_does_not_panic() {
140        let p = Progress::new(100, true);
141        p.tick_by(40);
142        p.tick_by(60);
143        p.finish();
144    }
145
146    #[test]
147    fn concurrent_tick_from_multiple_threads_reaches_correct_total() {
148        use std::sync::Arc;
149        let progress = Arc::new(Progress::new(1000, true));
150        let handles: Vec<_> = (0..10)
151            .map(|_| {
152                let p = Arc::clone(&progress);
153                std::thread::spawn(move || {
154                    for _ in 0..100 {
155                        p.tick();
156                    }
157                })
158            })
159            .collect();
160        for h in handles {
161            h.join().unwrap();
162        }
163        assert_eq!(progress.done.load(Ordering::Relaxed), 1000);
164    }
165}