pdk-cache-lib 1.10.0

// Copyright (c) 2026, Salesforce, Inc.,
// All rights reserved.
// For full license text, see the LICENSE.txt file

//! Implementation of the [`crate::Cache`] using an LRU (Least Recently Used) approach.
//!
//! # Sharding Strategy
//!
//! To optimize serialization/deserialization overhead for large caches, this implementation
//! uses a sharding strategy where the LRU state is split across multiple fixed shards.
//!
//! Each shard maintains its own independent LRU queue, reducing the serialization overhead
//! from O(N) to O(N/S) where S is the number of shards. Keys are deterministically mapped
//! to shards using FNV-1a hashing.
//!
//! ## Trade-offs
//!
//! - **Performance**: 15-16x serialization reduction for large caches (1000+ entries)
//! - **LRU Accuracy**: Per-shard LRU approximates global LRU with uniform key distribution
//! - **Eviction**: Happens independently within each shard's capacity limits
//!
//! # Storage
//!
//! Both the cached values and the per-shard LRU queues are persisted through
//! [`data_storage_lib`] local storage.
//!
//! - **Values** are written unconditionally with [`StoreMode::Always`]: the last writer wins,
//!   caching is best effort.
//! - **Shard queues** are updated with an optimistic read-modify-write loop: the current queue
//!   and its CAS token are read (or its absence detected), the mutation is applied, and the
//!   result is written back with [`StoreMode::Cas`] (or [`StoreMode::Absent`] when the shard did
//!   not exist yet). On a CAS mismatch the loop retries until the CAS succeeds (or the read/write
//!   errors out), since caching is best effort.

use data_storage_lib::{BlockingDataStorage, DataStorageError, LocalDataStorage, StoreMode};
use indexmap::IndexSet;
use log;
use serde::{Deserialize, Serialize};
use std::hash::RandomState;

type LruQueue = IndexSet<String, RandomState>;

use super::{error::CacheError, Cache};

/// Per-shard LRU state.
///
/// Each shard maintains its own independent LRU queue containing only the keys
/// that hash to this shard. This reduces serialization overhead by a factor of
/// the number of shards (typically 15-16x for 16 shards).
#[derive(Debug, Clone, Serialize, Deserialize)]
struct ShardData {
    shard_id: u8,
    lru_queue: LruQueue,
}

impl ShardData {
    fn new(shard_id: u8) -> Self {
        Self {
            shard_id,
            lru_queue: LruQueue::default(),
        }
    }
}

/// Moves `key` to the back (most recently used position) of the queue if present.
/// Returns whether the queue order changed.
fn move_to_back(queue: &mut LruQueue, key: &str) -> bool {
    if let Some(current_index) = queue.get_index_of(key) {
        let last_index = queue.len() - 1;
        if current_index != last_index {
            queue.move_index(current_index, last_index);
            return true;
        }
    }
    false
}

/// Inserts `key` at the back (most recently used) or moves it there if it already
/// exists. Returns whether the queue changed.
fn insert_or_touch(queue: &mut LruQueue, key: &str) -> bool {
    if queue.insert(key.to_string()) {
        true
    } else {
        move_to_back(queue, key)
    }
}

/// An LRU cache that shares global state along all incoming requests of a policy.
///
/// As the name suggests, an LRU (Least Recently Used) eviction strategy is used, meaning that
/// when the maximum number of entries is reached, the least recently accessed cache entries
/// will be the first to be evicted. Both `get()` and `save()` operations update the access time.
///
/// ## Sharding
///
/// Internally, the cache uses a sharding strategy to optimize performance for large caches.
/// The LRU state is split across multiple shards (default: 16), with each shard maintaining
/// its own independent LRU queue. This reduces serialization overhead significantly while
/// maintaining approximate LRU ordering across the cache.
pub(crate) struct LruCache<S: BlockingDataStorage = LocalDataStorage> {
    cache_id: String,
    storage: S,
    max_entries_size: usize,
    num_shards: u8,
}

impl<S: BlockingDataStorage> LruCache<S> {
    pub(crate) fn new(
        cache_id: String,
        max_entries_size: usize,
        storage: S,
        num_shards: u8,
    ) -> Self {
        Self {
            cache_id,
            storage,
            max_entries_size,
            num_shards,
        }
    }

