Skip to main content

clawft_kernel/
mesh_swarm_rate.rs

1//! Bandwidth caps for the artifact swarm (ADR-099 section 6, card
2//! mesh-placement-25).
3//!
4//! [`RateLimiter`] paces byte transfers to a fixed rate. It is a pacing
5//! reservation, not a bursty bucket: each [`RateLimiter::acquire`] books the
6//! next slice of time after the previous booking and waits until that
7//! slice ends, so `n` bytes never complete faster than `n / rate`. It is
8//! driven by `tokio::time`, so tests can run under paused time and measure
9//! throughput deterministically.
10
11use std::sync::Mutex;
12use std::time::Duration;
13
14use tokio::time::Instant;
15
16/// Paces byte transfers to `bytes_per_sec`.
17#[derive(Debug)]
18pub struct RateLimiter {
19    bytes_per_sec: u64,
20    next_free: Mutex<Option<Instant>>,
21}
22
23impl RateLimiter {
24    /// Limiter at `bytes_per_sec` (clamped to at least 1).
25    pub fn new(bytes_per_sec: u64) -> Self {
26        Self {
27            bytes_per_sec: bytes_per_sec.max(1),
28            next_free: Mutex::new(None),
29        }
30    }
31
32    /// The configured rate.
33    pub fn bytes_per_sec(&self) -> u64 {
34        self.bytes_per_sec
35    }
36
37    /// Time `n` bytes take at the configured rate.
38    pub fn cost(&self, n: u64) -> Duration {
39        Duration::from_secs_f64(n as f64 / self.bytes_per_sec as f64)
40    }
41
42    /// Reserve `n` bytes and wait until their slice has elapsed.
43    pub async fn acquire(&self, n: u64) {
44        let until = {
45            let mut next = self.next_free.lock().unwrap_or_else(|p| p.into_inner());
46            let start = next.map_or_else(Instant::now, |t| t.max(Instant::now()));
47            let end = start + self.cost(n);
48            *next = Some(end);
49            end
50        };
51        tokio::time::sleep_until(until).await;
52    }
53}
54
55/// A node's upload and download caps (either may be unlimited).
56#[derive(Debug, Default)]
57pub struct Bandwidth {
58    /// Cap on bytes sent to peers.
59    pub up: Option<RateLimiter>,
60    /// Cap on bytes received from peers.
61    pub down: Option<RateLimiter>,
62}
63
64impl Bandwidth {
65    /// Caps from bytes-per-second settings (`None` = unlimited).
66    pub fn new(up: Option<u64>, down: Option<u64>) -> Self {
67        Self {
68            up: up.map(RateLimiter::new),
69            down: down.map(RateLimiter::new),
70        }
71    }
72
73    /// Wait for `n` bytes of upload budget.
74    pub async fn upload(&self, n: usize) {
75        if let Some(l) = &self.up {
76            l.acquire(n as u64).await;
77        }
78    }
79
80    /// Wait for `n` bytes of download budget.
81    pub async fn download(&self, n: usize) {
82        if let Some(l) = &self.down {
83            l.acquire(n as u64).await;
84        }
85    }
86}
87
88#[cfg(test)]
89mod tests {
90    use super::*;
91
92    #[tokio::test(start_paused = true)]
93    async fn transfers_never_beat_the_rate_and_concurrent_callers_share_it() {
94        let l = RateLimiter::new(1_000);
95        let t0 = Instant::now();
96        l.acquire(2_000).await;
97        assert_eq!(t0.elapsed(), Duration::from_secs(2));
98        // Two callers at once still share one 1000 B/s budget.
99        let t1 = Instant::now();
100        tokio::join!(l.acquire(500), l.acquire(500), l.acquire(1_000));
101        assert_eq!(t1.elapsed(), Duration::from_secs(2));
102    }
103
104    #[tokio::test(start_paused = true)]
105    async fn unlimited_does_not_wait_and_idle_time_is_not_banked() {
106        let b = Bandwidth::new(None, Some(100));
107        let t0 = Instant::now();
108        b.upload(1 << 30).await;
109        assert_eq!(t0.elapsed(), Duration::ZERO);
110        tokio::time::sleep(Duration::from_secs(60)).await;
111        let t1 = Instant::now();
112        b.download(100).await;
113        assert_eq!(t1.elapsed(), Duration::from_secs(1), "no burst credit from idling");
114    }
115}