stateset_nsr/cache/

local.rs

//! # Caching Layer for Performance Optimization
//!
//! Provides thread-safe caching for expensive computations:
//! - Embedding computations
//! - Inference results
//! - Knowledge base queries
//!
//! ## Features
//!
//! - **LRU Eviction**: Least-recently-used eviction policy
//! - **TTL Support**: Time-to-live for cache entries
//! - **Thread-Safe**: Lock-free reads with DashMap
//! - **Metrics**: Cache hit/miss statistics
//!
//! ## Example
//!
//! ```rust
//! use stateset_nsr::cache::{Cache, CacheConfig};
//!
//! let cache = Cache::<String, Vec<f32>>::new(CacheConfig {
//!     max_entries: 10000,
//!     ttl_seconds: 3600,
//!     ..Default::default()
//! });
//!
//! // Store embedding
//! cache.insert("hello".to_string(), vec![0.1, 0.2, 0.3]);
//!
//! // Retrieve (cache hit)
//! if let Some(embedding) = cache.get(&"hello".to_string()) {
//!     println!("Got cached embedding: {:?}", embedding);
//! }
//! ```

use dashmap::DashMap;
use serde::{Deserialize, Serialize};
use std::hash::Hash;
use std::sync::atomic::{AtomicU64, Ordering};
use std::sync::Arc;
use std::time::{Duration, Instant};

/// Cache configuration
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct CacheConfig {
    /// Maximum number of entries in the cache
    pub max_entries: usize,
    /// Time-to-live in seconds (0 = no expiration)
    pub ttl_seconds: u64,
    /// Whether to enable cache statistics
    pub enable_stats: bool,
}

impl Default for CacheConfig {
    fn default() -> Self {
        Self {
            max_entries: 10000,
            ttl_seconds: 3600, // 1 hour
            enable_stats: true,
        }
    }
}

/// Cache entry with metadata
struct CacheEntry<V> {
    value: V,
    created_at: Instant,
    last_accessed: Instant,
    access_count: u64,
}

impl<V: Clone> CacheEntry<V> {
    fn new(value: V) -> Self {
        let now = Instant::now();
        Self {
            value,
            created_at: now,
            last_accessed: now,
            access_count: 1,
        }
    }

    fn is_expired(&self, ttl: Duration) -> bool {
        if ttl.is_zero() {
            return false;
        }
        self.created_at.elapsed() > ttl
    }
}

/// Cache statistics
#[derive(Debug, Clone, Default, Serialize, Deserialize)]
pub struct CacheStats {
    /// Number of cache hits
    pub hits: u64,
    /// Number of cache misses
    pub misses: u64,
    /// Number of entries currently in cache
    pub entries: usize,
    /// Number of evictions
    pub evictions: u64,
    /// Cache hit rate (0.0-1.0)
    pub hit_rate: f64,
}

/// Thread-safe LRU cache with TTL support
pub struct Cache<K, V> {
    data: DashMap<K, CacheEntry<V>>,
    config: CacheConfig,
    ttl: Duration,
    hits: AtomicU64,
    misses: AtomicU64,
    evictions: AtomicU64,
}

impl<K, V> Cache<K, V>
where
    K: Eq + Hash + Clone,
    V: Clone,
{
    /// Create a new cache with the given configuration
    pub fn new(config: CacheConfig) -> Self {
        let ttl = Duration::from_secs(config.ttl_seconds);
        Self {
            data: DashMap::with_capacity(config.max_entries),
            config,
            ttl,
            hits: AtomicU64::new(0),
            misses: AtomicU64::new(0),
            evictions: AtomicU64::new(0),
        }
    }

    /// Get a value from the cache
    pub fn get(&self, key: &K) -> Option<V> {
        if let Some(mut entry) = self.data.get_mut(key) {
            // Check TTL
            if entry.is_expired(self.ttl) {
                drop(entry);
                self.data.remove(key);
                self.misses.fetch_add(1, Ordering::Relaxed);
                return None;
            }

            // Update access metadata
            entry.last_accessed = Instant::now();
            entry.access_count += 1;
            self.hits.fetch_add(1, Ordering::Relaxed);
            Some(entry.value.clone())
        } else {
            self.misses.fetch_add(1, Ordering::Relaxed);
            None
        }
    }

    /// Insert a value into the cache
    pub fn insert(&self, key: K, value: V) {
        // Check if we need to evict
        if self.data.len() >= self.config.max_entries {
            self.evict_lru();
        }

        self.data.insert(key, CacheEntry::new(value));
    }

    /// Get or insert a value using a closure
    pub fn get_or_insert_with<F>(&self, key: K, f: F) -> V
    where
        F: FnOnce() -> V,
    {
        if let Some(value) = self.get(&key) {
            return value;
        }

        let value = f();
        self.insert(key, value.clone());
        value
    }

    /// Remove a value from the cache
    pub fn remove(&self, key: &K) -> Option<V> {
        self.data.remove(key).map(|(_, entry)| entry.value)
    }

    /// Clear all entries from the cache
    pub fn clear(&self) {
        self.data.clear();
    }

    /// Get cache statistics
    pub fn stats(&self) -> CacheStats {
        let hits = self.hits.load(Ordering::Relaxed);
        let misses = self.misses.load(Ordering::Relaxed);
        let total = hits + misses;

        CacheStats {
            hits,
            misses,
            entries: self.data.len(),
            evictions: self.evictions.load(Ordering::Relaxed),
            hit_rate: if total > 0 {
                hits as f64 / total as f64
            } else {
                0.0
            },
        }
    }

    /// Check if cache contains a key
    pub fn contains(&self, key: &K) -> bool {
        if let Some(entry) = self.data.get(key) {
            !entry.is_expired(self.ttl)
        } else {
            false
        }
    }

    /// Get the number of entries in the cache
    pub fn len(&self) -> usize {
        self.data.len()
    }

    /// Check if cache is empty
    pub fn is_empty(&self) -> bool {
        self.data.is_empty()
    }

    /// Evict the least recently used entry
    fn evict_lru(&self) {
        // Find the LRU entry
        let mut lru_key: Option<K> = None;
        let mut oldest_access = Instant::now();

        for entry in self.data.iter() {
            if entry.last_accessed < oldest_access {
                oldest_access = entry.last_accessed;
                lru_key = Some(entry.key().clone());
            }
        }

        // Remove LRU entry
        if let Some(key) = lru_key {
            self.data.remove(&key);
            self.evictions.fetch_add(1, Ordering::Relaxed);
        }
    }

    /// Remove all expired entries
    pub fn cleanup_expired(&self) {
        if self.ttl.is_zero() {
            return;
        }

        let keys_to_remove: Vec<K> = self
            .data
            .iter()
            .filter(|entry| entry.is_expired(self.ttl))
            .map(|entry| entry.key().clone())
            .collect();

        for key in keys_to_remove {
            self.data.remove(&key);
            self.evictions.fetch_add(1, Ordering::Relaxed);
        }
    }
}

