Skip to main content

rskit_cache/adapters/memory/
cache.rs

1use std::collections::BTreeMap;
2use std::sync::Arc;
3use std::time::{Duration, Instant};
4
5use parking_lot::Mutex;
6use rskit_errors::{AppError, AppResult, ErrorCode};
7
8use crate::registry::CacheStore;
9
10#[derive(Clone)]
11pub(crate) struct Entry {
12    value: String,
13    expires_at: Option<Instant>,
14}
15
16impl Entry {
17    fn is_expired(&self, now: Instant) -> bool {
18        self.expires_at.is_some_and(|expires_at| expires_at <= now)
19    }
20}
21
22/// Lean in-process cache store adapter for local development and tests.
23pub struct MemoryCache {
24    prefix: Option<String>,
25    max_entries: Option<usize>,
26    pub(crate) entries: Mutex<BTreeMap<String, Entry>>,
27    clock: Arc<dyn Fn() -> Instant + Send + Sync>,
28}
29
30impl MemoryCache {
31    /// Create an empty in-memory cache.
32    #[must_use]
33    pub fn new(prefix: Option<String>, max_entries: Option<usize>) -> Self {
34        Self::new_with_clock(prefix, max_entries, Instant::now)
35    }
36
37    /// Create an in-memory cache with an injected clock.
38    ///
39    /// This is primarily useful for deterministic tests
40    /// and simulations that need to advance cache expiry without sleeping.
41    #[must_use]
42    pub fn new_with_clock(
43        prefix: Option<String>,
44        max_entries: Option<usize>,
45        clock: impl Fn() -> Instant + Send + Sync + 'static,
46    ) -> Self {
47        Self {
48            prefix,
49            max_entries: max_entries.filter(|entries| *entries > 0),
50            entries: Mutex::new(BTreeMap::new()),
51            clock: Arc::new(clock),
52        }
53    }
54
55    fn now(&self) -> Instant {
56        (self.clock)()
57    }
58
59    fn key(&self, key: &str) -> String {
60        self.prefix
61            .as_ref()
62            .map_or_else(|| key.to_owned(), |prefix| format!("{prefix}:{key}"))
63    }
64
65    fn prune_expired(entries: &mut BTreeMap<String, Entry>, now: Instant) {
66        entries.retain(|_, entry| !entry.is_expired(now));
67    }
68
69    fn expires_at(now: Instant, ttl: Option<Duration>) -> AppResult<Option<Instant>> {
70        let Some(duration) = ttl else {
71            return Ok(None);
72        };
73        now.checked_add(duration).map(Some).ok_or_else(|| {
74            AppError::new(
75                ErrorCode::InvalidInput,
76                "cache TTL is too large to represent safely",
77            )
78        })
79    }
80}
81
82#[async_trait::async_trait]
83impl CacheStore for MemoryCache {
84    async fn get(&self, key: &str) -> AppResult<Option<String>> {
85        let mut entries = self.entries.lock();
86        Self::prune_expired(&mut entries, self.now());
87        let full_key = self.key(key);
88        Ok(entries.get(&full_key).map(|entry| entry.value.clone()))
89    }
90
91    async fn set(&self, key: &str, val: &str, ttl: Option<Duration>) -> AppResult<()> {
92        let mut entries = self.entries.lock();
93        let now = self.now();
94        Self::prune_expired(&mut entries, now);
95        let key = self.key(key);
96
97        if ttl.is_some_and(|ttl| ttl.is_zero()) {
98            return Err(AppError::new(
99                ErrorCode::InvalidInput,
100                "cache TTL must be greater than zero",
101            ));
102        }
103        let expires_at = Self::expires_at(now, ttl)?;
104
105        if let Some(max_entries) = self.max_entries
106            && entries.len() >= max_entries
107            && !entries.contains_key(&key)
108            && let Some(first_key) = entries.keys().next().cloned()
109        {
110            entries.remove(&first_key);
111        }
112
113        entries.insert(
114            key,
115            Entry {
116                value: val.to_owned(),
117                expires_at,
118            },
119        );
120        Ok(())
121    }
122
123    async fn delete(&self, key: &str) -> AppResult<bool> {
124        let mut entries = self.entries.lock();
125        Self::prune_expired(&mut entries, self.now());
126        let key = self.key(key);
127        Ok(entries.remove(&key).is_some())
128    }
129
130    async fn exists(&self, key: &str) -> AppResult<bool> {
131        self.get(key).await.map(|value| value.is_some())
132    }
133}
134
135impl Default for MemoryCache {
136    fn default() -> Self {
137        Self::new(None, None)
138    }
139}