Skip to main content

aa_core/storage/
rate_limit_counter.rs

1//! [`RateLimitCounter`] — read-modify-write counters for rate limiting.
2
3use super::Result;
4use async_trait::async_trait;
5
6/// Atomic counters keyed by an arbitrary string, used for rate limiting.
7///
8/// The defining operation is [`increment`](RateLimitCounter::increment): it must
9/// apply the read-modify-write **atomically** so two concurrent callers can never
10/// observe or commit the same pre-increment value. Counters are scoped to a
11/// fixed-length window; `window_secs` lets the backend bucket and expire counts
12/// without the caller tracking wall-clock time.
13///
14/// # Example
15///
16/// ```
17/// use aa_core::storage::{RateLimitCounter, Result};
18/// use async_trait::async_trait;
19///
20/// /// A counter that always reports a single hit (a stand-in for a real backend).
21/// struct AlwaysOne;
22///
23/// #[async_trait]
24/// impl RateLimitCounter for AlwaysOne {
25///     async fn increment(&self, _key: &str, _amount: u64, _window_secs: u64) -> Result<u64> {
26///         Ok(1)
27///     }
28///
29///     async fn current(&self, _key: &str) -> Result<u64> {
30///         Ok(1)
31///     }
32///
33///     async fn reset(&self, _key: &str) -> Result<()> {
34///         Ok(())
35///     }
36/// }
37/// ```
38#[async_trait]
39pub trait RateLimitCounter: Send + Sync {
40    /// Atomically add `amount` to the counter for `key` within the window of
41    /// length `window_secs`, returning the new total for the current window.
42    ///
43    /// The read-modify-write is atomic with respect to concurrent callers.
44    async fn increment(&self, key: &str, amount: u64, window_secs: u64) -> Result<u64>;
45
46    /// Return the current total for `key` without modifying it.
47    ///
48    /// Returns `0` for a key that has never been incremented (or whose window
49    /// has expired).
50    async fn current(&self, key: &str) -> Result<u64>;
51
52    /// Reset the counter for `key` back to zero.
53    ///
54    /// Idempotent: resetting an absent key succeeds.
55    async fn reset(&self, key: &str) -> Result<()>;
56}