skp-ratelimit 0.1.2

Advanced, modular, extensible rate limiting library with GCRA, per-route quotas, and composite keys
Documentation 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115

//! Fixed Window rate limiting algorithm.

use std::time::Duration;

use crate::algorithm::{current_timestamp_ms, timestamp_to_instant, Algorithm};
use crate::decision::{Decision, RateLimitInfo};
use crate::error::Result;
use crate::quota::Quota;
use crate::storage::Storage;

/// Fixed Window rate limiting algorithm.
///
/// Simple counter that resets at fixed intervals.
/// Fast but has "boundary burst" problem.
#[derive(Debug, Clone, Default)]
pub struct FixedWindow;

impl FixedWindow {
    /// Create a new Fixed Window algorithm instance.
    pub fn new() -> Self {
        Self
    }

    /// Calculate the current window start.
    fn window_start(&self, now: u64, window_ms: u64) -> u64 {
        (now / window_ms) * window_ms
    }
}

impl Algorithm for FixedWindow {
    fn name(&self) -> &'static str {
        "fixed_window"
    }

    async fn check_and_record<S: Storage>(
        &self,
        storage: &S,
        key: &str,
        quota: &Quota,
    ) -> Result<Decision> {
        let now = current_timestamp_ms();
        let window_ms = quota.window().as_millis() as u64;
        let window_start = self.window_start(now, window_ms);
        let ttl = Duration::from_millis(window_ms * 2);

        let count = storage.increment(key, 1, window_start, ttl).await?;

        let limit = quota.max_requests();
        let remaining = limit.saturating_sub(count);
        let reset_at = timestamp_to_instant(window_start + window_ms);
        let window_start_instant = timestamp_to_instant(window_start);

        let info = RateLimitInfo::new(limit, remaining, reset_at, window_start_instant)
            .with_algorithm("fixed_window");

        Ok(if count <= limit {
            Decision::allowed(info)
        } else {
            let retry_after = Duration::from_millis(window_start + window_ms - now);
            Decision::denied(info.with_retry_after(retry_after))
        })
    }

    async fn check<S: Storage>(
        &self,
        storage: &S,
        key: &str,
        quota: &Quota,
    ) -> Result<Decision> {
        let now = current_timestamp_ms();
        let window_ms = quota.window().as_millis() as u64;
        let window_start = self.window_start(now, window_ms);

        let entry = storage.get(key).await?;
        let count = entry
            .filter(|e| e.window_start == window_start)
            .map(|e| e.count)
            .unwrap_or(0);

        let limit = quota.max_requests();
        let remaining = limit.saturating_sub(count);
        let reset_at = timestamp_to_instant(window_start + window_ms);

        let info = RateLimitInfo::new(limit, remaining, reset_at, timestamp_to_instant(window_start))
            .with_algorithm("fixed_window");

        Ok(if count < limit {
            Decision::allowed(info)
        } else {
            let retry_after = Duration::from_millis(window_start + window_ms - now);
            Decision::denied(info.with_retry_after(retry_after))
        })
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use crate::storage::MemoryStorage;

    #[tokio::test]
    async fn test_fixed_window_basic() {
        let algorithm = FixedWindow::new();
        let storage = MemoryStorage::new();
        let quota = Quota::per_minute(5);

        for i in 1..=5 {
            let decision = algorithm.check_and_record(&storage, "user:1", &quota).await.unwrap();
            assert!(decision.is_allowed(), "Request {} should be allowed", i);
        }

        let decision = algorithm.check_and_record(&storage, "user:1", &quota).await.unwrap();
        assert!(decision.is_denied());
    }
}