saluki_common/
rate.rs

1//! Rate limiting primitives.
2
3use std::time::Instant;
4
5/// Token bucket rate limiter.
6///
7/// Provides a `rate` tokens-per-second refill up to `capacity`, and allows consuming one token at a
8/// time via [`allow`][TokenBucket::allow]. Mirrors `golang.org/x/time/rate.Limiter`.
9pub struct TokenBucket {
10    capacity: f64,
11    tokens: f64,
12    last_refill: Instant,
13    rate: f64,
14}
15
16impl TokenBucket {
17    /// Creates a new `TokenBucket` with the given rate (tokens per second) and burst capacity.
18    ///
19    /// The bucket starts full.
20    pub fn new(rate: f64, burst: usize) -> Self {
21        Self {
22            capacity: burst as f64,
23            tokens: burst as f64,
24            last_refill: Instant::now(),
25            rate,
26        }
27    }
28
29    /// Attempt to consume one token. Returns `true` if a token was available.
30    pub fn allow(&mut self) -> bool {
31        let now = Instant::now();
32        saluki_antithesis::always_or_unreachable!(
33            now >= self.last_refill,
34            "token-bucket refill clock did not move backward"
35        );
36        let elapsed = now.duration_since(self.last_refill).as_secs_f64();
37        self.tokens = (self.tokens + elapsed * self.rate).min(self.capacity);
38        self.last_refill = now.max(self.last_refill);
39        if self.tokens >= 1.0 {
40            self.tokens -= 1.0;
41            true
42        } else {
43            false
44        }
45    }
46}
47
48#[cfg(test)]
49mod tests {
50    use std::time::Duration;
51
52    use super::TokenBucket;
53
54    #[test]
55    fn full_bucket_allows_up_to_burst() {
56        let burst = 5;
57        let mut bucket = TokenBucket::new(1.0, burst);
58        for _ in 0..burst {
59            assert!(bucket.allow());
60        }
61        assert!(!bucket.allow());
62    }
63
64    #[test]
65    fn empty_bucket_refills_over_time() {
66        let mut bucket = TokenBucket::new(100.0, 1);
67        assert!(bucket.allow()); // consume the initial token
68        assert!(!bucket.allow()); // empty
69
70        std::thread::sleep(Duration::from_millis(20)); // ~2 tokens at 100 TPS
71        assert!(bucket.allow());
72    }
73
74    #[test]
75    fn refill_does_not_exceed_capacity() {
76        let burst = 3;
77        let mut bucket = TokenBucket::new(1000.0, burst);
78        assert!(bucket.allow());
79        assert!(bucket.allow());
80        assert!(bucket.allow());
81        assert!(!bucket.allow());
82
83        std::thread::sleep(Duration::from_millis(50)); // would add 50 tokens at 1000 TPS, capped at burst
84        for _ in 0..burst {
85            assert!(bucket.allow());
86        }
87        assert!(!bucket.allow());
88    }
89
90    #[test]
91    fn zero_rate_never_refills() {
92        let mut bucket = TokenBucket::new(0.0, 1);
93        assert!(bucket.allow()); // initial token
94        assert!(!bucket.allow());
95        std::thread::sleep(Duration::from_millis(20));
96        assert!(!bucket.allow()); // still empty, no refill
97    }
98}
saluki_common/rate.rs

saluki_common/
rate.rs
