warpdrive_proxy/cache/

memory.rs

//! In-memory LRU cache implementation
//!
//! This module provides a high-performance in-memory cache with:
//! - LRU eviction using randomized sampling
//! - TTL-based expiration
//! - Thread-safe concurrent access via RwLock
//! - Configurable capacity and max item size
//!
//! # Eviction Strategy
//!
//! The cache uses a probabilistic LRU approximation:
//! - When eviction is needed, randomly sample 5 entries
//! - Evict the oldest one (by last_accessed_at)
//! - If an expired entry is found during sampling, evict it immediately
//!
//! This approach is O(5) = O(1) instead of O(n) for scanning all entries,
//! while still evicting items in the oldest ~20% on average.
//!
//! # Performance Characteristics
//!
//! - Get: O(1) hash lookup + O(1) expiry check
//! - Set: O(1) hash insertion + O(5) eviction sampling if needed
//! - Thread-safe: Uses parking_lot::RwLock for better performance than std::Mutex
//! - Zero-copy when possible: Returns references where applicable
//!
//! # Example
//!
//! ```no_run
//! use warpdrive::cache::memory::MemoryCache;
//! use warpdrive::cache::Cache;
//!
//! # async fn example() -> anyhow::Result<()> {
//! // 32 MB cache, 1 MB max item size
//! let cache = MemoryCache::new(32 * 1024 * 1024, 1024 * 1024);
//!
//! // Store data with 60 second TTL
//! cache.set("user:123", b"user_data", 60).await?;
//!
//! // Retrieve data
//! if let Some(data) = cache.get("user:123").await? {
//!     println!("Found {} bytes", data.len());
//! }
//! # Ok(())
//! # }
//! ```

use async_trait::async_trait;
use parking_lot::RwLock;
use std::collections::HashMap;
use std::sync::Arc;
use std::time::{Duration, SystemTime};

use crate::cache::Cache;

/// Hash key type for cache entries
///
/// We use u64 hashing for efficient key storage. The hash is computed
/// from the string key using a fast hashing algorithm.
type CacheKey = u64;

/// Cache entry with metadata
///
/// Stores the value along with access time and expiration metadata
/// required for LRU eviction and TTL support.
#[derive(Clone)]
struct MemoryCacheEntry {
    /// Last time this entry was accessed (for LRU)
    last_accessed_at: SystemTime,

    /// Absolute time when this entry expires
    expires_at: SystemTime,

    /// The cached value (binary data)
    value: Vec<u8>,
}

/// Inner cache state protected by RwLock
///
/// All mutable state is contained within this struct, which is
/// protected by a RwLock for thread-safe concurrent access.
struct MemoryCacheInner {
    /// Maximum total size in bytes
    capacity: usize,

    /// Maximum size of a single item in bytes
    max_item_size: usize,

    /// Current total size of all cached items
    size: usize,

    /// List of cache keys (for random sampling during eviction)
    keys: Vec<CacheKey>,

    /// Map of cache key to entry
    items: HashMap<CacheKey, MemoryCacheEntry>,
}

/// In-memory LRU cache with TTL support
///
/// This cache provides ultra-fast local caching for frequently accessed data.
/// It complements distributed caches (like Redis) by providing:
/// - Sub-microsecond latency (no network round trip)
/// - Automatic memory management via LRU eviction
/// - TTL-based expiration
/// - Thread-safe concurrent access
///
/// # Thread Safety
///
/// The cache uses `Arc<RwLock<...>>` for thread-safe access:
/// - Multiple concurrent reads (via read lock)
/// - Exclusive writes (via write lock)
/// - Clone is cheap (just increments Arc reference count)
///
/// # Size Management
///
/// The cache tracks size in bytes and enforces two limits:
/// - Total capacity: Sum of all cached values
/// - Max item size: Individual value size limit
///
/// When capacity is exceeded, items are evicted using probabilistic LRU.
#[derive(Clone)]
pub struct MemoryCache {
    inner: Arc<RwLock<MemoryCacheInner>>,
}

