memlink_runtime/

cache.rs

//! Request/response caching for modules.
//!
//! Caches responses for identical requests to improve performance.

use std::collections::hash_map::RandomState;
use std::hash::{BuildHasher, Hash, Hasher};
use std::sync::atomic::{AtomicU64, AtomicUsize, Ordering};
use std::time::{Duration, Instant};

use dashmap::DashMap;

/// Cache entry with TTL.
#[derive(Debug)]
pub struct CacheEntry {
    /// Cached response data.
    pub data: Vec<u8>,
    /// Creation timestamp.
    pub created_at: Instant,
    /// Time-to-live.
    pub ttl: Duration,
    /// Access count.
    pub access_count: AtomicU64,
    /// Last access timestamp.
    pub last_access: AtomicU64,
}

impl CacheEntry {
    /// Creates a new cache entry.
    pub fn new(data: Vec<u8>, ttl: Duration) -> Self {
        let now = Instant::now().duration_since(Instant::now()).as_secs();
        Self {
            data,
            created_at: Instant::now(),
            ttl,
            access_count: AtomicU64::new(0),
            last_access: AtomicU64::new(now),
        }
    }

    /// Returns whether the entry is expired.
    pub fn is_expired(&self) -> bool {
        self.created_at.elapsed() > self.ttl
    }

    /// Records an access.
    pub fn record_access(&self) {
        self.access_count.fetch_add(1, Ordering::Relaxed);
        let now = Instant::now().duration_since(Instant::now()).as_secs();
        self.last_access.store(now, Ordering::Relaxed);
    }
}

/// Cache configuration.
#[derive(Debug, Clone)]
pub struct CacheConfig {
    /// Default time-to-live for entries.
    pub default_ttl: Duration,
    /// Maximum number of entries.
    pub max_entries: usize,
    /// Maximum memory usage (bytes).
    pub max_memory_bytes: usize,
}

impl Default for CacheConfig {
    fn default() -> Self {
        Self {
            default_ttl: Duration::from_secs(60),
            max_entries: 10000,
            max_memory_bytes: 256 * 1024 * 1024, // 256 MB
        }
    }
}

/// Cache statistics.
#[derive(Debug, Clone)]
pub struct CacheStats {
    pub hits: u64,
    pub misses: u64,
    pub evictions: u64,
    pub entries: usize,
    pub memory_bytes: usize,
    pub hit_rate: f64,
}

/// Request cache for a module.
#[derive(Debug)]
pub struct RequestCache {
    /// Cached entries.
    entries: DashMap<u64, CacheEntry>,
    /// Configuration.
    config: CacheConfig,
    /// Cache hits.
    hits: AtomicU64,
    /// Cache misses.
    misses: AtomicU64,
    /// Evictions.
    evictions: AtomicU64,
    /// Current memory usage.
    memory_bytes: AtomicUsize,
}

impl RequestCache {
    /// Creates a new request cache.
    pub fn new(config: CacheConfig) -> Self {
        Self {
            entries: DashMap::new(),
            config,
            hits: AtomicU64::new(0),
            misses: AtomicU64::new(0),
            evictions: AtomicU64::new(0),
            memory_bytes: AtomicUsize::new(0),
        }
    }

    /// Generates a cache key from method and arguments.
    pub fn generate_key(method: &str, args: &[u8]) -> u64 {
        let hasher = RandomState::new().build_hasher();
        let mut hasher = hasher;
        method.hash(&mut hasher);
        args.hash(&mut hasher);
        hasher.finish()
    }

    /// Gets a cached response.
    pub fn get(&self, key: u64) -> Option<Vec<u8>> {
        let entry = self.entries.get(&key)?;

        if entry.is_expired() {
            drop(entry);
            self.remove(key);
            return None;
        }

        entry.record_access();
        self.hits.fetch_add(1, Ordering::Relaxed);
        Some(entry.data.clone())
    }

