xmtp_common/
rate_limit.rs

1//! A token bucket shared by stream admission and client update scheduling.
2
3use crate::time::{Duration, Instant};
4
5/// A burst allowance that refills at a fixed number of tokens per second.
6pub struct Bucket {
7    rate: f64,
8    capacity: f64,
9    tokens: f64,
10    updated: Instant,
11}
12
13impl Bucket {
14    /// Start with the full burst allowance. A zero rate disables refill.
15    pub fn new(rate: u32, burst: u32) -> Self {
16        Self {
17            rate: rate as f64,
18            capacity: burst as f64,
19            tokens: burst as f64,
20            updated: Instant::now(),
21        }
22    }
23
24    /// Add a delay margin to each token's refill interval. A zero rate stays disabled.
25    pub fn with_margin(rate: u32, burst: u32, margin: Duration) -> Self {
26        let mut bucket = Self::new(rate, burst);
27        if rate > 0 {
28            bucket.rate = 1.0 / (1.0 / f64::from(rate) + margin.as_secs_f64());
29        }
30        bucket
31    }
32
33    fn refill(&mut self) {
34        let now = Instant::now();
35        self.tokens = (self.tokens + now.duration_since(self.updated).as_secs_f64() * self.rate)
36            .min(self.capacity);
37        self.updated = now;
38    }
39
40    /// Consume one available token without waiting.
41    pub fn take(&mut self) -> bool {
42        self.refill();
43        if self.tokens < 1.0 {
44            false
45        } else {
46            self.tokens -= 1.0;
47            true
48        }
49    }
50
51    /// Return a token when the transport did not accept the frame.
52    pub fn refund(&mut self) {
53        self.tokens = (self.tokens + 1.0).min(self.capacity);
54    }
55
56    /// Time until one token is available. A disabled bucket never refills.
57    pub fn wait(&mut self) -> Duration {
58        self.refill();
59        if self.tokens >= 1.0 {
60            Duration::ZERO
61        } else if self.rate == 0.0 || self.capacity < 1.0 {
62            Duration::MAX
63        } else {
64            Duration::from_secs_f64((1.0 - self.tokens) / self.rate)
65        }
66    }
67}
xmtp_common/rate_limit.rs

xmtp_common/
rate_limit.rs
