bitfold_protocol/
bandwidth.rs

1use std::time::{Duration, Instant};
2
3/// Bandwidth manager for controlling incoming/outgoing data rates.
4#[derive(Debug, Clone)]
5pub struct BandwidthManager {
6    /// Maximum bytes per second for outgoing traffic (0 = unlimited)
7    outgoing_bandwidth: u32,
8    /// Maximum bytes per second for incoming traffic (0 = unlimited)
9    incoming_bandwidth: u32,
10    /// Bytes sent in current window
11    bytes_sent_this_window: u32,
12    /// Bytes received in current window
13    bytes_recv_this_window: u32,
14    /// Start of current measurement window
15    window_start: Instant,
16    /// Window duration for bandwidth measurement
17    window_duration: Duration,
18}
19
20impl BandwidthManager {
21    /// Creates a new bandwidth manager with specified limits (bytes/sec).
22    pub fn new(outgoing_bw: u32, incoming_bw: u32, window_duration: Duration) -> Self {
23        Self {
24            outgoing_bandwidth: outgoing_bw,
25            incoming_bandwidth: incoming_bw,
26            bytes_sent_this_window: 0,
27            bytes_recv_this_window: 0,
28            window_start: Instant::now(),
29            window_duration,
30        }
31    }
32
33    /// Creates an unlimited bandwidth manager (no throttling).
34    pub fn unlimited() -> Self {
35        Self::new(0, 0, Duration::from_secs(1))
36    }
37
38    /// Sets the outgoing bandwidth limit in bytes per second.
39    pub fn set_outgoing_bandwidth(&mut self, bytes_per_sec: u32) {
40        self.outgoing_bandwidth = bytes_per_sec;
41    }
42
43    /// Sets the incoming bandwidth limit in bytes per second.
44    pub fn set_incoming_bandwidth(&mut self, bytes_per_sec: u32) {
45        self.incoming_bandwidth = bytes_per_sec;
46    }
47
48    /// Updates the bandwidth window if needed.
49    pub fn update_window(&mut self, now: Instant) {
50        if now.duration_since(self.window_start) >= self.window_duration {
51            self.bytes_sent_this_window = 0;
52            self.bytes_recv_this_window = 0;
53            self.window_start = now;
54        }
55    }
56
57    /// Checks if sending the given number of bytes is allowed.
58    /// Returns true if allowed, false if would exceed bandwidth limit.
59    pub fn can_send_outgoing(&self, byte_count: usize) -> bool {
60        self.update_window_if_needed();
61
62        if self.outgoing_bandwidth == 0 {
63            return true; // Unlimited
64        }
65
66        self.bytes_sent_this_window + byte_count as u32 <= self.outgoing_bandwidth
67    }
68
69    /// Records that the given number of bytes were sent.
70    pub fn record_sent(&mut self, byte_count: usize) {
71        self.bytes_sent_this_window += byte_count as u32;
72    }
73
74    /// Checks if receiving the given number of bytes is allowed.
75    pub fn can_receive_incoming(&self, byte_count: usize) -> bool {
76        self.update_window_if_needed();
77
78        if self.incoming_bandwidth == 0 {
79            return true; // Unlimited
80        }
81
82        self.bytes_recv_this_window + byte_count as u32 <= self.incoming_bandwidth
83    }
84
85    /// Records that the given number of bytes were received.
86    pub fn record_received(&mut self, byte_count: usize) {
87        self.bytes_recv_this_window += byte_count as u32;
88    }
89
90    /// Returns the current outgoing bandwidth utilization (0.0 to 1.0+).
91    pub fn outgoing_utilization(&self) -> f32 {
92        if self.outgoing_bandwidth == 0 {
93            return 0.0; // Unlimited
94        }
95        self.bytes_sent_this_window as f32 / self.outgoing_bandwidth as f32
96    }
97
98    /// Returns the current incoming bandwidth utilization (0.0 to 1.0+).
99    pub fn incoming_utilization(&self) -> f32 {
100        if self.incoming_bandwidth == 0 {
101            return 0.0; // Unlimited
102        }
103        self.bytes_recv_this_window as f32 / self.incoming_bandwidth as f32
104    }
105
106    fn update_window_if_needed(&self) {
107        // This is a const method check - actual update happens via update_window()
108    }
109}
110
111#[cfg(test)]
112mod tests {
113    use std::thread::sleep;
114
115    use super::*;
116
117    #[test]
118    fn test_unlimited_bandwidth() {
119        let bw = BandwidthManager::unlimited();
120
121        // Should allow any amount
122        assert!(bw.can_send_outgoing(1_000_000));
123        assert!(bw.can_receive_incoming(1_000_000));
124        assert_eq!(bw.outgoing_utilization(), 0.0);
125        assert_eq!(bw.incoming_utilization(), 0.0);
126    }
127
128    #[test]
129    fn test_utilization() {
130        let mut bw = BandwidthManager::new(1000, 2000, Duration::from_secs(1));
131
132        bw.record_sent(500);
133        assert_eq!(bw.outgoing_utilization(), 0.5);
134
135        bw.record_received(1000);
136        assert_eq!(bw.incoming_utilization(), 0.5);
137    }
138
139    #[test]
140    fn test_outgoing_bandwidth_limit() {
141        let mut bw = BandwidthManager::new(1000, 0, Duration::from_secs(1));
142
143        // Should allow within limit
144        assert!(bw.can_send_outgoing(500));
145        bw.record_sent(500);
146
147        assert!(bw.can_send_outgoing(500));
148        bw.record_sent(500);
149
150        // Should deny exceeding limit
151        assert!(!bw.can_send_outgoing(100));
152    }
153
154    #[test]
155    fn test_bandwidth_window_reset() {
156        let mut bw = BandwidthManager::new(1000, 0, Duration::from_millis(10));
157
158        bw.record_sent(1000);
159        assert!(!bw.can_send_outgoing(100));
160
161        // Wait for window to reset
162        sleep(Duration::from_millis(15));
163        bw.update_window(Instant::now());
164
165        // Should allow again after window reset
166        assert!(bw.can_send_outgoing(500));
167    }
168}
bitfold_protocol/bandwidth.rs

bitfold_protocol/
bandwidth.rs
