clawft_kernel/
mesh_swarm_rate.rs1use std::sync::Mutex;
12use std::time::Duration;
13
14use tokio::time::Instant;
15
16#[derive(Debug)]
18pub struct RateLimiter {
19 bytes_per_sec: u64,
20 next_free: Mutex<Option<Instant>>,
21}
22
23impl RateLimiter {
24 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 pub fn bytes_per_sec(&self) -> u64 {
34 self.bytes_per_sec
35 }
36
37 pub fn cost(&self, n: u64) -> Duration {
39 Duration::from_secs_f64(n as f64 / self.bytes_per_sec as f64)
40 }
41
42 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#[derive(Debug, Default)]
57pub struct Bandwidth {
58 pub up: Option<RateLimiter>,
60 pub down: Option<RateLimiter>,
62}
63
64impl Bandwidth {
65 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 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 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 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}