    /// Returns cache ID.
    fn cache_id(&self) -> &str {
        self.cache_id.as_str()
    }

    /// Storage key under which a cached value is stored.
    fn value_key(&self, item_key: &str) -> String {
        format!("value-{item_key}")
    }

    /// Storage key under which a shard's LRU queue is stored.
    fn shard_key(&self, shard_id: u8) -> String {
        format!("shard-{shard_id}")
    }

    /// Calculate shard ID from an item key using FNV-1a hash function.
    ///
    /// FNV-1a provides fast, deterministic hashing with good distribution for string keys.
    /// The result is modulo num_shards to distribute keys across shards.
    fn get_shard_id(&self, item_key: &str) -> u8 {
        const FNV_OFFSET_BASIS: u64 = 0xcbf29ce484222325;
        const FNV_PRIME: u64 = 0x100000001b3;

        let mut hash = FNV_OFFSET_BASIS;
        for byte in item_key.as_bytes() {
            hash ^= *byte as u64;
            hash = hash.wrapping_mul(FNV_PRIME);
        }
        (hash % self.num_shards as u64) as u8
    }

    /// Per-shard capacity: the ceiling of `max_entries / num_shards`.
    ///
    /// Rounding up guarantees every shard has capacity for at least one entry
    /// (no zero-size shards silently dropping keys) and spreads the remainder
    /// across all shards. Total capacity may slightly exceed `max_entries`.
    fn get_shard_max_entries(&self) -> usize {
        self.max_entries_size.div_ceil(self.num_shards as usize)
    }

    /// Applies `mutate` to the shard's LRU queue with an optimistic read-modify-write loop.
    ///
    /// The current shard queue and its CAS token are read (or its absence detected), `mutate`
    /// is applied, the shard is trimmed to its capacity (evicting the least recently used
    /// entries), and the result is written back conditionally. On a CAS mismatch the loop
    /// re-reads and retries until the CAS succeeds (or the read/write errors out), since
    /// caching is best effort.
    ///
    /// Values corresponding to evicted keys are deleted only after the shard write commits.
    fn update_shard<M: Fn(&mut LruQueue) -> bool>(&self, shard_id: u8, mutate: M) {
        let shard_key = self.shard_key(shard_id);
        let capacity = self.get_shard_max_entries();

        loop {
            let (mut shard, mode) = match self.storage.get::<ShardData>(&shard_key) {
                Ok(Some((shard, cas))) => (shard, StoreMode::Cas(cas)),
                Ok(None) => (ShardData::new(shard_id), StoreMode::Absent),
                Err(e) => {
                    log::debug!(
                        "Failed to read shard {} for cache {}: {:?}. Skipping shard update.",
                        shard_id,
                        self.cache_id(),
                        e
                    );
                    return;
                }
            };

            let changed = mutate(&mut shard.lru_queue);
            if !changed {
                return;
            }

            // Trim the shard to its capacity, evicting the least recently used entries (front).
            let mut evicted: Vec<String> = Vec::new();
            while shard.lru_queue.len() > capacity {
                if let Some(least_recently_used) = shard.lru_queue.shift_remove_index(0) {
                    evicted.push(least_recently_used);
                } else {
                    break;
                }
            }

            match self.storage.store(&shard_key, &mode, &shard) {
                Ok(()) => {
                    // Only delete evicted values once the shard state is committed.
                    for evicted_key in evicted {
                        if let Err(e) = self.storage.delete(&self.value_key(&evicted_key)) {
                            log::debug!(
                                "Failed to delete evicted value {} for cache {}: {:?}",
                                evicted_key,
                                self.cache_id(),
                                e
                            );
                        }
                    }
                    return;
                }
                Err(DataStorageError::CasMismatch) => {
                    log::trace!(
                        "CAS mismatch updating shard {} for cache {}, retrying",
                        shard_id,
                        self.cache_id()
                    );
                    continue;
                }
                Err(e) => {
                    log::debug!(
                        "Failed to store shard {} for cache {}: {:?}. Skipping shard update.",
                        shard_id,
                        self.cache_id(),
                        e
                    );
                    return;
                }
            }
        }
    }
}

