Skip to main content

laterite_core/cache/
memory.rs

1//! The `memory` store: values in this process, each with its own lifetime,
2//! the least recently used evicted when the capacity is reached.
3
4use std::sync::Arc;
5use std::time::{Duration, Instant};
6
7use async_trait::async_trait;
8use moka::future::Cache as Moka;
9use moka::Expiry;
10
11use super::{Entry, ValueStore};
12use crate::error::CoreResult;
13
14/// Sixty-four MiB, the capacity [`MemoryStore::new`] takes.
15pub const DEFAULT_CAPACITY: u64 = 64 * 1024 * 1024;
16
17struct Held {
18    entry: Entry,
19    ttl: Duration,
20}
21
22struct PerEntry;
23
24impl Expiry<String, Arc<Held>> for PerEntry {
25    fn expire_after_create(
26        &self,
27        _key: &String,
28        held: &Arc<Held>,
29        _at: Instant,
30    ) -> Option<Duration> {
31        Some(held.ttl)
32    }
33}
34
35/// Values in this process's memory.
36#[derive(Clone)]
37pub struct MemoryStore {
38    inner: Moka<String, Arc<Held>>,
39}
40
41impl Default for MemoryStore {
42    fn default() -> Self {
43        Self::new()
44    }
45}
46
47impl MemoryStore {
48    /// A store of [`DEFAULT_CAPACITY`] bytes.
49    pub fn new() -> Self {
50        Self::with_capacity(DEFAULT_CAPACITY)
51    }
52
53    /// A store holding up to `bytes` of values (keys and overhead counted
54    /// roughly).
55    pub fn with_capacity(bytes: u64) -> Self {
56        let inner = Moka::builder()
57            .max_capacity(bytes.max(1))
58            .weigher(|key: &String, held: &Arc<Held>| {
59                let size = key.len() + held.entry.bytes.len() + 64;
60                u32::try_from(size).unwrap_or(u32::MAX)
61            })
62            .expire_after(PerEntry)
63            .build();
64        Self { inner }
65    }
66
67    /// How many values are held (an estimate while writes are settling).
68    pub fn len(&self) -> u64 {
69        self.inner.entry_count()
70    }
71
72    pub fn is_empty(&self) -> bool {
73        self.len() == 0
74    }
75}
76
77#[async_trait]
78impl ValueStore for MemoryStore {
79    fn name(&self) -> &'static str {
80        "memory"
81    }
82    async fn get(&self, key: &str) -> CoreResult<Option<Entry>> {
83        Ok(self.inner.get(key).await.map(|h| h.entry.clone()))
84    }
85    async fn put(&self, key: &str, entry: Entry, ttl: Duration) -> CoreResult<()> {
86        self.inner
87            .insert(key.to_string(), Arc::new(Held { entry, ttl }))
88            .await;
89        Ok(())
90    }
91    async fn add(&self, key: &str, entry: Entry, ttl: Duration) -> CoreResult<bool> {
92        let mut stored = false;
93        self.inner
94            .entry(key.to_string())
95            .or_insert_with(async {
96                stored = true;
97                Arc::new(Held { entry, ttl })
98            })
99            .await;
100        Ok(stored)
101    }
102    async fn forget(&self, key: &str) -> CoreResult<bool> {
103        Ok(self.inner.remove(key).await.is_some())
104    }
105    async fn flush(&self) -> CoreResult<()> {
106        self.inner.invalidate_all();
107        self.inner.run_pending_tasks().await;
108        Ok(())
109    }
110}