skp_ratelimit/
lib.rs

1//! Advanced, modular rate limiting library for Rust.
2//!
3//! `skp_ratelimit` provides a comprehensive rate limiting solution with:
4//!
5//! - **Multiple Algorithms**: GCRA, Token Bucket, Leaky Bucket, Sliding Log, and more
6//! - **Pluggable Storage**: In-memory with GC, Redis with connection pooling
7//! - **Per-Route Quotas**: Different limits for different endpoints
8//! - **Composite Keys**: Rate limit by IP + Path, User + API Key, etc.
9//! - **Framework Integration**: Axum and Actix-web middleware
10//!
11//! # Quick Start
12//!
13//! ```ignore
14//! use skp_ratelimit::{GCRA, Quota, MemoryStorage, Algorithm};
15//! use std::time::Duration;
16//!
17//! #[tokio::main]
18//! async fn main() {
19//!     // Create storage and algorithm
20//!     let storage = MemoryStorage::new();
21//!     let algorithm = GCRA::new();
22//!     let quota = Quota::per_second(10).with_burst(15);
23//!
24//!     // Check and record a request
25//!     let decision = algorithm.check_and_record(&storage, "user:123", &quota).await.unwrap();
26//!
27//!     if decision.is_allowed() {
28//!         println!("Request allowed! {} remaining", decision.info().remaining);
29//!     } else {
30//!         println!("Rate limited! Retry after {:?}", decision.info().retry_after);
31//!     }
32//! }
33//! ```
34//!
35//! # Algorithms
36//!
37//! | Algorithm | Best For | Memory | Feature Flag |
38//! |-----------|----------|--------|--------------|
39//! | GCRA | Precise rate control | Low | `gcra` |
40//! | Token Bucket | Bursty traffic | Low | default |
41//! | Leaky Bucket | Smooth output | Low | `leaky-bucket` |
42//! | Sliding Log | Precision critical | High | `sliding-log` |
43//! | Sliding Window | General purpose | Low | default |
44//! | Fixed Window | Simple use cases | Low | default |
45//! | Concurrent | Limit parallelism | Low | `concurrent` |
46//!
47//! # Feature Flags
48//!
49//! - `memory` (default): In-memory storage with garbage collection
50//! - `redis`: Redis storage backend
51//! - `axum`: Axum middleware integration
52//! - `gcra`: GCRA algorithm
53//! - `leaky-bucket`: Leaky Bucket algorithm
54//! - `sliding-log`: Sliding Log algorithm
55//! - `concurrent`: Concurrent request limiter
56
57pub mod algorithm;
58pub mod decision;
59pub mod error;
60pub mod extensions;
61pub mod headers;
62pub mod key;
63pub mod manager;
64pub mod policy;
65pub mod quota;
66pub mod storage;
67
68#[cfg(feature = "axum")]
69pub mod middleware;
70
71// Re-export main types
72pub use algorithm::Algorithm;
73pub use decision::{Decision, DecisionMetadata, RateLimitInfo};
74pub use error::{ConfigError, ConnectionError, RateLimitError, Result, StorageError};
75pub use key::{CompositeKey, FnKey, GlobalKey, Key, StaticKey};
76pub use manager::{RateLimitManager, RateLimitManagerBuilder, RouteConfig};
77pub use quota::{Quota, QuotaBuilder};
78pub use storage::{Storage, StorageEntry};
79
80// Re-export policy types
81pub use policy::{CompositePolicy, CreditPolicy, DefaultPolicy, PenaltyPolicy, Policy};
82
83// Re-export extensions and headers
84pub use extensions::{RateLimitExt, RateLimitResponse};
85pub use headers::RateLimitHeaders;
86
87// Re-export algorithms
88pub use algorithm::{FixedWindow, SlidingWindow, TokenBucket};
89
90#[cfg(feature = "gcra")]
91pub use algorithm::GCRA;
92
93#[cfg(feature = "leaky-bucket")]
94pub use algorithm::LeakyBucket;
95
96#[cfg(feature = "sliding-log")]
97pub use algorithm::SlidingLog;
98
99#[cfg(feature = "concurrent")]
100pub use algorithm::ConcurrentLimiter;
101
102// Re-export storage types
103#[cfg(feature = "memory")]
104pub use storage::{GcConfig, GcInterval, MemoryStorage};
105
106/// Prelude module for convenient imports.
107pub mod prelude {
108    pub use crate::algorithm::Algorithm;
109    pub use crate::decision::{Decision, RateLimitInfo};
110    pub use crate::error::{RateLimitError, Result};
111    pub use crate::quota::Quota;
112    pub use crate::storage::Storage;
113
114    pub use crate::algorithm::{FixedWindow, SlidingWindow, TokenBucket};
115
116    #[cfg(feature = "gcra")]
117    pub use crate::algorithm::GCRA;
118
119    #[cfg(feature = "leaky-bucket")]
120    pub use crate::algorithm::LeakyBucket;
121
122    #[cfg(feature = "sliding-log")]
123    pub use crate::algorithm::SlidingLog;
124
125    #[cfg(feature = "concurrent")]
126    pub use crate::algorithm::ConcurrentLimiter;
127
128    #[cfg(feature = "memory")]
129    pub use crate::storage::{GcConfig, GcInterval, MemoryStorage};
130}
131
132#[cfg(test)]
133mod tests {
134    use super::*;
135
136    #[cfg(feature = "memory")]
137    #[tokio::test]
138    async fn test_integration_gcra() {
139        use crate::prelude::*;
140
141        let storage = MemoryStorage::new();
142        let algorithm = GCRA::new();
143        let quota = Quota::per_second(10).with_burst(5);
144
145        // Should allow burst
146        for i in 1..=5 {
147            let decision = algorithm
148                .check_and_record(&storage, "user:1", &quota)
149                .await
150                .unwrap();
151            assert!(decision.is_allowed(), "Request {} should be allowed", i);
152        }
153
154        // Should deny after burst
155        let decision = algorithm
156            .check_and_record(&storage, "user:1", &quota)
157            .await
158            .unwrap();
159        assert!(decision.is_denied());
160        assert!(decision.info().retry_after.is_some());
161    }
162
163    #[cfg(feature = "memory")]
164    #[tokio::test]
165    async fn test_integration_token_bucket() {
166        let storage = MemoryStorage::new();
167        let algorithm = TokenBucket::new();
168        let quota = Quota::per_minute(60).with_burst(10);
169
170        let decision = algorithm
171            .check_and_record(&storage, "user:1", &quota)
172            .await
173            .unwrap();
174
175        assert!(decision.is_allowed());
176        assert_eq!(decision.info().remaining, 9);
177        assert_eq!(decision.info().algorithm, Some("token_bucket"));
178    }
179
180    #[cfg(feature = "memory")]
181    #[tokio::test]
182    async fn test_integration_headers() {
183        let storage = MemoryStorage::new();
184        let algorithm = FixedWindow::new();
185        let quota = Quota::per_minute(100);
186
187        let decision = algorithm
188            .check_and_record(&storage, "user:1", &quota)
189            .await
190            .unwrap();
191
192        let headers = decision.info().to_headers();
193        assert!(headers.iter().any(|(k, _)| *k == "X-RateLimit-Limit"));
194        assert!(headers.iter().any(|(k, _)| *k == "X-RateLimit-Remaining"));
195        assert!(headers.iter().any(|(k, _)| *k == "X-RateLimit-Reset"));
196    }
197
198    #[cfg(all(feature = "memory", feature = "concurrent"))]
199    #[tokio::test]
200    async fn test_integration_concurrent() {
201        let limiter = ConcurrentLimiter::new(2);
202
203        let _permit1 = limiter.try_acquire("user:1").unwrap();
204        let _permit2 = limiter.try_acquire("user:1").unwrap();
205
206        assert!(limiter.try_acquire("user:1").is_none());
207        assert_eq!(limiter.remaining("user:1"), 0);
208    }
209}
skp_ratelimit/lib.rs

skp_ratelimit/
lib.rs
