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.entries.entry(key.to_string()).or_insert_with(|| Entry {
50            last_check: Instant::now(),
51            remaining: burst,
52        });
53
54        let now = Instant::now();
55        let elapsed = now.duration_since(entry.last_check);
56
57        // GCRA: compute how many tokens have been replenished.
58        let replenished = (elapsed.as_nanos() as u64 / interval.as_nanos() as u64) as u32;
59        if replenished > 0 {
60            entry.remaining = (entry.remaining + replenished).min(burst);
61            let consumption =
62                Duration::from_nanos(replenished as u64 * interval.as_nanos() as u64);
63            entry.last_check += consumption;
64        }
65
66        let allowed = entry.remaining > 0;
67        if allowed {
68            entry.remaining -= 1;
69        }
70
71        let remaining = entry.remaining;
72        let reset_at = entry.last_check + interval;
73
74        RateLimitResult {
75            allowed,
76            remaining: remaining as u64,
77            reset_at,
78            limit: burst as u64,
79            retry_after: if allowed {
80                None
81            } else {
82                Some(interval)
83            },
84        }
85    }
86}
throttle_kit/backend.rs

throttle_kit/
backend.rs
