Skip to main content

virustotal_rs/iterator_utils/
progress.rs

1//! Progress tracking utilities for iterators
2
3use std::sync::Arc;
4use std::time::{Duration, SystemTime};
5use tokio::sync::RwLock;
6
7/// Progress tracker for iterator operations
8pub struct ProgressTracker {
9    pub name: String,
10    pub items_processed: Arc<RwLock<u64>>,
11    pub batches_processed: Arc<RwLock<u32>>,
12    pub start_time: SystemTime,
13    pub last_update: Arc<RwLock<SystemTime>>,
14    pub update_interval: Duration,
15}
16
17impl ProgressTracker {
18    /// Create a new progress tracker
19    pub fn new(name: impl Into<String>) -> Self {
20        let now = SystemTime::now();
21        Self {
22            name: name.into(),
23            items_processed: Arc::new(RwLock::new(0)),
24            batches_processed: Arc::new(RwLock::new(0)),
25            start_time: now,
26            last_update: Arc::new(RwLock::new(now)),
27            update_interval: Duration::from_secs(1),
28        }
29    }
30
31    /// Set the update interval for progress reports
32    pub fn with_update_interval(mut self, interval: Duration) -> Self {
33        self.update_interval = interval;
34        self
35    }
36
37    /// Update progress with new batch
38    pub async fn update_batch(&self, batch_size: usize) {
39        let mut items = self.items_processed.write().await;
40        let mut batches = self.batches_processed.write().await;
41        let mut last_update = self.last_update.write().await;
42
43        *items += batch_size as u64;
44        *batches += 1;
45
46        let now = SystemTime::now();
47        if now.duration_since(*last_update).unwrap_or(Duration::ZERO) >= self.update_interval {
48            let elapsed = now
49                .duration_since(self.start_time)
50                .unwrap_or(Duration::ZERO);
51            let rate = if elapsed.as_secs_f64() > 0.0 {
52                *items as f64 / elapsed.as_secs_f64()
53            } else {
54                0.0
55            };
56
57            println!(
58                "[{}] Processed {} items in {} batches ({:.2} items/sec)",
59                self.name, *items, *batches, rate
60            );
61
62            *last_update = now;
63        }
64    }
65
66    /// Get current statistics
67    pub async fn stats(&self) -> ProgressStats {
68        let items = *self.items_processed.read().await;
69        let batches = *self.batches_processed.read().await;
70        let elapsed = SystemTime::now()
71            .duration_since(self.start_time)
72            .unwrap_or(Duration::ZERO);
73        let rate = if elapsed.as_secs_f64() > 0.0 {
74            items as f64 / elapsed.as_secs_f64()
75        } else {
76            0.0
77        };
78
79        ProgressStats {
80            items_processed: items,
81            batches_processed: batches,
82            elapsed,
83            items_per_second: rate,
84        }
85    }
86
87    /// Reset progress statistics
88    pub async fn reset(&self) {
89        *self.items_processed.write().await = 0;
90        *self.batches_processed.write().await = 0;
91        let now = SystemTime::now();
92        *self.last_update.write().await = now;
93    }
94}
95
96/// Progress statistics
97#[derive(Debug, Clone)]
98pub struct ProgressStats {
99    pub items_processed: u64,
100    pub batches_processed: u32,
101    pub elapsed: Duration,
102    pub items_per_second: f64,
103}
104
105impl ProgressStats {
106    /// Create empty progress stats
107    pub fn new() -> Self {
108        Self {
109            items_processed: 0,
110            batches_processed: 0,
111            elapsed: Duration::ZERO,
112            items_per_second: 0.0,
113        }
114    }
115
116    /// Get estimated time remaining based on current rate and target
117    pub fn estimated_remaining(&self, target_items: u64) -> Duration {
118        if self.items_per_second <= 0.0 || self.items_processed >= target_items {
119            return Duration::ZERO;
120        }
121
122        let remaining_items = target_items - self.items_processed;
123        let seconds = remaining_items as f64 / self.items_per_second;
124        Duration::from_secs_f64(seconds)
125    }
126}
127
128impl Default for ProgressStats {
129    fn default() -> Self {
130        Self::new()
131    }
132}