rok-rate-limit 0.3.0

Rate limiting Tower middleware and programmatic Limiter API for the rok ecosystem
Documentation
use std::{
    sync::{
        atomic::{AtomicU64, Ordering},
        Arc,
    },
    time::{Duration, SystemTime, UNIX_EPOCH},
};

use dashmap::DashMap;

struct WindowEntry {
    count: AtomicU64,
    reset_epoch: u64,
}

/// Fixed-window in-memory rate limit backend.
///
/// Each key+window is stored as an entry in a `DashMap`. Old windows are
/// naturally isolated by including the window-start timestamp in the map key;
/// they accumulate until the map is explicitly pruned (not yet implemented for
/// V2 — memory use is bounded to one entry per active key per window).
#[derive(Clone)]
pub struct MemoryBackend {
    windows: Arc<DashMap<String, Arc<WindowEntry>>>,
}

impl Default for MemoryBackend {
    fn default() -> Self {
        Self::new()
    }
}

impl MemoryBackend {
    pub fn new() -> Self {
        Self {
            windows: Arc::new(DashMap::new()),
        }
    }

    /// Increment the counter for `key` inside the current fixed window of
    /// `window_secs` seconds.
    ///
    /// Returns `(new_count, reset_epoch_secs)` where `reset_epoch_secs` is the
    /// Unix timestamp (seconds) at which the window resets.
    pub fn increment(&self, key: &str, window_secs: u64) -> (u64, u64) {
        let now_epoch = epoch_secs();
        let window_start = (now_epoch / window_secs) * window_secs;
        let reset_epoch = window_start + window_secs;
        let map_key = format!("{}:{}", key, window_start);

        let entry = self.windows.entry(map_key).or_insert_with(|| {
            Arc::new(WindowEntry {
                count: AtomicU64::new(0),
                reset_epoch,
            })
        });

        let count = entry.count.fetch_add(1, Ordering::Relaxed) + 1;
        (count, entry.reset_epoch)
    }
}

fn epoch_secs() -> u64 {
    SystemTime::now()
        .duration_since(UNIX_EPOCH)
        .unwrap_or(Duration::ZERO)
        .as_secs()
}