unistore-cache 0.1.0

In-memory cache capability for UniStore
Documentation
//! 【缓存核心】- LRU 缓存实现
//!
//! 职责:
//! - 实现 LRU(最近最少使用)淘汰策略
//! - 管理缓存条目的存取
//! - 处理 TTL 过期

use crate::config::CacheConfig;
use crate::deps::*;
use crate::entry::CacheEntry;
use crate::stats::CacheStats;
use tracing::debug;

/// LRU 缓存
pub struct Cache<K, V>
where
    K: Eq + Hash + Clone,
{
    /// 缓存数据
    data: RwLock<HashMap<K, CacheEntry<V>>>,
    /// 配置
    config: CacheConfig,
    /// 统计
    stats: CacheStats,
}

impl<K, V> Cache<K, V>
where
    K: Eq + Hash + Clone,
{
    /// 创建新缓存
    pub fn new(config: CacheConfig) -> Self {
        Self {
            data: RwLock::new(HashMap::with_capacity(config.max_capacity)),
            config,
            stats: CacheStats::new(),
        }
    }

    /// 使用默认配置创建
    pub fn with_defaults() -> Self {
        Self::new(CacheConfig::default())
    }

    /// 插入条目
    pub fn insert(&self, key: K, value: V) {
        self.insert_with_ttl(key, value, self.config.default_ttl)
    }

    /// 插入条目(自定义 TTL)
    pub fn insert_with_ttl(&self, key: K, value: V, ttl: Option<Duration>) {
        let entry = CacheEntry::new(value, ttl);
        let mut data = self.data.write();

        // 检查是否需要淘汰
        if data.len() >= self.config.max_capacity && !data.contains_key(&key) {
            self.evict_lru(&mut data);
        }

        data.insert(key, entry);

        if self.config.enable_stats {
            self.stats.record_insert();
        }
    }

    /// 获取值(如果存在且未过期)
    pub fn get(&self, key: &K) -> Option<V>
    where
        V: Clone,
    {
        let mut data = self.data.write();

        if let Some(entry) = data.get_mut(key) {
            if entry.is_expired() {
                data.remove(key);
                if self.config.enable_stats {
                    self.stats.record_expiration();
                    self.stats.record_miss();
                }
                return None;
            }

            entry.touch();
            if self.config.enable_stats {
                self.stats.record_hit();
            }
            return Some(entry.value().clone());
        }

        if self.config.enable_stats {
            self.stats.record_miss();
        }
        None
    }

    /// 获取值的引用(通过闭包访问)
    pub fn with_value<R, F>(&self, key: &K, f: F) -> Option<R>
    where
        F: FnOnce(&V) -> R,
    {
        let mut data = self.data.write();

        if let Some(entry) = data.get_mut(key) {
            if entry.is_expired() {
                data.remove(key);
                if self.config.enable_stats {
                    self.stats.record_expiration();
                    self.stats.record_miss();
                }
                return None;
            }

            entry.touch();
            if self.config.enable_stats {
                self.stats.record_hit();
            }
            return Some(f(entry.value()));
        }

        if self.config.enable_stats {
            self.stats.record_miss();
        }
        None
    }

    /// 检查键是否存在(不更新访问时间)
    pub fn contains_key(&self, key: &K) -> bool {
        let data = self.data.read();
        if let Some(entry) = data.get(key) {
            !entry.is_expired()
        } else {
            false
        }
    }

    /// 删除条目
    pub fn remove(&self, key: &K) -> Option<V> {
        self.data.write().remove(key).map(|e| e.into_value())
    }

    /// 获取或插入
    pub fn get_or_insert_with<F>(&self, key: K, f: F) -> V
    where
        V: Clone,
        F: FnOnce() -> V,
    {
        // 先尝试读取
        if let Some(value) = self.get(&key) {
            return value;
        }

        // 需要插入
        let value = f();
        self.insert(key, value.clone());
        value
    }

    /// 清空缓存
    pub fn clear(&self) {
        self.data.write().clear();
    }

    /// 获取当前大小
    pub fn len(&self) -> usize {
        self.data.read().len()
    }

    /// 检查是否为空
    pub fn is_empty(&self) -> bool {
        self.data.read().is_empty()
    }

    /// 获取统计快照
    pub fn stats(&self) -> crate::stats::CacheStatsSnapshot {
        self.stats.snapshot()
    }

    /// 重置统计
    pub fn reset_stats(&self) {
        self.stats.reset();
    }

    /// 手动清理过期条目
    pub fn cleanup_expired(&self) -> usize {
        let mut data = self.data.write();
        let before = data.len();

        data.retain(|_, entry| {
            let expired = entry.is_expired();
            if expired && self.config.enable_stats {
                self.stats.record_expiration();
            }
            !expired
        });

        let removed = before - data.len();
        if removed > 0 {
            debug!(removed = removed, "清理过期缓存条目");
        }
        removed
    }

    /// LRU 淘汰(内部方法)
    fn evict_lru(&self, data: &mut HashMap<K, CacheEntry<V>>) {
        // 找到最旧的条目
        let to_evict: Vec<K> = data
            .iter()
            .filter(|(_, entry)| entry.is_expired())
            .map(|(k, _)| k.clone())
            .take(self.config.eviction_batch_size)
            .collect();

        // 如果过期条目不够,按 LRU 淘汰
        let mut to_evict = to_evict;
        if to_evict.len() < self.config.eviction_batch_size {
            let needed = self.config.eviction_batch_size - to_evict.len();
            let mut entries: Vec<_> = data
                .iter()
                .filter(|(k, _)| !to_evict.contains(k))
                .map(|(k, e)| (k.clone(), e.last_accessed()))
                .collect();

            entries.sort_by_key(|(_, t)| *t);
            to_evict.extend(entries.into_iter().take(needed).map(|(k, _)| k));
        }

        // 执行淘汰
        for key in to_evict {
            data.remove(&key);
            if self.config.enable_stats {
                self.stats.record_eviction();
            }
        }
    }

    /// 获取所有键
    pub fn keys(&self) -> Vec<K> {
        self.data.read().keys().cloned().collect()
    }

    /// 批量插入
    pub fn insert_many<I>(&self, items: I)
    where
        I: IntoIterator<Item = (K, V)>,
    {
        let mut data = self.data.write();
        for (key, value) in items {
            let entry = CacheEntry::new(value, self.config.default_ttl);
            data.insert(key, entry);
            if self.config.enable_stats {
                self.stats.record_insert();
            }
        }
    }

    /// 更新已有条目的值(如果存在)
    pub fn update<F>(&self, key: &K, f: F) -> bool
    where
        F: FnOnce(&mut V),
    {
        let mut data = self.data.write();
        if let Some(entry) = data.get_mut(key) {
            if entry.is_expired() {
                data.remove(key);
                if self.config.enable_stats {
                    self.stats.record_expiration();
                }
                return false;
            }
            f(entry.value_mut());
            entry.touch();
            true
        } else {
            false
        }
    }
}