    /// Caches a response.
    pub fn set(&self, key: u64, data: Vec<u8>, ttl: Option<Duration>) {
        // Check limits
        if self.entries.len() >= self.config.max_entries {
            self.evict_one();
        }

        let entry_size = data.len();
        if self.memory_bytes.load(Ordering::Relaxed) + entry_size > self.config.max_memory_bytes {
            self.evict_until_fits(entry_size);
        }

        let ttl = ttl.unwrap_or(self.config.default_ttl);
        let entry = CacheEntry::new(data, ttl);
        let entry_size = entry.data.len();

        self.entries.insert(key, entry);
        self.memory_bytes.fetch_add(entry_size, Ordering::Relaxed);
    }

    /// Removes a cached entry.
    pub fn remove(&self, key: u64) -> Option<CacheEntry> {
        if let Some((_, entry)) = self.entries.remove(&key) {
            self.memory_bytes.fetch_sub(entry.data.len(), Ordering::Relaxed);
            self.evictions.fetch_add(1, Ordering::Relaxed);
            Some(entry)
        } else {
            None
        }
    }

    /// Clears all entries.
    pub fn clear(&self) {
        self.entries.clear();
        self.memory_bytes.store(0, Ordering::Relaxed);
    }

    /// Evicts one entry (oldest by last access).
    fn evict_one(&self) {
        let mut oldest: Option<(u64, u64)> = None;

        for entry in self.entries.iter() {
            let last_access = entry.value().last_access.load(Ordering::Relaxed);
            if oldest.is_none() || last_access < oldest.unwrap().1 {
                oldest = Some((*entry.key(), last_access));
            }
        }

        if let Some((key, _)) = oldest {
            self.remove(key);
        }
    }

    /// Evicts entries until there's room for the given size.
    fn evict_until_fits(&self, size: usize) {
        while self.memory_bytes.load(Ordering::Relaxed) + size > self.config.max_memory_bytes {
            self.evict_one();
            if self.entries.is_empty() {
                break;
            }
        }
    }

    /// Returns 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,
            evictions: self.evictions.load(Ordering::Relaxed),
            entries: self.entries.len(),
            memory_bytes: self.memory_bytes.load(Ordering::Relaxed),
            hit_rate: if total > 0 { hits as f64 / total as f64 } else { 0.0 },
        }
    }

    /// Records a cache miss.
    pub fn record_miss(&self) {
        self.misses.fetch_add(1, Ordering::Relaxed);
    }
}

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

    #[test]
    fn test_cache_set_get() {
        let cache = RequestCache::new(CacheConfig::default());
        let key = RequestCache::generate_key("method", b"args");

        cache.set(key, b"response".to_vec(), None);

        let result = cache.get(key);
        assert_eq!(result, Some(b"response".to_vec()));
    }

    #[test]
    fn test_cache_miss() {
        let cache = RequestCache::new(CacheConfig::default());

        let result = cache.get(12345);
        assert_eq!(result, None);

        // Note: record_miss is not automatically called by get()
        // It should be called by the caller when they need to fetch from origin
        cache.record_miss();

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

    #[test]
    fn test_cache_expiration() {
        let config = CacheConfig {
            default_ttl: Duration::from_millis(100),
            ..CacheConfig::default()
        };
        let cache = RequestCache::new(config);
        let key = RequestCache::generate_key("method", b"args");

        cache.set(key, b"response".to_vec(), None);

        // Wait for expiration
        std::thread::sleep(Duration::from_millis(150));

        let result = cache.get(key);
        assert_eq!(result, None);
    }

    #[test]
    fn test_cache_clear() {
        let cache = RequestCache::new(CacheConfig::default());

        for i in 0..10 {
            cache.set(i, vec![i as u8], None);
        }

        assert_eq!(cache.stats().entries, 10);

        cache.clear();

        assert_eq!(cache.stats().entries, 0);
    }
}