impl MemoryCache {
    /// Create a new in-memory cache
    ///
    /// # Parameters
    ///
    /// - `capacity`: Maximum total size in bytes (sum of all values)
    /// - `max_item_size`: Maximum size of a single item in bytes
    ///
    /// # Example
    ///
    /// ```
    /// use warpdrive::cache::memory::MemoryCache;
    ///
    /// // 32 MB cache with 1 MB max item size
    /// let cache = MemoryCache::new(32 * 1024 * 1024, 1024 * 1024);
    /// ```
    pub fn new(capacity: usize, max_item_size: usize) -> Self {
        MemoryCache {
            inner: Arc::new(RwLock::new(MemoryCacheInner {
                capacity,
                max_item_size,
                size: 0,
                keys: Vec::new(),
                items: HashMap::new(),
            })),
        }
    }

    /// Hash a string key to a u64
    ///
    /// Uses the default hash algorithm (SipHash 1-3) for security and speed.
    #[inline]
    fn hash_key(key: &str) -> CacheKey {
        use std::collections::hash_map::DefaultHasher;
        use std::hash::{Hash, Hasher};

        let mut hasher = DefaultHasher::new();
        key.hash(&mut hasher);
        hasher.finish()
    }

    /// Get current time (used internally)
    ///
    /// Isolated to a function for potential testing/mocking in the future.
    #[inline]
    fn current_time() -> SystemTime {
        SystemTime::now()
    }

    /// Internal get implementation with write lock
    ///
    /// This is used by the public get() method and requires a write lock
    /// because we update last_accessed_at on cache hits.
    fn get_internal(&self, key: CacheKey) -> Option<Vec<u8>> {
        let mut inner = self.inner.write();
        let now = Self::current_time();

        // Check if entry exists
        let entry = inner.items.get_mut(&key)?;

        // Check if expired
        if entry.expires_at < now {
            tracing::debug!("Cache: entry expired");
            return None;
        }

        // Update last accessed time (for LRU)
        entry.last_accessed_at = now;

        // Return cloned value
        Some(entry.value.clone())
    }

    /// Internal set implementation
    ///
    /// Handles size validation, eviction, and insertion.
    fn set_internal(&self, key: CacheKey, value: Vec<u8>, expires_at: SystemTime) {
        let mut inner = self.inner.write();

        let item_size = value.len();

        // Reject items that are too large
        if item_size > inner.max_item_size || item_size > inner.capacity {
            tracing::debug!(
                "Cache: item is too large to store, len={}, max_item_size={}, capacity={}",
                item_size,
                inner.max_item_size,
                inner.capacity
            );
            return;
        }

        // Evict items until we have enough space
        let limit = inner.capacity - item_size;
        while inner.size > limit {
            tracing::debug!(
                "Cache: evicting item to make space, current_size={}, need_size={}",
                inner.size,
                limit
            );
            Self::evict_oldest_item(&mut inner);
        }

        // If key exists, subtract old value size
        if let Some(existing) = inner.items.get(&key) {
            inner.size -= existing.value.len();
        } else {
            // New key, add to keys list
            inner.keys.push(key);
        }

        // Insert new entry
        inner.items.insert(
            key,
            MemoryCacheEntry {
                last_accessed_at: Self::current_time(),
                expires_at,
                value,
            },
        );

        // Update total size
        inner.size += item_size;

        tracing::debug!(
            "Cache: added item, key={}, size={}, expires_at={:?}",
            key,
            item_size,
            expires_at
        );
    }

