pdk-cache-lib 1.10.0

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

//! Primitives for building a cache
//!
//! The entrypoint of this crate is the [`CacheBuilder`] struct.

use super::fifo_cache::FifoCache;
use super::lru_cache::LruCache;
use super::ttl_cache::TtlCache;
use super::Cache;
use data_storage_lib::DataStorageBuilder;
use pdk_core::classy::extract::context::ConfigureContext;
use pdk_core::classy::extract::{Extract, FromContext};
use pdk_core::classy::{Clock, SharedData};
use pdk_core::policy_context::api::Metadata;
use std::convert::Infallible;
use std::rc::Rc;
use std::time::Duration;
cfg_if::cfg_if! {
    if #[cfg(test)] {
        use crate::mocks::MockConcurrentSharedData as ConcurrentSharedData;
    } else {
        use pdk_core::classy::shared_data::concurrent_shared_data::ConcurrentSharedData;
    }
}

const DEFAULT_MAX_ENTRIES: usize = 100;

/// Default number of shards used by the sharded LRU strategy.
const DEFAULT_LRU_SHARDS: u8 = 16;

/// Eviction strategy for the cache.
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
#[non_exhaustive]
pub enum EvictionStrategy {
    /// First-In-First-Out: evicts the oldest entries first (by insertion time).
    Fifo,
    /// Sharded, approximate Least Recently Used eviction. See [`LruStrategy`].
    ///
    /// This is **not** a globally-exact LRU. To reduce serialization overhead on
    /// large caches, the entries are split across independent shards (see
    /// [`LruStrategy::shards`]) and eviction happens per shard: when a shard is
    /// full, that shard evicts its own least recently accessed entry. With keys
    /// spread across shards, the globally least recently used entry is only
    /// evicted approximately, since eviction only ever considers one shard's
    /// entries at a time. Both `get()` and `save()` update the access time.
    ///
    /// With a single shard ([`LruStrategy::shards`]`(1)`) the whole cache lives in
    /// one queue, so eviction is exact, global LRU at the cost of no sharding.
    Lru(LruStrategy),
}

/// Configuration for the sharded LRU [`EvictionStrategy`].
///
/// Use [`LruStrategy::default`] for the default shard count, or override it with
/// [`LruStrategy::with_shards`]. The shard count controls how the LRU state is partitioned;
/// see [`EvictionStrategy::Lru`] for how sharding affects eviction accuracy.
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub struct LruStrategy {
    shards: u8,
}

impl Default for LruStrategy {
    fn default() -> Self {
        Self {
            shards: DEFAULT_LRU_SHARDS,
        }
    }
}

impl LruStrategy {
    /// Sets the number of shards, rounded to the nearest power of two
    /// (minimum 1). Rounding keeps hash-based key distribution well-balanced
    /// without rejecting arbitrary counts.
    pub fn with_shards(mut self, shards: u8) -> Self {
        self.shards = nearest_power_of_two(shards);
        self
    }

    /// Number of shards this strategy partitions its state into.
    pub fn shards(&self) -> u8 {
        self.shards
    }
}

/// Rounds `count` to the nearest power of two (ties round down), minimum 1.
fn nearest_power_of_two(count: u8) -> u8 {
    if count <= 1 {
        return 1;
    }
    // The next power of two above `count` overflows u8 for counts > 128, and 128 is
    // the largest power of two that fits, so those counts clamp to 128.
    let Some(higher) = count.checked_next_power_of_two() else {
        return 128;
    };
    if higher == count {
        return count;
    }
    let lower = higher >> 1;
    if count - lower <= higher - count {
        lower
    } else {
        higher
    }
}

/// Internal enum that wraps different cache implementations.
/// This allows returning a concrete type instead of Box<dyn Cache>.
enum CacheImpl {
    Fifo(FifoCache),
    Lru(LruCache),
    Ttl(Box<TtlCache<CacheImpl>>),
}

impl Cache for CacheImpl {
    fn save(&self, key: &str, value: Vec<u8>) -> Result<(), super::error::CacheError> {
        match self {
            CacheImpl::Fifo(cache) => cache.save(key, value),
            CacheImpl::Lru(cache) => cache.save(key, value),
            CacheImpl::Ttl(cache) => cache.save(key, value),
        }
    }

    fn get(&self, key: &str) -> Option<Vec<u8>> {
        match self {
            CacheImpl::Fifo(cache) => cache.get(key),
            CacheImpl::Lru(cache) => cache.get(key),
            CacheImpl::Ttl(cache) => cache.get(key),
        }
    }

