turbulence/
bandwidth_limiter.rs

1use std::time::Duration;
2
3use crate::Runtime;
4
5pub struct BandwidthLimiter<R: Runtime> {
6    runtime: R,
7    bandwidth: u32,
8    burst_bandwidth: u32,
9    bytes_available: f64,
10    last_calculation: R::Instant,
11}
12
13impl<R: Runtime> BandwidthLimiter<R> {
14    /// The `burst_bandwidth` is the maximum amount of bandwidth credit that can accumulate.
15    pub fn new(runtime: R, bandwidth: u32, burst_bandwidth: u32) -> BandwidthLimiter<R> {
16        let last_calculation = runtime.now();
17        BandwidthLimiter {
18            runtime,
19            bandwidth,
20            burst_bandwidth,
21            bytes_available: burst_bandwidth as f64,
22            last_calculation,
23        }
24    }
25
26    /// Delay until a time where there will be bandwidth available.
27    pub async fn delay_until_available(&self) {
28        if self.bytes_available < 0. {
29            self.runtime
30                .sleep(Duration::from_secs_f64(
31                    (-self.bytes_available) / self.bandwidth as f64,
32                ))
33                .await;
34        }
35    }
36
37    /// Actually update the amount of available bandwidth.  Additional available bytes are not added
38    /// until this method is called to add them.
39    pub fn update_available(&mut self) {
40        let now = self.runtime.now();
41        self.bytes_available += self
42            .runtime
43            .duration_between(self.last_calculation, now)
44            .as_secs_f64()
45            * self.bandwidth as f64;
46        self.bytes_available = self.bytes_available.min(self.burst_bandwidth as f64);
47        self.last_calculation = now;
48    }
49
50    /// The bandwidth limiter only needs to limit outgoing packets being sent at all, not their
51    /// size, so this returns true if a non-negative amount of bytes is available.  If a packet is
52    /// sent that is larger than the available bytes, the available bytes will go negative and this
53    /// will no longer return true.
54    pub fn bytes_available(&self) -> bool {
55        self.bytes_available >= 0.
56    }
57
58    /// Record that bytes were sent, possibly going into bandwidth debt.
59    pub fn take_bytes(&mut self, bytes: u32) {
60        self.bytes_available -= bytes as f64
61    }
62}