litellm_rs/core/cache/

memory.rs

//! In-memory cache implementation
//!
//! This module provides a high-performance in-memory cache using DashMap
//! for lock-free concurrent access with LRU eviction support.
//!
//! Hot-path eviction metadata is maintained with per-entry atomics in sharded
//! indexes so cache hits and writes do not serialize on a global LRU lock.

use super::types::{AtomicCacheStats, CacheEntry, CacheKey, DualCacheConfig, EvictionPolicy};
use dashmap::DashMap;
use lru::LruCache;
use std::sync::atomic::{AtomicBool, AtomicU64, AtomicUsize, Ordering};
use std::sync::{Arc, Mutex, MutexGuard};
use std::time::{Duration, Instant};
use tokio::sync::watch;
use tracing::{debug, trace};

const EVICTION_SAMPLE_SIZE: usize = 64;
const MIN_ACCESS_SHARDS: usize = 4;
const MAX_ACCESS_SHARDS: usize = 64;

#[derive(Debug)]
struct CacheAccessMeta {
    last_access_tick: AtomicU64,
    access_count: AtomicU64,
}

impl CacheAccessMeta {
    fn new(insert_tick: u64) -> Self {
        Self {
            last_access_tick: AtomicU64::new(insert_tick),
            access_count: AtomicU64::new(0),
        }
    }

    fn record_access(&self, tick: u64) -> u64 {
        self.last_access_tick.store(tick, Ordering::Relaxed);
        self.access_count.fetch_add(1, Ordering::Relaxed) + 1
    }

    fn reset_for_insert(&self, tick: u64) {
        self.last_access_tick.store(tick, Ordering::Relaxed);
        self.access_count.store(0, Ordering::Relaxed);
    }

    fn snapshot(&self) -> (u64, u64) {
        (
            self.last_access_tick.load(Ordering::Relaxed),
            self.access_count.load(Ordering::Relaxed),
        )
    }
}

#[derive(Debug)]
struct EvictionCandidate {
    key: CacheKey,
    last_access_tick: u64,
    access_count: u64,
    created_at: Instant,
    remaining_ttl: Option<Duration>,
}

/// In-memory cache with LRU eviction and TTL expiration
pub struct InMemoryCache<T> {
    /// Main cache storage using DashMap for lock-free access
    cache: Arc<DashMap<CacheKey, CacheEntry<T>>>,
    /// Sharded eviction metadata. Shard count defaults to available CPU
    /// parallelism rounded to a bounded power of two.
    access_meta: Arc<Vec<DashMap<CacheKey, CacheAccessMeta>>>,
    /// Per-shard unique candidate indexes, with one slot per tracked key.
    /// Hits only update atomics; inserts/removals maintain these indexes in O(1).
    /// The LRU order is only a rotating sample order, not the eviction policy.
    access_queue: Arc<Vec<Mutex<LruCache<CacheKey, ()>>>>,
    /// Monotonic logical clock for access ordering.
    access_clock: AtomicU64,
    /// Rotates the first shard inspected by sampled eviction.
    eviction_cursor: AtomicUsize,
    /// Configuration
    config: DualCacheConfig,
    /// Statistics
    stats: Arc<AtomicCacheStats>,
    /// Cleanup has one owner, even if startup is requested more than once.
    cleanup_started: AtomicBool,
    /// Persistent shutdown state, including shutdown before the task first polls.
    shutdown: watch::Sender<bool>,
}

