cratestack_core/store/
ratelimit.rs

1//! Rate limiting store trait and configuration types.
2
3use async_trait::async_trait;
4
5use crate::CoolError;
6
7/// Configuration for a single bucket: capacity (max burst) and refill rate
8/// in tokens per second. Banks running high-frequency back-office traffic
9/// pick large bursts; consumer-facing channels use small bursts to dampen
10/// abuse.
11#[derive(Debug, Clone, Copy)]
12pub struct RateLimitConfig {
13    pub burst: u32,
14    pub refill_per_second: f64,
15}
16
17impl RateLimitConfig {
18    pub fn new(burst: u32, refill_per_second: f64) -> Self {
19        Self {
20            burst,
21            refill_per_second,
22        }
23    }
24}
25
26/// Result of attempting to consume a token. `Allowed` carries the number
27/// of tokens left after consumption; `Throttled` carries seconds the
28/// caller should wait before retrying.
29#[derive(Debug, Clone, Copy, PartialEq)]
30pub enum RateLimitDecision {
31    Allowed { remaining: u32 },
32    Throttled { retry_after_secs: u32 },
33}
34
35/// Sleep helper for tests — exposes the bucket's wall-clock refill model so
36/// the integration tests can exercise both the burst and the throttle path
37/// without depending on real time.
38#[doc(hidden)]
39pub fn _bucket_capacity_for(config: RateLimitConfig) -> u32 {
40    config.burst
41}
42
43/// Pluggable storage for token-bucket state. Implementations must be safe
44/// to share across tasks (use a Mutex internally, or rely on the backing
45/// store's atomicity).
46#[async_trait]
47pub trait RateLimitStore: Send + Sync + 'static {
48    /// Atomically consume one token for `key`. Returns the decision based
49    /// on the bucket state after the consumption attempt.
50    async fn consume(
51        &self,
52        key: &str,
53        config: RateLimitConfig,
54    ) -> Result<RateLimitDecision, CoolError>;
55}
56
57#[cfg(test)]
58mod tests {
59    use super::*;
60
61    #[test]
62    fn rate_limit_config_new_creates_correct_values() {
63        let config = RateLimitConfig::new(100, 10.5);
64        assert_eq!(config.burst, 100);
65        assert_eq!(config.refill_per_second, 10.5);
66    }
67
68    #[test]
69    fn rate_limit_decision_allowed_equality() {
70        let d1 = RateLimitDecision::Allowed { remaining: 42 };
71        let d2 = RateLimitDecision::Allowed { remaining: 42 };
72        assert_eq!(d1, d2);
73    }
74
75    #[test]
76    fn rate_limit_decision_throttled_equality() {
77        let d1 = RateLimitDecision::Throttled {
78            retry_after_secs: 5,
79        };
80        let d2 = RateLimitDecision::Throttled {
81            retry_after_secs: 5,
82        };
83        assert_eq!(d1, d2);
84    }
85
86    #[test]
87    fn bucket_capacity_for_returns_burst() {
88        let config = RateLimitConfig::new(42, 1.0);
89        assert_eq!(_bucket_capacity_for(config), 42);
90    }
91}
cratestack_core/store/ratelimit.rs

cratestack_core/store/
ratelimit.rs
