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    fn refill(&mut self) {
25        let now = Instant::now();
26        self.tokens = (self.tokens + now.duration_since(self.updated).as_secs_f64() * self.rate)
27            .min(self.capacity);
28        self.updated = now;
29    }
30
31    /// Consume one available token without waiting.
32    pub fn take(&mut self) -> bool {
33        self.refill();
34        if self.tokens < 1.0 {
35            false
36        } else {
37            self.tokens -= 1.0;
38            true
39        }
40    }
41
42    /// Return a token when the transport did not accept the frame.
43    pub fn refund(&mut self) {
44        self.tokens = (self.tokens + 1.0).min(self.capacity);
45    }
46
47    /// Time until one token is available. A disabled bucket never refills.
48    pub fn wait(&mut self) -> Duration {
49        self.refill();
50        if self.tokens >= 1.0 {
51            Duration::ZERO
52        } else if self.rate == 0.0 || self.capacity < 1.0 {
53            Duration::MAX
54        } else {
55            Duration::from_secs_f64((1.0 - self.tokens) / self.rate)
56        }
57    }
58}
xmtp_common/rate_limit.rs

xmtp_common/
rate_limit.rs