impl<T: Clone + Send + Sync + 'static> InMemoryCache<T> {
    /// Create a new in-memory cache with the given configuration
    pub fn new(config: DualCacheConfig) -> Self {
        Self::with_stats(config, Arc::new(AtomicCacheStats::new()))
    }

    /// Create a new in-memory cache with shared statistics
    pub fn with_stats(config: DualCacheConfig, stats: Arc<AtomicCacheStats>) -> Self {
        let cache = Arc::new(DashMap::with_capacity(config.max_size));
        let (shutdown, _) = watch::channel(false);
        let access_shards = default_access_shard_count();
        let access_meta: Arc<Vec<DashMap<CacheKey, CacheAccessMeta>>> =
            Arc::new((0..access_shards).map(|_| DashMap::new()).collect());
        let access_queue: Arc<Vec<Mutex<LruCache<CacheKey, ()>>>> = Arc::new(
            (0..access_shards)
                .map(|_| Mutex::new(LruCache::unbounded()))
                .collect(),
        );

        Self {
            cache,
            access_meta,
            access_queue,
            access_clock: AtomicU64::new(0),
            eviction_cursor: AtomicUsize::new(0),
            config,
            stats,
            cleanup_started: AtomicBool::new(false),
            shutdown,
        }
    }

    /// Create with default configuration
    pub fn with_defaults() -> Self {
        Self::new(DualCacheConfig::memory_only())
    }

    /// Start the background cleanup task
    pub fn start_cleanup_task(self: &Arc<Self>) {
        if *self.shutdown.borrow() || self.cleanup_started.swap(true, Ordering::AcqRel) {
            return;
        }
        // Background maintenance must not keep a discarded runtime candidate
        // alive. Upgrade only while actually cleaning an externally owned cache.
        let cache = Arc::downgrade(self);
        let mut shutdown = self.shutdown.subscribe();
        let interval = self.config.cleanup_interval;

        tokio::spawn(async move {
            loop {
                if *shutdown.borrow() {
                    break;
                }
                tokio::select! {
                    _ = tokio::time::sleep(interval) => {
                        let Some(cache) = cache.upgrade() else { break };
                        cache.cleanup_expired().await;
                    }
                    _ = shutdown.changed() => {
                        debug!("In-memory cache cleanup task shutting down");
                        break;
                    }
                }
            }
        });
    }

    /// Get a value from the cache
    pub async fn get(&self, key: &CacheKey) -> Option<T> {
        // Atomically remove expired entries to avoid TOCTOU race
        if let Some((_, removed)) = self.cache.remove_if(key, |_k, v| v.is_expired()) {
            self.remove_access_meta(key);
            self.stats.sub_total_size(removed.size_bytes);
            self.stats.set_entry_count(self.cache.len());
            self.stats.record_memory_miss();
            trace!(key = %key, "Cache entry expired");
            return None;
        }

        if let Some(entry) = self.cache.get(key) {
            let value = entry.value.clone();
            drop(entry);
            self.record_access(key);
            self.stats.record_memory_hit();
            trace!(key = %key, "Cache hit");
            Some(value)
        } else {
            self.stats.record_memory_miss();
            trace!(key = %key, "Cache miss");
            None
        }
    }

    /// Peek a value without recording hit/miss statistics.
    ///
    /// Used for under-lock Dual L1 rechecks after an accounting miss so a single
    /// client lookup is not double-counted. Still updates eviction access meta on
    /// a live hit and still removes expired entries.
    pub async fn peek(&self, key: &CacheKey) -> Option<T> {
        if let Some((_, removed)) = self.cache.remove_if(key, |_k, v| v.is_expired()) {
            self.remove_access_meta(key);
            self.stats.sub_total_size(removed.size_bytes);
            self.stats.set_entry_count(self.cache.len());
            return None;
        }

        if let Some(entry) = self.cache.get(key) {
            let value = entry.value.clone();
            drop(entry);
            self.record_access(key);
            Some(value)
        } else {
            None
        }
    }

    /// Get an entry with metadata from the cache
    pub async fn get_entry(&self, key: &CacheKey) -> Option<CacheEntry<T>> {
        // Atomically remove expired entries to avoid TOCTOU race
        if let Some((_, removed)) = self.cache.remove_if(key, |_k, v| v.is_expired()) {
            self.remove_access_meta(key);
            self.stats.sub_total_size(removed.size_bytes);
            self.stats.set_entry_count(self.cache.len());
            self.stats.record_memory_miss();
            return None;
        }

        if let Some(entry) = self.cache.get(key) {
            let mut snapshot = entry.clone();
            drop(entry);
            let access_count = self.record_access(key);
            snapshot.access_count = access_count;
            snapshot.last_accessed = Instant::now();
            self.stats.record_memory_hit();
            Some(snapshot)
        } else {
            self.stats.record_memory_miss();
            None
        }
    }

    /// Peek an entry without recording hit/miss statistics (see [`Self::peek`]).
    pub async fn peek_entry(&self, key: &CacheKey) -> Option<CacheEntry<T>> {
        if let Some((_, removed)) = self.cache.remove_if(key, |_k, v| v.is_expired()) {
            self.remove_access_meta(key);
            self.stats.sub_total_size(removed.size_bytes);
            self.stats.set_entry_count(self.cache.len());
            return None;
        }

        if let Some(entry) = self.cache.get(key) {
            let mut snapshot = entry.clone();
            drop(entry);
            let access_count = self.record_access(key);
            snapshot.access_count = access_count;
            snapshot.last_accessed = Instant::now();
            Some(snapshot)
        } else {
            None
        }
    }

    /// Set a value in the cache with the default TTL
    pub async fn set(&self, key: CacheKey, value: T) {
        self.set_with_ttl(key, value, self.config.default_ttl).await;
    }

    /// Set a value in the cache with a specific TTL
    pub async fn set_with_ttl(&self, key: CacheKey, value: T, ttl: Duration) {
        // Check if we need to evict entries
        if self.cache.len() >= self.config.max_size {
            self.evict_one().await;
        }

        let entry = CacheEntry::new(value, ttl);
        let new_size = entry.size_bytes;
        // Atomic insert returns the old entry if key existed (no TOCTOU gap)
        let old = self.cache.insert(key.clone(), entry);
        self.reset_access_meta_for_insert(&key);
        self.stats.record_write();

        if let Some(old_entry) = old {
            self.stats.sub_total_size(old_entry.size_bytes);
        }

        self.stats.add_total_size(new_size);
        self.stats.set_entry_count(self.cache.len());
        trace!(key = %key, ttl_secs = ttl.as_secs(), "Cache set");
    }

    /// Set a value with size tracking
    pub async fn set_with_size(&self, key: CacheKey, value: T, ttl: Duration, size_bytes: usize) {
        if self.cache.len() >= self.config.max_size {
            self.evict_one().await;
        }

        let entry = CacheEntry::with_size(value, ttl, size_bytes);
        let new_size = entry.size_bytes;
        // Atomic insert returns the old entry if key existed (no TOCTOU gap)
        let old = self.cache.insert(key.clone(), entry);
        self.reset_access_meta_for_insert(&key);
        self.stats.record_write();

        if let Some(old_entry) = old {
            self.stats.sub_total_size(old_entry.size_bytes);
        }

        self.stats.add_total_size(new_size);
        self.stats.set_entry_count(self.cache.len());
    }

    /// Delete a value from the cache
    pub async fn delete(&self, key: &CacheKey) -> bool {
        if let Some((_, removed)) = self.cache.remove(key) {
            self.remove_access_meta(key);
            self.stats.record_deletion();
            self.stats.sub_total_size(removed.size_bytes);
            self.stats.set_entry_count(self.cache.len());
            trace!(key = %key, "Cache delete");
            true
        } else {
            false
        }
    }

    /// Atomically delete `key` only while the live (non-expired) value matches `predicate`.
    ///
    /// Uses DashMap `remove_if` so a concurrent replacement under the same key is
    /// not removed after a stale match decision. Bookkeeping removal checks live
    /// state while holding the candidate-index lock, so a replacement inserted
    /// after `remove_if` keeps its eviction metadata and candidate slot.
    ///
    /// Returns the removed value when a matching entry was deleted.
    pub async fn delete_if<F>(&self, key: &CacheKey, predicate: F) -> Option<T>
    where
        F: Fn(&T) -> bool,
    {
        let removed = self.cache.remove_if(key, |_k, entry| {
            !entry.is_expired() && predicate(&entry.value)
        });
        if let Some((_, removed)) = removed {
            self.remove_access_meta(key);
            self.stats.record_deletion();
            self.stats.sub_total_size(removed.size_bytes);
            self.stats.set_entry_count(self.cache.len());
            trace!(key = %key, "Cache conditional delete");
            Some(removed.value)
        } else {
            None
        }
    }

    /// Check if a key exists in the cache
    pub async fn exists(&self, key: &CacheKey) -> bool {
        // Atomically remove expired entries to avoid TOCTOU race
        if self.cache.remove_if(key, |_k, v| v.is_expired()).is_some() {
            self.remove_access_meta(key);
            self.stats.set_entry_count(self.cache.len());
            return false;
        }
        self.cache.contains_key(key)
    }

    /// Get the remaining TTL for a key
    pub fn ttl(&self, key: &CacheKey) -> Option<Duration> {
        if let Some(entry) = self.cache.get(key) {
            entry.remaining_ttl()
        } else {
            None
        }
    }

    /// Clear all entries from the cache
    pub async fn clear(&self) {
        self.cache.clear();
        for (index, shard) in self.access_meta.iter().enumerate() {
            let mut queue = lock_queue(&self.access_queue[index]);
            queue.clear();
            // A concurrent writer may already have installed a new value after
            // cache.clear(). Preserve its metadata, or let its pending reset
            // install both metadata and a candidate after this lock is released.
            shard.retain(|key, _| self.cache.contains_key(key));
            for meta in shard.iter() {
                queue.put(meta.key().clone(), ());
            }
        }
        // Do not rewind the logical clock while concurrent writers can use it.
        self.eviction_cursor.store(0, Ordering::Relaxed);
        self.stats.reset();
        debug!("Cache cleared");
    }

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

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

    /// Get cache statistics
    pub fn stats(&self) -> Arc<AtomicCacheStats> {
        Arc::clone(&self.stats)
    }

    /// Get all keys in the cache
    pub fn keys(&self) -> Vec<CacheKey> {
        self.cache.iter().map(|r| r.key().clone()).collect()
    }

    /// Shutdown the cache and cleanup task
    pub fn shutdown(&self) {
        self.shutdown.send_replace(true);
    }

    // ==================== Private Methods ====================

    fn next_access_tick(&self) -> u64 {
        self.access_clock.fetch_add(1, Ordering::Relaxed) + 1
    }

    fn access_shard_index(&self, key: &CacheKey) -> usize {
        key.hash_value() as usize % self.access_meta.len()
    }

    fn access_shard(&self, key: &CacheKey) -> &DashMap<CacheKey, CacheAccessMeta> {
        &self.access_meta[self.access_shard_index(key)]
    }

    fn reset_access_meta_for_insert(&self, key: &CacheKey) {
        let index = self.access_shard_index(key);
        let mut queue = lock_queue(&self.access_queue[index]);
        // The value may have been removed while this writer waited for the
        // index. Never recreate bookkeeping for an already absent key.
        if !self.cache.contains_key(key) {
            return;
        }
        let tick = self.next_access_tick();
        let shard = &self.access_meta[index];
        if let Some(meta) = shard.get(key) {
            meta.reset_for_insert(tick);
        } else {
            shard.insert(key.clone(), CacheAccessMeta::new(tick));
        }
        queue.put(key.clone(), ());
    }

    fn record_access(&self, key: &CacheKey) -> u64 {
        let tick = self.next_access_tick();
        // A read can finish after its entry was deleted. Only an insertion may
        // create metadata, so such a read cannot leave an orphan candidate.
        self.access_shard(key)
            .get(key)
            .map_or(0, |meta| meta.record_access(tick))
    }

    fn remove_access_meta(&self, key: &CacheKey) {
        let index = self.access_shard_index(key);
        let mut queue = lock_queue(&self.access_queue[index]);
        // Serialize with insert bookkeeping, then validate current live state.
        // A replacement keeps its slot; a later insertion will create its own.
        if !self.cache.contains_key(key) {
            self.access_meta[index].remove(key);
            queue.pop(key);
        }
    }

    fn eviction_candidates(&self) -> Vec<EvictionCandidate> {
        let target = self.cache.len().min(EVICTION_SAMPLE_SIZE);
        let mut candidates = Vec::with_capacity(target);
        if self.cache.is_empty() {
            return candidates;
        }

        let shard_count = self.access_meta.len();
        let start = self.eviction_cursor.fetch_add(1, Ordering::Relaxed) % shard_count;

        for offset in 0..shard_count {
            if candidates.len() >= target {
                break;
            }

            let shard_index = (start + offset) % shard_count;
            let sampled = self.sample_eviction_keys(shard_index, target - candidates.len());
            if sampled.is_empty() {
                continue;
            }
            let shard = &self.access_meta[(start + offset) % shard_count];

            for key in sampled {
                let snapshot = shard.get(&key).map(|meta| meta.snapshot());
                if let Some((last_access_tick, access_count)) = snapshot {
                    if let Some(entry) = self.cache.get(&key) {
                        if candidates.len() < target {
                            candidates.push(EvictionCandidate {
                                key: key.clone(),
                                last_access_tick,
                                access_count,
                                created_at: entry.created_at,
                                remaining_ttl: entry.remaining_ttl(),
                            });
                        }
                    } else {
                        self.remove_access_meta(&key);
                    }
                }
            }
        }

        candidates
    }

    fn expired_eviction_candidate(&self) -> Option<EvictionCandidate> {
        let shard_count = self.access_meta.len();
        let start = self.eviction_cursor.fetch_add(1, Ordering::Relaxed) % shard_count;
        let mut inspected = 0;

        for offset in 0..shard_count {
            if inspected >= EVICTION_SAMPLE_SIZE {
                break;
            }

            let shard_index = (start + offset) % shard_count;
            let sampled = self.sample_eviction_keys(shard_index, EVICTION_SAMPLE_SIZE - inspected);
            inspected += sampled.len();
            if sampled.is_empty() {
                continue;
            }

            let shard = &self.access_meta[shard_index];
            let mut expired = None;

            for key in sampled {
                let snapshot = shard.get(&key).map(|meta| meta.snapshot());
                if let Some((last_access_tick, access_count)) = snapshot {
                    if let Some(entry) = self.cache.get(&key) {
                        if entry.is_expired() && expired.is_none() {
                            expired = Some(EvictionCandidate {
                                key: key.clone(),
                                last_access_tick,
                                access_count,
                                created_at: entry.created_at,
                                remaining_ttl: entry.remaining_ttl(),
                            });
                        }
                    } else {
                        self.remove_access_meta(&key);
                    }
                }
            }

            if expired.is_some() {
                return expired;
            }
        }

        None
    }

    fn sample_eviction_keys(&self, shard_index: usize, budget: usize) -> Vec<CacheKey> {
        if budget == 0 {
            return Vec::new();
        }

        let mut queue = lock_queue(&self.access_queue[shard_index]);
        let attempts = budget
            .saturating_mul(2)
            .min(EVICTION_SAMPLE_SIZE)
            .min(queue.len());
        let mut sampled = Vec::with_capacity(attempts);
        for _ in 0..attempts {
            let Some((key, ())) = queue.peek_lru() else {
                break;
            };
            let key = key.clone();
            // Rotate under the same lock without allocating another index node.
            // A delete cannot interleave and resurrect a removed candidate.
            queue.promote(&key);
            sampled.push(key);
        }
        sampled
    }

    /// Evict one entry based on the eviction policy
    async fn evict_one(&self) -> bool {
        for _ in 0..EVICTION_SAMPLE_SIZE {
            let removed = match self.config.eviction_policy {
                EvictionPolicy::LRU => self.evict_lru().await,
                EvictionPolicy::LFU => self.evict_lfu().await,
                EvictionPolicy::TTL => self.evict_ttl().await,
                EvictionPolicy::FIFO => self.evict_fifo().await,
            };
            if removed || self.cache.len() < self.config.max_size {
                return removed;
            }
        }

        false
    }

    /// Evict the least recently used entry
    async fn evict_lru(&self) -> bool {
        let candidate = self
            .eviction_candidates()
            .into_iter()
            .min_by_key(|candidate| candidate.last_access_tick);

        if let Some(candidate) = candidate {
            return self.evict_candidate(candidate, "LRU");
        }

        false
    }

    /// Evict the least frequently used entry
    async fn evict_lfu(&self) -> bool {
        let candidate = self
            .eviction_candidates()
            .into_iter()
            .min_by_key(|candidate| (candidate.access_count, candidate.last_access_tick));

        if let Some(candidate) = candidate {
            return self.evict_candidate(candidate, "LFU");
        }

        false
    }

    /// Evict entry with shortest remaining TTL
    async fn evict_ttl(&self) -> bool {
        if let Some(candidate) = self.expired_eviction_candidate() {
            return self.evict_candidate(candidate, "TTL");
        }

        let candidate = self
            .eviction_candidates()
            .into_iter()
            .min_by_key(|candidate| candidate.remaining_ttl.unwrap_or(Duration::ZERO));

        if let Some(candidate) = candidate {
            return self.evict_candidate(candidate, "TTL");
        }

        false
    }

    /// Evict the oldest entry (FIFO)
    async fn evict_fifo(&self) -> bool {
        let candidate = self
            .eviction_candidates()
            .into_iter()
            .min_by_key(|candidate| candidate.created_at);

        if let Some(candidate) = candidate {
            return self.evict_candidate(candidate, "FIFO");
        }

        false
    }

    fn evict_candidate(&self, candidate: EvictionCandidate, policy: &'static str) -> bool {
        if let Some((_, removed)) = self.cache.remove_if(&candidate.key, |_key, entry| {
            entry.created_at == candidate.created_at
        }) {
            self.stats.sub_total_size(removed.size_bytes);
            self.stats.record_eviction();
            self.stats.set_entry_count(self.cache.len());
            trace!(key = %candidate.key, policy = policy, "Cache eviction");
            self.remove_access_meta(&candidate.key);
            return true;
        }

        self.remove_access_meta(&candidate.key);
        false
    }

    /// Clean up expired entries
    async fn cleanup_expired(&self) {
        let mut expired_keys = Vec::new();

        for entry in self.cache.iter() {
            if entry.value().is_expired() {
                expired_keys.push(entry.key().clone());
            }
        }

        let mut count = 0;
        for key in expired_keys {
            if let Some((_, removed)) = self.cache.remove_if(&key, |_key, entry| entry.is_expired())
            {
                self.stats.sub_total_size(removed.size_bytes);
                self.remove_access_meta(&key);
                self.stats.record_eviction();
                count += 1;
            }
        }

        if count > 0 {
            debug!(count = count, "Cleaned up expired entries");
            self.stats.set_entry_count(self.cache.len());
        }
    }
}

fn default_access_shard_count() -> usize {
    std::thread::available_parallelism()
        .map(usize::from)
        .unwrap_or(MIN_ACCESS_SHARDS)
        .next_power_of_two()
        .clamp(MIN_ACCESS_SHARDS, MAX_ACCESS_SHARDS)
}

fn lock_queue(queue: &Mutex<LruCache<CacheKey, ()>>) -> MutexGuard<'_, LruCache<CacheKey, ()>> {
    queue
        .lock()
        .unwrap_or_else(|poisoned| poisoned.into_inner())
}

impl<T> Drop for InMemoryCache<T> {
    fn drop(&mut self) {
        self.shutdown.send_replace(true);
    }
}

#[cfg(test)]
#[path = "memory_tests.rs"]
mod tests;