use lru::LruCache;
use moka::sync::Cache as MokaCache;
use std::num::NonZeroUsize;
use std::sync::atomic::{AtomicU64, Ordering};
use std::sync::{Arc, RwLock};
#[derive(Debug, Clone, Copy, PartialEq, Eq, Default)]
pub enum CacheStrategy {
Lru,
#[default]
Moka,
}
#[derive(Debug, Clone)]
pub struct CacheStats {
pub hits: u64,
pub misses: u64,
pub entry_count: u64,
pub capacity: u64,
}
impl CacheStats {
pub fn hit_rate(&self) -> f64 {
let total = self.hits + self.misses;
if total == 0 {
0.0
} else {
(self.hits as f64 / total as f64) * 100.0
}
}
}
pub(crate) trait Cache: Send + Sync {
fn get(&self, key: &str) -> Option<u64>;
fn put(&self, key: String, value: u64);
fn stats(&self) -> CacheStats;
}
pub(crate) struct LruCacheWrapper {
cache: Arc<RwLock<LruCache<String, u64>>>,
capacity: usize,
hits: Arc<AtomicU64>,
misses: Arc<AtomicU64>,
}
impl LruCacheWrapper {
pub fn new(capacity: usize) -> Self {
let capacity_nz = NonZeroUsize::new(capacity).expect("Cache capacity must be non-zero");
Self {
cache: Arc::new(RwLock::new(LruCache::new(capacity_nz))),
capacity,
hits: Arc::new(AtomicU64::new(0)),
misses: Arc::new(AtomicU64::new(0)),
}
}
}
impl Cache for LruCacheWrapper {
fn get(&self, key: &str) -> Option<u64> {
if let Ok(mut cache) = self.cache.write() {
let result = cache.get(key).copied();
if result.is_some() {
self.hits.fetch_add(1, Ordering::Relaxed);
} else {
self.misses.fetch_add(1, Ordering::Relaxed);
}
result
} else {
None
}
}
fn put(&self, key: String, value: u64) {
if let Ok(mut cache) = self.cache.write() {
cache.put(key, value);
}
}
fn stats(&self) -> CacheStats {
let entry_count = if let Ok(cache) = self.cache.read() {
cache.len() as u64
} else {
0
};
CacheStats {
hits: self.hits.load(Ordering::Relaxed),
misses: self.misses.load(Ordering::Relaxed),
entry_count,
capacity: self.capacity as u64,
}
}
}
pub(crate) struct MokaCacheWrapper {
cache: MokaCache<String, u64>,
capacity: usize,
hits: Arc<AtomicU64>,
misses: Arc<AtomicU64>,
}
impl MokaCacheWrapper {
pub fn new(capacity: usize) -> Self {
Self {
cache: MokaCache::builder().max_capacity(capacity as u64).build(),
capacity,
hits: Arc::new(AtomicU64::new(0)),
misses: Arc::new(AtomicU64::new(0)),
}
}
}
impl Cache for MokaCacheWrapper {
fn get(&self, key: &str) -> Option<u64> {
let result = self.cache.get(key);
if result.is_some() {
self.hits.fetch_add(1, Ordering::Relaxed);
} else {
self.misses.fetch_add(1, Ordering::Relaxed);
}
result
}
fn put(&self, key: String, value: u64) {
self.cache.insert(key, value);
}
fn stats(&self) -> CacheStats {
let hits = self.hits.load(Ordering::Relaxed);
let misses = self.misses.load(Ordering::Relaxed);
let entry_count = self.cache.entry_count();
CacheStats {
hits,
misses,
entry_count,
capacity: self.capacity as u64,
}
}
}
pub(crate) fn create_cache(strategy: CacheStrategy, capacity: usize) -> Arc<dyn Cache> {
match strategy {
CacheStrategy::Lru => Arc::new(LruCacheWrapper::new(capacity)),
CacheStrategy::Moka => Arc::new(MokaCacheWrapper::new(capacity)),
}
}