impl<K, V> Default for Cache<K, V>
where
    K: Eq + Hash + Clone,
{
    fn default() -> Self {
        Self::with_defaults()
    }
}

#[cfg(test)]
mod tests {
    use super::*;

    #[test]
    fn test_insert_and_get() {
        let cache: Cache<&str, i32> = Cache::with_defaults();
        cache.insert("key1", 42);

        assert_eq!(cache.get(&"key1"), Some(42));
        assert_eq!(cache.get(&"key2"), None);
    }

    #[test]
    fn test_ttl_expiration() {
        let config = CacheConfig::default().default_ttl(Duration::from_millis(10));
        let cache: Cache<&str, i32> = Cache::new(config);

        cache.insert("key1", 42);
        assert!(cache.contains_key(&"key1"));

        std::thread::sleep(Duration::from_millis(20));
        assert!(!cache.contains_key(&"key1"));
    }

    #[test]
    fn test_remove() {
        let cache: Cache<&str, i32> = Cache::with_defaults();
        cache.insert("key1", 42);

        assert_eq!(cache.remove(&"key1"), Some(42));
        assert!(!cache.contains_key(&"key1"));
    }

    #[test]
    fn test_get_or_insert() {
        let cache: Cache<&str, i32> = Cache::with_defaults();

        let value = cache.get_or_insert_with("key1", || 42);
        assert_eq!(value, 42);

        let value = cache.get_or_insert_with("key1", || 100);
        assert_eq!(value, 42); // 应该返回已有值
    }

    #[test]
    fn test_stats() {
        let cache: Cache<&str, i32> = Cache::with_defaults();

        cache.insert("key1", 42);
        cache.get(&"key1");
        cache.get(&"key2");

        let stats = cache.stats();
        assert_eq!(stats.hits, 1);
        assert_eq!(stats.misses, 1);
        assert_eq!(stats.inserts, 1);
    }

    #[test]
    fn test_lru_eviction() {
        let config = CacheConfig::default()
            .max_capacity(3)
            .no_ttl()
            .eviction_batch_size(1);
        let cache: Cache<i32, i32> = Cache::new(config);

        cache.insert(1, 10);
        std::thread::sleep(Duration::from_millis(1));
        cache.insert(2, 20);
        std::thread::sleep(Duration::from_millis(1));
        cache.insert(3, 30);
        std::thread::sleep(Duration::from_millis(1));

        // 访问 key 1,使其成为最近使用
        cache.get(&1);
        std::thread::sleep(Duration::from_millis(1));

        // 插入新条目,应该淘汰 key 2(最久未使用)
        cache.insert(4, 40);

        assert!(cache.contains_key(&1));
        assert!(!cache.contains_key(&2)); // 被淘汰
        assert!(cache.contains_key(&3));
        assert!(cache.contains_key(&4));
    }
}