axum_limit/policy/
token_bucket.rs

1use super::{PolicyState, RateLimitPolicy};
2use crate::codec::{decode_json, encode_json, CodecError};
3use crate::quota::Quota;
4use crate::snapshot::RateLimitSnapshot;
5use crate::time::saturating_sub_ms;
6use serde::{Deserialize, Serialize};
7
8/// Token bucket rate limiting policy.
9///
10/// Allows bursts up to the configured capacity while replenishing tokens smoothly
11/// at `quota.max` requests per `quota.per`.
12#[derive(Debug, Clone, Copy, Default)]
13pub struct TokenBucketPolicy;
14
15impl RateLimitPolicy for TokenBucketPolicy {
16    type State = TokenBucketState;
17    const STATE_ID: &'static str = "token_bucket";
18}
19
20/// Token bucket state for a single key.
21#[derive(Debug, Clone, Serialize, Deserialize)]
22pub struct TokenBucketState {
23    tokens: usize,
24    max_tokens: usize,
25    last_refill_time_ms: u64,
26    ns_per_token: u64,
27}
28
29impl TokenBucketState {
30    pub(crate) fn new_at(quota: Quota, now_ms: u64) -> Self {
31        let max_tokens = quota.burst().max(1);
32        let rate = quota.max.max(1);
33        let ns_per_token = (quota.per_ms.saturating_mul(1_000_000) / rate as u64).max(1);
34
35        Self {
36            tokens: max_tokens,
37            max_tokens,
38            last_refill_time_ms: now_ms,
39            ns_per_token,
40        }
41    }
42
43    fn refill(&mut self, now_ms: u64) {
44        let elapsed_ms = saturating_sub_ms(now_ms, self.last_refill_time_ms);
45        let elapsed_ns = elapsed_ms.saturating_mul(1_000_000);
46
47        if elapsed_ns >= self.ns_per_token {
48            let new_tokens = (elapsed_ns / self.ns_per_token) as usize;
49            self.tokens = (self.tokens + new_tokens).min(self.max_tokens);
50            let remainder_ns = elapsed_ns % self.ns_per_token;
51            self.last_refill_time_ms = now_ms.saturating_sub(remainder_ns / 1_000_000);
52        }
53    }
54
55    fn next_token_at_ms(&self, now_ms: u64) -> u64 {
56        if self.tokens > 0 {
57            now_ms
58        } else {
59            now_ms
60                .saturating_add(self.ns_per_token / 1_000_000)
61                .max(now_ms + 1)
62        }
63    }
64}
65
66impl PolicyState for TokenBucketState {
67    fn try_acquire(&mut self, now_ms: u64) -> RateLimitSnapshot {
68        self.refill(now_ms);
69        let limit = self.max_tokens;
70
71        if self.tokens > 0 {
72            self.tokens -= 1;
73            RateLimitSnapshot {
74                allowed: true,
75                limit,
76                remaining: self.tokens,
77                reset_at_ms: self.next_token_at_ms(now_ms),
78            }
79        } else {
80            RateLimitSnapshot {
81                allowed: false,
82                limit,
83                remaining: 0,
84                reset_at_ms: self.next_token_at_ms(now_ms),
85            }
86        }
87    }
88
89    fn encode(&self) -> Result<Vec<u8>, CodecError> {
90        encode_json(self)
91    }
92
93    fn decode(bytes: &[u8], _quota: Quota) -> Result<Self, CodecError> {
94        decode_json(bytes)
95    }
96
97    fn create(quota: Quota, now_ms: u64) -> Self {
98        Self::new_at(quota, now_ms)
99    }
100}
101
102#[cfg(test)]
103mod tests {
104    use super::*;
105
106    const BASE_MS: u64 = 1_000_000;
107
108    fn at(offset_ms: u64) -> u64 {
109        BASE_MS + offset_ms
110    }
111
112    fn bucket(quota: Quota) -> TokenBucketState {
113        TokenBucketState::new_at(quota, BASE_MS)
114    }
115
116    #[test]
117    fn allows_burst_up_to_capacity() {
118        let mut bucket = bucket(Quota::with_burst(5, 1000, 10));
119
120        for _ in 0..10 {
121            assert!(bucket.try_acquire(at(0)).allowed);
122        }
123        assert!(!bucket.try_acquire(at(0)).allowed);
124    }
125
126    #[test]
127    fn refills_at_sustained_rate() {
128        let mut bucket = bucket(Quota::new(5, 1000));
129
130        for _ in 0..5 {
131            assert!(bucket.try_acquire(at(0)).allowed);
132        }
133        assert!(!bucket.try_acquire(at(0)).allowed);
134
135        assert!(!bucket.try_acquire(at(199)).allowed);
136        assert!(bucket.try_acquire(at(200)).allowed);
137        assert!(!bucket.try_acquire(at(250)).allowed);
138    }
139
140    #[test]
141    fn refills_full_quota_after_period() {
142        let mut bucket = bucket(Quota::new(5, 1000));
143
144        for _ in 0..5 {
145            assert!(bucket.try_acquire(at(0)).allowed);
146        }
147        assert!(!bucket.try_acquire(at(0)).allowed);
148
149        for _ in 0..5 {
150            assert!(bucket.try_acquire(at(1000)).allowed);
151        }
152        assert!(!bucket.try_acquire(at(1000)).allowed);
153    }
154
155    #[test]
156    fn round_trips_through_json() {
157        let state = bucket(Quota::new(5, 1000));
158        let encoded = state.encode().expect("encode");
159        let decoded = TokenBucketState::decode(&encoded, Quota::new(5, 1000)).expect("decode");
160        assert_eq!(decoded.tokens, state.tokens);
161    }
162}
axum_limit/policy/token_bucket.rs

axum_limit/policy/
token_bucket.rs
