rskit_cache/adapters/memory/
cache.rs1use 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
22pub 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 #[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 #[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}