1use std::collections::VecDeque;
27
28use serde::{Deserialize, Serialize};
29
30pub const MIN_SAMPLES: usize = 3;
32pub const MIN_SPAN_MS: u64 = 3_000;
34pub const WINDOW_MS: u64 = 30_000;
37const MAX_SAMPLES: usize = 256;
42
43#[derive(Debug, Clone, Copy, PartialEq, Eq)]
44struct Sample {
45 at_ms: u64,
46 bytes: u64,
47}
48
49#[derive(Debug, Clone, Copy, PartialEq, Serialize, Deserialize)]
52pub struct RateReport {
53 pub stable: bool,
55 #[serde(skip_serializing_if = "Option::is_none")]
56 pub bytes_per_second: Option<f64>,
57 #[serde(skip_serializing_if = "Option::is_none")]
60 pub eta_seconds: Option<f64>,
61 pub samples: usize,
64}
65
66impl RateReport {
67 fn warming(samples: usize) -> Self {
68 RateReport {
69 stable: false,
70 bytes_per_second: None,
71 eta_seconds: None,
72 samples,
73 }
74 }
75}
76
77#[derive(Debug, Default, Clone)]
79pub struct RateEstimator {
80 window: VecDeque<Sample>,
81}
82
83impl RateEstimator {
84 pub fn new() -> Self {
85 Self::default()
86 }
87
88 pub fn observe(&mut self, at_ms: u64, bytes: u64) {
98 if let Some(last) = self.window.back() {
99 if at_ms < last.at_ms {
100 return;
101 }
102 if bytes < last.bytes {
103 self.window.clear();
104 }
105 }
106 self.window.push_back(Sample { at_ms, bytes });
107
108 let cutoff = at_ms.saturating_sub(WINDOW_MS);
109 while self.window.len() > 1 && self.window.front().is_some_and(|s| s.at_ms < cutoff) {
110 self.window.pop_front();
111 }
112 while self.window.len() > MAX_SAMPLES {
113 self.window.remove(1);
114 }
115 }
116
117 pub fn reset(&mut self) {
121 self.window.clear();
122 }
123
124 pub fn bytes_done(&self) -> Option<u64> {
126 self.window.back().map(|s| s.bytes)
127 }
128
129 pub fn report(&self, total_bytes: Option<u64>) -> RateReport {
133 let (Some(first), Some(last)) = (self.window.front(), self.window.back()) else {
134 return RateReport::warming(0);
135 };
136 let span_ms = last.at_ms - first.at_ms;
137 if self.window.len() < MIN_SAMPLES || span_ms < MIN_SPAN_MS {
138 return RateReport::warming(self.window.len());
139 }
140
141 let bytes = last.bytes.saturating_sub(first.bytes) as f64;
142 let rate = bytes / (span_ms as f64 / 1000.0);
143 let eta = total_bytes.filter(|_| rate > 0.0).map(|total| {
144 total.saturating_sub(last.bytes) as f64 / rate
148 });
149 RateReport {
150 stable: true,
151 bytes_per_second: Some(rate),
152 eta_seconds: eta,
153 samples: self.window.len(),
154 }
155 }
156}
157
158#[cfg(test)]
159mod tests {
160 use super::*;
161
162 #[test]
163 fn first_tick_reports_no_rate_at_all() {
164 let mut est = RateEstimator::new();
165 est.observe(0, 0);
166 est.observe(2, 8 * 1024 * 1024); let report = est.report(Some(1 << 30));
168 assert!(!report.stable);
169 assert_eq!(report.bytes_per_second, None);
170 assert_eq!(report.eta_seconds, None);
171 }
172
173 #[test]
174 fn three_samples_are_not_enough_without_three_seconds() {
175 let mut est = RateEstimator::new();
176 for i in 0..5 {
177 est.observe(i * 100, i * 1_000_000);
178 }
179 assert!(!est.report(None).stable);
180 }
181
182 #[test]
183 fn three_seconds_are_not_enough_without_three_samples() {
184 let mut est = RateEstimator::new();
185 est.observe(0, 0);
186 est.observe(5_000, 5_000_000);
187 assert!(!est.report(None).stable);
188 }
189
190 #[test]
191 fn reports_a_stable_rate_and_eta_once_the_window_qualifies() {
192 let mut est = RateEstimator::new();
193 for i in 0..=4u64 {
195 est.observe(i * 1000, i * 1_000_000);
196 }
197 let report = est.report(Some(10_000_000));
198 assert!(report.stable);
199 assert_eq!(report.bytes_per_second, Some(1_000_000.0));
200 assert_eq!(report.eta_seconds, Some(6.0));
202 }
203
204 #[test]
205 fn a_restarted_transfer_clears_the_window_instead_of_going_negative() {
206 let mut est = RateEstimator::new();
207 for i in 0..=4u64 {
208 est.observe(i * 1000, i * 1_000_000);
209 }
210 assert!(est.report(None).stable);
211 est.observe(5_000, 0); let report = est.report(None);
213 assert!(!report.stable);
214 assert_eq!(report.samples, 1);
215 assert_eq!(est.bytes_done(), Some(0));
216 }
217
218 #[test]
219 fn eta_is_clamped_at_zero_when_the_total_is_wrong() {
220 let mut est = RateEstimator::new();
221 for i in 0..=4u64 {
222 est.observe(i * 1000, i * 1_000_000);
223 }
224 assert_eq!(est.report(Some(1_000_000)).eta_seconds, Some(0.0));
226 }
227
228 #[test]
229 fn no_total_means_a_rate_but_no_eta() {
230 let mut est = RateEstimator::new();
231 for i in 0..=4u64 {
232 est.observe(i * 1000, i * 1_000_000);
233 }
234 let report = est.report(None);
235 assert!(report.stable);
236 assert!(report.bytes_per_second.is_some());
237 assert_eq!(report.eta_seconds, None);
238 }
239
240 #[test]
241 fn a_stalled_transfer_reports_zero_rather_than_an_eta() {
242 let mut est = RateEstimator::new();
243 for i in 0..=4u64 {
244 est.observe(i * 1000, 1_000_000);
245 }
246 let report = est.report(Some(2_000_000));
247 assert_eq!(report.bytes_per_second, Some(0.0));
248 assert_eq!(report.eta_seconds, None);
250 }
251
252 #[test]
253 fn samples_older_than_the_window_are_dropped() {
254 let mut est = RateEstimator::new();
255 est.observe(0, 0);
256 for i in 0..=4u64 {
257 est.observe(WINDOW_MS + i * 1000, 1_000_000 + i * 1_000_000);
258 }
259 assert_eq!(est.report(None).bytes_per_second, Some(1_000_000.0));
261 }
262
263 #[test]
264 fn a_fast_ticking_job_still_becomes_stable() {
265 let mut est = RateEstimator::new();
268 for i in 0..=1000u64 {
269 est.observe(i * 10, i * 10_000);
270 }
271 let report = est.report(None);
272 assert!(report.stable, "{report:?}");
273 assert_eq!(report.bytes_per_second, Some(1_000_000.0));
274 }
275
276 #[test]
277 fn a_backwards_clock_sample_is_ignored() {
278 let mut est = RateEstimator::new();
279 for i in 0..=4u64 {
280 est.observe(i * 1000, i * 1_000_000);
281 }
282 est.observe(500, 9_000_000);
283 assert_eq!(est.bytes_done(), Some(4_000_000));
284 assert_eq!(est.report(None).bytes_per_second, Some(1_000_000.0));
285 }
286
287 #[test]
288 fn reset_returns_to_warming() {
289 let mut est = RateEstimator::new();
290 for i in 0..=4u64 {
291 est.observe(i * 1000, i * 1_000_000);
292 }
293 est.reset();
294 let report = est.report(None);
295 assert!(!report.stable);
296 assert_eq!(report.samples, 0);
297 }
298
299 #[test]
300 fn warming_reports_omit_the_absent_numbers_rather_than_nulling_them() {
301 let json = serde_json::to_string(&RateReport::warming(1)).unwrap();
302 assert_eq!(json, "{\"stable\":false,\"samples\":1}");
303 }
304}