use std::collections::HashMap;
use std::hash::Hash;
use std::sync::{RwLock, RwLockReadGuard, RwLockWriteGuard};
use std::time::Duration;
use tokio::time::Instant;
use tracing::{debug, warn};
#[derive(Debug, Clone)]
struct CacheEntry<V> {
value: V,
inserted_at: Instant,
ttl: Duration,
}
impl<V> CacheEntry<V> {
fn new(value: V, ttl: Duration) -> Self {
Self {
value,
inserted_at: Instant::now(),
ttl,
}
}
fn is_expired(&self) -> bool {
self.inserted_at.elapsed() > self.ttl
}
fn age(&self) -> Duration {
self.inserted_at.elapsed()
}
}
pub struct TtlCache<K, V> {
entries: RwLock<HashMap<K, CacheEntry<V>>>,
default_ttl: Duration,
max_capacity: usize,
}
const DEFAULT_MAX_CAPACITY: usize = 1024;
impl<K, V> TtlCache<K, V>
where
K: Eq + Hash + Clone + std::fmt::Debug,
V: Clone,
{
pub fn new(default_ttl: Duration) -> Self {
Self {
entries: RwLock::new(HashMap::new()),
default_ttl,
max_capacity: DEFAULT_MAX_CAPACITY,
}
}
pub fn with_max_capacity(default_ttl: Duration, max_capacity: usize) -> Self {
Self {
entries: RwLock::new(HashMap::new()),
default_ttl,
max_capacity,
}
}
fn read(&self) -> RwLockReadGuard<'_, HashMap<K, CacheEntry<V>>> {
self.entries.read().unwrap_or_else(|poisoned| {
warn!("Cache read lock poisoned, recovering");
poisoned.into_inner()
})
}
fn write(&self) -> RwLockWriteGuard<'_, HashMap<K, CacheEntry<V>>> {
self.entries.write().unwrap_or_else(|poisoned| {
warn!("Cache write lock poisoned, recovering");
poisoned.into_inner()
})
}
pub fn get(&self, key: &K) -> Option<V> {
let entries = self.read();
let entry = entries.get(key)?;
if entry.is_expired() {
debug!(
hit = false,
?key,
age_secs = entry.age().as_secs(),
"cache lookup (expired)"
);
None
} else {
debug!(hit = true, ?key, "cache lookup");
Some(entry.value.clone())
}
}
pub fn insert(&self, key: K, value: V) {
self.insert_with_ttl(key, value, self.default_ttl);
}
pub fn insert_with_ttl(&self, key: K, value: V, ttl: Duration) {
let mut entries = self.write();
if entries.len() >= self.max_capacity && !entries.contains_key(&key) {
let before = entries.len();
entries.retain(|_, entry| !entry.is_expired());
let removed = before - entries.len();
if removed > 0 {
debug!(removed, "Evicted expired entries to make room");
}
if entries.len() >= self.max_capacity {
let low_water = self.max_capacity.saturating_mul(9) / 10;
let to_remove = entries.len().saturating_sub(low_water).max(1);
let mut aged: Vec<(K, Duration)> =
entries.iter().map(|(k, e)| (k.clone(), e.age())).collect();
let take = to_remove.min(aged.len());
if take < aged.len() {
aged.select_nth_unstable_by(take - 1, |a, b| b.1.cmp(&a.1));
}
for (k, _) in aged.into_iter().take(take) {
entries.remove(&k);
}
debug!(evicted = take, "Batch-evicted oldest entries to make room");
}
}
debug!(?key, ttl_secs = ttl.as_secs(), "Inserting cache entry");
entries.insert(key, CacheEntry::new(value, ttl));
}
pub fn remove(&self, key: &K) -> Option<V> {
self.write().remove(key).map(|e| e.value)
}
pub fn len(&self) -> usize {
self.read().len()
}
pub fn is_empty(&self) -> bool {
self.len() == 0
}
pub fn clear(&self) {
self.write().clear();
}
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn test_cache_insert_and_get() {
let cache: TtlCache<String, String> = TtlCache::new(Duration::from_secs(3600));
cache.insert("key".to_string(), "value".to_string());
assert_eq!(cache.get(&"key".to_string()), Some("value".to_string()));
}
#[test]
fn test_cache_get_missing_key() {
let cache: TtlCache<String, String> = TtlCache::new(Duration::from_secs(3600));
assert_eq!(cache.get(&"missing".to_string()), None);
}
#[test]
fn test_cache_expiration() {
let cache: TtlCache<String, String> = TtlCache::new(Duration::from_millis(10));
cache.insert("key".to_string(), "value".to_string());
assert_eq!(cache.get(&"key".to_string()), Some("value".to_string()));
std::thread::sleep(Duration::from_millis(20));
assert_eq!(cache.get(&"key".to_string()), None);
}
#[test]
fn test_cache_remove() {
let cache: TtlCache<String, String> = TtlCache::new(Duration::from_secs(3600));
cache.insert("key".to_string(), "value".to_string());
assert!(cache.get(&"key".to_string()).is_some());
cache.remove(&"key".to_string());
assert!(cache.get(&"key".to_string()).is_none());
}
#[test]
fn capacity_eviction_batches_to_low_water_mark() {
let cap = 100;
let cache: TtlCache<u32, u32> = TtlCache::with_max_capacity(Duration::from_secs(3600), cap);
for i in 0..=cap as u32 {
cache.insert(i, i);
}
assert!(
cache.len() <= (cap * 9 / 10) + 1,
"batch eviction should drop to ~90% low-water, got len {}",
cache.len()
);
assert!(
cache.len() < cap,
"must be below capacity after a batch evict"
);
assert_eq!(cache.get(&(cap as u32)), Some(cap as u32));
assert_eq!(cache.get(&0), None);
}
#[test]
fn capacity_eviction_never_exceeds_capacity_under_churn() {
let cap = 50;
let cache: TtlCache<u32, u32> = TtlCache::with_max_capacity(Duration::from_secs(3600), cap);
for i in 0..1000u32 {
cache.insert(i, i);
assert!(cache.len() <= cap, "len {} exceeded cap {cap}", cache.len());
}
}
#[test]
fn test_cache_clear() {
let cache: TtlCache<String, String> = TtlCache::new(Duration::from_secs(3600));
cache.insert("key1".to_string(), "value1".to_string());
cache.insert("key2".to_string(), "value2".to_string());
assert_eq!(cache.len(), 2);
cache.clear();
assert_eq!(cache.len(), 0);
assert!(cache.is_empty());
}
#[tokio::test(start_paused = true)]
async fn test_entry_expires_after_ttl() {
let cache: TtlCache<String, String> = TtlCache::new(Duration::from_secs(1));
cache.insert("key".to_string(), "value".to_string());
tokio::time::advance(Duration::from_millis(800)).await;
assert!(
cache.get(&"key".to_string()).is_some(),
"entry must not be expired at t=800ms (< 1000ms TTL)"
);
tokio::time::advance(Duration::from_millis(300)).await;
assert!(
cache.get(&"key".to_string()).is_none(),
"entry must be expired at t=1100ms (> 1000ms TTL)"
);
}
}