envoy/rate_limit/
state.rs

1//! Rate limit state types — token bucket, configuration, per-agent state.
2
3use std::time::Duration;
4
5/// Token bucket for rate limiting.
6///
7/// Tokens replenish over time based on `replenish_rate` (tokens/second).
8/// Capacity is capped at `max_tokens`.
9#[derive(Debug, Clone)]
10pub struct TokenBucket {
11    pub tokens: u64,
12    pub max_tokens: u64,
13    pub replenish_rate: u64,
14    pub last_replenish: std::time::Instant,
15}
16
17impl TokenBucket {
18    pub fn new(max_tokens: u64, replenish_rate: u64) -> Self {
19        Self {
20            tokens: max_tokens,
21            max_tokens,
22            replenish_rate,
23            last_replenish: std::time::Instant::now(),
24        }
25    }
26
27    /// Replenish tokens based on elapsed time.
28    pub fn replenish(&mut self, elapsed: Duration) {
29        let secs = elapsed.as_secs_f64();
30        let tokens_to_add = (secs * self.replenish_rate as f64) as u64;
31        self.tokens = (self.tokens + tokens_to_add).min(self.max_tokens);
32        self.last_replenish = std::time::Instant::now();
33    }
34
35    /// Consume tokens if available.
36    pub fn consume(&mut self, amount: u64) {
37        self.tokens = self.tokens.saturating_sub(amount);
38    }
39
40    /// Try to consume tokens, returning error if insufficient.
41    pub fn try_consume(&mut self, amount: u64) -> Result<(), u64> {
42        if self.tokens >= amount {
43            self.tokens -= amount;
44            Ok(())
45        } else {
46            Err(self.tokens)
47        }
48    }
49}
50
51/// Rate limit decision result.
52#[derive(Debug, Clone, PartialEq)]
53pub struct RateLimitDecision {
54    pub allowed: bool,
55    pub retry_after: Option<Duration>,
56}
57
58/// Per-agent rate limit state.
59#[derive(Debug, Clone)]
60pub struct RateLimitState {
61    pub agent_id: String,
62    bucket: TokenBucket,
63}
64
65impl RateLimitState {
66    pub fn new(agent_id: &str, max_tokens: u64, replenish_rate: u64) -> Self {
67        Self {
68            agent_id: agent_id.to_string(),
69            bucket: TokenBucket::new(max_tokens, replenish_rate),
70        }
71    }
72
73    /// Create from an existing bucket (used by RateLimitStore).
74    pub(crate) fn from_bucket(agent_id: String, bucket: TokenBucket) -> Self {
75        Self { agent_id, bucket }
76    }
77
78    /// Check if a request is allowed, consuming tokens if so.
79    pub fn check(&mut self, cost: u64) -> RateLimitDecision {
80        if self.bucket.tokens >= cost {
81            self.bucket.consume(cost);
82            RateLimitDecision {
83                allowed: true,
84                retry_after: None,
85            }
86        } else {
87            let deficit = cost - self.bucket.tokens;
88            let retry_after =
89                Duration::from_secs_f64(deficit as f64 / self.bucket.replenish_rate as f64);
90            RateLimitDecision {
91                allowed: false,
92                retry_after: Some(retry_after),
93            }
94        }
95    }
96
97    /// Replenish tokens based on elapsed time.
98    pub fn replenish(&mut self, elapsed: Duration) {
99        self.bucket.replenish(elapsed);
100    }
101
102    /// Access the underlying bucket (for testing).
103    pub fn bucket(&self) -> &TokenBucket {
104        &self.bucket
105    }
106
107    /// Access the underlying bucket mutably (for testing).
108    pub fn bucket_mut(&mut self) -> &mut TokenBucket {
109        &mut self.bucket
110    }
111}
112
113#[cfg(test)]
114mod tests {
115    use super::*;
116
117    #[test]
118    fn test_token_bucket_initial_state() {
119        let bucket = TokenBucket::new(100, 10);
120        assert_eq!(bucket.tokens, 100);
121        assert_eq!(bucket.max_tokens, 100);
122        assert_eq!(bucket.replenish_rate, 10);
123    }
124}
envoy/rate_limit/state.rs

envoy/rate_limit/
state.rs