    fn delete(&self, key: &str) -> Option<Vec<u8>> {
        match self {
            CacheImpl::Fifo(cache) => cache.delete(key),
            CacheImpl::Lru(cache) => cache.delete(key),
            CacheImpl::Ttl(cache) => cache.delete(key),
        }
    }

    fn purge(&self) {
        match self {
            CacheImpl::Fifo(cache) => cache.purge(),
            CacheImpl::Lru(cache) => cache.purge(),
            CacheImpl::Ttl(cache) => cache.purge(),
        }
    }
}

/// Represents the base builder for a [`Cache`] instance.
pub struct CacheBuilder {
    clock: Rc<dyn Clock>,
    shared_data: Rc<dyn SharedData>,
    storage_builder: Option<DataStorageBuilder>,
    max_entries: usize,
    prefix: String,
    eviction_strategy: EvictionStrategy,
    ttl: Option<Duration>,
}

/// CacheBuilder can be injected in your configuration function.
/// ```rust
/// #[entrypoint]
/// async fn configure(
///     launcher: Launcher,
///     cache_builder: CacheBuilder,
///     Configuration(configuration): Configuration,
/// ) -> anyhow::Result<()> {
///     // configure the cache here.
/// }
/// ```
impl FromContext<ConfigureContext> for CacheBuilder {
    type Error = Infallible;

    fn from_context(context: &ConfigureContext) -> Result<Self, Self::Error> {
        let clock: Rc<dyn Clock> = context.extract_always();
        let shared_data: Rc<dyn SharedData> = context.extract_always();
        let metadata: Metadata = context.extract_always();
        let storage_builder: Option<DataStorageBuilder> = context.extract_always();
        Ok(CacheBuilder {
            clock,
            shared_data,
            storage_builder,
            max_entries: DEFAULT_MAX_ENTRIES,
            prefix: format!(
                "isolated-cache_id-{}-{}",
                metadata.policy_metadata.policy_name, metadata.policy_metadata.policy_namespace
            ),
            eviction_strategy: EvictionStrategy::Fifo,
            ttl: None,
        })
    }
}

impl CacheBuilder {
    /// Returns a [`CacheBuilderInstance`] with an ID to uniquely identify the cache.
    ///
    /// _Note: if 2 caches are assigned the same ID in a same policy, they will share the state.
    /// This behavior is extended to other policies when using the
    /// [shared cache option](CacheBuilderInstance::shared)._
    #[allow(clippy::new_ret_no_self)]
    pub fn new(&self, cache_id: String) -> CacheBuilderInstance {
        CacheBuilderInstance {
            clock: Rc::clone(&self.clock),
            shared_data: Rc::clone(&self.shared_data),
            storage_builder: self.storage_builder.clone(),
            max_entries: self.max_entries,
            prefix: self.prefix.clone(),
            key: cache_id,
            eviction_strategy: self.eviction_strategy,
            ttl: self.ttl,
        }
    }

    /// Sets the eviction strategy for all caches created by this builder.
    /// Default is [`EvictionStrategy::Fifo`].
    pub fn eviction_strategy(mut self, strategy: EvictionStrategy) -> Self {
        self.eviction_strategy = strategy;
        self
    }
}

/// Represents a particular instance of a [`CacheBuilder`]. It includes extra metadata of the
/// policy, required to successfully work in all scenarios involving multiple policies.
pub struct CacheBuilderInstance {
    clock: Rc<dyn Clock>,
    shared_data: Rc<dyn SharedData>,
    storage_builder: Option<DataStorageBuilder>,
    max_entries: usize,
    prefix: String,
    key: String,
    eviction_strategy: EvictionStrategy,
    ttl: Option<Duration>,
}

impl CacheBuilderInstance {
    /// Sets the maximum number of entries available for the cache.
    pub fn max_entries(mut self, max_entries: usize) -> Self {
        self.max_entries = max_entries;
        self
    }

    /// Indicates that the cache should not be isolated to the policy. The state will be shared
    /// across different policy instances and different policies that use the same cache ID.
    pub fn shared(mut self) -> Self {
        self.prefix = "shared-cache_id".to_string();
        self.storage_builder = self.storage_builder.map(|b| b.shared());
        self
    }

    /// Sets the eviction strategy for this specific cache instance.
    /// This overrides the strategy set on the builder.
    pub fn eviction_strategy(mut self, strategy: EvictionStrategy) -> Self {
        self.eviction_strategy = strategy;
        self
    }

