load_balancer/

rate_limit.rs

use dashmap::DashMap;
use dashmap::Entry;
use std::hash::Hash;
use std::sync::Arc;
use std::time::{Duration, Instant};
use tokio::sync::Mutex;
use tokio::task::yield_now;
use tokio::time::sleep;

/// A token-bucket rate limiter.
///
/// Limits the rate of operations to `rate` per `interval`. When the
/// rate is exceeded, further requests are denied until the next window.
#[derive(Clone)]
pub struct RateLimiter {
    rate: u64,
    interval: Duration,
    state: Arc<Mutex<Inner>>,
}

struct Inner {
    count: u64,
    last_reset: Instant,
}

impl RateLimiter {
    /// Create a new rate limiter allowing `rate` operations per second.
    pub fn new_rate(rate: u64) -> Self {
        Self::new_rate_interval(rate, Duration::from_secs(1))
    }

    /// Create a new rate limiter with a custom rate and interval.
    pub fn new_rate_interval(rate: u64, interval: Duration) -> Self {
        RateLimiter {
            rate,
            interval,
            state: Arc::new(Mutex::new(Inner {
                count: 0,
                last_reset: Instant::now(),
            })),
        }
    }

    /// Check whether an operation is allowed right now.
    ///
    /// Returns `true` if within the rate limit, `false` if the bucket is exhausted.
    pub async fn check(&self) -> bool {
        let mut state = self.state.lock().await;
        let now = Instant::now();

        if now >= state.last_reset + self.interval {
            state.count = 0;
            state.last_reset = now;
        }

        if state.count < self.rate {
            state.count += 1;
            true
        } else {
            false
        }
    }

    /// Returns the duration until the next token becomes available.
    ///
    /// Returns `Duration::ZERO` if tokens are immediately available,
    /// `Duration::MAX` if the rate is zero.
    pub async fn time_until_available(&self) -> Duration {
        if self.rate == 0 {
            return Duration::MAX;
        }

        let state = self.state.lock().await;
        let now = Instant::now();
        let window_end = state.last_reset + self.interval;

        if now >= window_end || state.count < self.rate {
            Duration::ZERO
        } else {
            window_end - now
        }
    }

    /// Block until a token is available, then consume it.
    ///
    /// Sleeps for the remaining window duration when the bucket is exhausted,
    /// then retries. Returns immediately if tokens are available.
    pub async fn wait(&self) {
        loop {
            let mut state = self.state.lock().await;
            let now = Instant::now();

            if now >= state.last_reset + self.interval {
                state.count = 0;
                state.last_reset = now;
            }

            if state.count < self.rate {
                state.count += 1;
                return;
            }

            let wait = (state.last_reset + self.interval).saturating_duration_since(now);

            drop(state);

            if wait > Duration::ZERO {
                sleep(wait).await;
            } else {
                yield_now().await;
            }
        }
    }
}

/// A per-key rate limiter map.
///
/// Each key can be assigned its own rate and interval via [`insert_rate`](Self::insert_rate).
/// Keys that have not been inserted are rejected by default — call
/// [`insert_rate`](Self::insert_rate) or [`insert_rate_interval`](Self::insert_rate_interval) first.
///
/// # Example
///
/// ```rust
/// use load_balancer::rate_limit::RateLimiterMap;
/// use std::time::Duration;
///
/// #[tokio::main]
/// async fn main() {
///     let map = RateLimiterMap::new();
///
///     // alice: 3 req/s,  bob: 10 req/s
///     map.insert_rate("alice", 3);
///     map.insert_rate_interval("bob", 10, Duration::from_secs(1));
///
///     assert!(map.check(&"alice").await);
///     assert!(map.check(&"bob").await);
///     assert!(!map.check(&"charlie").await); // not inserted → rejected
/// }
/// ```
#[derive(Clone)]
pub struct RateLimiterMap<K> {
    map: Arc<DashMap<K, RateLimiter>>,
}

impl<K> RateLimiterMap<K>
where
    K: Hash + Eq + Clone + Send + Sync + 'static,
{
    /// Create an empty map.
    pub fn new() -> Self {
        RateLimiterMap {
            map: Arc::new(DashMap::new()),
        }
    }

    /// Insert a rate limiter for `key` with the given `rate` and a 1-second interval.
    ///
    /// Overwrites any existing limiter for the same key.
    pub fn insert_rate(&self, key: K, rate: u64) {
        self.insert_rate_interval(key, rate, Duration::from_secs(1))
    }

    /// Insert a rate limiter for `key` with the given `rate` and `interval`.
    pub fn insert_rate_interval(&self, key: K, rate: u64, interval: Duration) {
        let limiter = RateLimiter::new_rate_interval(rate, interval);
        self.map.insert(key, limiter);
    }

    /// Remove the rate limiter for `key`.
    ///
    /// After removal the key is unknown and will be rejected by
    /// [`check`](Self::check) (same as a key that was never inserted).
    pub fn remove(&self, key: &K) {
        self.map.remove(key);
    }

    /// Returns the number of keys currently managed.
    pub fn len(&self) -> usize {
        self.map.len()
    }

    /// Returns `true` if no keys are being rate-limited.
    pub fn is_empty(&self) -> bool {
        self.map.is_empty()
    }

    /// Returns `true` if `key` has an active rate limiter.
    pub fn contains_key(&self, key: &K) -> bool {
        self.map.contains_key(key)
    }

    /// Remove all rate limiters.
    pub fn clear(&self) {
        self.map.clear();
    }

    /// Retain only the keys for which `f` returns `true`.
    ///
    /// Each remaining key's limiter is unchanged.
    pub fn retain(&self, f: impl FnMut(&K, &mut RateLimiter) -> bool) {
        self.map.retain(f);
    }

    /// Get an [`Entry`] for `key`, allowing in-place insertion or modification.
    ///
    /// Use `entry(key).or_insert(limiter)` to insert a limiter only if the key
    /// is absent, or `entry(key).and_modify(...)` to adjust an existing one.
    pub fn entry(&self, key: K) -> Entry<'_, K, RateLimiter> {
        self.map.entry(key)
    }

    /// Check whether `key` is allowed right now.
    ///
    /// Returns `true` if the key's limiter has available tokens.
    /// Returns `false` if the key has not been inserted (unknown keys are
    /// rejected by default).
    pub async fn check(&self, key: &K) -> bool {
        match self.map.get(key).map(|r| r.clone()) {
            Some(limiter) => limiter.check().await,
            None => false,
        }
    }

    /// Check `key` and distinguish "unknown" from "rate-limited".
    ///
    /// Returns `Some(true)` if allowed, `Some(false)` if rate-limited,
    /// `None` if the key has not been inserted.
    pub async fn check_opt(&self, key: &K) -> Option<bool> {
        match self.map.get(key).map(|r| r.clone()) {
            Some(limiter) => Some(limiter.check().await),
            None => None,
        }
    }

    /// Returns the duration until `key` becomes available again.
    ///
    /// Returns `Duration::MAX` if the key is not inserted or the rate is zero,
    /// `Duration::ZERO` if tokens are immediately available.
    pub async fn time_until_available(&self, key: &K) -> Duration {
        match self.map.get(key).map(|r| r.clone()) {
            Some(limiter) => limiter.time_until_available().await,
            None => Duration::MAX,
        }
    }

    /// Block until `key` is allowed, then consume a token.
    ///
    /// Returns immediately if the key is not inserted (no-op).
    pub async fn wait(&self, key: &K) {
        if let Some(limiter) = self.map.get(key).map(|r| r.clone()) {
            limiter.wait().await;
        }
    }
}