Skip to main content

rustfs_io_core/
progress.rs

1// Copyright 2024 RustFS Team
2//
3// Licensed under the Apache License, Version 2.0 (the "License");
4// you may not use this file except in compliance with the License.
5// You may obtain a copy of the License at
6//
7//     http://www.apache.org/licenses/LICENSE-2.0
8//
9// Unless required by applicable law or agreed to in writing, software
10// distributed under the License is distributed on an "AS IS" BASIS,
11// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12// See the License for the specific language governing permissions and
13// limitations under the License.
14
15//! Progress tracking for long-running I/O operations.
16//!
17//! Re-exported as `rustfs_concurrency::OperationProgress` for the storage
18//! timeout implementation, which uses `is_stale` to tell a slow transfer
19//! apart from a stalled one.
20
21use std::sync::atomic::{AtomicU64, Ordering};
22use std::time::{Duration, Instant};
23
24/// Operation progress tracker.
25#[derive(Debug)]
26pub struct OperationProgress {
27    /// Total size (if known).
28    pub total_size: Option<u64>,
29    /// Bytes processed.
30    bytes_processed: AtomicU64,
31    /// Last update time.
32    last_update: std::sync::Mutex<Instant>,
33    /// Stale timeout.
34    stale_timeout: Duration,
35    /// Start time for transfer rate calculation.
36    start_time: Instant,
37}
38
39impl OperationProgress {
40    /// Create new operation progress.
41    pub fn new(total_size: Option<u64>, stale_timeout: Duration) -> Self {
42        Self {
43            total_size,
44            bytes_processed: AtomicU64::new(0),
45            last_update: std::sync::Mutex::new(Instant::now()),
46            stale_timeout,
47            start_time: Instant::now(),
48        }
49    }
50
51    /// Update progress.
52    pub fn update(&self, bytes: u64) {
53        self.bytes_processed.store(bytes, Ordering::Relaxed);
54        if let Ok(mut last) = self.last_update.lock() {
55            *last = Instant::now();
56        }
57    }
58
59    /// Add to progress.
60    pub fn add(&self, bytes: u64) {
61        self.bytes_processed.fetch_add(bytes, Ordering::Relaxed);
62        if let Ok(mut last) = self.last_update.lock() {
63            *last = Instant::now();
64        }
65    }
66
67    /// Get current progress.
68    pub fn current(&self) -> u64 {
69        self.bytes_processed.load(Ordering::Relaxed)
70    }
71
72    /// Check if progress is stale.
73    pub fn is_stale(&self) -> bool {
74        if let Ok(last) = self.last_update.lock() {
75            last.elapsed() > self.stale_timeout
76        } else {
77            false
78        }
79    }
80
81    /// Get progress percentage.
82    pub fn progress_percent(&self) -> Option<f64> {
83        self.total_size.map(|total| {
84            if total == 0 {
85                100.0
86            } else {
87                let processed = self.bytes_processed.load(Ordering::Relaxed);
88                (processed as f64 / total as f64 * 100.0).min(100.0)
89            }
90        })
91    }
92
93    /// Get remaining bytes.
94    pub fn remaining(&self) -> Option<u64> {
95        self.total_size.map(|total| {
96            let processed = self.bytes_processed.load(Ordering::Relaxed);
97            total.saturating_sub(processed)
98        })
99    }
100
101    /// Calculate transfer rate in bytes per second.
102    ///
103    /// Returns 0 if no time has elapsed or no data transferred.
104    pub fn transfer_rate(&self) -> u64 {
105        let processed = self.bytes_processed.load(Ordering::Relaxed);
106        if processed == 0 {
107            return 0;
108        }
109
110        let elapsed = self.start_time.elapsed().as_secs_f64();
111        if elapsed > 0.0 {
112            (processed as f64 / elapsed) as u64
113        } else {
114            0
115        }
116    }
117}
118
119#[cfg(test)]
120mod tests {
121    use super::*;
122
123    #[test]
124    fn test_operation_progress() {
125        let progress = OperationProgress::new(Some(1000), Duration::from_secs(5));
126
127        assert_eq!(progress.current(), 0);
128        assert_eq!(progress.progress_percent(), Some(0.0));
129
130        progress.update(500);
131        assert_eq!(progress.current(), 500);
132        assert_eq!(progress.progress_percent(), Some(50.0));
133
134        progress.add(300);
135        assert_eq!(progress.current(), 800);
136        assert_eq!(progress.remaining(), Some(200));
137    }
138}