    /// Sets the time-to-live for cached entries. After this duration, entries expire and are
    /// lazily removed on the next read.
    ///
    /// **Note:** When using shared caches, all policy instances sharing the cache ID must use
    /// the same TTL setting. A cache with TTL enabled serializes values with a TTL envelope;
    /// a cache without TTL (or reading with a different encoding) will fail to deserialize
    /// the entry and treat it as a cache miss.
    pub fn ttl(mut self, ttl: Duration) -> Self {
        self.ttl = Some(ttl);
        self
    }

    /// Builds a new [`Cache`] with the corresponding specifications.
    pub fn build(self) -> impl Cache + 'static {
        let cache_id = format!("{}-{}", self.prefix, self.key);
        let ttl = self.ttl;
        let clock = Rc::clone(&self.clock);

        let inner: CacheImpl = match self.eviction_strategy {
            EvictionStrategy::Fifo => {
                let concurrent_shared_data =
                    ConcurrentSharedData::new(self.clock, self.shared_data);
                CacheImpl::Fifo(FifoCache::new(
                    cache_id,
                    self.max_entries,
                    concurrent_shared_data,
                ))
            }
            EvictionStrategy::Lru(lru) => {
                // The LRU strategy is backed by data-storage-lib local storage. The cache_id
                // already encodes policy isolation (or the shared prefix), so it is used as the
                // storage key to keep every LRU cache's entries and shard queues namespaced.
                let storage_builder = self.storage_builder.expect("EvictionStrategy::Lru requires local data storage, which is unavailable in this context");
                let storage = storage_builder.local(cache_id.clone());
                CacheImpl::Lru(LruCache::new(
                    cache_id,
                    self.max_entries,
                    storage,
                    lru.shards(),
                ))
            }
        };

        match ttl {
            Some(ttl) => CacheImpl::Ttl(Box::new(TtlCache::new(inner, clock, ttl))),
            None => inner,
        }
    }
}

#[cfg(test)]
mod tests {
    use super::{
        nearest_power_of_two, CacheBuilder, EvictionStrategy, LruStrategy, DEFAULT_LRU_SHARDS,
    };
    use crate::fifo_cache::{stateful_shared_data_mock, SharedValueStore};
    use crate::mocks::MockConcurrentSharedData;
    use crate::Cache;
    use pdk_core::classy::proxy_wasm::types::{Bytes, Status};
    use pdk_core::classy::{Clock, SharedData};
    use std::cell::RefCell;
    use std::collections::HashMap;
    use std::rc::Rc;
    use std::time::{Duration, SystemTime};

    // Clock whose time can be advanced manually, to drive TTL expiry.
    struct ManualClock {
        current: RefCell<SystemTime>,
    }

    impl ManualClock {
        fn new() -> Self {
            Self {
                current: RefCell::new(SystemTime::now()),
            }
        }

        fn move_forward(&self, duration: Duration) {
            let new_time = self.get_current_time() + duration;
            *self.current.borrow_mut() = new_time;
        }
    }

    impl Clock for ManualClock {
        fn get_current_time(&self) -> SystemTime {
            *self.current.borrow()
        }
    }

    // The build() Fifo path passes shared data into the mock's static `new`,
    // which ignores it, so a no-op implementation suffices.
    struct DummySharedData;

    impl SharedData for DummySharedData {
        fn shared_data_get(&self, _key: &str) -> (Option<Bytes>, Option<u32>) {
            (None, None)
        }

        fn shared_data_set(
            &self,
            _key: &str,
            _value: &[u8],
            _cas: Option<u32>,
        ) -> Result<(), Status> {
            Ok(())
        }

        fn shared_data_remove(
            &self,
            _key: &str,
            _version: Option<u32>,
        ) -> Result<Option<Bytes>, Status> {
            Ok(None)
        }

        fn shared_data_keys(&self) -> Vec<String> {
            Vec::new()
        }
    }

