laterite_core/cache/
memory.rs1use 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
14pub 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#[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 pub fn new() -> Self {
50 Self::with_capacity(DEFAULT_CAPACITY)
51 }
52
53 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 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}