skp_ratelimit/algorithm/
fixed_window.rs1use std::time::Duration;
4
5use crate::algorithm::{current_timestamp_ms, timestamp_to_instant, Algorithm};
6use crate::decision::{Decision, RateLimitInfo};
7use crate::error::Result;
8use crate::quota::Quota;
9use crate::storage::Storage;
10
11#[derive(Debug, Clone, Default)]
16pub struct FixedWindow;
17
18impl FixedWindow {
19 pub fn new() -> Self {
21 Self
22 }
23
24 fn window_start(&self, now: u64, window_ms: u64) -> u64 {
26 (now / window_ms) * window_ms
27 }
28}
29
30impl Algorithm for FixedWindow {
31 fn name(&self) -> &'static str {
32 "fixed_window"
33 }
34
35 async fn check_and_record<S: Storage>(
36 &self,
37 storage: &S,
38 key: &str,
39 quota: &Quota,
40 ) -> Result<Decision> {
41 let now = current_timestamp_ms();
42 let window_ms = quota.window().as_millis() as u64;
43 let window_start = self.window_start(now, window_ms);
44 let ttl = Duration::from_millis(window_ms * 2);
45
46 let count = storage.increment(key, 1, window_start, ttl).await?;
47
48 let limit = quota.max_requests();
49 let remaining = limit.saturating_sub(count);
50 let reset_at = timestamp_to_instant(window_start + window_ms);
51 let window_start_instant = timestamp_to_instant(window_start);
52
53 let info = RateLimitInfo::new(limit, remaining, reset_at, window_start_instant)
54 .with_algorithm("fixed_window");
55
56 Ok(if count <= limit {
57 Decision::allowed(info)
58 } else {
59 let retry_after = Duration::from_millis(window_start + window_ms - now);
60 Decision::denied(info.with_retry_after(retry_after))
61 })
62 }
63
64 async fn check<S: Storage>(
65 &self,
66 storage: &S,
67 key: &str,
68 quota: &Quota,
69 ) -> Result<Decision> {
70 let now = current_timestamp_ms();
71 let window_ms = quota.window().as_millis() as u64;
72 let window_start = self.window_start(now, window_ms);
73
74 let entry = storage.get(key).await?;
75 let count = entry
76 .filter(|e| e.window_start == window_start)
77 .map(|e| e.count)
78 .unwrap_or(0);
79
80 let limit = quota.max_requests();
81 let remaining = limit.saturating_sub(count);
82 let reset_at = timestamp_to_instant(window_start + window_ms);
83
84 let info = RateLimitInfo::new(limit, remaining, reset_at, timestamp_to_instant(window_start))
85 .with_algorithm("fixed_window");
86
87 Ok(if count < limit {
88 Decision::allowed(info)
89 } else {
90 let retry_after = Duration::from_millis(window_start + window_ms - now);
91 Decision::denied(info.with_retry_after(retry_after))
92 })
93 }
94}
95
96#[cfg(test)]
97mod tests {
98 use super::*;
99 use crate::storage::MemoryStorage;
100
101 #[tokio::test]
102 async fn test_fixed_window_basic() {
103 let algorithm = FixedWindow::new();
104 let storage = MemoryStorage::new();
105 let quota = Quota::per_minute(5);
106
107 for i in 1..=5 {
108 let decision = algorithm.check_and_record(&storage, "user:1", "a).await.unwrap();
109 assert!(decision.is_allowed(), "Request {} should be allowed", i);
110 }
111
112 let decision = algorithm.check_and_record(&storage, "user:1", "a).await.unwrap();
113 assert!(decision.is_denied());
114 }
115}