pdk_cache_lib/
builder.rs

1// Copyright (c) 2026, Salesforce, Inc.,
2// All rights reserved.
3// For full license text, see the LICENSE.txt file
4
5//! Primitives for building a cache
6//!
7//! The entrypoint of this crate is the [`CacheBuilder`] struct.
8
9use super::fifo_cache::FifoCache;
10use super::lru_cache::LruCache;
11use super::ttl_cache::TtlCache;
12use super::Cache;
13use data_storage_lib::DataStorageBuilder;
14use pdk_core::classy::extract::context::ConfigureContext;
15use pdk_core::classy::extract::{Extract, FromContext};
16use pdk_core::classy::{Clock, SharedData};
17use pdk_core::policy_context::api::Metadata;
18use std::convert::Infallible;
19use std::rc::Rc;
20use std::time::Duration;
21cfg_if::cfg_if! {
22    if #[cfg(test)] {
23        use crate::mocks::MockConcurrentSharedData as ConcurrentSharedData;
24    } else {
25        use pdk_core::classy::shared_data::concurrent_shared_data::ConcurrentSharedData;
26    }
27}
28
29const DEFAULT_MAX_ENTRIES: usize = 100;
30
31/// Default number of shards used by the sharded LRU strategy.
32const DEFAULT_LRU_SHARDS: u8 = 16;
33
34/// Eviction strategy for the cache.
35#[derive(Debug, Clone, Copy, PartialEq, Eq)]
36#[non_exhaustive]
37pub enum EvictionStrategy {
38    /// First-In-First-Out: evicts the oldest entries first (by insertion time).
39    Fifo,
40    /// Sharded, approximate Least Recently Used eviction. See [`LruStrategy`].
41    ///
42    /// This is **not** a globally-exact LRU. To reduce serialization overhead on
43    /// large caches, the entries are split across independent shards (see
44    /// [`LruStrategy::shards`]) and eviction happens per shard: when a shard is
45    /// full, that shard evicts its own least recently accessed entry. With keys
46    /// spread across shards, the globally least recently used entry is only
47    /// evicted approximately, since eviction only ever considers one shard's
48    /// entries at a time. Both `get()` and `save()` update the access time.
49    ///
50    /// With a single shard ([`LruStrategy::shards`]`(1)`) the whole cache lives in
51    /// one queue, so eviction is exact, global LRU at the cost of no sharding.
52    Lru(LruStrategy),
53}
54
55/// Configuration for the sharded LRU [`EvictionStrategy`].
56///
57/// Use [`LruStrategy::default`] for the default shard count, or override it with
58/// [`LruStrategy::with_shards`]. The shard count controls how the LRU state is partitioned;
59/// see [`EvictionStrategy::Lru`] for how sharding affects eviction accuracy.
60#[derive(Debug, Clone, Copy, PartialEq, Eq)]
61pub struct LruStrategy {
62    shards: u8,
63}
64
65impl Default for LruStrategy {
66    fn default() -> Self {
67        Self {
68            shards: DEFAULT_LRU_SHARDS,
69        }
70    }
71}
72
73impl LruStrategy {
74    /// Sets the number of shards, rounded to the nearest power of two
75    /// (minimum 1). Rounding keeps hash-based key distribution well-balanced
76    /// without rejecting arbitrary counts.
77    pub fn with_shards(mut self, shards: u8) -> Self {
78        self.shards = nearest_power_of_two(shards);
79        self
80    }
81
82    /// Number of shards this strategy partitions its state into.
83    pub fn shards(&self) -> u8 {
84        self.shards
85    }
86}
87
88/// Rounds `count` to the nearest power of two (ties round down), minimum 1.
89fn nearest_power_of_two(count: u8) -> u8 {
90    if count <= 1 {
91        return 1;
92    }
93    // The next power of two above `count` overflows u8 for counts > 128, and 128 is
94    // the largest power of two that fits, so those counts clamp to 128.
95    let Some(higher) = count.checked_next_power_of_two() else {
96        return 128;
97    };
98    if higher == count {
99        return count;
100    }
101    let lower = higher >> 1;
102    if count - lower <= higher - count {
103        lower
104    } else {
105        higher
106    }
107}
108
109/// Internal enum that wraps different cache implementations.
110/// This allows returning a concrete type instead of Box<dyn Cache>.
111enum CacheImpl {
112    Fifo(FifoCache),
113    Lru(LruCache),
114    Ttl(Box<TtlCache<CacheImpl>>),
115}
116
117impl Cache for CacheImpl {
118    fn save(&self, key: &str, value: Vec<u8>) -> Result<(), super::error::CacheError> {
119        match self {
120            CacheImpl::Fifo(cache) => cache.save(key, value),
121            CacheImpl::Lru(cache) => cache.save(key, value),
122            CacheImpl::Ttl(cache) => cache.save(key, value),
123        }
124    }
125
126    fn get(&self, key: &str) -> Option<Vec<u8>> {
127        match self {
128            CacheImpl::Fifo(cache) => cache.get(key),
129            CacheImpl::Lru(cache) => cache.get(key),
130            CacheImpl::Ttl(cache) => cache.get(key),
131        }
132    }
133
134    fn delete(&self, key: &str) -> Option<Vec<u8>> {
135        match self {
136            CacheImpl::Fifo(cache) => cache.delete(key),
137            CacheImpl::Lru(cache) => cache.delete(key),
138            CacheImpl::Ttl(cache) => cache.delete(key),
139        }
140    }
141
142    fn purge(&self) {
143        match self {
144            CacheImpl::Fifo(cache) => cache.purge(),
145            CacheImpl::Lru(cache) => cache.purge(),
146            CacheImpl::Ttl(cache) => cache.purge(),
147        }
148    }
149}
150
151/// Represents the base builder for a [`Cache`] instance.
152pub struct CacheBuilder {
153    clock: Rc<dyn Clock>,
154    shared_data: Rc<dyn SharedData>,
155    storage_builder: Option<DataStorageBuilder>,
156    max_entries: usize,
157    prefix: String,
158    eviction_strategy: EvictionStrategy,
159    ttl: Option<Duration>,
160}
161
162/// CacheBuilder can be injected in your configuration function.
163/// ```rust
164/// #[entrypoint]
165/// async fn configure(
166///     launcher: Launcher,
167///     cache_builder: CacheBuilder,
168///     Configuration(configuration): Configuration,
169/// ) -> anyhow::Result<()> {
170///     // configure the cache here.
171/// }
172/// ```
173impl FromContext<ConfigureContext> for CacheBuilder {
174    type Error = Infallible;
175
176    fn from_context(context: &ConfigureContext) -> Result<Self, Self::Error> {
177        let clock: Rc<dyn Clock> = context.extract_always();
178        let shared_data: Rc<dyn SharedData> = context.extract_always();
179        let metadata: Metadata = context.extract_always();
180        let storage_builder: Option<DataStorageBuilder> = context.extract_always();
181        Ok(CacheBuilder {
182            clock,
183            shared_data,
184            storage_builder,
185            max_entries: DEFAULT_MAX_ENTRIES,
186            prefix: format!(
187                "isolated-cache_id-{}-{}",
188                metadata.policy_metadata.policy_name, metadata.policy_metadata.policy_namespace
189            ),
190            eviction_strategy: EvictionStrategy::Fifo,
191            ttl: None,
192        })
193    }
194}
195
196impl CacheBuilder {
197    /// Returns a [`CacheBuilderInstance`] with an ID to uniquely identify the cache.
198    ///
199    /// _Note: if 2 caches are assigned the same ID in a same policy, they will share the state.
200    /// This behavior is extended to other policies when using the
201    /// [shared cache option](CacheBuilderInstance::shared)._
202    #[allow(clippy::new_ret_no_self)]
203    pub fn new(&self, cache_id: String) -> CacheBuilderInstance {
204        CacheBuilderInstance {
205            clock: Rc::clone(&self.clock),
206            shared_data: Rc::clone(&self.shared_data),
207            storage_builder: self.storage_builder.clone(),
208            max_entries: self.max_entries,
209            prefix: self.prefix.clone(),
210            key: cache_id,
211            eviction_strategy: self.eviction_strategy,
212            ttl: self.ttl,
213        }
214    }
215
216    /// Sets the eviction strategy for all caches created by this builder.
217    /// Default is [`EvictionStrategy::Fifo`].
218    pub fn eviction_strategy(mut self, strategy: EvictionStrategy) -> Self {
219        self.eviction_strategy = strategy;
220        self
221    }
222}
223
224/// Represents a particular instance of a [`CacheBuilder`]. It includes extra metadata of the
225/// policy, required to successfully work in all scenarios involving multiple policies.
226pub struct CacheBuilderInstance {
227    clock: Rc<dyn Clock>,
228    shared_data: Rc<dyn SharedData>,
229    storage_builder: Option<DataStorageBuilder>,
230    max_entries: usize,
231    prefix: String,
232    key: String,
233    eviction_strategy: EvictionStrategy,
234    ttl: Option<Duration>,
235}
236
237impl CacheBuilderInstance {
238    /// Sets the maximum number of entries available for the cache.
239    pub fn max_entries(mut self, max_entries: usize) -> Self {
240        self.max_entries = max_entries;
241        self
242    }
243
244    /// Indicates that the cache should not be isolated to the policy. The state will be shared
245    /// across different policy instances and different policies that use the same cache ID.
246    pub fn shared(mut self) -> Self {
247        self.prefix = "shared-cache_id".to_string();
248        self.storage_builder = self.storage_builder.map(|b| b.shared());
249        self
250    }
251
252    /// Sets the eviction strategy for this specific cache instance.
253    /// This overrides the strategy set on the builder.
254    pub fn eviction_strategy(mut self, strategy: EvictionStrategy) -> Self {
255        self.eviction_strategy = strategy;
256        self
257    }
258
259    /// Sets the time-to-live for cached entries. After this duration, entries expire and are
260    /// lazily removed on the next read.
261    ///
262    /// **Note:** When using shared caches, all policy instances sharing the cache ID must use
263    /// the same TTL setting. A cache with TTL enabled serializes values with a TTL envelope;
264    /// a cache without TTL (or reading with a different encoding) will fail to deserialize
265    /// the entry and treat it as a cache miss.
266    pub fn ttl(mut self, ttl: Duration) -> Self {
267        self.ttl = Some(ttl);
268        self
269    }
270
271    /// Builds a new [`Cache`] with the corresponding specifications.
272    pub fn build(self) -> impl Cache + 'static {
273        let cache_id = format!("{}-{}", self.prefix, self.key);
274        let ttl = self.ttl;
275        let clock = Rc::clone(&self.clock);
276
277        let inner: CacheImpl = match self.eviction_strategy {
278            EvictionStrategy::Fifo => {
279                let concurrent_shared_data =
280                    ConcurrentSharedData::new(self.clock, self.shared_data);
281                CacheImpl::Fifo(FifoCache::new(
282                    cache_id,
283                    self.max_entries,
284                    concurrent_shared_data,
285                ))
286            }
287            EvictionStrategy::Lru(lru) => {
288                // The LRU strategy is backed by data-storage-lib local storage. The cache_id
289                // already encodes policy isolation (or the shared prefix), so it is used as the
290                // storage key to keep every LRU cache's entries and shard queues namespaced.
291                let storage_builder = self.storage_builder.expect("EvictionStrategy::Lru requires local data storage, which is unavailable in this context");
292                let storage = storage_builder.local(cache_id.clone());
293                CacheImpl::Lru(LruCache::new(
294                    cache_id,
295                    self.max_entries,
296                    storage,
297                    lru.shards(),
298                ))
299            }
300        };
301
302        match ttl {
303            Some(ttl) => CacheImpl::Ttl(Box::new(TtlCache::new(inner, clock, ttl))),
304            None => inner,
305        }
306    }
307}
308
309#[cfg(test)]
310mod tests {
311    use super::{
312        nearest_power_of_two, CacheBuilder, EvictionStrategy, LruStrategy, DEFAULT_LRU_SHARDS,
313    };
314    use crate::fifo_cache::{stateful_shared_data_mock, SharedValueStore};
315    use crate::mocks::MockConcurrentSharedData;
316    use crate::Cache;
317    use pdk_core::classy::proxy_wasm::types::{Bytes, Status};
318    use pdk_core::classy::{Clock, SharedData};
319    use std::cell::RefCell;
320    use std::collections::HashMap;
321    use std::rc::Rc;
322    use std::time::{Duration, SystemTime};
323
324    // Clock whose time can be advanced manually, to drive TTL expiry.
325    struct ManualClock {
326        current: RefCell<SystemTime>,
327    }
328
329    impl ManualClock {
330        fn new() -> Self {
331            Self {
332                current: RefCell::new(SystemTime::now()),
333            }
334        }
335
336        fn move_forward(&self, duration: Duration) {
337            let new_time = self.get_current_time() + duration;
338            *self.current.borrow_mut() = new_time;
339        }
340    }
341
342    impl Clock for ManualClock {
343        fn get_current_time(&self) -> SystemTime {
344            *self.current.borrow()
345        }
346    }
347
348    // The build() Fifo path passes shared data into the mock's static `new`,
349    // which ignores it, so a no-op implementation suffices.
350    struct DummySharedData;
351
352    impl SharedData for DummySharedData {
353        fn shared_data_get(&self, _key: &str) -> (Option<Bytes>, Option<u32>) {
354            (None, None)
355        }
356
357        fn shared_data_set(
358            &self,
359            _key: &str,
360            _value: &[u8],
361            _cas: Option<u32>,
362        ) -> Result<(), Status> {
363            Ok(())
364        }
365
366        fn shared_data_remove(
367            &self,
368            _key: &str,
369            _version: Option<u32>,
370        ) -> Result<Option<Bytes>, Status> {
371            Ok(None)
372        }
373
374        fn shared_data_keys(&self) -> Vec<String> {
375            Vec::new()
376        }
377    }
378
379    // Exercises the real `.ttl(..).build()` path end to end: build() wraps the
380    // FIFO cache in `CacheImpl::Ttl`, so save/get/delete/purge all dispatch
381    // through the TTL decorator over a genuine FifoCache backend. The FifoCache
382    // is backed by a stateful mock (a shared HashMap) so the TTL-framed bytes
383    // survive a real save/get roundtrip, and the injected ManualClock drives
384    // expiry deterministically.
385    #[test]
386    fn build_with_ttl_expires_entry_through_builder() {
387        // build()'s Fifo branch constructs the shared data via the mock's static
388        // `new`; intercept it to return the stateful mock backed by `store`.
389        let store: SharedValueStore = Rc::new(RefCell::new(HashMap::new()));
390        let new_ctx = MockConcurrentSharedData::new_context();
391        let store_for_ctx = Rc::clone(&store);
392        new_ctx.expect().returning_st(move |_clock, _shared| {
393            stateful_shared_data_mock(Rc::clone(&store_for_ctx))
394        });
395
396        let manual_clock = Rc::new(ManualClock::new());
397        let clock: Rc<dyn Clock> = Rc::clone(&manual_clock) as Rc<dyn Clock>;
398
399        let builder = CacheBuilder {
400            clock,
401            shared_data: Rc::new(DummySharedData),
402            storage_builder: None,
403            max_entries: 100,
404            prefix: "isolated-cache_id-test".to_string(),
405            eviction_strategy: EvictionStrategy::Fifo,
406            ttl: None,
407        };
408
409        let cache = builder
410            .new("ttl-cache".to_string())
411            .ttl(Duration::from_millis(50))
412            .build();
413
414        cache.save("key1", b"value1".to_vec()).unwrap();
415
416        // Within the TTL window the value is present through the TTL decorator.
417        manual_clock.move_forward(Duration::from_millis(20));
418        assert_eq!(cache.get("key1"), Some(b"value1".to_vec()));
419
420        // Past the TTL, get() reports expiry and lazily deletes from the backend.
421        manual_clock.move_forward(Duration::from_millis(40));
422        assert!(cache.get("key1").is_none());
423        assert!(store
424            .borrow()
425            .get("isolated-cache_id-test-ttl-cache-key1")
426            .is_none());
427
428        // purge() dispatches through the TTL arm without panicking.
429        cache.purge();
430    }
431
432    #[test]
433    fn nearest_power_of_two_rounds_expected() {
434        // Zero and one collapse to the minimum of 1.
435        assert_eq!(nearest_power_of_two(0), 1);
436        assert_eq!(nearest_power_of_two(1), 1);
437        // Exact powers of two are unchanged.
438        assert_eq!(nearest_power_of_two(2), 2);
439        assert_eq!(nearest_power_of_two(4), 4);
440        assert_eq!(nearest_power_of_two(16), 16);
441        assert_eq!(nearest_power_of_two(128), 128);
442        // Ties round down (equidistant between two powers of two).
443        assert_eq!(nearest_power_of_two(3), 2); // between 2 and 4, tie -> 2
444        assert_eq!(nearest_power_of_two(6), 4); // between 4 and 8, tie -> 4
445                                                // Non-ties round to the closer power of two.
446        assert_eq!(nearest_power_of_two(5), 4);
447        assert_eq!(nearest_power_of_two(7), 8);
448        assert_eq!(nearest_power_of_two(9), 8);
449        assert_eq!(nearest_power_of_two(15), 16);
450        // Counts above the largest u8 power of two clamp to 128.
451        assert_eq!(nearest_power_of_two(200), 128);
452        assert_eq!(nearest_power_of_two(255), 128);
453    }
454
455    #[test]
456    fn with_shards_applies_rounding() {
457        assert_eq!(LruStrategy::default().with_shards(0).shards(), 1);
458        assert_eq!(LruStrategy::default().with_shards(1).shards(), 1);
459        assert_eq!(LruStrategy::default().with_shards(3).shards(), 2);
460        assert_eq!(LruStrategy::default().with_shards(5).shards(), 4);
461        assert_eq!(LruStrategy::default().with_shards(7).shards(), 8);
462        assert_eq!(LruStrategy::default().with_shards(16).shards(), 16);
463    }
464
465    #[test]
466    fn default_strategy_uses_default_shards() {
467        assert_eq!(LruStrategy::default().shards(), DEFAULT_LRU_SHARDS);
468    }
469
470    #[test]
471    fn eviction_strategy_wraps_configured_strategy() {
472        let strategy = LruStrategy::default().with_shards(4);
473        assert_eq!(
474            EvictionStrategy::Lru(strategy),
475            EvictionStrategy::Lru(strategy)
476        );
477        assert_ne!(
478            EvictionStrategy::Lru(LruStrategy::default().with_shards(1)),
479            EvictionStrategy::Lru(LruStrategy::default().with_shards(16)),
480        );
481    }
482
483    // Note: the genuine "TtlCache wrapping a real FifoCache" orthogonality
484    // tests live in fifo_cache.rs, alongside the FifoCache internals and mock
485    // they need (the FIFO internal-state type is private to that module). See
486    // `ttl_wrapping_real_fifo_cache_present_before_ttl` /
487    // `ttl_wrapping_real_fifo_cache_gone_after_ttl` there, mirroring the
488    // real-LruCache orthogonality tests in ttl_cache.rs.
489}
pdk_cache_lib/builder.rs

pdk_cache_lib/
builder.rs
