1use crate::error::CacheError;
6use std::collections::HashMap;
7use std::sync::{Arc, RwLock};
8use std::time::{Duration, Instant};
9
10pub trait Cache: Send + Sync {
12 fn get(&self, key: &str) -> Result<Option<Vec<u8>>, CacheError>;
14 fn set(&self, key: &str, value: Vec<u8>, ttl: Option<Duration>) -> Result<(), CacheError>;
16 fn delete(&self, key: &str) -> Result<(), CacheError>;
18 fn clear(&self) -> Result<(), CacheError>;
20 fn exists(&self, key: &str) -> Result<bool, CacheError>;
22 fn expire(&self, key: &str, ttl: Duration) -> Result<(), CacheError>;
24 fn ttl(&self, key: &str) -> Result<Option<Duration>, CacheError>;
26}
27
28#[derive(Clone)]
30pub struct MemoryCache {
31 data: Arc<RwLock<HashMap<String, CacheEntry>>>,
32 default_ttl: Option<Duration>,
33}
34
35struct CacheEntry {
36 value: Vec<u8>,
37 expires_at: Option<Instant>,
38}
39
40impl MemoryCache {
41 pub fn new() -> Self {
43 Self {
44 data: Arc::new(RwLock::new(HashMap::new())),
45 default_ttl: None,
46 }
47 }
48
49 pub fn with_ttl(ttl: Duration) -> Self {
51 Self {
52 data: Arc::new(RwLock::new(HashMap::new())),
53 default_ttl: Some(ttl),
54 }
55 }
56}
57
58impl Default for MemoryCache {
59 fn default() -> Self {
60 Self::new()
61 }
62}
63
64impl Cache for MemoryCache {
65 fn get(&self, key: &str) -> Result<Option<Vec<u8>>, CacheError> {
66 let data = self.data.read()?;
67 if let Some(entry) = data.get(key) {
68 if let Some(expires_at) = entry.expires_at {
69 if expires_at <= Instant::now() {
70 return Ok(None);
71 }
72 }
73 Ok(Some(entry.value.clone()))
74 } else {
75 Ok(None)
76 }
77 }
78
79 fn set(&self, key: &str, value: Vec<u8>, ttl: Option<Duration>) -> Result<(), CacheError> {
80 let expires_at = ttl.or(self.default_ttl).map(|d| Instant::now() + d);
81 let mut data = self.data.write()?;
82 data.insert(key.to_string(), CacheEntry { value, expires_at });
83 Ok(())
84 }
85
86 fn delete(&self, key: &str) -> Result<(), CacheError> {
87 let mut data = self.data.write()?;
88 data.remove(key);
89 Ok(())
90 }
91
92 fn clear(&self) -> Result<(), CacheError> {
93 let mut data = self.data.write()?;
94 data.clear();
95 Ok(())
96 }
97
98 fn exists(&self, key: &str) -> Result<bool, CacheError> {
99 let data = self.data.read()?;
100 if let Some(entry) = data.get(key) {
101 if let Some(expires_at) = entry.expires_at {
102 if expires_at <= Instant::now() {
103 return Ok(false);
104 }
105 }
106 Ok(true)
107 } else {
108 Ok(false)
109 }
110 }
111
112 fn expire(&self, key: &str, ttl: Duration) -> Result<(), CacheError> {
113 let mut data = self.data.write()?;
114 if let Some(entry) = data.get_mut(key) {
115 entry.expires_at = Some(Instant::now() + ttl);
116 Ok(())
117 } else {
118 Err(CacheError::NotFound(key.to_string()))
119 }
120 }
121
122 fn ttl(&self, key: &str) -> Result<Option<Duration>, CacheError> {
123 let data = self.data.read()?;
124 if let Some(entry) = data.get(key) {
125 if let Some(expires_at) = entry.expires_at {
126 if expires_at <= Instant::now() {
127 return Ok(None);
128 }
129 let remaining = expires_at.duration_since(Instant::now());
130 Ok(Some(remaining))
131 } else {
132 Ok(None)
133 }
134 } else {
135 Err(CacheError::NotFound(key.to_string()))
136 }
137 }
138}
139
140pub struct MultiLevelCache {
142 caches: Vec<Box<dyn Cache>>,
143}
144
145impl MultiLevelCache {
146 pub fn new() -> Self {
148 Self { caches: Vec::new() }
149 }
150
151 pub fn add_cache(mut self, cache: Box<dyn Cache>) -> Self {
153 self.caches.push(cache);
154 self
155 }
156}
157
158impl Default for MultiLevelCache {
159 fn default() -> Self {
160 Self::new()
161 }
162}
163
164impl Cache for MultiLevelCache {
165 fn get(&self, key: &str) -> Result<Option<Vec<u8>>, CacheError> {
166 for (i, cache) in self.caches.iter().enumerate() {
167 if let Ok(Some(value)) = cache.get(key) {
168 let ttl = cache.ttl(key).ok().flatten();
170 for j in 0..i {
171 let _ = self.caches[j].set(key, value.clone(), ttl);
172 }
173 return Ok(Some(value));
174 }
175 }
176 Ok(None)
177 }
178
179 fn set(&self, key: &str, value: Vec<u8>, ttl: Option<Duration>) -> Result<(), CacheError> {
180 for cache in &self.caches {
181 cache.set(key, value.clone(), ttl)?;
182 }
183 Ok(())
184 }
185
186 fn delete(&self, key: &str) -> Result<(), CacheError> {
187 for cache in &self.caches {
188 cache.delete(key)?;
189 }
190 Ok(())
191 }
192
193 fn clear(&self) -> Result<(), CacheError> {
194 for cache in &self.caches {
195 cache.clear()?;
196 }
197 Ok(())
198 }
199
200 fn exists(&self, key: &str) -> Result<bool, CacheError> {
201 for cache in &self.caches {
202 if cache.exists(key)? {
203 return Ok(true);
204 }
205 }
206 Ok(false)
207 }
208
209 fn expire(&self, key: &str, ttl: Duration) -> Result<(), CacheError> {
210 for cache in &self.caches {
211 cache.expire(key, ttl)?;
212 }
213 Ok(())
214 }
215
216 fn ttl(&self, key: &str) -> Result<Option<Duration>, CacheError> {
217 if let Some(cache) = self.caches.first() {
218 cache.ttl(key)
219 } else {
220 Err(CacheError::NotFound("No caches configured".to_string()))
221 }
222 }
223}
224
225#[derive(Debug, Clone, Default)]
227pub struct CacheStats {
228 pub hits: u64,
230 pub misses: u64,
232 pub sets: u64,
234 pub deletes: u64,
236}
237
238pub fn read_through<F>(
253 cache: &dyn Cache,
254 key: &str,
255 ttl: Option<Duration>,
256 loader: F,
257) -> Result<Option<Vec<u8>>, CacheError>
258where
259 F: FnOnce() -> Result<Option<Vec<u8>>, CacheError>,
260{
261 if let Some(v) = cache.get(key)? {
263 return Ok(Some(v));
264 }
265 let value = loader()?;
267 if let Some(ref v) = value {
269 cache.set(key, v.clone(), ttl)?;
270 }
271 Ok(value)
272}
273
274pub async fn read_through_async<F, Fut>(
278 cache: &dyn Cache,
279 key: &str,
280 ttl: Option<Duration>,
281 loader: F,
282) -> Result<Option<Vec<u8>>, CacheError>
283where
284 F: FnOnce() -> Fut,
285 Fut: std::future::Future<Output = Result<Option<Vec<u8>>, CacheError>>,
286{
287 if let Some(v) = cache.get(key)? {
289 return Ok(Some(v));
290 }
291 let value = loader().await?;
293 if let Some(ref v) = value {
295 cache.set(key, v.clone(), ttl)?;
296 }
297 Ok(value)
298}
299
300pub fn write_through<F>(
311 cache: &dyn Cache,
312 key: &str,
313 value: Vec<u8>,
314 ttl: Option<Duration>,
315 writer: F,
316) -> Result<(), CacheError>
317where
318 F: FnOnce(&str, &[u8]) -> Result<(), CacheError>,
319{
320 writer(key, &value)?;
322 cache.set(key, value, ttl)?;
324 Ok(())
325}
326
327pub async fn write_through_async<F, Fut>(
331 cache: &dyn Cache,
332 key: &str,
333 value: Vec<u8>,
334 ttl: Option<Duration>,
335 writer: F,
336) -> Result<(), CacheError>
337where
338 F: FnOnce(&str, Vec<u8>) -> Fut,
339 Fut: std::future::Future<Output = Result<Vec<u8>, CacheError>>,
340{
341 let stored = writer(key, value).await?;
343 cache.set(key, stored, ttl)?;
345 Ok(())
346}
347
348pub fn write_around<F>(cache: &dyn Cache, key: &str, writer: F) -> Result<(), CacheError>
352where
353 F: FnOnce(&str) -> Result<(), CacheError>,
354{
355 writer(key)?;
356 cache.delete(key)?;
357 Ok(())
358}
359
360#[derive(Debug, Clone, PartialEq, Eq)]
367pub enum CacheLookup {
368 Found(Vec<u8>),
370 Miss,
372 NotFound,
374}
375
376pub struct NegativeCache<C: Cache> {
383 inner: C,
385 negatives: RwLock<HashMap<String, Instant>>,
387 negative_ttl: Duration,
389}
390
391impl<C: Cache> NegativeCache<C> {
392 pub fn new(inner: C) -> Self {
394 Self {
395 inner,
396 negatives: RwLock::new(HashMap::new()),
397 negative_ttl: Duration::from_secs(60),
398 }
399 }
400
401 pub fn with_negative_ttl(mut self, ttl: Duration) -> Self {
403 self.negative_ttl = ttl;
404 self
405 }
406
407 pub fn get_or_negative(&self, key: &str) -> Result<CacheLookup, CacheError> {
414 let now = Instant::now();
416 {
417 let neg = self.negatives.read()?;
418 if let Some(expires_at) = neg.get(key) {
419 if *expires_at > now {
420 return Ok(CacheLookup::NotFound);
421 }
422 }
423 }
424 match self.inner.get(key)? {
426 Some(v) => Ok(CacheLookup::Found(v)),
427 None => {
428 let mut neg = self.negatives.write()?;
430 neg.insert(key.to_string(), now + self.negative_ttl);
431 Ok(CacheLookup::Miss)
432 }
433 }
434 }
435
436 pub fn purge_expired(&self) -> Result<usize, CacheError> {
438 let now = Instant::now();
439 let mut neg = self.negatives.write()?;
440 let before = neg.len();
441 neg.retain(|_, expires_at| *expires_at > now);
442 Ok(before - neg.len())
443 }
444}
445
446impl<C: Cache> Cache for NegativeCache<C> {
447 fn get(&self, key: &str) -> Result<Option<Vec<u8>>, CacheError> {
448 match self.get_or_negative(key)? {
449 CacheLookup::Found(v) => Ok(Some(v)),
450 CacheLookup::Miss | CacheLookup::NotFound => Ok(None),
451 }
452 }
453
454 fn set(&self, key: &str, value: Vec<u8>, ttl: Option<Duration>) -> Result<(), CacheError> {
455 {
457 let mut neg = self.negatives.write()?;
458 neg.remove(key);
459 }
460 self.inner.set(key, value, ttl)
461 }
462
463 fn delete(&self, key: &str) -> Result<(), CacheError> {
464 {
465 let mut neg = self.negatives.write()?;
466 neg.remove(key);
467 }
468 self.inner.delete(key)
469 }
470
471 fn clear(&self) -> Result<(), CacheError> {
472 {
473 let mut neg = self.negatives.write()?;
474 neg.clear();
475 }
476 self.inner.clear()
477 }
478
479 fn exists(&self, key: &str) -> Result<bool, CacheError> {
480 self.inner.exists(key)
481 }
482
483 fn expire(&self, key: &str, ttl: Duration) -> Result<(), CacheError> {
484 self.inner.expire(key, ttl)
485 }
486
487 fn ttl(&self, key: &str) -> Result<Option<Duration>, CacheError> {
488 self.inner.ttl(key)
489 }
490}
491
492#[cfg(test)]
493mod tests {
494 use super::*;
495
496 #[test]
497 fn test_memory_cache_set_get() {
498 let cache = MemoryCache::new();
499 cache.set("key1", b"value1".to_vec(), None).unwrap();
500 let val = cache.get("key1").unwrap();
501 assert_eq!(val, Some(b"value1".to_vec()));
502 }
503
504 #[test]
505 fn test_memory_cache_delete() {
506 let cache = MemoryCache::new();
507 cache.set("key1", b"value1".to_vec(), None).unwrap();
508 cache.delete("key1").unwrap();
509 let val = cache.get("key1").unwrap();
510 assert_eq!(val, None);
511 }
512
513 #[test]
514 fn test_memory_cache_exists() {
515 let cache = MemoryCache::new();
516 cache.set("key1", b"value1".to_vec(), None).unwrap();
517 let exists = cache.exists("key1").unwrap();
518 assert!(exists);
519 let exists2 = cache.exists("nonexistent").unwrap();
520 assert!(!exists2);
521 }
522
523 #[test]
524 fn test_memory_cache_clear() {
525 let cache = MemoryCache::new();
526 cache.set("key1", b"value1".to_vec(), None).unwrap();
527 cache.set("key2", b"value2".to_vec(), None).unwrap();
528 cache.clear().unwrap();
529 let val = cache.get("key1").unwrap();
530 assert_eq!(val, None);
531 }
532
533 #[test]
534 fn test_memory_cache_with_ttl() {
535 let cache = MemoryCache::with_ttl(Duration::from_secs(1));
536 cache.set("key1", b"value1".to_vec(), None).unwrap();
537 let val = cache.get("key1").unwrap();
538 assert!(val.is_some());
539 }
540
541 #[test]
542 fn test_cache_stats() {
543 let stats = CacheStats::default();
544 assert_eq!(stats.hits, 0);
545 assert_eq!(stats.misses, 0);
546 }
547
548 #[test]
549 fn test_multi_level_cache() {
550 let cache1 = MemoryCache::new();
551 let cache2 = MemoryCache::new();
552 let multi = MultiLevelCache::new()
553 .add_cache(Box::new(cache1))
554 .add_cache(Box::new(cache2));
555
556 multi.set("key1", b"value1".to_vec(), None).unwrap();
557 let val = multi.get("key1").unwrap();
558 assert_eq!(val, Some(b"value1".to_vec()));
559
560 multi.delete("key1").unwrap();
561 let val = multi.get("key1").unwrap();
562 assert_eq!(val, None);
563 }
564
565 #[test]
567 fn test_negative_cache_found() {
568 let inner = MemoryCache::new();
569 inner.set("key1", b"value1".to_vec(), None).unwrap();
570 let cache = NegativeCache::new(inner);
571 let result = cache.get_or_negative("key1").unwrap();
572 assert_eq!(result, CacheLookup::Found(b"value1".to_vec()));
573 }
574
575 #[test]
576 fn test_negative_cache_miss_then_not_found() {
577 let inner = MemoryCache::new();
578 let cache = NegativeCache::new(inner);
579 let result = cache.get_or_negative("missing").unwrap();
581 assert_eq!(result, CacheLookup::Miss);
582 let result = cache.get_or_negative("missing").unwrap();
584 assert_eq!(result, CacheLookup::NotFound);
585 }
586
587 #[test]
588 fn test_negative_cache_set_clears_negative() {
589 let inner = MemoryCache::new();
590 let cache = NegativeCache::new(inner);
591 cache.get_or_negative("key1").unwrap();
593 assert_eq!(
594 cache.get_or_negative("key1").unwrap(),
595 CacheLookup::NotFound
596 );
597 cache.set("key1", b"value1".to_vec(), None).unwrap();
599 assert_eq!(
600 cache.get_or_negative("key1").unwrap(),
601 CacheLookup::Found(b"value1".to_vec())
602 );
603 }
604
605 #[test]
606 fn test_negative_cache_delete_clears_negative() {
607 let inner = MemoryCache::new();
608 inner.set("key1", b"value1".to_vec(), None).unwrap();
609 let cache = NegativeCache::new(inner);
610 cache.delete("key1").unwrap();
612 let result = cache.get_or_negative("key1").unwrap();
613 assert_eq!(result, CacheLookup::Miss);
614 }
615
616 #[test]
617 fn test_negative_cache_ttl_expiry() {
618 let inner = MemoryCache::new();
619 let cache = NegativeCache::new(inner).with_negative_ttl(Duration::from_millis(50));
620 cache.get_or_negative("key1").unwrap();
622 assert_eq!(
623 cache.get_or_negative("key1").unwrap(),
624 CacheLookup::NotFound
625 );
626 std::thread::sleep(Duration::from_millis(60));
628 let result = cache.get_or_negative("key1").unwrap();
630 assert_eq!(result, CacheLookup::Miss);
631 }
632
633 #[test]
634 fn test_negative_cache_clear() {
635 let inner = MemoryCache::new();
636 let cache = NegativeCache::new(inner);
637 cache.get_or_negative("key1").unwrap();
638 cache.get_or_negative("key2").unwrap();
639 cache.clear().unwrap();
640 assert_eq!(cache.get_or_negative("key1").unwrap(), CacheLookup::Miss);
642 assert_eq!(cache.get_or_negative("key2").unwrap(), CacheLookup::Miss);
643 }
644
645 #[test]
646 fn test_negative_cache_purge_expired() {
647 let inner = MemoryCache::new();
648 let cache = NegativeCache::new(inner).with_negative_ttl(Duration::from_millis(50));
649 cache.get_or_negative("key1").unwrap();
650 cache.get_or_negative("key2").unwrap();
651 std::thread::sleep(Duration::from_millis(60));
652 let purged = cache.purge_expired().unwrap();
653 assert_eq!(purged, 2);
654 }
655
656 #[test]
657 fn test_negative_cache_as_cache_trait() {
658 let inner = MemoryCache::new();
659 let cache = NegativeCache::new(inner);
660 let val = cache.get("missing").unwrap();
662 assert_eq!(val, None);
663 let val = cache.get("missing").unwrap();
665 assert_eq!(val, None);
666 cache.set("key1", b"value1".to_vec(), None).unwrap();
668 let val = cache.get("key1").unwrap();
669 assert_eq!(val, Some(b"value1".to_vec()));
670 }
671}