    /// Evict the oldest item using randomized sampling
    ///
    /// This implements a probabilistic LRU eviction strategy:
    /// 1. Pick 5 random items from the cache
    /// 2. Find the one with the oldest last_accessed_at
    /// 3. If any expired item is found, evict it immediately
    /// 4. Remove the selected item from cache
    ///
    /// This is O(5) = O(1) instead of O(n), while still providing
    /// good approximation of LRU behavior (evicts from oldest ~20%).
    fn evict_oldest_item(inner: &mut MemoryCacheInner) {
        use rand::Rng;

        if inner.keys.is_empty() {
            return;
        }

        let mut oldest_key: CacheKey = 0;
        let mut oldest_index: usize = 0;
        let mut oldest_time = SystemTime::UNIX_EPOCH;

        let now = Self::current_time();
        let mut rng = rand::rng();

        // Sample up to 5 random items (or fewer if cache is small)
        let sample_size = std::cmp::min(5, inner.keys.len());

        for _ in 0..sample_size {
            let index = rng.random_range(0..inner.keys.len());
            let key = inner.keys[index];

            if let Some(entry) = inner.items.get(&key) {
                // If expired, evict immediately
                if entry.expires_at < now {
                    oldest_key = key;
                    oldest_index = index;
                    break;
                }

                // Track oldest by last_accessed_at
                if oldest_time == SystemTime::UNIX_EPOCH || entry.last_accessed_at < oldest_time {
                    oldest_time = entry.last_accessed_at;
                    oldest_key = key;
                    oldest_index = index;
                }
            }
        }

        // Remove from keys list (swap with last element, then pop)
        let keys_len = inner.keys.len();
        if oldest_index < keys_len {
            inner.keys.swap(oldest_index, keys_len - 1);
            inner.keys.pop();
        }

        // Remove from items map and update size
        if let Some(entry) = inner.items.remove(&oldest_key) {
            inner.size -= entry.value.len();
            tracing::debug!(
                "Cache: evicted item, key={}, size={}",
                oldest_key,
                entry.value.len()
            );
        }
    }

    /// Get cache statistics (for debugging/monitoring)
    ///
    /// Returns current size, capacity, and item count.
    ///
    /// # Example
    ///
    /// ```
    /// # use warpdrive::cache::memory::MemoryCache;
    /// let cache = MemoryCache::new(1024, 256);
    /// let (size, capacity, count) = cache.stats();
    /// println!("Cache: {}/{} bytes, {} items", size, capacity, count);
    /// ```
    pub fn stats(&self) -> (usize, usize, usize) {
        let inner = self.inner.read();
        (inner.size, inner.capacity, inner.items.len())
    }

    /// Clear all entries from the cache
    ///
    /// Resets the cache to empty state.
    ///
    /// # Example
    ///
    /// ```
    /// # use warpdrive::cache::memory::MemoryCache;
    /// let cache = MemoryCache::new(1024, 256);
    /// // ... use cache ...
    /// cache.clear();
    /// ```
    pub fn clear(&self) {
        let mut inner = self.inner.write();
        inner.keys.clear();
        inner.items.clear();
        inner.size = 0;
        tracing::debug!("Cache: cleared all entries");
    }
}

#[async_trait]
impl Cache for MemoryCache {
    async fn get(&self, key: &str) -> anyhow::Result<Option<Vec<u8>>> {
        let hash = Self::hash_key(key);
        Ok(self.get_internal(hash))
    }

    async fn set(&self, key: &str, value: &[u8], ttl_seconds: u64) -> anyhow::Result<()> {
        let hash = Self::hash_key(key);
        let expires_at = if ttl_seconds > 0 {
            Self::current_time() + Duration::from_secs(ttl_seconds)
        } else {
            // No TTL = expire far in future (100 years)
            Self::current_time() + Duration::from_secs(100 * 365 * 24 * 3600)
        };

        self.set_internal(hash, value.to_vec(), expires_at);
        Ok(())
    }

