throttle_kit/
backend.rs

1#[cfg(feature = "in-memory")]
2use std::time::{Duration, Instant};
3
4use crate::metrics::RateLimitResult;
5use crate::quota::Quota;
6
7/// Backend trait for rate limit state storage.
8#[async_trait::async_trait]
9pub trait RateLimitBackend: Send + Sync + 'static {
10    /// Check if a request for `key` is allowed under the given `quota`.
11    async fn check(&self, key: &str, quota: &Quota) -> RateLimitResult;
12}
13
14/// In-memory backend backed by a `DashMap`.
15///
16/// Suitable for single-node deployments. Each key tracks its last-allowed
17/// timestamp and a token-bucket counter.
18#[cfg(feature = "in-memory")]
19#[derive(Clone)]
20pub struct InMemoryBackend {
21    entries: dashmap::DashMap<String, Entry>,
22}
23
24#[cfg(feature = "in-memory")]
25#[derive(Clone)]
26struct Entry {
27    last_check: Instant,
28    remaining: u32,
29}
30
31#[cfg(feature = "in-memory")]
32impl InMemoryBackend {
33    /// Create a new in-memory backend.
34    pub fn new() -> Self {
35        Self {
36            entries: dashmap::DashMap::new(),
37        }
38    }
39}
40
41#[cfg(feature = "in-memory")]
42impl Default for InMemoryBackend {
43    fn default() -> Self {
44        Self::new()
45    }
46}
47
48#[cfg(feature = "in-memory")]
49#[async_trait::async_trait]
50impl RateLimitBackend for InMemoryBackend {
51    async fn check(&self, key: &str, quota: &Quota) -> RateLimitResult {
52        let interval = quota.interval();
53        let burst = quota.burst;
54
55        let mut entry = self
56            .entries
57            .entry(key.to_string())
58            .or_insert_with(|| Entry {
59                last_check: Instant::now(),
60                remaining: burst,
61            });
62
63        let now = Instant::now();
64        let elapsed = now.duration_since(entry.last_check);
65
66        // GCRA: compute how many tokens have been replenished.
67        let replenished = (elapsed.as_nanos() as u64 / interval.as_nanos() as u64) as u32;
68        if replenished > 0 {
69            entry.remaining = (entry.remaining + replenished).min(burst);
70            let consumption = Duration::from_nanos(replenished as u64 * interval.as_nanos() as u64);
71            entry.last_check += consumption;
72        }
73
74        let allowed = entry.remaining > 0;
75        if allowed {
76            entry.remaining -= 1;
77        }
78
79        let remaining = entry.remaining;
80        let reset_at = entry.last_check + interval;
81
82        #[cfg(feature = "metrics")]
83        {
84            let utilization = (burst - remaining) as f64 / burst as f64;
85            metrics::histogram!("ratelimit_utilization").record(utilization);
86            if allowed {
87                metrics::counter!("ratelimit_allowed_total", "key" => key.to_string()).increment(1);
88            } else {
89                metrics::counter!("ratelimit_rejected_total", "key" => key.to_string())
90                    .increment(1);
91            }
92        }
93
94        RateLimitResult {
95            allowed,
96            remaining: remaining as u64,
97            reset_at,
98            limit: burst as u64,
99            retry_after: if allowed { None } else { Some(interval) },
100        }
101    }
102}
throttle_kit/backend.rs

throttle_kit/
backend.rs