    // Exercises the real `.ttl(..).build()` path end to end: build() wraps the
    // FIFO cache in `CacheImpl::Ttl`, so save/get/delete/purge all dispatch
    // through the TTL decorator over a genuine FifoCache backend. The FifoCache
    // is backed by a stateful mock (a shared HashMap) so the TTL-framed bytes
    // survive a real save/get roundtrip, and the injected ManualClock drives
    // expiry deterministically.
    #[test]
    fn build_with_ttl_expires_entry_through_builder() {
        // build()'s Fifo branch constructs the shared data via the mock's static
        // `new`; intercept it to return the stateful mock backed by `store`.
        let store: SharedValueStore = Rc::new(RefCell::new(HashMap::new()));
        let new_ctx = MockConcurrentSharedData::new_context();
        let store_for_ctx = Rc::clone(&store);
        new_ctx.expect().returning_st(move |_clock, _shared| {
            stateful_shared_data_mock(Rc::clone(&store_for_ctx))
        });

        let manual_clock = Rc::new(ManualClock::new());
        let clock: Rc<dyn Clock> = Rc::clone(&manual_clock) as Rc<dyn Clock>;

        let builder = CacheBuilder {
            clock,
            shared_data: Rc::new(DummySharedData),
            storage_builder: None,
            max_entries: 100,
            prefix: "isolated-cache_id-test".to_string(),
            eviction_strategy: EvictionStrategy::Fifo,
            ttl: None,
        };

        let cache = builder
            .new("ttl-cache".to_string())
            .ttl(Duration::from_millis(50))
            .build();

        cache.save("key1", b"value1".to_vec()).unwrap();

        // Within the TTL window the value is present through the TTL decorator.
        manual_clock.move_forward(Duration::from_millis(20));
        assert_eq!(cache.get("key1"), Some(b"value1".to_vec()));

        // Past the TTL, get() reports expiry and lazily deletes from the backend.
        manual_clock.move_forward(Duration::from_millis(40));
        assert!(cache.get("key1").is_none());
        assert!(store
            .borrow()
            .get("isolated-cache_id-test-ttl-cache-key1")
            .is_none());

        // purge() dispatches through the TTL arm without panicking.
        cache.purge();
    }

    #[test]
    fn nearest_power_of_two_rounds_expected() {
        // Zero and one collapse to the minimum of 1.
        assert_eq!(nearest_power_of_two(0), 1);
        assert_eq!(nearest_power_of_two(1), 1);
        // Exact powers of two are unchanged.
        assert_eq!(nearest_power_of_two(2), 2);
        assert_eq!(nearest_power_of_two(4), 4);
        assert_eq!(nearest_power_of_two(16), 16);
        assert_eq!(nearest_power_of_two(128), 128);
        // Ties round down (equidistant between two powers of two).
        assert_eq!(nearest_power_of_two(3), 2); // between 2 and 4, tie -> 2
        assert_eq!(nearest_power_of_two(6), 4); // between 4 and 8, tie -> 4
                                                // Non-ties round to the closer power of two.
        assert_eq!(nearest_power_of_two(5), 4);
        assert_eq!(nearest_power_of_two(7), 8);
        assert_eq!(nearest_power_of_two(9), 8);
        assert_eq!(nearest_power_of_two(15), 16);
        // Counts above the largest u8 power of two clamp to 128.
        assert_eq!(nearest_power_of_two(200), 128);
        assert_eq!(nearest_power_of_two(255), 128);
    }

    #[test]
    fn with_shards_applies_rounding() {
        assert_eq!(LruStrategy::default().with_shards(0).shards(), 1);
        assert_eq!(LruStrategy::default().with_shards(1).shards(), 1);
        assert_eq!(LruStrategy::default().with_shards(3).shards(), 2);
        assert_eq!(LruStrategy::default().with_shards(5).shards(), 4);
        assert_eq!(LruStrategy::default().with_shards(7).shards(), 8);
        assert_eq!(LruStrategy::default().with_shards(16).shards(), 16);
    }

    #[test]
    fn default_strategy_uses_default_shards() {
        assert_eq!(LruStrategy::default().shards(), DEFAULT_LRU_SHARDS);
    }

    #[test]
    fn eviction_strategy_wraps_configured_strategy() {
        let strategy = LruStrategy::default().with_shards(4);
        assert_eq!(
            EvictionStrategy::Lru(strategy),
            EvictionStrategy::Lru(strategy)
        );
        assert_ne!(
            EvictionStrategy::Lru(LruStrategy::default().with_shards(1)),
            EvictionStrategy::Lru(LruStrategy::default().with_shards(16)),
        );
    }

    // Note: the genuine "TtlCache wrapping a real FifoCache" orthogonality
    // tests live in fifo_cache.rs, alongside the FifoCache internals and mock
    // they need (the FIFO internal-state type is private to that module). See
    // `ttl_wrapping_real_fifo_cache_present_before_ttl` /
    // `ttl_wrapping_real_fifo_cache_gone_after_ttl` there, mirroring the
    // real-LruCache orthogonality tests in ttl_cache.rs.
}