Skip to main content

api_scanner/
progress_tracker.rs

1//! Universal Progress Tracker
2//!
3//! A reusable progress tracking utility that automatically detects terminal capabilities
4//! and provides appropriate output formatting for both TTY and non-TTY environments.
5//!
6//! ## Features
7//! - Automatic TTY detection
8//! - Thread-safe progress updates
9//! - Configurable update frequency
10//! - Clean output for both interactive and CI/CD environments
11//!
12//! ## Usage
13//!
14//! ```rust
15//! use api_scanner::progress_tracker::ProgressTracker;
16//!
17//! #[tokio::main]
18//! async fn main() {
19//!     let tracker = ProgressTracker::new(100); // 100 total items
20//!
21//!     for _ in 0..100 {
22//!         // Do work...
23//!         tracker.increment(Some("Processing item")).await;
24//!     }
25//!
26//!     tracker.finish().await;
27//! }
28//! ```
29
30use std::io::{IsTerminal, Write};
31use std::sync::Arc;
32use std::time::Instant;
33use tokio::sync::Mutex;
34
35/// Configuration for progress tracking behavior
36#[derive(Clone)]
37pub struct ProgressConfig {
38    /// Total number of items to process
39    pub total: usize,
40    /// Update frequency (every N items) for TTY mode
41    pub tty_update_frequency: usize,
42    /// Update frequency (every N items) for non-TTY mode
43    pub non_tty_update_frequency: usize,
44    /// Whether to show elapsed time
45    pub show_elapsed: bool,
46    /// Whether to show ETA
47    pub show_eta: bool,
48    /// Whether to show rate
49    pub show_rate: bool,
50    /// Custom prefix for progress messages
51    pub prefix: String,
52    /// Whether to show detailed messages (URLs, findings, etc.)
53    pub show_details: bool,
54}
55
56impl Default for ProgressConfig {
57    fn default() -> Self {
58        Self {
59            total: 0,
60            tty_update_frequency: 5,
61            non_tty_update_frequency: 1,
62            show_elapsed: true,
63            show_eta: true,
64            show_rate: true,
65            prefix: String::new(),
66            show_details: true,
67        }
68    }
69}
70
71/// Thread-safe progress tracker
72pub struct ProgressTracker {
73    config: ProgressConfig,
74    progress: Arc<Mutex<usize>>,
75    start_time: Instant,
76    is_tty: bool,
77}
78
79impl ProgressTracker {
80    /// Create a new progress tracker with default configuration
81    pub fn new(total: usize) -> Self {
82        Self::with_config(ProgressConfig {
83            total,
84            ..Default::default()
85        })
86    }
87
88    /// Create a new progress tracker with custom configuration
89    pub fn with_config(config: ProgressConfig) -> Self {
90        Self {
91            config,
92            progress: Arc::new(Mutex::new(0)),
93            start_time: Instant::now(),
94            is_tty: std::io::stderr().is_terminal(),
95        }
96    }
97
98    /// Get a cloneable handle for use in async tasks
99    pub fn handle(&self) -> ProgressHandle {
100        ProgressHandle {
101            config: self.config.clone(),
102            progress: Arc::clone(&self.progress),
103            start_time: self.start_time,
104            is_tty: self.is_tty,
105        }
106    }
107
108    /// Increment progress by 1 and optionally display a message
109    pub async fn increment(&self, message: Option<&str>) {
110        let mut p = self.progress.lock().await;
111        *p += 1;
112        let current = *p;
113        drop(p);
114
115        self.display_progress(current, message).await;
116    }
117
118    /// Set progress to a specific value
119    pub async fn set(&self, value: usize, message: Option<&str>) {
120        let mut p = self.progress.lock().await;
121        *p = value;
122        drop(p);
123
124        self.display_progress(value, message).await;
125    }
126
127    /// Get current progress value
128    pub async fn current(&self) -> usize {
129        *self.progress.lock().await
130    }
131
132    /// Display progress based on TTY detection
133    async fn display_progress(&self, current: usize, message: Option<&str>) {
134        let update_freq = if self.is_tty {
135            self.config.tty_update_frequency
136        } else {
137            self.config.non_tty_update_frequency
138        };
139
140        // Only display at specified frequency or when complete
141        let is_multiple = update_freq != 0 && current % update_freq == 0;
142        if !is_multiple && current != self.config.total {
143            return;
144        }
145
146        let elapsed = self.start_time.elapsed().as_secs();
147        let percentage = if self.config.total > 0 {
148            (current as f64 / self.config.total as f64) * 100.0
149        } else {
150            0.0
151        };
152
153        if self.is_tty {
154            // TTY mode: use carriage return for same-line updates
155            let mut output = format!(
156                "\r{}{}/{} ({:.1}%)",
157                self.config.prefix, current, self.config.total, percentage
158            );
159
160            if self.config.show_rate && elapsed > 0 {
161                let rate = current as f64 / elapsed as f64;
162                output.push_str(&format!(" | {:.1}/s", rate));
163            }
164
165            if self.config.show_eta && elapsed > 0 && current > 0 {
166                let rate = current as f64 / elapsed as f64;
167                let remaining = self.config.total.saturating_sub(current);
168                let eta_secs = (remaining as f64 / rate) as u64;
169                let eta_mins = eta_secs / 60;
170                output.push_str(&format!(" | ETA: {}m{}s", eta_mins, eta_secs % 60));
171            }
172
173            if self.config.show_elapsed {
174                output.push_str(&format!(" | Elapsed: {}s", elapsed));
175            }
176
177            if self.config.show_details {
178                if let Some(msg) = message {
179                    output.push_str(&format!(" | {}", msg));
180                }
181            }
182
183            output.push_str("   "); // Clear any leftover characters
184
185            eprint!("{}", output);
186            std::io::stderr().flush().ok();
187        } else {
188            // Non-TTY mode: new line for each update
189            let mut output = format!("[{}/{}] ({:.1}%)", current, self.config.total, percentage);
190
191            if !self.config.prefix.is_empty() {
192                output = format!("{} {}", self.config.prefix, output);
193            }
194
195            if self.config.show_elapsed {
196                output.push_str(&format!(" | Elapsed: {}s", elapsed));
197            }
198
199            if self.config.show_details {
200                if let Some(msg) = message {
201                    output.push_str(&format!(" | {}", msg));
202                }
203            }
204
205            eprintln!("{}", output);
206        }
207    }
208
209    /// Finish progress tracking and clear the line (TTY mode)
210    pub async fn finish(&self) {
211        if self.is_tty {
212            eprintln!(); // Move to next line
213        }
214        let elapsed = self.start_time.elapsed();
215        eprintln!(
216            "✅ Completed {} items in {:.2}s",
217            self.config.total,
218            elapsed.as_secs_f64()
219        );
220    }
221}
222
223/// Cloneable handle for use in async tasks
224#[derive(Clone)]
225pub struct ProgressHandle {
226    config: ProgressConfig,
227    progress: Arc<Mutex<usize>>,
228    start_time: Instant,
229    is_tty: bool,
230}
231
232impl ProgressHandle {
233    /// Increment progress by 1
234    pub async fn increment(&self, message: Option<&str>) {
235        let mut p = self.progress.lock().await;
236        *p += 1;
237        let current = *p;
238        drop(p);
239
240        self.display_progress(current, message).await;
241    }
242
243    /// Get current progress value
244    pub async fn current(&self) -> usize {
245        *self.progress.lock().await
246    }
247
248    async fn display_progress(&self, current: usize, message: Option<&str>) {
249        let update_freq = if self.is_tty {
250            self.config.tty_update_frequency
251        } else {
252            self.config.non_tty_update_frequency
253        };
254
255        let is_multiple = update_freq != 0 && current % update_freq == 0;
256        if !is_multiple && current != self.config.total {
257            return;
258        }
259
260        let elapsed = self.start_time.elapsed().as_secs();
261        let percentage = if self.config.total > 0 {
262            (current as f64 / self.config.total as f64) * 100.0
263        } else {
264            0.0
265        };
266
267        if self.is_tty {
268            let mut output = format!(
269                "\r{}{}/{} ({:.1}%)",
270                self.config.prefix, current, self.config.total, percentage
271            );
272
273            if self.config.show_rate && elapsed > 0 {
274                let rate = current as f64 / elapsed as f64;
275                output.push_str(&format!(" | {:.1}/s", rate));
276            }
277
278            if self.config.show_eta && elapsed > 0 && current > 0 {
279                let rate = current as f64 / elapsed as f64;
280                let remaining = self.config.total.saturating_sub(current);
281                let eta_secs = (remaining as f64 / rate) as u64;
282                let eta_mins = eta_secs / 60;
283                output.push_str(&format!(" | ETA: {}m{}s", eta_mins, eta_secs % 60));
284            }
285
286            if self.config.show_elapsed {
287                output.push_str(&format!(" | Elapsed: {}s", elapsed));
288            }
289
290            if self.config.show_details {
291                if let Some(msg) = message {
292                    output.push_str(&format!(" | {}", msg));
293                }
294            }
295
296            output.push_str("   ");
297
298            eprint!("{}", output);
299            std::io::stderr().flush().ok();
300        } else {
301            let mut output = format!("[{}/{}] ({:.1}%)", current, self.config.total, percentage);
302
303            if !self.config.prefix.is_empty() {
304                output = format!("{} {}", self.config.prefix, output);
305            }
306
307            if self.config.show_elapsed {
308                output.push_str(&format!(" | Elapsed: {}s", elapsed));
309            }
310
311            if self.config.show_details {
312                if let Some(msg) = message {
313                    output.push_str(&format!(" | {}", msg));
314                }
315            }
316
317            eprintln!("{}", output);
318        }
319    }
320}