use std::cell::RefCell;
use std::collections::{HashMap, VecDeque};
use std::fmt::Debug;
use std::thread::LocalKey;
use crate::{CacheEntry, EvictionPolicy};
#[cfg(feature = "stats")]
use crate::CacheStats;
use crate::utils::{
find_arc_eviction_key, find_min_frequency_key, find_tlru_eviction_key, move_key_to_end,
remove_key_from_cache_local,
};
pub struct ThreadLocalCache<R: 'static> {
pub cache: &'static LocalKey<RefCell<HashMap<String, CacheEntry<R>>>>,
pub order: &'static LocalKey<RefCell<VecDeque<String>>>,
pub limit: Option<usize>,
pub max_memory: Option<usize>,
pub policy: EvictionPolicy,
pub ttl: Option<u64>,
pub frequency_weight: Option<f64>,
pub window_ratio: Option<f64>,
pub sketch_width: Option<usize>,
pub sketch_depth: Option<usize>,
pub decay_interval: Option<u64>,
#[cfg(feature = "stats")]
pub stats: CacheStats,
}
impl<R: Clone + 'static> ThreadLocalCache<R> {
pub fn new(
cache: &'static LocalKey<RefCell<HashMap<String, CacheEntry<R>>>>,
order: &'static LocalKey<RefCell<VecDeque<String>>>,
limit: Option<usize>,
max_memory: Option<usize>,
policy: EvictionPolicy,
ttl: Option<u64>,
frequency_weight: Option<f64>,
window_ratio: Option<f64>,
sketch_width: Option<usize>,
sketch_depth: Option<usize>,
decay_interval: Option<u64>,
) -> Self {
Self {
cache,
order,
limit,
max_memory,
policy,
ttl,
frequency_weight,
window_ratio,
sketch_width,
sketch_depth,
decay_interval,
#[cfg(feature = "stats")]
stats: CacheStats::new(),
}
}
pub fn get(&self, key: &str) -> Option<R> {
let mut expired = false;
let val = self.cache.with(|c| {
let c = c.borrow();
if let Some(entry) = c.get(key) {
if entry.is_expired(self.ttl) {
expired = true;
return None;
}
Some(entry.value.clone())
} else {
None
}
});
if expired {
self.remove_key(key);
#[cfg(feature = "stats")]
self.stats.record_miss();
return None;
}
#[cfg(feature = "stats")]
{
if val.is_some() {
self.stats.record_hit();
} else {
self.stats.record_miss();
}
}
if val.is_some() {
match self.policy {
EvictionPolicy::LRU => {
self.move_to_end(key);
}
EvictionPolicy::LFU => {
self.increment_frequency(key);
}
EvictionPolicy::ARC => {
self.move_to_end(key);
self.increment_frequency(key);
}
EvictionPolicy::TLRU => {
self.move_to_end(key);
self.increment_frequency(key);
}
EvictionPolicy::WTinyLFU => {
self.move_to_end(key);
self.increment_frequency(key);
}
EvictionPolicy::FIFO | EvictionPolicy::Random => {
}
}
}
val
}
fn move_to_end(&self, key: &str) {
self.order.with(|o| {
let mut o = o.borrow_mut();
move_key_to_end(&mut o, key);
});
}
fn increment_frequency(&self, key: &str) {
self.cache.with(|c| {
let mut c = c.borrow_mut();
if let Some(entry) = c.get_mut(key) {
entry.increment_frequency();
}
});
}
pub fn insert(&self, key: &str, value: R) {
let key = key.to_string();
let entry = CacheEntry::new(value);
self.cache.with(|c| {
c.borrow_mut().insert(key.clone(), entry);
});
self.order.with(|o| {
let mut order = o.borrow_mut();
if let Some(pos) = order.iter().position(|k| *k == key) {
order.remove(pos);
}
order.push_back(key.clone());
self.handle_entry_limit_eviction(&mut order);
});
}
#[cfg(feature = "stats")]
pub fn stats(&self) -> &CacheStats {
&self.stats
}
fn remove_key(&self, key: &str) {
self.cache.with(|c| {
self.order.with(|o| {
remove_key_from_cache_local(&mut c.borrow_mut(), &mut o.borrow_mut(), key);
});
});
}
fn handle_entry_limit_eviction(&self, order: &mut VecDeque<String>) {
if let Some(limit) = self.limit {
if order.len() > limit {
match self.policy {
EvictionPolicy::LFU => {
let min_freq_key = self
.cache
.with(|c| find_min_frequency_key(&c.borrow(), order));
if let Some(evict_key) = min_freq_key {
self.remove_key(&evict_key);
}
}
EvictionPolicy::ARC => {
let evict_key = self
.cache
.with(|c| find_arc_eviction_key(&c.borrow(), order.iter().enumerate()));
if let Some(key) = evict_key {
self.remove_key(&key);
}
}
EvictionPolicy::TLRU => {
let evict_key = self.cache.with(|c| {
find_tlru_eviction_key(
&c.borrow(),
order.iter().enumerate(),
self.ttl,
self.frequency_weight,
)
});
if let Some(key) = evict_key {
self.remove_key(&key);
}
}
EvictionPolicy::WTinyLFU => {
let window_ratio = self.window_ratio.unwrap_or(0.20); let window_size = crate::utils::calculate_window_size(limit, window_ratio);
if order.len() <= window_size {
while let Some(evict_key) = order.pop_front() {
let mut removed = false;
self.cache.with(|c| {
let mut cache = c.borrow_mut();
if cache.contains_key(&evict_key) {
cache.remove(&evict_key);
removed = true;
}
});
if removed {
break;
}
}
} else {
let mut evicted = false;
for i in 0..window_size.min(order.len()) {
if let Some(evict_key) = order.get(i) {
let mut removed = false;
self.cache.with(|c| {
let mut cache = c.borrow_mut();
if cache.contains_key(evict_key) {
cache.remove(evict_key);
removed = true;
}
});
if removed {
order.remove(i);
evicted = true;
break;
}
}
}
if !evicted {
let protected_keys: VecDeque<String> =
order.iter().skip(window_size).cloned().collect();
let evict_key = self
.cache
.with(|c| find_min_frequency_key(&c.borrow(), &protected_keys));
if let Some(key) = evict_key {
self.remove_key(&key);
}
}
}
}
EvictionPolicy::Random => {
if !order.is_empty() {
let pos = fastrand::usize(..order.len());
if let Some(evict_key) = order.remove(pos) {
self.cache.with(|c| {
c.borrow_mut().remove(&evict_key);
});
}
}
}
EvictionPolicy::FIFO | EvictionPolicy::LRU => {
while let Some(evict_key) = order.pop_front() {
let mut removed = false;
self.cache.with(|c| {
let mut cache = c.borrow_mut();
if cache.contains_key(&evict_key) {
cache.remove(&evict_key);
removed = true;
}
});
if removed {
break;
}
}
}
}
}
}
}
}
impl<R: Clone + 'static + crate::MemoryEstimator> ThreadLocalCache<R> {
pub fn insert_with_memory(&self, key: &str, value: R) {
let key = key.to_string();
let entry = CacheEntry::new(value);
self.cache.with(|c| {
c.borrow_mut().insert(key.clone(), entry);
});
self.order.with(|o| {
let mut order = o.borrow_mut();
if let Some(pos) = order.iter().position(|k| *k == key) {
order.remove(pos);
}
order.push_back(key.clone());
if let Some(max_mem) = self.max_memory {
let new_value_size = self.cache.with(|c| {
c.borrow()
.get(&key)
.map(|e| e.value.estimate_memory())
.unwrap_or(0)
});
if new_value_size > max_mem {
self.cache.with(|c| {
c.borrow_mut().remove(&key);
});
order.pop_back(); return;
}
loop {
let current_mem = self.cache.with(|c| {
let cache = c.borrow();
cache
.values()
.map(|e| e.value.estimate_memory())
.sum::<usize>()
});
if current_mem <= max_mem {
break;
}
let evicted = match self.policy {
EvictionPolicy::LFU => {
let min_freq_key = self
.cache
.with(|c| find_min_frequency_key(&c.borrow(), &order));
if let Some(evict_key) = min_freq_key {
self.remove_key(&evict_key);
true
} else {
false
}
}
EvictionPolicy::ARC => {
let evict_key = self.cache.with(|c| {
find_arc_eviction_key(&c.borrow(), order.iter().enumerate())
});
if let Some(key) = evict_key {
self.remove_key(&key);
true
} else {
false
}
}
EvictionPolicy::TLRU => {
let evict_key = self.cache.with(|c| {
find_tlru_eviction_key(
&c.borrow(),
order.iter().enumerate(),
self.ttl,
self.frequency_weight,
)
});
if let Some(key) = evict_key {
self.remove_key(&key);
true
} else {
false
}
}
EvictionPolicy::WTinyLFU => {
let evict_key = self
.cache
.with(|c| find_min_frequency_key(&c.borrow(), &order));
if let Some(key) = evict_key {
self.remove_key(&key);
true
} else {
false
}
}
EvictionPolicy::Random => {
if !order.is_empty() {
let pos = fastrand::usize(..order.len());
if let Some(evict_key) = order.remove(pos) {
self.cache.with(|c| {
c.borrow_mut().remove(&evict_key);
});
true
} else {
false
}
} else {
false
}
}
EvictionPolicy::FIFO | EvictionPolicy::LRU => {
if let Some(evict_key) = order.pop_front() {
self.cache.with(|c| {
c.borrow_mut().remove(&evict_key);
});
true
} else {
false
}
}
};
if !evicted {
break; }
}
}
self.handle_entry_limit_eviction(&mut order);
});
}
}
impl<T: Clone + Debug + 'static, E: Clone + Debug + 'static> ThreadLocalCache<Result<T, E>> {
pub fn insert_result(&self, key: &str, value: &Result<T, E>) {
if let Ok(val) = value {
self.insert(key, Ok(val.clone()));
}
}
}
impl<
T: Clone + Debug + 'static + crate::MemoryEstimator,
E: Clone + Debug + 'static + crate::MemoryEstimator,
> ThreadLocalCache<Result<T, E>>
{
pub fn insert_result_with_memory(&self, key: &str, value: &Result<T, E>) {
if let Ok(val) = value {
self.insert_with_memory(key, Ok(val.clone()));
}
}
}
#[cfg(test)]
mod tests {
use super::*;
thread_local! {
static TEST_CACHE: RefCell<HashMap<String, CacheEntry<i32>>> = RefCell::new(HashMap::new());
static TEST_ORDER: RefCell<VecDeque<String>> = RefCell::new(VecDeque::new());
}
fn setup_cache(
limit: Option<usize>,
policy: EvictionPolicy,
ttl: Option<u64>,
) -> ThreadLocalCache<i32> {
TEST_CACHE.with(|c| c.borrow_mut().clear());
TEST_ORDER.with(|o| o.borrow_mut().clear());
ThreadLocalCache::new(
&TEST_CACHE,
&TEST_ORDER,
limit,
None,
policy,
ttl,
None,
None,
None,
None,
None,
)
}
fn setup_cache_with_weight(
limit: Option<usize>,
policy: EvictionPolicy,
ttl: Option<u64>,
frequency_weight: Option<f64>,
) -> ThreadLocalCache<i32> {
TEST_CACHE.with(|c| c.borrow_mut().clear());
TEST_ORDER.with(|o| o.borrow_mut().clear());
ThreadLocalCache::new(
&TEST_CACHE,
&TEST_ORDER,
limit,
None,
policy,
ttl,
frequency_weight,
None,
None,
None,
None,
)
}
#[test]
fn test_basic_insert_get() {
let cache = setup_cache(None, EvictionPolicy::FIFO, None);
cache.insert("key1", 42);
assert_eq!(cache.get("key1"), Some(42));
}
#[test]
fn test_missing_key() {
let cache = setup_cache(None, EvictionPolicy::FIFO, None);
assert_eq!(cache.get("missing"), None);
}
#[test]
fn test_update_existing_key() {
let cache = setup_cache(None, EvictionPolicy::FIFO, None);
cache.insert("key", 1);
cache.insert("key", 2);
assert_eq!(cache.get("key"), Some(2));
}
#[test]
fn test_fifo_eviction() {
let cache = setup_cache(Some(2), EvictionPolicy::FIFO, None);
cache.insert("k1", 1);
cache.insert("k2", 2);
cache.insert("k3", 3);
assert_eq!(cache.get("k1"), None);
assert_eq!(cache.get("k2"), Some(2));
assert_eq!(cache.get("k3"), Some(3));
}
#[test]
fn test_lru_eviction() {
let cache = setup_cache(Some(2), EvictionPolicy::LRU, None);
cache.insert("k1", 1);
cache.insert("k2", 2);
let _ = cache.get("k1"); cache.insert("k3", 3);
assert_eq!(cache.get("k1"), Some(1));
assert_eq!(cache.get("k2"), None);
assert_eq!(cache.get("k3"), Some(3));
}
#[test]
fn test_lru_access_updates_order() {
let cache = setup_cache(Some(3), EvictionPolicy::LRU, None);
cache.insert("k1", 1);
cache.insert("k2", 2);
cache.insert("k3", 3);
let _ = cache.get("k1");
let _ = cache.get("k1");
cache.insert("k4", 4);
assert_eq!(cache.get("k1"), Some(1));
assert_eq!(cache.get("k2"), None);
assert_eq!(cache.get("k3"), Some(3));
assert_eq!(cache.get("k4"), Some(4));
}
#[test]
fn test_result_caching_ok() {
thread_local! {
static RES_CACHE: RefCell<HashMap<String, CacheEntry<Result<i32, String>>>> = RefCell::new(HashMap::new());
static RES_ORDER: RefCell<VecDeque<String>> = RefCell::new(VecDeque::new());
}
let cache = ThreadLocalCache::new(
&RES_CACHE,
&RES_ORDER,
None,
None,
EvictionPolicy::FIFO,
None,
None,
None,
None,
None,
None,
);
let ok_result = Ok(100);
cache.insert_result("success", &ok_result);
assert_eq!(cache.get("success"), Some(Ok(100)));
}
#[test]
fn test_result_caching_err() {
thread_local! {
static RES_CACHE: RefCell<HashMap<String, CacheEntry<Result<i32, String>>>> = RefCell::new(HashMap::new());
static RES_ORDER: RefCell<VecDeque<String>> = RefCell::new(VecDeque::new());
}
let cache = ThreadLocalCache::new(
&RES_CACHE,
&RES_ORDER,
None,
None,
EvictionPolicy::FIFO,
None,
None,
None,
None,
None,
None,
);
let err_result: Result<i32, String> = Err("error".to_string());
cache.insert_result("failure", &err_result);
assert_eq!(cache.get("failure"), None); }
#[test]
fn test_ttl_expiration() {
use std::thread;
use std::time::Duration;
let cache = setup_cache(None, EvictionPolicy::FIFO, Some(1));
cache.insert("expires", 999);
assert_eq!(cache.get("expires"), Some(999));
thread::sleep(Duration::from_secs(2));
assert_eq!(cache.get("expires"), None);
}
#[test]
fn test_no_limit() {
let cache = setup_cache(None, EvictionPolicy::FIFO, None);
for i in 0..1000 {
cache.insert(&format!("key{}", i), i);
}
for i in 0..1000 {
assert_eq!(cache.get(&format!("key{}", i)), Some(i));
}
}
#[test]
#[cfg(feature = "stats")]
fn test_stats_basic() {
let cache = setup_cache(None, EvictionPolicy::FIFO, None);
cache.insert("k1", 1);
cache.insert("k2", 2);
let _ = cache.get("k1"); let _ = cache.get("k2"); let _ = cache.get("k3");
let stats = cache.stats();
assert_eq!(stats.hits(), 2);
assert_eq!(stats.misses(), 1);
assert_eq!(stats.total_accesses(), 3);
assert!((stats.hit_rate() - 0.6666).abs() < 0.001);
}
#[test]
#[cfg(feature = "stats")]
fn test_stats_expired_counts_as_miss() {
use std::thread;
use std::time::Duration;
let cache = setup_cache(None, EvictionPolicy::FIFO, Some(1));
cache.insert("expires", 999);
let _ = cache.get("expires");
assert_eq!(cache.stats().hits(), 1);
assert_eq!(cache.stats().misses(), 0);
thread::sleep(Duration::from_secs(2));
let _ = cache.get("expires");
assert_eq!(cache.stats().hits(), 1);
assert_eq!(cache.stats().misses(), 1);
}
#[test]
#[cfg(feature = "stats")]
fn test_stats_reset() {
let cache = setup_cache(None, EvictionPolicy::FIFO, None);
cache.insert("k1", 1);
let _ = cache.get("k1");
let _ = cache.get("k2");
let stats = cache.stats();
assert_eq!(stats.hits(), 1);
assert_eq!(stats.misses(), 1);
stats.reset();
assert_eq!(stats.hits(), 0);
assert_eq!(stats.misses(), 0);
}
#[test]
#[cfg(feature = "stats")]
fn test_stats_all_hits() {
let cache = setup_cache(None, EvictionPolicy::FIFO, None);
cache.insert("k1", 1);
cache.insert("k2", 2);
for _ in 0..10 {
let _ = cache.get("k1");
let _ = cache.get("k2");
}
let stats = cache.stats();
assert_eq!(stats.hits(), 20);
assert_eq!(stats.misses(), 0);
assert_eq!(stats.hit_rate(), 1.0);
assert_eq!(stats.miss_rate(), 0.0);
}
#[test]
#[cfg(feature = "stats")]
fn test_stats_all_misses() {
let cache = setup_cache(None, EvictionPolicy::FIFO, None);
for i in 0..10 {
let _ = cache.get(&format!("k{}", i));
}
let stats = cache.stats();
assert_eq!(stats.hits(), 0);
assert_eq!(stats.misses(), 10);
assert_eq!(stats.hit_rate(), 0.0);
assert_eq!(stats.miss_rate(), 1.0);
}
#[test]
fn test_tlru_with_frequency_weight_basic() {
let cache = setup_cache_with_weight(Some(10), EvictionPolicy::TLRU, Some(10), Some(1.5));
cache.insert("k1", 1);
cache.insert("k2", 2);
cache.insert("k3", 3);
for _ in 0..5 {
assert_eq!(cache.get("k1"), Some(1));
}
assert_eq!(cache.get("k1"), Some(1));
assert_eq!(cache.get("k2"), Some(2));
assert_eq!(cache.get("k3"), Some(3));
}
#[test]
fn test_tlru_default_frequency_weight_basic() {
let cache = setup_cache_with_weight(Some(10), EvictionPolicy::TLRU, Some(5), None);
cache.insert("k1", 1);
cache.insert("k2", 2);
for _ in 0..3 {
let _ = cache.get("k1");
}
assert_eq!(cache.get("k1"), Some(1));
assert_eq!(cache.get("k2"), Some(2));
}
#[test]
fn test_tlru_no_ttl_with_frequency_weight() {
let cache = setup_cache_with_weight(Some(10), EvictionPolicy::TLRU, None, Some(1.5));
cache.insert("k1", 1);
cache.insert("k2", 2);
cache.insert("k3", 3);
for _ in 0..10 {
let _ = cache.get("k1");
}
assert_eq!(cache.get("k1"), Some(1));
assert_eq!(cache.get("k2"), Some(2));
assert_eq!(cache.get("k3"), Some(3));
}
#[test]
fn test_tlru_frequency_tracking() {
let cache = setup_cache_with_weight(Some(10), EvictionPolicy::TLRU, Some(10), Some(1.0));
cache.insert("k1", 1);
cache.insert("k2", 2);
for _ in 0..5 {
assert_eq!(cache.get("k1"), Some(1));
}
assert_eq!(cache.get("k2"), Some(2));
assert_eq!(cache.get("k1"), Some(1));
assert_eq!(cache.get("k2"), Some(2));
}
#[test]
fn test_tlru_with_different_weights() {
let cache_low =
setup_cache_with_weight(Some(10), EvictionPolicy::TLRU, Some(10), Some(0.3));
let cache_high =
setup_cache_with_weight(Some(10), EvictionPolicy::TLRU, Some(10), Some(2.0));
cache_low.insert("k1", 1);
cache_high.insert("k1", 1);
assert_eq!(cache_low.get("k1"), Some(1));
assert_eq!(cache_high.get("k1"), Some(1));
}
}