/// Embedding cache specialized for text-to-embedding lookups
pub type EmbeddingCache = Cache<String, Vec<f32>>;

/// Inference result cache
pub type InferenceCache = Cache<String, String>;

/// Query result cache for knowledge base queries
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct QueryResult {
    pub entities: Vec<String>,
    pub confidence: f64,
}

pub type QueryCache = Cache<String, QueryResult>;

/// Create a shared embedding cache
pub fn create_embedding_cache(max_size: usize) -> Arc<EmbeddingCache> {
    Arc::new(EmbeddingCache::new(CacheConfig {
        max_entries: max_size,
        ttl_seconds: 86400, // 24 hours for embeddings
        enable_stats: true,
    }))
}

/// Create a shared inference cache
pub fn create_inference_cache(max_size: usize) -> Arc<InferenceCache> {
    Arc::new(InferenceCache::new(CacheConfig {
        max_entries: max_size,
        ttl_seconds: 300, // 5 minutes for inference results
        enable_stats: true,
    }))
}

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

    #[test]
    fn test_cache_basic_operations() {
        let cache = Cache::<String, i32>::new(CacheConfig::default());

        // Insert and get
        cache.insert("key1".to_string(), 42);
        assert_eq!(cache.get(&"key1".to_string()), Some(42));
        assert_eq!(cache.get(&"key2".to_string()), None);

        // Stats
        let stats = cache.stats();
        assert_eq!(stats.hits, 1);
        assert_eq!(stats.misses, 1);
        assert!(stats.hit_rate > 0.4 && stats.hit_rate < 0.6);
    }

    #[test]
    fn test_cache_eviction() {
        let cache = Cache::<i32, i32>::new(CacheConfig {
            max_entries: 3,
            ttl_seconds: 0,
            enable_stats: true,
        });

        // Fill cache
        cache.insert(1, 100);
        cache.insert(2, 200);
        cache.insert(3, 300);

        // Access key 1 to make it recently used
        cache.get(&1);

        // Insert new entry, should evict LRU (key 2 or 3)
        cache.insert(4, 400);

        assert_eq!(cache.len(), 3);
        assert!(cache.contains(&1)); // Recently accessed
        assert!(cache.contains(&4)); // Just inserted
    }

    #[test]
    fn test_get_or_insert_with() {
        let cache = Cache::<String, i32>::new(CacheConfig::default());
        let mut computed = false;

        // First call computes
        let v1 = cache.get_or_insert_with("key".to_string(), || {
            computed = true;
            42
        });
        assert_eq!(v1, 42);
        assert!(computed);

        // Second call uses cache
        computed = false;
        let v2 = cache.get_or_insert_with("key".to_string(), || {
            computed = true;
            99
        });
        assert_eq!(v2, 42);
        assert!(!computed);
    }

    #[test]
    fn test_cache_ttl() {
        let cache = Cache::<String, i32>::new(CacheConfig {
            max_entries: 100,
            ttl_seconds: 0, // Immediate expiration for testing would need time manipulation
            enable_stats: true,
        });

        cache.insert("key".to_string(), 42);
        // With ttl_seconds = 0, entries don't expire
        assert_eq!(cache.get(&"key".to_string()), Some(42));
    }
}