db_mover/
progress.rs

1use std::{
2    fmt::Display,
3    sync::atomic,
4    time::{Duration, Instant},
5};
6
7use num_format::{Locale, ToFormattedString};
8use tracing::info;
9
10struct RateLimiter {
11    last_usage: atomic::AtomicU64,
12    min_duration_between: Duration,
13    started: Instant,
14}
15
16impl RateLimiter {
17    fn new(min_secs_between: u64) -> Self {
18        return Self {
19            last_usage: atomic::AtomicU64::new(0),
20            min_duration_between: Duration::from_secs(min_secs_between),
21            started: Instant::now(),
22        };
23    }
24
25    fn get_token(&self) -> Result<(), ()> {
26        let last_usage = Duration::from_secs(self.last_usage.load(atomic::Ordering::Relaxed));
27        if self.started.elapsed() - last_usage < self.min_duration_between {
28            return Err(());
29        }
30        let new_usage = self.started.elapsed();
31
32        return match self.last_usage.compare_exchange_weak(
33            last_usage.as_secs(),
34            new_usage.as_secs(),
35            atomic::Ordering::Relaxed,
36            atomic::Ordering::Relaxed,
37        ) {
38            Ok(_) => Ok(()),
39            Err(_) => Err(()),
40        };
41    }
42}
43
44#[derive(Debug)]
45pub struct FormattedDuration(pub Duration);
46
47impl Display for FormattedDuration {
48    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
49        let mut t = self.0.as_secs();
50        let seconds = t % 60;
51        t /= 60;
52        let minutes = t % 60;
53        t /= 60;
54        let hours = t % 24;
55        t /= 24;
56        if t > 0 {
57            let days = t;
58            write!(f, "{days}d {hours:02}:{minutes:02}:{seconds:02}")
59        } else {
60            write!(f, "{hours:02}:{minutes:02}:{seconds:02}")
61        }
62    }
63}
64
65struct ProgressTracker {
66    total: Option<u64>,
67    current: atomic::AtomicU64,
68    finished: atomic::AtomicBool,
69    started: Instant,
70}
71
72impl ProgressTracker {
73    fn new(total: Option<u64>) -> Self {
74        return Self {
75            total,
76            current: atomic::AtomicU64::new(0),
77            finished: atomic::AtomicBool::new(false),
78            started: Instant::now(),
79        };
80    }
81
82    fn inc(&self, value: u64) {
83        self.current.fetch_add(value, atomic::Ordering::Relaxed);
84    }
85
86    fn current(&self) -> u64 {
87        return self.current.load(atomic::Ordering::Relaxed);
88    }
89
90    fn finished(&self) -> bool {
91        return self.finished.load(atomic::Ordering::Relaxed);
92    }
93
94    fn finish(&self) {
95        self.finished.store(true, atomic::Ordering::Relaxed);
96    }
97}
98
99impl Display for ProgressTracker {
100    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
101        let current = self.current();
102        f.write_fmt(format_args!(
103            "[{}] ",
104            FormattedDuration(self.started.elapsed()),
105        ))?;
106        let finished = self.finished();
107        if finished {
108            f.write_str("Finished ")?;
109        }
110        // Prevent zero division
111        let per_sec = current / std::cmp::max(self.started.elapsed().as_secs(), 1);
112
113        if let Some(total) = self.total {
114            let mut percent = 100;
115            let mut percent_remainder = 0;
116            if total > 0 {
117                percent = current * 100 / total;
118                percent_remainder = (current * 10000 / total) % 100;
119            }
120            f.write_fmt(format_args!(
121                "Processed: {percent}.{percent_remainder:02}% ({}/{}) ",
122                current.to_formatted_string(&Locale::en),
123                total.to_formatted_string(&Locale::en),
124            ))?;
125            if !finished && per_sec > 0 {
126                let eta = FormattedDuration(Duration::from_secs((total - current) / per_sec));
127                f.write_fmt(format_args!("ETA: {eta} "))?;
128            }
129        } else {
130            f.write_fmt(format_args!("Processed: {current} "))?;
131        }
132
133        if !finished {
134            f.write_fmt(format_args!("Rows per sec: {per_sec} "))?;
135        }
136        return Ok(());
137    }
138}
139
140pub struct TableMigrationProgress {
141    table: String,
142    reader: ProgressTracker,
143    writer: ProgressTracker,
144    limiter: RateLimiter,
145}
146
147impl TableMigrationProgress {
148    pub fn inc_reader(&self, value: u64) {
149        self.reader.inc(value);
150        self.log_with_limit();
151    }
152
153    pub fn inc_writer(&self, value: u64) {
154        self.writer.inc(value);
155        self.log_with_limit();
156    }
157
158    pub fn finish_reader(&self) {
159        self.reader.finish();
160    }
161
162    pub fn finish_writer(&self) {
163        self.writer.finish();
164    }
165
166    pub fn reader_processed(&self) -> u64 {
167        return self.reader.current();
168    }
169
170    pub fn writer_processed(&self) -> u64 {
171        return self.writer.current();
172    }
173
174    fn log_with_limit(&self) {
175        if self.limiter.get_token().is_ok() {
176            self.log();
177        }
178    }
179
180    fn log(&self) {
181        info!("Reading table \"{}\" {}", self.table, self.reader);
182        info!("Writing table \"{}\" {}", self.table, self.writer);
183    }
184}
185
186impl Drop for TableMigrationProgress {
187    fn drop(&mut self) {
188        if self.reader.current() > 0 || self.writer.current() > 0 {
189            self.log();
190        }
191    }
192}
193
194impl TableMigrationProgress {
195    pub fn new(table: &str, num_rows: Option<u64>) -> Self {
196        return Self {
197            table: table.to_string(),
198            reader: ProgressTracker::new(num_rows),
199            writer: ProgressTracker::new(num_rows),
200            limiter: RateLimiter::new(1),
201        };
202    }
203}
db_mover/progress.rs

db_mover/
progress.rs
