1use crate::{
2 BoxFuture, Container, Injectable, InternalServerErrorException, Module, ModuleMetadata,
3 PayloadTooLargeException, Result,
4};
5use serde::{Serialize, de::DeserializeOwned};
6use serde_json::Value;
7use std::{
8 collections::HashMap,
9 sync::{Arc, RwLock},
10 time::{Duration, Instant},
11};
12
13const DEFAULT_MAX_ENTRIES: usize = 1024;
14const DEFAULT_MAX_VALUE_BYTES: usize = 1024 * 1024;
15
16pub trait CacheStore: Send + Sync + 'static {
18 fn get(&self, key: String) -> BoxFuture<'_, Result<Option<Value>>>;
20 fn set(&self, key: String, value: Value, ttl: Option<Duration>) -> BoxFuture<'_, Result<()>>;
22 fn delete(&self, key: String) -> BoxFuture<'_, Result<()>>;
24 fn clear(&self) -> BoxFuture<'_, Result<()>>;
26}
27
28#[derive(Clone)]
29struct CacheEntry {
30 value: Value,
31 inserted_at: Instant,
32 expires_at: Option<Instant>,
33}
34
35impl CacheEntry {
36 fn new(value: Value, ttl: Option<Duration>) -> Self {
37 let now = Instant::now();
38 Self {
39 value,
40 inserted_at: now,
41 expires_at: ttl.and_then(|ttl| now.checked_add(ttl)),
42 }
43 }
44
45 fn is_expired(&self) -> bool {
46 self.expires_at
47 .is_some_and(|expires_at| Instant::now() >= expires_at)
48 }
49}
50
51#[derive(Clone, Debug)]
52pub struct MemoryCacheOptions {
54 pub max_entries: usize,
56 pub max_value_bytes: usize,
58 pub default_ttl: Option<Duration>,
60}
61
62impl Default for MemoryCacheOptions {
63 fn default() -> Self {
64 Self {
65 max_entries: DEFAULT_MAX_ENTRIES,
66 max_value_bytes: DEFAULT_MAX_VALUE_BYTES,
67 default_ttl: None,
68 }
69 }
70}
71
72pub struct MemoryCacheStore {
74 options: MemoryCacheOptions,
75 entries: RwLock<HashMap<String, CacheEntry>>,
76}
77
78impl MemoryCacheStore {
79 pub fn new() -> Self {
81 Self::with_options(MemoryCacheOptions::default())
82 }
83
84 pub fn with_options(options: MemoryCacheOptions) -> Self {
86 Self {
87 options,
88 entries: RwLock::new(HashMap::new()),
89 }
90 }
91
92 pub fn len(&self) -> usize {
94 self.entries.read().map_or(0, |entries| entries.len())
95 }
96
97 pub fn is_empty(&self) -> bool {
99 self.len() == 0
100 }
101
102 fn remove_expired(entries: &mut HashMap<String, CacheEntry>) {
103 entries.retain(|_, entry| !entry.is_expired());
104 }
105
106 fn evict_to_capacity(&self, entries: &mut HashMap<String, CacheEntry>) {
107 if self.options.max_entries == 0 {
108 entries.clear();
109 return;
110 }
111
112 while entries.len() > self.options.max_entries {
113 let Some(oldest_key) = entries
114 .iter()
115 .min_by_key(|(_, entry)| entry.inserted_at)
116 .map(|(key, _)| key.clone())
117 else {
118 return;
119 };
120
121 entries.remove(&oldest_key);
122 }
123 }
124}
125
126impl Default for MemoryCacheStore {
127 fn default() -> Self {
128 Self::new()
129 }
130}
131
132impl Injectable for MemoryCacheStore {
133 fn dependencies() -> Vec<crate::ProviderDependency> {
134 crate::provider_dependencies![]
135 }
136
137 fn create(_container: &Container) -> BoxFuture<'_, crate::Result<Self>> {
138 Box::pin(async { Ok(Self::new()) })
139 }
140}
141
142impl CacheStore for MemoryCacheStore {
143 fn get(&self, key: String) -> BoxFuture<'_, crate::Result<Option<Value>>> {
144 Box::pin(async move {
145 let mut entries = self.entries.write().map_err(|_| {
146 InternalServerErrorException::new(anyhow::anyhow!("cache store lock poisoned"))
147 })?;
148 let Some(entry) = entries.get(&key) else {
149 return Ok(None);
150 };
151
152 if entry.is_expired() {
153 entries.remove(&key);
154 return Ok(None);
155 }
156
157 Ok(Some(entry.value.clone()))
158 })
159 }
160
161 fn set(
162 &self,
163 key: String,
164 value: Value,
165 ttl: Option<Duration>,
166 ) -> BoxFuture<'_, crate::Result<()>> {
167 Box::pin(async move {
168 let value_size = serde_json::to_vec(&value)
169 .map_err(InternalServerErrorException::new)?
170 .len();
171
172 if value_size > self.options.max_value_bytes {
173 return Err(PayloadTooLargeException::new(format!(
174 "cache value exceeds the configured limit of {} bytes",
175 self.options.max_value_bytes
176 )));
177 }
178
179 let ttl = ttl.or(self.options.default_ttl);
180 let mut entries = self.entries.write().map_err(|_| {
181 InternalServerErrorException::new(anyhow::anyhow!("cache store lock poisoned"))
182 })?;
183 Self::remove_expired(&mut entries);
184 entries.insert(key, CacheEntry::new(value, ttl));
185 self.evict_to_capacity(&mut entries);
186 Ok(())
187 })
188 }
189
190 fn delete(&self, key: String) -> BoxFuture<'_, crate::Result<()>> {
191 Box::pin(async move {
192 self.entries
193 .write()
194 .map_err(|_| {
195 InternalServerErrorException::new(anyhow::anyhow!("cache store lock poisoned"))
196 })?
197 .remove(&key);
198 Ok(())
199 })
200 }
201
202 fn clear(&self) -> BoxFuture<'_, crate::Result<()>> {
203 Box::pin(async move {
204 self.entries
205 .write()
206 .map_err(|_| {
207 InternalServerErrorException::new(anyhow::anyhow!("cache store lock poisoned"))
208 })?
209 .clear();
210 Ok(())
211 })
212 }
213}
214
215pub struct Cache {
217 store: Arc<dyn CacheStore>,
218}
219
220impl Cache {
221 pub fn new(store: Arc<dyn CacheStore>) -> Self {
223 Self { store }
224 }
225
226 pub async fn get<T>(&self, key: impl Into<String>) -> Result<Option<T>>
228 where
229 T: DeserializeOwned,
230 {
231 let Some(value) = self.store.get(key.into()).await? else {
232 return Ok(None);
233 };
234
235 serde_json::from_value(value)
236 .map(Some)
237 .map_err(InternalServerErrorException::new)
238 }
239
240 pub async fn set<T>(&self, key: impl Into<String>, value: T) -> Result<()>
242 where
243 T: Serialize,
244 {
245 self.set_with_optional_ttl(key, value, None).await
246 }
247
248 pub async fn set_with_ttl<T>(
250 &self,
251 key: impl Into<String>,
252 value: T,
253 ttl: Duration,
254 ) -> Result<()>
255 where
256 T: Serialize,
257 {
258 self.set_with_optional_ttl(key, value, Some(ttl)).await
259 }
260
261 pub async fn set_with_optional_ttl<T>(
263 &self,
264 key: impl Into<String>,
265 value: T,
266 ttl: Option<Duration>,
267 ) -> Result<()>
268 where
269 T: Serialize,
270 {
271 let value = serde_json::to_value(value).map_err(InternalServerErrorException::new)?;
272 self.store.set(key.into(), value, ttl).await
273 }
274
275 pub async fn delete(&self, key: impl Into<String>) -> Result<()> {
277 self.store.delete(key.into()).await
278 }
279
280 pub async fn clear(&self) -> Result<()> {
282 self.store.clear().await
283 }
284}
285
286impl Injectable for Cache {
287 fn dependencies() -> Vec<crate::ProviderDependency> {
288 crate::provider_dependencies![MemoryCacheStore]
289 }
290
291 fn create(container: &Container) -> BoxFuture<'_, crate::Result<Self>> {
292 Box::pin(async move {
293 let store = container.resolve::<MemoryCacheStore>()? as Arc<dyn CacheStore>;
294 Ok(Self::new(store))
295 })
296 }
297}
298
299pub struct CacheModule;
301
302impl Module for CacheModule {
303 fn register() -> ModuleMetadata {
304 ModuleMetadata::new()
305 .provider::<MemoryCacheStore>()
306 .provider::<Cache>()
307 .export::<MemoryCacheStore>()
308 .export::<Cache>()
309 }
310}