    async fn delete(&self, key: &str) -> anyhow::Result<()> {
        let hash = Self::hash_key(key);
        let mut inner = self.inner.write();

        // Remove from items map
        if let Some(entry) = inner.items.remove(&hash) {
            inner.size -= entry.value.len();

            // Remove from keys list
            if let Some(pos) = inner.keys.iter().position(|&k| k == hash) {
                let keys_len = inner.keys.len();
                inner.keys.swap(pos, keys_len - 1);
                inner.keys.pop();
            }

            tracing::debug!("Cache: deleted item, key={}", hash);
        }

        Ok(())
    }

    async fn exists(&self, key: &str) -> anyhow::Result<bool> {
        let hash = Self::hash_key(key);
        let inner = self.inner.read();
        let now = Self::current_time();

        if let Some(entry) = inner.items.get(&hash) {
            // Check if not expired
            Ok(entry.expires_at >= now)
        } else {
            Ok(false)
        }
    }
}

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

    const KB: usize = 1024;
    const MB: usize = 1024 * KB;

    #[tokio::test]
    async fn test_store_and_retrieve() {
        let cache = MemoryCache::new(32 * MB, 1 * MB);
        cache
            .set("test_key", b"hello world", 30)
            .await
            .expect("set failed");

        let value = cache.get("test_key").await.expect("get failed");
        assert_eq!(value, Some(b"hello world".to_vec()));
    }

    #[tokio::test]
    async fn test_storing_updates_existing_value() {
        let cache = MemoryCache::new(32 * MB, 1 * MB);
        cache.set("key", b"first", 30).await.expect("set failed");
        cache.set("key", b"second", 30).await.expect("set failed");

        let value = cache.get("key").await.expect("get failed");
        assert_eq!(value, Some(b"second".to_vec()));
    }

    #[tokio::test]
    async fn test_storing_existing_value_keeps_size_correct() {
        let cache = MemoryCache::new(32 * MB, 1 * MB);
        cache.set("key", b"first", 30).await.expect("set failed");
        cache.set("key", b"second", 30).await.expect("set failed");

        let (size, _, count) = cache.stats();
        assert_eq!(count, 1);
        assert_eq!(size, 6); // "second".len() == 6
    }

    #[tokio::test]
    async fn test_expiry() {
        let cache = MemoryCache::new(32 * MB, 1 * MB);

        // Set with 1 second TTL
        cache
            .set("key", b"hello world", 1)
            .await
            .expect("set failed");

        // Should be available immediately
        let value = cache.get("key").await.expect("get failed");
        assert_eq!(value, Some(b"hello world".to_vec()));

        // Wait for expiration
        tokio::time::sleep(Duration::from_secs(2)).await;

        // Should be expired
        let value = cache.get("key").await.expect("get failed");
        assert_eq!(value, None);
    }

    #[tokio::test]
    async fn test_does_not_store_items_over_cache_limit() {
        let cache = MemoryCache::new(3 * KB, 50 * KB);

        let payload = vec![0u8; 10 * KB];
        cache.set("key", &payload, 3600).await.expect("set failed");

        let value = cache.get("key").await.expect("get failed");
        assert_eq!(value, None);
    }

    #[tokio::test]
    async fn test_cache_of_size_zero_does_not_store_items() {
        let cache = MemoryCache::new(0, 1 * KB);

        cache
            .set("key", b"no storage", 3600)
            .await
            .expect("set failed");

        let value = cache.get("key").await.expect("get failed");
        assert_eq!(value, None);
    }

    #[tokio::test]
    async fn test_items_are_evicted_to_make_space() {
        let max_cache_size = 10 * KB;
        let cache = MemoryCache::new(max_cache_size, 1 * KB);

        // Fill cache with 20 items of 1KB each
        for i in 0u32..20 {
            let key = format!("key_{}", i);
            let payload = vec![i as u8; 1 * KB];
            cache.set(&key, &payload, 3600).await.expect("set failed");

            // Most recent item should always be available
            let value = cache.get(&key).await.expect("get failed");
            assert_eq!(value, Some(payload));
        }

        // Cache should be at capacity
        let (size, capacity, _) = cache.stats();
        assert_eq!(size, capacity);
        assert_eq!(size, max_cache_size);
    }

    #[tokio::test]
    async fn test_does_not_store_items_over_item_limit() {
        let cache = MemoryCache::new(50 * KB, 3 * KB);

        let payload = vec![0u8; 10 * KB];
        cache.set("key", &payload, 3600).await.expect("set failed");

        let value = cache.get("key").await.expect("get failed");
        assert_eq!(value, None);
    }

    #[tokio::test]
    async fn test_delete() {
        let cache = MemoryCache::new(32 * MB, 1 * MB);
        cache.set("key", b"value", 30).await.expect("set failed");

        // Verify exists
        let exists = cache.exists("key").await.expect("exists failed");
        assert!(exists);

        // Delete
        cache.delete("key").await.expect("delete failed");

        // Verify deleted
        let exists = cache.exists("key").await.expect("exists failed");
        assert!(!exists);

        let value = cache.get("key").await.expect("get failed");
        assert_eq!(value, None);
    }

    #[tokio::test]
    async fn test_exists() {
        let cache = MemoryCache::new(32 * MB, 1 * MB);

        // Non-existent key
        let exists = cache.exists("missing").await.expect("exists failed");
        assert!(!exists);

        // Add key
        cache.set("key", b"value", 30).await.expect("set failed");

        // Should exist
        let exists = cache.exists("key").await.expect("exists failed");
        assert!(exists);
    }

    #[tokio::test]
    async fn test_exists_returns_false_for_expired() {
        let cache = MemoryCache::new(32 * MB, 1 * MB);

        // Set with 1 second TTL
        cache.set("key", b"value", 1).await.expect("set failed");

        // Should exist initially
        let exists = cache.exists("key").await.expect("exists failed");
        assert!(exists);

        // Wait for expiration
        tokio::time::sleep(Duration::from_secs(2)).await;

        // Should not exist after expiration
        let exists = cache.exists("key").await.expect("exists failed");
        assert!(!exists);
    }

    #[tokio::test]
    async fn test_clear() {
        let cache = MemoryCache::new(32 * MB, 1 * MB);

        // Add multiple items
        for i in 0..10 {
            cache
                .set(&format!("key_{}", i), b"value", 30)
                .await
                .expect("set failed");
        }

        let (_, _, count_before) = cache.stats();
        assert_eq!(count_before, 10);

        // Clear cache
        cache.clear();

        let (size, _, count_after) = cache.stats();
        assert_eq!(count_after, 0);
        assert_eq!(size, 0);
    }

    #[tokio::test]
    async fn test_concurrent_access() {
        use std::sync::Arc;
        use tokio::task;

        let cache = Arc::new(MemoryCache::new(32 * MB, 1 * MB));

        // Spawn multiple tasks doing concurrent reads/writes
        let mut handles = vec![];

        for i in 0..10 {
            let cache = Arc::clone(&cache);
            let handle = task::spawn(async move {
                for j in 0..100 {
                    let key = format!("key_{}_{}", i, j);
                    let value = format!("value_{}_{}", i, j);

                    cache.set(&key, value.as_bytes(), 30).await.unwrap();

                    let retrieved = cache.get(&key).await.unwrap();
                    assert_eq!(retrieved, Some(value.as_bytes().to_vec()));
                }
            });
            handles.push(handle);
        }

        // Wait for all tasks to complete
        for handle in handles {
            handle.await.expect("task failed");
        }

        // Verify cache has items
        let (_, _, count) = cache.stats();
        assert!(count > 0);
    }

    #[test]
    fn test_hash_key_consistency() {
        let key1 = "test_key";
        let key2 = "test_key";
        let key3 = "different_key";

        assert_eq!(MemoryCache::hash_key(key1), MemoryCache::hash_key(key2));
        assert_ne!(MemoryCache::hash_key(key1), MemoryCache::hash_key(key3));
    }
}