throttle_kit/
backend.rs

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

throttle_kit/
backend.rs
