use crate::value::Value;
use std::collections::HashMap;
use std::sync::RwLock;
use std::time::{Duration, Instant};
#[derive(Debug, Clone, PartialEq, Eq, Hash)]
pub struct CacheKey {
pub table: String,
pub kind: CacheKeyKind,
pub identifier: String,
}
#[derive(Debug, Clone, PartialEq, Eq, Hash)]
pub enum CacheKeyKind {
ByPk,
ByQuery,
ByRelation,
}
impl CacheKey {
pub fn by_pk(table: impl Into<String>, pk: impl std::fmt::Display) -> Self {
Self {
table: table.into(),
kind: CacheKeyKind::ByPk,
identifier: pk.to_string(),
}
}
pub fn by_query(table: impl Into<String>, query_hash: impl std::fmt::Display) -> Self {
Self {
table: table.into(),
kind: CacheKeyKind::ByQuery,
identifier: query_hash.to_string(),
}
}
pub fn by_relation(table: impl Into<String>, relation: impl std::fmt::Display) -> Self {
Self {
table: table.into(),
kind: CacheKeyKind::ByRelation,
identifier: relation.to_string(),
}
}
pub fn to_string_key(&self) -> String {
let kind_str = match self.kind {
CacheKeyKind::ByPk => "pk",
CacheKeyKind::ByQuery => "q",
CacheKeyKind::ByRelation => "rel",
};
format!("l2:{}:{}:{}", self.table, kind_str, self.identifier)
}
}
impl std::fmt::Display for CacheKey {
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
write!(f, "{}", self.to_string_key())
}
}
#[derive(Debug, Clone, Default)]
pub struct L2CacheStats {
pub hits: u64,
pub misses: u64,
pub sets: u64,
pub evictions: u64,
pub size: usize,
}
impl L2CacheStats {
pub fn total_lookups(&self) -> u64 {
self.hits + self.misses
}
pub fn hit_rate(&self) -> f64 {
let total = self.total_lookups();
if total == 0 {
0.0
} else {
self.hits as f64 / total as f64
}
}
pub fn miss_rate(&self) -> f64 {
1.0 - self.hit_rate()
}
pub fn merge(&mut self, other: &L2CacheStats) {
self.hits += other.hits;
self.misses += other.misses;
self.sets += other.sets;
self.evictions += other.evictions;
self.size += other.size;
}
}
#[derive(Debug, Clone)]
struct CacheEntry {
value: Value,
expires_at: Option<Instant>,
}
impl CacheEntry {
fn new(value: Value, ttl: Option<Duration>) -> Self {
let expires_at = ttl.and_then(|d| {
if d == Duration::MAX {
None
} else {
Some(Instant::now() + d)
}
});
Self { value, expires_at }
}
fn is_expired(&self) -> bool {
self.expires_at
.map(|t| t <= Instant::now())
.unwrap_or(false)
}
}
pub struct L2Cache {
data: RwLock<HashMap<String, CacheEntry>>,
table_index: RwLock<HashMap<String, Vec<String>>>,
access_order: RwLock<Vec<String>>,
stats: RwLock<L2CacheStats>,
default_ttl: Option<Duration>,
max_size: usize,
}
impl Default for L2Cache {
fn default() -> Self {
Self::new()
}
}
impl L2Cache {
pub fn new() -> Self {
Self {
data: RwLock::new(HashMap::new()),
table_index: RwLock::new(HashMap::new()),
access_order: RwLock::new(Vec::new()),
stats: RwLock::new(L2CacheStats::default()),
default_ttl: None,
max_size: 10_000,
}
}
pub fn with_default_ttl(mut self, ttl: Duration) -> Self {
self.default_ttl = Some(ttl);
self
}
pub fn with_max_size(mut self, max_size: usize) -> Self {
self.max_size = max_size;
self
}
pub fn put(&self, key: &CacheKey, value: Value, ttl: Option<Duration>) {
let actual_ttl = ttl.or(self.default_ttl);
let entry = CacheEntry::new(value, actual_ttl);
let key_str = key.to_string_key();
let is_new_key = {
let mut data = self.data.write().unwrap();
let exists = data.contains_key(&key_str);
if !exists && data.len() >= self.max_size {
let victim = {
let order = self.access_order.read().unwrap();
order
.iter()
.find(|k| data.get(*k).map(|e| e.is_expired()).unwrap_or(false))
.cloned()
.or_else(|| order.first().cloned())
};
if let Some(victim) = victim {
data.remove(&victim);
let mut order = self.access_order.write().unwrap();
order.retain(|k| k != &victim);
}
}
data.insert(key_str.clone(), entry);
!exists
};
{
let mut order = self.access_order.write().unwrap();
if is_new_key {
order.push(key_str.clone());
} else {
order.retain(|k| k != &key_str);
order.push(key_str.clone());
}
}
{
let mut idx = self.table_index.write().unwrap();
let keys = idx.entry(key.table.clone()).or_default();
if !keys.contains(&key_str) {
keys.push(key_str);
}
}
{
let mut stats = self.stats.write().unwrap();
stats.sets += 1;
}
}
pub fn get(&self, key: &CacheKey) -> Option<Value> {
let key_str = key.to_string_key();
let result = {
let data = self.data.read().ok()?;
if let Some(entry) = data.get(&key_str) {
if entry.is_expired() {
None
} else {
Some(entry.value.clone())
}
} else {
None
}
};
if result.is_some() {
let mut order = self.access_order.write().unwrap();
order.retain(|k| k != &key_str);
order.push(key_str);
}
if let Ok(mut stats) = self.stats.write() {
if result.is_some() {
stats.hits += 1;
} else {
stats.misses += 1;
}
}
result
}
pub fn invalidate(&self, key: &CacheKey) {
let key_str = key.to_string_key();
let removed = {
let mut data = self.data.write().unwrap();
data.remove(&key_str).is_some()
};
if removed {
let mut order = self.access_order.write().unwrap();
order.retain(|k| k != &key_str);
}
if removed {
let mut stats = self.stats.write().unwrap();
stats.evictions += 1;
}
}
pub fn invalidate_table(&self, table: &str) {
let keys_to_remove: Vec<String> = {
let idx = match self.table_index.read() {
Ok(i) => i,
Err(_) => return,
};
idx.get(table).cloned().unwrap_or_default()
};
let mut actually_removed: usize = 0;
{
let mut data = self.data.write().unwrap();
for k in &keys_to_remove {
if data.remove(k).is_some() {
actually_removed += 1;
}
}
}
if actually_removed > 0 {
let mut order = self.access_order.write().unwrap();
order.retain(|k| !keys_to_remove.contains(k));
}
if let Ok(mut idx) = self.table_index.write() {
idx.remove(table);
}
if actually_removed > 0 {
let mut stats = self.stats.write().unwrap();
stats.evictions += actually_removed as u64;
}
}
pub fn clear(&self) {
let removed = {
let mut data = self.data.write().unwrap();
let n = data.len();
data.clear();
n
};
if let Ok(mut order) = self.access_order.write() {
order.clear();
}
if let Ok(mut idx) = self.table_index.write() {
idx.clear();
}
if removed > 0 {
let mut stats = self.stats.write().unwrap();
stats.evictions += removed as u64;
stats.size = 0;
}
}
pub fn size(&self) -> usize {
self.data.read().map(|d| d.len()).unwrap_or(0)
}
pub fn stats(&self) -> L2CacheStats {
let mut s = self.stats.read().map(|s| s.clone()).unwrap_or_default();
s.size = self.size();
s
}
pub fn reset_stats(&self) {
if let Ok(mut stats) = self.stats.write() {
*stats = L2CacheStats::default();
}
}
pub fn contains(&self, key: &CacheKey) -> bool {
let key_str = key.to_string_key();
self.data
.read()
.map(|d| d.get(&key_str).map(|e| !e.is_expired()).unwrap_or(false))
.unwrap_or(false)
}
pub fn evict_expired(&self) -> usize {
let expired_keys: Vec<String> = {
let data = self.data.read().unwrap();
data.iter()
.filter(|(_, e)| e.is_expired())
.map(|(k, _)| k.clone())
.collect()
};
let mut removed = 0;
if !expired_keys.is_empty() {
let mut data = self.data.write().unwrap();
for k in &expired_keys {
if data.remove(k).is_some() {
removed += 1;
}
}
}
if removed > 0 {
let mut order = self.access_order.write().unwrap();
order.retain(|k| !expired_keys.contains(k));
let mut stats = self.stats.write().unwrap();
stats.evictions += removed as u64;
}
removed
}
}
#[cfg(test)]
mod tests {
use super::*;
use crate::Value;
use std::thread;
use std::time::Duration;
#[test]
fn test_cache_key_by_pk() {
let key = CacheKey::by_pk("users", 1);
assert_eq!(key.table, "users");
assert_eq!(key.kind, CacheKeyKind::ByPk);
assert_eq!(key.identifier, "1");
assert_eq!(key.to_string_key(), "l2:users:pk:1");
}
#[test]
fn test_cache_key_by_query() {
let key = CacheKey::by_query("orders", "abc123");
assert_eq!(key.kind, CacheKeyKind::ByQuery);
assert_eq!(key.to_string_key(), "l2:orders:q:abc123");
}
#[test]
fn test_cache_key_by_relation() {
let key = CacheKey::by_relation("users", "posts:1");
assert_eq!(key.kind, CacheKeyKind::ByRelation);
assert_eq!(key.to_string_key(), "l2:users:rel:posts:1");
}
#[test]
fn test_cache_key_equality() {
let k1 = CacheKey::by_pk("users", 1);
let k2 = CacheKey::by_pk("users", 1);
let k3 = CacheKey::by_pk("users", 2);
assert_eq!(k1, k2);
assert_ne!(k1, k3);
}
#[test]
fn test_cache_key_display() {
let key = CacheKey::by_pk("users", 42);
assert_eq!(format!("{}", key), "l2:users:pk:42");
}
#[test]
fn test_stats_hit_rate_empty() {
let stats = L2CacheStats::default();
assert_eq!(stats.hit_rate(), 0.0);
assert_eq!(stats.total_lookups(), 0);
}
#[test]
fn test_stats_hit_rate_calculation() {
let stats = L2CacheStats {
hits: 80,
misses: 20,
..Default::default()
};
assert_eq!(stats.total_lookups(), 100);
assert!((stats.hit_rate() - 0.8).abs() < 0.001);
assert!((stats.miss_rate() - 0.2).abs() < 0.001);
}
#[test]
fn test_stats_merge() {
let mut s1 = L2CacheStats {
hits: 10,
misses: 5,
sets: 15,
evictions: 2,
size: 100,
};
let s2 = L2CacheStats {
hits: 20,
misses: 10,
sets: 30,
evictions: 5,
size: 200,
};
s1.merge(&s2);
assert_eq!(s1.hits, 30);
assert_eq!(s1.misses, 15);
assert_eq!(s1.sets, 45);
assert_eq!(s1.evictions, 7);
assert_eq!(s1.size, 300);
}
#[test]
fn test_put_and_get() {
let cache = L2Cache::new();
let key = CacheKey::by_pk("users", 1);
cache.put(&key, Value::String("Alice".to_string()), None);
let val = cache.get(&key);
assert_eq!(val, Some(Value::String("Alice".to_string())));
}
#[test]
fn test_get_missing_returns_none() {
let cache = L2Cache::new();
let key = CacheKey::by_pk("users", 999);
assert_eq!(cache.get(&key), None);
}
#[test]
fn test_overwrite_existing_key() {
let cache = L2Cache::new();
let key = CacheKey::by_pk("users", 1);
cache.put(&key, Value::String("Alice".to_string()), None);
cache.put(&key, Value::String("Bob".to_string()), None);
assert_eq!(cache.get(&key), Some(Value::String("Bob".to_string())));
}
#[test]
fn test_invalidate_single_key() {
let cache = L2Cache::new();
let key = CacheKey::by_pk("users", 1);
cache.put(&key, Value::I64(42), None);
assert!(cache.get(&key).is_some());
cache.invalidate(&key);
assert!(cache.get(&key).is_none());
}
#[test]
fn test_invalidate_table_removes_all_entries_for_table() {
let cache = L2Cache::new();
let k1 = CacheKey::by_pk("users", 1);
let k2 = CacheKey::by_pk("users", 2);
let k3 = CacheKey::by_query("users", "hash1");
let k4 = CacheKey::by_pk("orders", 1);
cache.put(&k1, Value::I64(1), None);
cache.put(&k2, Value::I64(2), None);
cache.put(&k3, Value::I64(3), None);
cache.put(&k4, Value::I64(4), None);
cache.invalidate_table("users");
assert!(cache.get(&k1).is_none());
assert!(cache.get(&k2).is_none());
assert!(cache.get(&k3).is_none());
assert!(cache.get(&k4).is_some());
}
#[test]
fn test_invalidate_table_no_op_for_unknown_table() {
let cache = L2Cache::new();
let k1 = CacheKey::by_pk("users", 1);
cache.put(&k1, Value::I64(1), None);
cache.invalidate_table("nonexistent");
assert!(cache.get(&k1).is_some());
}
#[test]
fn test_ttl_expiration() {
let cache = L2Cache::new();
let key = CacheKey::by_pk("users", 1);
cache.put(&key, Value::I64(42), Some(Duration::from_millis(50)));
assert!(cache.get(&key).is_some());
thread::sleep(Duration::from_millis(100));
assert!(cache.get(&key).is_none());
}
#[test]
fn test_default_ttl_applied_when_no_explicit_ttl() {
let cache = L2Cache::new().with_default_ttl(Duration::from_millis(50));
let key = CacheKey::by_pk("users", 1);
cache.put(&key, Value::I64(42), None); assert!(cache.get(&key).is_some());
thread::sleep(Duration::from_millis(100));
assert!(cache.get(&key).is_none());
}
#[test]
fn test_explicit_ttl_overrides_default() {
let cache = L2Cache::new().with_default_ttl(Duration::from_millis(50));
let key = CacheKey::by_pk("users", 1);
cache.put(&key, Value::I64(42), Some(Duration::MAX));
thread::sleep(Duration::from_millis(100));
assert!(cache.get(&key).is_some());
}
#[test]
fn test_none_ttl_uses_default_ttl() {
let cache = L2Cache::new().with_default_ttl(Duration::from_millis(50));
let key = CacheKey::by_pk("users", 1);
cache.put(&key, Value::I64(42), None);
assert!(cache.get(&key).is_some());
thread::sleep(Duration::from_millis(100));
assert!(cache.get(&key).is_none());
}
#[test]
fn test_stats_tracks_hits_and_misses() {
let cache = L2Cache::new();
let k1 = CacheKey::by_pk("users", 1);
let k2 = CacheKey::by_pk("users", 2);
cache.put(&k1, Value::I64(1), None);
cache.get(&k1);
cache.get(&k2);
cache.get(&k2);
let stats = cache.stats();
assert_eq!(stats.hits, 1);
assert_eq!(stats.misses, 2);
assert_eq!(stats.sets, 1);
}
#[test]
fn test_stats_tracks_evictions() {
let cache = L2Cache::new();
let k1 = CacheKey::by_pk("users", 1);
let k2 = CacheKey::by_pk("users", 2);
cache.put(&k1, Value::I64(1), None);
cache.put(&k2, Value::I64(2), None);
cache.invalidate(&k1); cache.invalidate_table("users");
let stats = cache.stats();
assert_eq!(stats.evictions, 2);
}
#[test]
fn test_stats_reset() {
let cache = L2Cache::new();
let k1 = CacheKey::by_pk("users", 1);
cache.put(&k1, Value::I64(1), None);
cache.get(&k1);
cache.get(&k1);
let stats_before = cache.stats();
assert!(stats_before.hits > 0);
cache.reset_stats();
let stats_after = cache.stats();
assert_eq!(stats_after.hits, 0);
assert_eq!(stats_after.misses, 0);
assert_eq!(stats_after.sets, 0);
}
#[test]
fn test_max_size_eviction() {
let cache = L2Cache::new().with_max_size(3);
for i in 0..5 {
let k = CacheKey::by_pk("users", i);
cache.put(&k, Value::I64(i), None);
}
let size = cache.size();
assert_eq!(
size, 3,
"size should be exactly max_size after LRU eviction, got {}",
size
);
}
#[test]
fn test_lru_eviction_order() {
let cache = L2Cache::new().with_max_size(3);
let k0 = CacheKey::by_pk("users", 0);
let k1 = CacheKey::by_pk("users", 1);
let k2 = CacheKey::by_pk("users", 2);
let k3 = CacheKey::by_pk("users", 3);
cache.put(&k0, Value::I64(0), None);
cache.put(&k1, Value::I64(1), None);
cache.put(&k2, Value::I64(2), None);
let _ = cache.get(&k0);
cache.put(&k3, Value::I64(3), None);
assert!(
cache.get(&k0).is_some(),
"k0 should survive (recently accessed)"
);
assert!(
cache.get(&k1).is_none(),
"k1 should be evicted (LRU victim)"
);
assert!(cache.get(&k2).is_some(), "k2 should survive");
assert!(
cache.get(&k3).is_some(),
"k3 should survive (just inserted)"
);
}
#[test]
fn test_clear_all() {
let cache = L2Cache::new();
cache.put(&CacheKey::by_pk("users", 1), Value::I64(1), None);
cache.put(&CacheKey::by_pk("users", 2), Value::I64(2), None);
cache.put(&CacheKey::by_pk("orders", 1), Value::I64(3), None);
assert_eq!(cache.size(), 3);
cache.clear();
assert_eq!(cache.size(), 0);
}
#[test]
fn test_contains_does_not_update_stats() {
let cache = L2Cache::new();
let k1 = CacheKey::by_pk("users", 1);
cache.put(&k1, Value::I64(1), None);
let exists = cache.contains(&k1);
assert!(exists);
let stats = cache.stats();
assert_eq!(stats.hits, 0);
assert_eq!(stats.misses, 0);
}
#[test]
fn test_contains_returns_false_for_missing() {
let cache = L2Cache::new();
let k = CacheKey::by_pk("users", 999);
assert!(!cache.contains(&k));
}
#[test]
fn test_contains_returns_false_for_expired() {
let cache = L2Cache::new();
let k = CacheKey::by_pk("users", 1);
cache.put(&k, Value::I64(1), Some(Duration::from_millis(10)));
thread::sleep(Duration::from_millis(50));
assert!(!cache.contains(&k));
}
#[test]
fn test_evict_expired_removes_only_expired_entries() {
let cache = L2Cache::new();
let k1 = CacheKey::by_pk("users", 1);
let k2 = CacheKey::by_pk("users", 2);
cache.put(&k1, Value::I64(1), Some(Duration::from_millis(10)));
cache.put(&k2, Value::I64(2), None);
thread::sleep(Duration::from_millis(50));
let removed = cache.evict_expired();
assert_eq!(removed, 1);
assert!(cache.get(&k1).is_none());
assert!(cache.get(&k2).is_some());
}
#[test]
fn test_evict_expired_returns_zero_if_no_expired() {
let cache = L2Cache::new();
let k1 = CacheKey::by_pk("users", 1);
cache.put(&k1, Value::I64(1), None);
let removed = cache.evict_expired();
assert_eq!(removed, 0);
}
#[test]
fn test_concurrent_access() {
let cache = std::sync::Arc::new(L2Cache::new());
let mut handles = Vec::new();
for i in 0..4 {
let c = cache.clone();
handles.push(thread::spawn(move || {
for j in 0..10 {
let k = CacheKey::by_pk("users", i * 10 + j);
c.put(&k, Value::I64(i * 10 + j), None);
}
}));
}
for h in handles {
h.join().unwrap();
}
assert_eq!(cache.size(), 40);
let mut handles = Vec::new();
for i in 0..4 {
let c = cache.clone();
handles.push(thread::spawn(move || {
for j in 0..10 {
let k = CacheKey::by_pk("users", i * 10 + j);
let v = c.get(&k);
assert!(v.is_some());
}
}));
}
for h in handles {
h.join().unwrap();
}
let stats = cache.stats();
assert_eq!(stats.hits, 40);
}
#[test]
fn test_default() {
let cache = L2Cache::default();
assert_eq!(cache.size(), 0);
}
#[test]
fn test_realistic_scenario() {
let cache = L2Cache::new();
for i in 1..=5 {
cache.put(
&CacheKey::by_pk("users", i),
Value::String(format!("user_{}", i)),
None,
);
}
cache.put(
&CacheKey::by_query("users", "active_users_hash"),
Value::I64(5),
None,
);
for i in 1..=10 {
let _ = cache.get(&CacheKey::by_pk("users", i));
}
let stats = cache.stats();
assert_eq!(stats.hits, 5); assert_eq!(stats.misses, 5); assert_eq!(stats.sets, 6);
cache.invalidate_table("users");
cache.reset_stats();
for i in 1..=5 {
let _ = cache.get(&CacheKey::by_pk("users", i));
}
let stats2 = cache.stats();
assert_eq!(stats2.hits, 0);
assert_eq!(stats2.misses, 5);
}
}