Skip to main content

netspeed_cli/
bandwidth_loop.rs

1//! Shared bandwidth measurement loop for download/upload tests.
2//!
3//! Eliminates duplication between `download.rs` and `upload.rs` by providing
4//! a unified state for:
5//! - Throttled speed sampling (20 Hz max)
6//! - Peak speed tracking
7//! - Progress bar updates
8//! - Atomic byte counting
9//!
10//! Each I/O operation (download chunk, upload round) calls `record_bytes()`
11//! to update shared state. Call `finish()` at the end to compute final results.
12
13use crate::common;
14use crate::progress::SpeedProgress;
15use std::sync::Arc;
16use std::sync::Mutex;
17use std::sync::atomic::{AtomicU64, Ordering};
18use std::time::Instant;
19
20/// Throttle interval for speed sampling (20 Hz max).
21const SAMPLE_INTERVAL_MS: u64 = 50;
22
23/// Shared state for a bandwidth test (download or upload).
24///
25/// All fields are thread-safe for use across multiple concurrent streams.
26pub struct BandwidthLoopState {
27    pub total_bytes: Arc<AtomicU64>,
28    pub peak_bps: Arc<AtomicU64>,
29    pub speed_samples: Arc<Mutex<Vec<f64>>>,
30    pub start: Instant,
31    pub last_sample_ms: Arc<AtomicU64>,
32    pub estimated_total: u64,
33    pub progress: Arc<SpeedProgress>,
34}
35
36/// Final result from a bandwidth test.
37pub struct BandwidthResult {
38    pub avg_bps: f64,
39    pub peak_bps: f64,
40    pub total_bytes: u64,
41    pub duration_secs: f64,
42    pub speed_samples: Vec<f64>,
43}
44
45impl BandwidthLoopState {
46    /// Create a new bandwidth measurement state.
47    #[must_use]
48    pub fn new(estimated_total: u64, progress: Arc<SpeedProgress>) -> Self {
49        Self {
50            total_bytes: Arc::new(AtomicU64::new(0)),
51            peak_bps: Arc::new(AtomicU64::new(0)),
52            speed_samples: Arc::new(Mutex::new(Vec::new())),
53            start: Instant::now(),
54            last_sample_ms: Arc::new(AtomicU64::new(0)),
55            estimated_total,
56            progress,
57        }
58    }
59
60    /// Record transferred bytes and update progress (throttled to 20 Hz).
61    ///
62    /// This is the single point where all expensive operations (bandwidth calc,
63    /// peak tracking, sample recording, progress update) are throttled.
64    pub fn record_bytes(&self, len: u64) {
65        self.total_bytes.fetch_add(len, Ordering::Relaxed);
66
67        let elapsed_ms = self.start.elapsed().as_millis() as u64;
68        let last_ms = self.last_sample_ms.load(Ordering::Relaxed);
69        let should_sample =
70            last_ms == 0 || elapsed_ms.saturating_sub(last_ms) >= SAMPLE_INTERVAL_MS;
71
72        if should_sample {
73            self.last_sample_ms.store(elapsed_ms, Ordering::Relaxed);
74            self.sample_now();
75        }
76    }
77
78    /// Take a speed sample and update progress (no throttle check — caller must gate).
79    fn sample_now(&self) {
80        let total = self.total_bytes.load(Ordering::Acquire);
81        let elapsed = self.start.elapsed().as_secs_f64();
82        let speed = common::calculate_bandwidth(total, elapsed);
83
84        let current_peak = self.peak_bps.load(Ordering::Relaxed);
85        if speed > current_peak as f64 {
86            self.peak_bps.store(speed as u64, Ordering::Relaxed);
87        }
88
89        if let Ok(mut samples) = self.speed_samples.lock() {
90            samples.push(speed);
91        }
92
93        let pct = (total as f64 / self.estimated_total as f64).min(1.0);
94        self.progress.update(speed / 1_000_000.0, pct, total);
95    }
96
97    /// Compute final results from accumulated state.
98    #[must_use]
99    pub fn finish(&self) -> BandwidthResult {
100        let total = self.total_bytes.load(Ordering::Relaxed);
101        let peak = self.peak_bps.load(Ordering::Relaxed) as f64;
102        let duration = self.start.elapsed().as_secs_f64();
103        let samples = self.speed_samples.lock().unwrap().to_vec();
104        let avg = common::calculate_bandwidth(total, duration);
105
106        BandwidthResult {
107            avg_bps: avg,
108            peak_bps: peak,
109            total_bytes: total,
110            duration_secs: duration,
111            speed_samples: samples,
112        }
113    }
114}