axum_limit/policy/
sliding_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/// Sliding window counter rate limiting policy.
9///
10/// Estimates request volume in a rolling window by weighting the previous window,
11/// avoiding the sharp reset spikes of fixed windows while remaining memory efficient.
12#[derive(Debug, Clone, Copy, Default)]
13pub struct SlidingWindowPolicy;
14
15impl RateLimitPolicy for SlidingWindowPolicy {
16    type State = SlidingWindowState;
17    const STATE_ID: &'static str = "sliding_window";
18}
19
20/// Sliding window counter state for a single key.
21#[derive(Debug, Clone, Serialize, Deserialize)]
22pub struct SlidingWindowState {
23    current_count: usize,
24    previous_count: usize,
25    window_start_ms: u64,
26    window_size_ms: u64,
27    max: usize,
28}
29
30impl SlidingWindowState {
31    pub(crate) fn new_at(quota: Quota, now_ms: u64) -> Self {
32        Self {
33            current_count: 0,
34            previous_count: 0,
35            window_start_ms: now_ms,
36            window_size_ms: quota.per_ms,
37            max: quota.max.max(1),
38        }
39    }
40
41    fn roll_window(&mut self, now_ms: u64) {
42        while saturating_sub_ms(now_ms, self.window_start_ms) >= self.window_size_ms {
43            self.previous_count = self.current_count;
44            self.current_count = 0;
45            self.window_start_ms = self.window_start_ms.saturating_add(self.window_size_ms);
46        }
47    }
48
49    fn estimated_count(&self, now_ms: u64) -> f64 {
50        let elapsed_ms = saturating_sub_ms(now_ms, self.window_start_ms);
51        let window_ms = self.window_size_ms.max(1) as f64;
52        let weight = 1.0 - (elapsed_ms as f64 / window_ms);
53        self.previous_count as f64 * weight + self.current_count as f64
54    }
55
56    fn reset_at_ms(&self) -> u64 {
57        self.window_start_ms.saturating_add(self.window_size_ms)
58    }
59}
60
61impl PolicyState for SlidingWindowState {
62    fn try_acquire(&mut self, now_ms: u64) -> RateLimitSnapshot {
63        self.roll_window(now_ms);
64        let limit = self.max;
65        let estimated = self.estimated_count(now_ms);
66
67        if estimated >= limit as f64 {
68            return RateLimitSnapshot {
69                allowed: false,
70                limit,
71                remaining: 0,
72                reset_at_ms: self.reset_at_ms(),
73            };
74        }
75
76        self.current_count += 1;
77        let remaining = limit.saturating_sub(self.estimated_count(now_ms).ceil() as usize);
78
79        RateLimitSnapshot {
80            allowed: true,
81            limit,
82            remaining,
83            reset_at_ms: self.reset_at_ms(),
84        }
85    }
86
87    fn encode(&self) -> Result<Vec<u8>, CodecError> {
88        encode_json(self)
89    }
90
91    fn decode(bytes: &[u8], _quota: Quota) -> Result<Self, CodecError> {
92        decode_json(bytes)
93    }
94
95    fn create(quota: Quota, now_ms: u64) -> Self {
96        Self::new_at(quota, now_ms)
97    }
98}
99
100#[cfg(test)]
101mod tests {
102    use super::*;
103
104    const BASE_MS: u64 = 3_000_000;
105
106    fn at(offset_ms: u64) -> u64 {
107        BASE_MS + offset_ms
108    }
109
110    fn window(quota: Quota) -> SlidingWindowState {
111        SlidingWindowState::new_at(quota, BASE_MS)
112    }
113
114    #[test]
115    fn allows_up_to_max_requests() {
116        let mut state = window(Quota::new(3, 1000));
117
118        for _ in 0..3 {
119            assert!(state.try_acquire(at(0)).allowed);
120        }
121        assert!(!state.try_acquire(at(0)).allowed);
122    }
123
124    #[test]
125    fn smooths_boundary_between_windows() {
126        let mut state = window(Quota::new(4, 1000));
127
128        for _ in 0..4 {
129            assert!(state.try_acquire(at(0)).allowed);
130        }
131        assert!(!state.try_acquire(at(100)).allowed);
132        assert!(!state.try_acquire(at(1000)).allowed);
133
134        for _ in 0..4 {
135            assert!(state.try_acquire(at(2000)).allowed);
136        }
137    }
138
139    #[test]
140    fn round_trips_through_json() {
141        let mut state = window(Quota::new(3, 1000));
142        assert!(state.try_acquire(at(0)).allowed);
143        assert!(state.try_acquire(at(100)).allowed);
144
145        let encoded = state.encode().expect("encode");
146        let mut decoded =
147            SlidingWindowState::decode(&encoded, Quota::new(3, 1000)).expect("decode");
148
149        assert!(decoded.try_acquire(at(200)).allowed);
150        assert!(!decoded.try_acquire(at(300)).allowed);
151    }
152}
axum_limit/policy/sliding_window.rs

axum_limit/policy/
sliding_window.rs