impl<S: BlockingDataStorage> Cache for LruCache<S> {
    fn save(&self, key: &str, value: Vec<u8>) -> Result<(), CacheError> {
        // Store the value unconditionally: last writer wins.
        if let Err(e) = self
            .storage
            .store(&self.value_key(key), &StoreMode::Always, &value)
        {
            return Err(CacheError::SaveError {
                status: format!("{e:?}"),
                cache_id: String::from(self.cache_id()),
            });
        }

        // Update ONLY the affected shard, inserting the key or moving it to most-recently-used.
        let shard_id = self.get_shard_id(key);
        let key = key.to_string();
        self.update_shard(shard_id, move |queue| insert_or_touch(queue, &key));

        Ok(())
    }

    fn get(&self, key: &str) -> Option<Vec<u8>> {
        match self.storage.get::<Vec<u8>>(&self.value_key(key)) {
            Ok(Some((value, _cas))) => {
                // For LRU, update access order when getting a value.
                let shard_id = self.get_shard_id(key);
                let key = key.to_string();
                self.update_shard(shard_id, move |queue| move_to_back(queue, &key));
                Some(value)
            }
            Ok(None) => None,
            Err(e) => {
                log::debug!(
                    "Failed to read value {} for cache {}: {:?}",
                    key,
                    self.cache_id(),
                    e
                );
                None
            }
        }
    }

    fn delete(&self, key: &str) -> Option<Vec<u8>> {
        log::trace!("Delete item '{key}' from cache");

        // Read the current value to return it, then delete it.
        let existing = match self.storage.get::<Vec<u8>>(&self.value_key(key)) {
            Ok(Some((value, _cas))) => Some(value),
            Ok(None) => None,
            Err(e) => {
                log::debug!(
                    "Failed to read value {} for cache {}: {:?}",
                    key,
                    self.cache_id(),
                    e
                );
                None
            }
        };

        if existing.is_some() {
            if let Err(e) = self.storage.delete(&self.value_key(key)) {
                log::debug!(
                    "Failed to delete value {} for cache {}: {:?}",
                    key,
                    self.cache_id(),
                    e
                );
            }
            let shard_id = self.get_shard_id(key);
            let key = key.to_string();
            self.update_shard(shard_id, move |queue| queue.shift_remove(&key));
        }

        existing
    }

    fn purge(&self) {
        if let Err(e) = self.storage.delete_all() {
            log::debug!("Failed to purge cache {}: {:?}", self.cache_id(), e);
        }
    }
}

#[cfg(test)]
mod tests {
    use std::cell::RefCell;
    use std::collections::HashMap;

    use data_storage_lib::{BlockingDataStorage, DataStorageError, StoreMode};
    use serde::{de::DeserializeOwned, Serialize};

    use super::{insert_or_touch, move_to_back, Cache, LruCache, LruQueue};

    /// In-memory fake of [`BlockingDataStorage`] backed by a HashMap, mirroring the
    /// serialize/CAS semantics of the real local storage so the LRU logic can be tested end to end.
    #[derive(Default)]
    struct FakeStore {
        entries: RefCell<HashMap<String, (Vec<u8>, u32)>>,
        cas_counter: RefCell<u32>,
    }

    impl FakeStore {
        fn next_cas(&self) -> u32 {
            let mut counter = self.cas_counter.borrow_mut();
            *counter += 1;
            *counter
        }

        fn contains(&self, key: &str) -> bool {
            self.entries.borrow().contains_key(key)
        }

        fn len(&self) -> usize {
            self.entries.borrow().len()
        }
    }

    impl BlockingDataStorage for FakeStore {
        fn get_keys(&self) -> Result<Vec<String>, DataStorageError> {
            Ok(self.entries.borrow().keys().cloned().collect())
        }

