shardline_cache/
memory.rs

1use std::{
2    collections::{BTreeMap, HashMap},
3    num::{NonZeroU64, NonZeroUsize},
4    sync::Arc,
5    time::Duration,
6};
7
8use tokio::time::Instant;
9
10use tokio::sync::{Notify, RwLock};
11
12use crate::{AsyncReconstructionCache, ReconstructionCacheFuture, ReconstructionCacheKey};
13
14#[derive(Debug, Clone)]
15struct MemoryEntry {
16    payload: Arc<Vec<u8>>,
17    expires_at: Instant,
18    inserted_at: Instant,
19    seq: u64,
20}
21
22#[derive(Debug, Clone, Eq, PartialEq)]
23struct EvictionKey(Instant, u64);
24
25impl Ord for EvictionKey {
26    fn cmp(&self, other: &Self) -> std::cmp::Ordering {
27        self.0.cmp(&other.0).then_with(|| self.1.cmp(&other.1))
28    }
29}
30
31impl PartialOrd for EvictionKey {
32    fn partial_cmp(&self, other: &Self) -> Option<std::cmp::Ordering> {
33        Some(self.cmp(other))
34    }
35}
36
37#[derive(Debug)]
38struct CacheInner {
39    entries: HashMap<ReconstructionCacheKey, MemoryEntry>,
40    eviction_order: BTreeMap<EvictionKey, ReconstructionCacheKey>,
41    next_seq: u64,
42    loading: HashMap<ReconstructionCacheKey, Arc<Notify>>,
43}
44
45impl CacheInner {
46    fn new() -> Self {
47        Self {
48            entries: HashMap::new(),
49            eviction_order: BTreeMap::new(),
50            next_seq: 0,
51            loading: HashMap::new(),
52        }
53    }
54
55    fn insert(&mut self, key: ReconstructionCacheKey, entry: MemoryEntry) {
56        let inserted_at = entry.inserted_at;
57        if let Some(old) = self.entries.insert(key.clone(), entry) {
58            self.eviction_order
59                .remove(&EvictionKey(old.inserted_at, old.seq));
60        }
61        let seq = self.next_seq;
62        self.next_seq = self.next_seq.wrapping_add(1);
63        self.eviction_order
64            .insert(EvictionKey(inserted_at, seq), key);
65    }
66
67    fn remove(&mut self, key: &ReconstructionCacheKey) -> Option<MemoryEntry> {
68        if let Some(entry) = self.entries.remove(key) {
69            self.eviction_order
70                .remove(&EvictionKey(entry.inserted_at, entry.seq));
71            Some(entry)
72        } else {
73            None
74        }
75    }
76
77    fn evict_oldest(&mut self) {
78        while let Some((_eviction_key, key)) = self.eviction_order.pop_first() {
79            if self.entries.contains_key(&key) {
80                self.entries.remove(&key);
81                return;
82            }
83        }
84    }
85}
86
87/// Bounded in-memory reconstruction cache adapter.
88#[derive(Debug, Clone)]
89pub struct MemoryReconstructionCache {
90    ttl: Duration,
91    max_entries: NonZeroUsize,
92    inner: Arc<RwLock<CacheInner>>,
93}
94
95impl MemoryReconstructionCache {
96    /// Creates a bounded in-memory reconstruction cache.
97    #[must_use]
98    pub fn new(ttl_seconds: NonZeroU64, max_entries: NonZeroUsize) -> Self {
99        Self {
100            ttl: Duration::from_secs(ttl_seconds.get()),
101            max_entries,
102            inner: Arc::new(RwLock::new(CacheInner::new())),
103        }
104    }
105}
106
107impl AsyncReconstructionCache for MemoryReconstructionCache {
108    fn ready(&self) -> ReconstructionCacheFuture<'_, ()> {
109        Box::pin(async { Ok(()) })
110    }
111
112    fn get<'operation>(
113        &'operation self,
114        key: &'operation ReconstructionCacheKey,
115    ) -> ReconstructionCacheFuture<'operation, Option<Vec<u8>>> {
116        Box::pin(async move {
117            let now = Instant::now();
118            {
119                let inner = self.inner.read().await;
120                if let Some(entry) = inner.entries.get(key) {
121                    if entry.expires_at > now {
122                        return Ok(Some(entry.payload.as_ref().clone()));
123                    }
124                } else if !inner.loading.contains_key(key) {
125                    return Ok(None);
126                }
127            }
128
129            let mut inner = self.inner.write().await;
130
131            if let Some(entry) = inner.entries.get(key)
132                && entry.expires_at > now
133            {
134                return Ok(Some(entry.payload.as_ref().clone()));
135            }
136
137            if let Some(notify) = inner.loading.get(key) {
138                let notify = Arc::clone(notify);
139                drop(inner);
140                notify.notified().await;
141
142                let read_inner = self.inner.read().await;
143                if let Some(entry) = read_inner.entries.get(key)
144                    && entry.expires_at > Instant::now()
145                {
146                    return Ok(Some(entry.payload.as_ref().clone()));
147                }
148                return Ok(None);
149            }
150
151            let notify = Arc::new(Notify::new());
152            inner.loading.insert(key.clone(), Arc::clone(&notify));
153
154            let should_remove = inner
155                .entries
156                .get(key)
157                .is_some_and(|entry| entry.expires_at <= now);
158            if should_remove {
159                inner.remove(key);
160            }
161            Ok(None)
162        })
163    }
164
165    fn put<'operation>(
166        &'operation self,
167        key: &'operation ReconstructionCacheKey,
168        payload: &'operation [u8],
169    ) -> ReconstructionCacheFuture<'operation, ()> {
170        Box::pin(async move {
171            let now = Instant::now();
172            let expires_at = now.checked_add(self.ttl).unwrap_or(now);
173            let mut inner = self.inner.write().await;
174            if !inner.entries.contains_key(key) && inner.entries.len() >= self.max_entries.get() {
175                inner.evict_oldest();
176            }
177            inner.insert(
178                key.clone(),
179                MemoryEntry {
180                    payload: Arc::new(payload.to_vec()),
181                    expires_at,
182                    inserted_at: now,
183                    seq: 0,
184                },
185            );
186            if let Some(notify) = inner.loading.remove(key) {
187                notify.notify_waiters();
188            }
189            Ok(())
190        })
191    }
192
193    fn delete<'operation>(
194        &'operation self,
195        key: &'operation ReconstructionCacheKey,
196    ) -> ReconstructionCacheFuture<'operation, bool> {
197        Box::pin(async move {
198            let mut inner = self.inner.write().await;
199            Ok(inner.remove(key).is_some())
200        })
201    }
202}
203
204#[cfg(test)]
205mod tests {
206    use std::{
207        num::{NonZeroU64, NonZeroUsize},
208        time::Duration,
209    };
210
211    use super::MemoryReconstructionCache;
212    use crate::{AsyncReconstructionCache, ReconstructionCacheKey};
213
214    #[tokio::test]
215    async fn memory_cache_roundtrips_one_payload() {
216        let cache = MemoryReconstructionCache::new(NonZeroU64::MIN, NonZeroUsize::MIN);
217        let key = ReconstructionCacheKey::latest("asset.bin", None);
218        let put = cache.put(&key, b"payload").await;
219        assert!(put.is_ok());
220
221        let value = cache.get(&key).await;
222
223        assert!(value.is_ok());
224        assert_eq!(value.ok(), Some(Some(b"payload".to_vec())));
225    }
226
227    #[tokio::test]
228    async fn memory_cache_evicts_oldest_entry_when_capacity_is_full() {
229        let max_entries = NonZeroUsize::new(2).unwrap_or(NonZeroUsize::MIN);
230        let ttl_seconds = NonZeroU64::new(60).unwrap_or(NonZeroU64::MIN);
231        let cache = MemoryReconstructionCache::new(ttl_seconds, max_entries);
232        let first = ReconstructionCacheKey::latest("asset-1.bin", None);
233        let second = ReconstructionCacheKey::latest("asset-2.bin", None);
234        let third = ReconstructionCacheKey::latest("asset-3.bin", None);
235
236        assert!(cache.put(&first, b"first").await.is_ok());
237        assert!(cache.put(&second, b"second").await.is_ok());
238        assert!(cache.put(&third, b"third").await.is_ok());
239
240        let first_value = cache.get(&first).await;
241        let second_value = cache.get(&second).await;
242        let third_value = cache.get(&third).await;
243
244        assert!(first_value.is_ok());
245        assert!(second_value.is_ok());
246        assert!(third_value.is_ok());
247        assert_eq!(first_value.ok(), Some(None));
248        assert_eq!(second_value.ok(), Some(Some(b"second".to_vec())));
249        assert_eq!(third_value.ok(), Some(Some(b"third".to_vec())));
250    }
251
252    #[tokio::test(start_paused = true)]
253    async fn memory_cache_expires_entries_after_ttl() {
254        let ttl_seconds = NonZeroU64::new(1).unwrap_or(NonZeroU64::MIN);
255        let cache = MemoryReconstructionCache::new(ttl_seconds, NonZeroUsize::MIN);
256        let key = ReconstructionCacheKey::latest("asset.bin", None);
257        assert!(cache.put(&key, b"payload").await.is_ok());
258
259        tokio::time::advance(Duration::from_secs(1)).await;
260        let value = cache.get(&key).await;
261
262        assert!(value.is_ok());
263        assert_eq!(value.ok(), Some(None));
264    }
265}
shardline_cache/memory.rs

shardline_cache/
memory.rs
