subms_rate_limiter/features/
distributed_backend.rs

1//! Pluggable backend for cross-process rate limiting state.
2//!
3//! The `Backend` trait abstracts the atomic INCR + EXPIRE primitive
4//! shared by Redis / Memcached / DynamoDB / any other store that can
5//! atomically bump a counter scoped to a fixed window. Real distributed
6//! backends are downstream user concerns; we ship a real
7//! `InMemoryBackend` so the trait is exercised + the in-process shape
8//! matches what a Redis-style impl would do over the wire.
9//!
10//! Algorithm: fixed-window counters. Each `(key, window_start_ns)`
11//! pair holds a count; `incr` returns the new count after bump, the
12//! backend collects garbage on window roll. Simpler than sliding-
13//! window + good enough for the cross-process case (the network
14//! round-trip cost dwarfs the windowing imprecision).
15
16use std::collections::HashMap;
17use std::sync::Mutex;
18
19use super::clock::{Clock, SystemClock};
20
21/// Cross-process state backend. Implementations bump a counter for
22/// `(key, window_start)` and return the new value after the bump.
23pub trait Backend: Send + Sync {
24    /// Increment the counter at `key` for `window_start_ns`. Returns
25    /// the new count after the bump. Must be atomic across concurrent
26    /// callers.
27    fn incr(&self, key: &str, window_start_ns: u64, ttl_ns: u64) -> u64;
28
29    /// Read the current counter without bumping. Returns 0 if the
30    /// (key, window) pair is unknown or expired.
31    fn read(&self, key: &str, window_start_ns: u64) -> u64;
32}
33
34/// Real in-process backend. Holds counters in a `HashMap<(key, window), u64>`
35/// guarded by a `Mutex`. Garbage-collects expired windows opportunistically
36/// on each `incr` call.
37pub struct InMemoryBackend {
38    inner: Mutex<Inner>,
39}
40
41struct Inner {
42    counters: HashMap<(String, u64), Cell>,
43}
44
45#[derive(Clone, Copy)]
46struct Cell {
47    count: u64,
48    expires_ns: u64,
49}
50
51impl InMemoryBackend {
52    pub fn new() -> Self {
53        Self {
54            inner: Mutex::new(Inner {
55                counters: HashMap::new(),
56            }),
57        }
58    }
59}
60
61impl Default for InMemoryBackend {
62    fn default() -> Self {
63        Self::new()
64    }
65}
66
67impl Backend for InMemoryBackend {
68    fn incr(&self, key: &str, window_start_ns: u64, ttl_ns: u64) -> u64 {
69        let mut g = self.inner.lock().unwrap();
70        // Opportunistic GC: drop windows whose TTL has expired by
71        // wall-clock-of-callsite. We accept this is imprecise vs
72        // actual now-ns; for the cross-process case the caller's
73        // clock IS the source of truth.
74        let now = window_start_ns;
75        g.counters.retain(|_, c| c.expires_ns > now);
76        let entry = g
77            .counters
78            .entry((key.to_string(), window_start_ns))
79            .or_insert(Cell {
80                count: 0,
81                expires_ns: window_start_ns.saturating_add(ttl_ns),
82            });
83        entry.count = entry.count.saturating_add(1);
84        entry.count
85    }
86
87    fn read(&self, key: &str, window_start_ns: u64) -> u64 {
88        let g = self.inner.lock().unwrap();
89        g.counters
90            .get(&(key.to_string(), window_start_ns))
91            .map(|c| c.count)
92            .unwrap_or(0)
93    }
94}
95
96/// Rate limiter backed by a pluggable `Backend`. Fixed-window
97/// algorithm: per `(key, window_size_ns)`, allow at most `limit`
98/// requests.
99pub struct DistributedLimiter {
100    backend: Box<dyn Backend>,
101    clock: Box<dyn Clock>,
102    limit: u64,
103    window_ns: u64,
104}
105
106impl DistributedLimiter {
107    pub fn new(backend: Box<dyn Backend>, limit: u64, window_ns: u64) -> Self {
108        Self::with_clock(backend, limit, window_ns, Box::new(SystemClock::new()))
109    }
110
111    pub fn with_clock(
112        backend: Box<dyn Backend>,
113        limit: u64,
114        window_ns: u64,
115        clock: Box<dyn Clock>,
116    ) -> Self {
117        Self {
118            backend,
119            clock,
120            limit: limit.max(1),
121            window_ns: window_ns.max(1),
122        }
123    }
124
125    /// Try to acquire one permit on `key`. Returns true if the post-bump
126    /// counter is within `limit`. The bump always happens (mirroring
127    /// the Redis INCR + EXPIRE shape) so contention races resolve
128    /// monotonically.
129    pub fn try_acquire(&self, key: &str) -> bool {
130        let now = self.clock.now_ns();
131        let window_start = now - (now % self.window_ns);
132        let count = self.backend.incr(key, window_start, self.window_ns);
133        count <= self.limit
134    }
135
136    pub fn limit(&self) -> u64 {
137        self.limit
138    }
139
140    pub fn window_ns(&self) -> u64 {
141        self.window_ns
142    }
143}
144
145#[cfg(test)]
146#[path = "distributed_backend_tests.rs"]
147mod tests;
subms_rate_limiter/features/distributed_backend.rs

subms_rate_limiter/features/
distributed_backend.rs