        fn store<T: Serialize>(
            &self,
            key: &str,
            mode: &StoreMode,
            item: &T,
        ) -> Result<(), DataStorageError> {
            let serialized = serde_fixint::to_vec(item)?;
            let mut entries = self.entries.borrow_mut();
            match mode {
                StoreMode::Always => {}
                StoreMode::Absent => {
                    if entries.contains_key(key) {
                        return Err(DataStorageError::CasMismatch);
                    }
                }
                StoreMode::Cas(expected) => {
                    let expected: u32 = expected.parse()?;
                    match entries.get(key) {
                        Some((_, cas)) if *cas == expected => {}
                        _ => return Err(DataStorageError::CasMismatch),
                    }
                }
            }
            let cas = self.next_cas();
            entries.insert(key.to_string(), (serialized, cas));
            Ok(())
        }

        fn get<T: DeserializeOwned>(
            &self,
            key: &str,
        ) -> Result<Option<(T, String)>, DataStorageError> {
            match self.entries.borrow().get(key) {
                Some((data, cas)) => {
                    let deserialized: T = serde_fixint::from_slice(data)?;
                    Ok(Some((deserialized, cas.to_string())))
                }
                None => Ok(None),
            }
        }

        fn delete(&self, key: &str) -> Result<(), DataStorageError> {
            self.entries.borrow_mut().remove(key);
            Ok(())
        }

        fn delete_all(&self) -> Result<(), DataStorageError> {
            self.entries.borrow_mut().clear();
            Ok(())
        }
    }

    fn new_cache(max_entries: usize, num_shards: u8) -> LruCache<FakeStore> {
        LruCache::new(
            String::from("test-cache"),
            max_entries,
            FakeStore::default(),
            num_shards,
        )
    }

    #[test]
    fn save_and_get_roundtrip() {
        let cache = new_cache(100, 16);
        assert!(cache.save("key1", b"value1".to_vec()).is_ok());
        assert_eq!(cache.get("key1"), Some(b"value1".to_vec()));
    }

    #[test]
    fn save_overwrites_existing_value() {
        let cache = new_cache(100, 16);
        assert!(cache.save("key1", b"value1".to_vec()).is_ok());
        assert!(cache.save("key1", b"value2".to_vec()).is_ok());
        assert_eq!(cache.get("key1"), Some(b"value2".to_vec()));
    }

    #[test]
    fn get_missing_returns_none() {
        let cache = new_cache(100, 16);
        assert!(cache.get("missing").is_none());
    }

    #[test]
    fn delete_returns_value_and_removes_it() {
        let cache = new_cache(100, 16);
        cache.save("key1", b"value1".to_vec()).unwrap();
        assert_eq!(cache.delete("key1"), Some(b"value1".to_vec()));
        assert!(cache.get("key1").is_none());
    }

    #[test]
    fn delete_missing_returns_none() {
        let cache = new_cache(100, 16);
        assert!(cache.delete("missing").is_none());
    }

    #[test]
    fn eviction_removes_least_recently_used_within_shard() {
        // Single shard, capacity 2: forces global exact LRU.
        let cache = new_cache(2, 1);

        cache.save("a", b"1".to_vec()).unwrap();
        cache.save("b", b"2".to_vec()).unwrap();
        // Access "a" so "b" becomes least recently used.
        assert_eq!(cache.get("a"), Some(b"1".to_vec()));
        // Insert "c": capacity exceeded, "b" (LRU) evicted.
        cache.save("c", b"3".to_vec()).unwrap();

        assert_eq!(cache.get("a"), Some(b"1".to_vec()));
        assert!(cache.get("b").is_none());
        assert_eq!(cache.get("c"), Some(b"3".to_vec()));
    }

    #[test]
    fn save_updates_lru_order_on_reinsert() {
        let cache = new_cache(2, 1);

        cache.save("a", b"1".to_vec()).unwrap();
        cache.save("b", b"2".to_vec()).unwrap();
        // Re-save "a" so it becomes most recently used; "b" is now LRU.
        cache.save("a", b"11".to_vec()).unwrap();
        // Insert "c": "b" evicted.
        cache.save("c", b"3".to_vec()).unwrap();

        assert_eq!(cache.get("a"), Some(b"11".to_vec()));
        assert!(cache.get("b").is_none());
        assert_eq!(cache.get("c"), Some(b"3".to_vec()));
    }

