axum_limit/policy/
fixed_window.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/// Fixed window rate limiting policy.
9///
10/// Allows up to `quota.max` requests per window of `quota.per`. The counter resets
11/// when a new window starts.
12#[derive(Debug, Clone, Copy, Default)]
13pub struct FixedWindowPolicy;
14
15impl RateLimitPolicy for FixedWindowPolicy {
16    type State = FixedWindowState;
17    const STATE_ID: &'static str = "fixed_window";
18}
19
20/// Fixed window state for a single key.
21#[derive(Debug, Clone, Serialize, Deserialize)]
22pub struct FixedWindowState {
23    count: usize,
24    window_start_ms: u64,
25    window_size_ms: u64,
26    max: usize,
27}
28
29impl FixedWindowState {
30    pub(crate) fn new_at(quota: Quota, now_ms: u64) -> Self {
31        Self {
32            count: 0,
33            window_start_ms: now_ms,
34            window_size_ms: quota.per_ms,
35            max: quota.max.max(1),
36        }
37    }
38
39    fn maybe_reset_window(&mut self, now_ms: u64) {
40        if saturating_sub_ms(now_ms, self.window_start_ms) >= self.window_size_ms {
41            self.count = 0;
42            self.window_start_ms = now_ms;
43        }
44    }
45
46    fn reset_at_ms(&self) -> u64 {
47        self.window_start_ms.saturating_add(self.window_size_ms)
48    }
49}
50
51impl PolicyState for FixedWindowState {
52    fn try_acquire(&mut self, now_ms: u64) -> RateLimitSnapshot {
53        self.maybe_reset_window(now_ms);
54        let limit = self.max;
55
56        if self.count < limit {
57            self.count += 1;
58            RateLimitSnapshot {
59                allowed: true,
60                limit,
61                remaining: limit.saturating_sub(self.count),
62                reset_at_ms: self.reset_at_ms(),
63            }
64        } else {
65            RateLimitSnapshot {
66                allowed: false,
67                limit,
68                remaining: 0,
69                reset_at_ms: self.reset_at_ms(),
70            }
71        }
72    }
73
74    fn encode(&self) -> Result<Vec<u8>, CodecError> {
75        encode_json(self)
76    }
77
78    fn decode(bytes: &[u8], _quota: Quota) -> Result<Self, CodecError> {
79        decode_json(bytes)
80    }
81
82    fn create(quota: Quota, now_ms: u64) -> Self {
83        Self::new_at(quota, now_ms)
84    }
85}
86
87#[cfg(test)]
88mod tests {
89    use super::*;
90
91    const BASE_MS: u64 = 2_000_000;
92
93    fn at(offset_ms: u64) -> u64 {
94        BASE_MS + offset_ms
95    }
96
97    fn window(quota: Quota) -> FixedWindowState {
98        FixedWindowState::new_at(quota, BASE_MS)
99    }
100
101    #[test]
102    fn allows_up_to_max_requests_in_window() {
103        let mut window = window(Quota::new(3, 1000));
104
105        assert!(window.try_acquire(at(0)).allowed);
106        assert!(window.try_acquire(at(100)).allowed);
107        assert!(window.try_acquire(at(500)).allowed);
108        assert!(!window.try_acquire(at(900)).allowed);
109    }
110
111    #[test]
112    fn resets_after_window_expires() {
113        let mut window = window(Quota::new(2, 1000));
114
115        assert!(window.try_acquire(at(0)).allowed);
116        assert!(window.try_acquire(at(100)).allowed);
117        assert!(!window.try_acquire(at(200)).allowed);
118
119        assert!(window.try_acquire(at(1000)).allowed);
120        assert!(window.try_acquire(at(1100)).allowed);
121        assert!(!window.try_acquire(at(1200)).allowed);
122    }
123
124    #[test]
125    fn round_trips_through_json() {
126        let mut state = window(Quota::new(3, 1000));
127        assert!(state.try_acquire(at(0)).allowed);
128        assert!(state.try_acquire(at(100)).allowed);
129
130        let encoded = state.encode().expect("encode");
131        let mut decoded = FixedWindowState::decode(&encoded, Quota::new(3, 1000)).expect("decode");
132
133        assert!(decoded.try_acquire(at(200)).allowed);
134        assert!(!decoded.try_acquire(at(300)).allowed);
135    }
136}
axum_limit/policy/fixed_window.rs

axum_limit/policy/
fixed_window.rs
