1
use std::sync::atomic::AtomicU64;
2
use std::sync::atomic::Ordering;
3
use std::time::Duration;
4
use std::time::Instant;
5

            
6
/// A time-based progress tracker that prints a message at most once per interval,
7
/// so a long-running procedure with many steps does not flood the log with
8
/// progress indications.
9
pub struct TimeProgress<T> {
10
    interval: Duration,
11
    /// Reference instant against which `last_update` is measured.
12
    start: Instant,
13
    /// Nanoseconds since `start` at the last printed update.
14
    last_update: AtomicU64,
15
    message: Box<dyn Fn(T) + Send + Sync>,
16
}
17

            
18
impl<T> TimeProgress<T> {
19
    /// Create a new time-based progress tracker with a given interval in seconds.
20
11621
    pub fn new<F>(message: F, interval_seconds: u64) -> TimeProgress<T>
21
11621
    where
22
11621
        F: Fn(T) + Send + Sync + 'static,
23
    {
24
11621
        TimeProgress {
25
11621
            message: Box::new(message),
26
11621
            interval: Duration::from_secs(interval_seconds),
27
11621
            start: Instant::now(),
28
11621
            last_update: AtomicU64::new(0),
29
11621
        }
30
11621
    }
31

            
32
    /// Increase the progress with the given amount, prints periodic progress
33
    /// messages based on time intervals.
34
    ///
35
    /// When called concurrently from several threads, a compare-exchange ensures
36
    /// exactly one caller claims the interval and emits the message, so no
37
    /// duplicate lines are printed.
38
9935950
    pub fn print(&self, object: T) {
39
9935950
        let elapsed = self.elapsed_nanos();
40
9935950
        let last = self.last_update.load(Ordering::Relaxed);
41
9935950
        if elapsed.saturating_sub(last) >= self.interval_nanos()
42
11184
            && self
43
11184
                .last_update
44
11184
                .compare_exchange(last, elapsed, Ordering::Relaxed, Ordering::Relaxed)
45
11184
                .is_ok()
46
11184
        {
47
11184
            (self.message)(object);
48
9929856
        }
49
9935950
    }
50

            
51
    /// Returns true iff the progress tracker is due for an update.
52
2103
    pub fn is_due(&self) -> bool {
53
2103
        self.elapsed_nanos()
54
2103
            .saturating_sub(self.last_update.load(Ordering::Relaxed))
55
2103
            >= self.interval_nanos()
56
2103
    }
57

            
58
    /// Nanoseconds elapsed since `start`, saturating at `u64::MAX` (~584 years).
59
9938053
    fn elapsed_nanos(&self) -> u64 {
60
9938053
        self.start.elapsed().as_nanos().min(u64::MAX as u128) as u64
61
9938053
    }
62

            
63
    /// The configured interval expressed in nanoseconds.
64
9938053
    fn interval_nanos(&self) -> u64 {
65
9938053
        self.interval.as_nanos().min(u64::MAX as u128) as u64
66
9938053
    }
67
}
68

            
69
#[cfg(test)]
70
mod tests {
71
    use super::TimeProgress;
72

            
73
    use std::sync::Arc;
74
    use std::sync::atomic::AtomicUsize;
75
    use std::sync::atomic::Ordering;
76

            
77
    use rand::RngExt;
78

            
79
    use merc_utilities::random_test;
80

            
81
    /// A `TimeProgress` whose closure counts how often it fired.
82
101
    fn counting_progress(interval_seconds: u64) -> (TimeProgress<()>, Arc<AtomicUsize>) {
83
101
        let count = Arc::new(AtomicUsize::new(0));
84
101
        let counter = count.clone();
85
101
        let progress = TimeProgress::new(
86
5190
            move |()| {
87
5190
                counter.fetch_add(1, Ordering::Relaxed);
88
5190
            },
89
101
            interval_seconds,
90
        );
91
101
        (progress, count)
92
101
    }
93

            
94
    #[test]
95
1
    fn test_zero_interval_fires_every_call() {
96
100
        random_test(100, |rng| {
97
100
            let n = rng.random_range(1..100);
98
100
            let (progress, count) = counting_progress(0);
99

            
100
100
            assert!(progress.is_due(), "a zero interval is always due");
101
5190
            for _ in 0..n {
102
5190
                progress.print(());
103
5190
            }
104
100
            assert_eq!(
105
100
                count.load(Ordering::Relaxed),
106
                n,
107
                "every print must fire at a zero interval"
108
            );
109
100
        });
110
1
    }
111

            
112
    #[test]
113
1
    fn test_large_interval_does_not_fire() {
114
1
        let (progress, count) = counting_progress(3600);
115

            
116
1
        assert!(!progress.is_due(), "a one-hour interval cannot be due immediately");
117
100
        for _ in 0..100 {
118
100
            progress.print(());
119
100
        }
120
1
        assert_eq!(
121
1
            count.load(Ordering::Relaxed),
122
            0,
123
            "no print should fire within the interval"
124
        );
125
1
    }
126
}