    #[test]
    fn purge_clears_everything() {
        let cache = new_cache(100, 16);
        cache.save("key1", b"value1".to_vec()).unwrap();
        cache.save("key2", b"value2".to_vec()).unwrap();

        cache.purge();

        assert!(cache.get("key1").is_none());
        assert!(cache.get("key2").is_none());
        assert_eq!(cache.storage.len(), 0);
    }

    #[test]
    fn save_persists_value_and_shard() {
        let cache = new_cache(100, 16);
        cache.save("key1", b"value1".to_vec()).unwrap();

        // The value and exactly one shard queue are persisted.
        assert!(cache.storage.contains(&cache.value_key("key1")));
        let shard_id = cache.get_shard_id("key1");
        assert!(cache.storage.contains(&cache.shard_key(shard_id)));
    }

    #[test]
    fn save_returns_error_when_value_store_fails() {
        // A store that always fails to persist values.
        struct FailingStore;
        impl BlockingDataStorage for FailingStore {
            fn get_keys(&self) -> Result<Vec<String>, DataStorageError> {
                Ok(vec![])
            }
            fn store<T: Serialize>(
                &self,
                _key: &str,
                _mode: &StoreMode,
                _item: &T,
            ) -> Result<(), DataStorageError> {
                Err(DataStorageError::Unexpected("boom".to_string()))
            }
            fn get<T: DeserializeOwned>(
                &self,
                _key: &str,
            ) -> Result<Option<(T, String)>, DataStorageError> {
                Ok(None)
            }
            fn delete(&self, _key: &str) -> Result<(), DataStorageError> {
                Ok(())
            }
            fn delete_all(&self) -> Result<(), DataStorageError> {
                Ok(())
            }
        }

        let cache = LruCache::new(String::from("test-cache"), 100, FailingStore, 16);
        assert!(cache.save("key1", b"value1".to_vec()).is_err());
    }

    #[test]
    fn shard_capacity_distribution() {
        let cache = new_cache(100, 16);
        // 100 / 16 ceiling = 7 for all shards.
        assert_eq!(cache.get_shard_max_entries(), 7);

        let cache = new_cache(64, 16);
        // 64 / 16 = 4 for all shards.
        assert_eq!(cache.get_shard_max_entries(), 4);
    }

    #[test]
    fn shard_capacity_small_cache() {
        let cache = new_cache(10, 16);
        // 10 / 16 ceiling = 1 for all shards (no zero-size shards).
        assert_eq!(cache.get_shard_max_entries(), 1);
    }

    #[test]
    fn shard_id_is_deterministic() {
        let cache = new_cache(100, 16);
        assert_eq!(
            cache.get_shard_id("test-key-123"),
            cache.get_shard_id("test-key-123")
        );
        assert!(cache.get_shard_id("test-key-456") < 16);
    }

    #[test]
    fn shard_id_distribution() {
        let cache = new_cache(100, 16);
        let mut shard_counts = vec![0; 16];
        for i in 0..1000 {
            let shard_id = cache.get_shard_id(&format!("test-key-{i}"));
            assert!(shard_id < 16);
            shard_counts[shard_id as usize] += 1;
        }
        for count in shard_counts {
            assert!(
                (30..=100).contains(&count),
                "Shard count {} outside expected range",
                count
            );
        }
    }

    #[test]
    fn move_to_back_reorders_queue() {
        let mut queue: LruQueue = ["a", "b", "c"].iter().map(|s| s.to_string()).collect();
        move_to_back(&mut queue, "a");
        assert_eq!(queue.last().unwrap(), "a");
        assert_eq!(queue.len(), 3);
    }

    #[test]
    fn insert_or_touch_inserts_new_and_touches_existing() {
        let mut queue: LruQueue = ["a", "b"].iter().map(|s| s.to_string()).collect();
        insert_or_touch(&mut queue, "c");
        assert_eq!(queue.last().unwrap(), "c");
        assert_eq!(queue.len(), 3);

        insert_or_touch(&mut queue, "a");
        assert_eq!(queue.last().unwrap(), "a");
        assert_eq!(queue.len(), 3);
    }
}