virustotal_rs/iterator_utils/
progress.rs1use std::sync::Arc;
4use std::time::{Duration, SystemTime};
5use tokio::sync::RwLock;
6
7pub 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 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 pub fn with_update_interval(mut self, interval: Duration) -> Self {
33 self.update_interval = interval;
34 self
35 }
36
37 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 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 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#[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 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 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}