use std::{
sync::{
atomic::{AtomicU64, Ordering},
Arc,
},
time::{Duration, SystemTime, UNIX_EPOCH},
};
use dashmap::DashMap;
struct WindowEntry {
count: AtomicU64,
reset_epoch: u64,
}
#[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()),
}
}
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()
}