use std::collections::hash_map::RandomState;
use std::hash::{BuildHasher, Hash, Hasher};
use std::sync::atomic::{AtomicU64, AtomicUsize, Ordering};
use std::time::{Duration, Instant};
use dashmap::DashMap;
#[derive(Debug)]
pub struct CacheEntry {
pub data: Vec<u8>,
pub created_at: Instant,
pub ttl: Duration,
pub access_count: AtomicU64,
pub last_access: AtomicU64,
}
impl CacheEntry {
pub fn new(data: Vec<u8>, ttl: Duration) -> Self {
let now = Instant::now().duration_since(Instant::now()).as_secs();
Self {
data,
created_at: Instant::now(),
ttl,
access_count: AtomicU64::new(0),
last_access: AtomicU64::new(now),
}
}
pub fn is_expired(&self) -> bool {
self.created_at.elapsed() > self.ttl
}
pub fn record_access(&self) {
self.access_count.fetch_add(1, Ordering::Relaxed);
let now = Instant::now().duration_since(Instant::now()).as_secs();
self.last_access.store(now, Ordering::Relaxed);
}
}
#[derive(Debug, Clone)]
pub struct CacheConfig {
pub default_ttl: Duration,
pub max_entries: usize,
pub max_memory_bytes: usize,
}
impl Default for CacheConfig {
fn default() -> Self {
Self {
default_ttl: Duration::from_secs(60),
max_entries: 10000,
max_memory_bytes: 256 * 1024 * 1024, }
}
}
#[derive(Debug, Clone)]
pub struct CacheStats {
pub hits: u64,
pub misses: u64,
pub evictions: u64,
pub entries: usize,
pub memory_bytes: usize,
pub hit_rate: f64,
}
#[derive(Debug)]
pub struct RequestCache {
entries: DashMap<u64, CacheEntry>,
config: CacheConfig,
hits: AtomicU64,
misses: AtomicU64,
evictions: AtomicU64,
memory_bytes: AtomicUsize,
}
impl RequestCache {
pub fn new(config: CacheConfig) -> Self {
Self {
entries: DashMap::new(),
config,
hits: AtomicU64::new(0),
misses: AtomicU64::new(0),
evictions: AtomicU64::new(0),
memory_bytes: AtomicUsize::new(0),
}
}
pub fn generate_key(method: &str, args: &[u8]) -> u64 {
let hasher = RandomState::new().build_hasher();
let mut hasher = hasher;
method.hash(&mut hasher);
args.hash(&mut hasher);
hasher.finish()
}
pub fn get(&self, key: u64) -> Option<Vec<u8>> {
let entry = self.entries.get(&key)?;
if entry.is_expired() {
drop(entry);
self.remove(key);
return None;
}
entry.record_access();
self.hits.fetch_add(1, Ordering::Relaxed);
Some(entry.data.clone())
}
pub fn set(&self, key: u64, data: Vec<u8>, ttl: Option<Duration>) {
if self.entries.len() >= self.config.max_entries {
self.evict_one();
}
let entry_size = data.len();
if self.memory_bytes.load(Ordering::Relaxed) + entry_size > self.config.max_memory_bytes {
self.evict_until_fits(entry_size);
}
let ttl = ttl.unwrap_or(self.config.default_ttl);
let entry = CacheEntry::new(data, ttl);
let entry_size = entry.data.len();
self.entries.insert(key, entry);
self.memory_bytes.fetch_add(entry_size, Ordering::Relaxed);
}
pub fn remove(&self, key: u64) -> Option<CacheEntry> {
if let Some((_, entry)) = self.entries.remove(&key) {
self.memory_bytes.fetch_sub(entry.data.len(), Ordering::Relaxed);
self.evictions.fetch_add(1, Ordering::Relaxed);
Some(entry)
} else {
None
}
}
pub fn clear(&self) {
self.entries.clear();
self.memory_bytes.store(0, Ordering::Relaxed);
}
fn evict_one(&self) {
let mut oldest: Option<(u64, u64)> = None;
for entry in self.entries.iter() {
let last_access = entry.value().last_access.load(Ordering::Relaxed);
if oldest.is_none() || last_access < oldest.unwrap().1 {
oldest = Some((*entry.key(), last_access));
}
}
if let Some((key, _)) = oldest {
self.remove(key);
}
}
fn evict_until_fits(&self, size: usize) {
while self.memory_bytes.load(Ordering::Relaxed) + size > self.config.max_memory_bytes {
self.evict_one();
if self.entries.is_empty() {
break;
}
}
}
pub fn stats(&self) -> CacheStats {
let hits = self.hits.load(Ordering::Relaxed);
let misses = self.misses.load(Ordering::Relaxed);
let total = hits + misses;
CacheStats {
hits,
misses,
evictions: self.evictions.load(Ordering::Relaxed),
entries: self.entries.len(),
memory_bytes: self.memory_bytes.load(Ordering::Relaxed),
hit_rate: if total > 0 { hits as f64 / total as f64 } else { 0.0 },
}
}
pub fn record_miss(&self) {
self.misses.fetch_add(1, Ordering::Relaxed);
}
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn test_cache_set_get() {
let cache = RequestCache::new(CacheConfig::default());
let key = RequestCache::generate_key("method", b"args");
cache.set(key, b"response".to_vec(), None);
let result = cache.get(key);
assert_eq!(result, Some(b"response".to_vec()));
}
#[test]
fn test_cache_miss() {
let cache = RequestCache::new(CacheConfig::default());
let result = cache.get(12345);
assert_eq!(result, None);
cache.record_miss();
let stats = cache.stats();
assert_eq!(stats.misses, 1);
}
#[test]
fn test_cache_expiration() {
let config = CacheConfig {
default_ttl: Duration::from_millis(100),
..CacheConfig::default()
};
let cache = RequestCache::new(config);
let key = RequestCache::generate_key("method", b"args");
cache.set(key, b"response".to_vec(), None);
std::thread::sleep(Duration::from_millis(150));
let result = cache.get(key);
assert_eq!(result, None);
}
#[test]
fn test_cache_clear() {
let cache = RequestCache::new(CacheConfig::default());
for i in 0..10 {
cache.set(i, vec![i as u8], None);
}
assert_eq!(cache.stats().entries, 10);
cache.clear();
assert_eq!(cache.stats().entries, 0);
}
}