Skip to main content

multi_tier_cache/backends/
moka_cache.rs

1use crate::error::CacheResult;
2use bytes::Bytes;
3use futures_util::future::BoxFuture;
4use moka::future::Cache;
5use std::any::Any;
6use std::sync::Arc;
7use std::sync::atomic::{AtomicU64, Ordering};
8use std::time::{Duration, Instant};
9use tracing::{debug, info};
10
11/// Cache entry with TTL information
12#[derive(Debug, Clone)]
13struct CacheEntry {
14    value: Bytes,
15    expires_at: Instant,
16}
17
18impl CacheEntry {
19    fn new(value: Bytes, ttl: Duration) -> Self {
20        Self {
21            value,
22            expires_at: Instant::now() + ttl,
23        }
24    }
25
26    fn is_expired(&self) -> bool {
27        Instant::now() > self.expires_at
28    }
29}
30
31/// Specialized entry for zero-cost deserialization on L1
32#[derive(Clone)]
33pub struct TypedCacheEntry {
34    pub value: Arc<dyn Any + Send + Sync>,
35    pub expires_at: Instant,
36}
37
38impl TypedCacheEntry {
39    pub fn new(value: Arc<dyn Any + Send + Sync>, ttl: Duration) -> Self {
40        Self {
41            value,
42            expires_at: Instant::now() + ttl,
43        }
44    }
45
46    #[must_use]
47    pub fn is_expired(&self) -> bool {
48        Instant::now() > self.expires_at
49    }
50}
51
52/// Configuration for `MokaCache`
53#[derive(Debug, Clone, Copy)]
54pub struct MokaCacheConfig {
55    /// Max capacity of the cache
56    pub max_capacity: u64,
57    /// Time to live for cache entries
58    pub time_to_live: Duration,
59    /// Time to idle for cache entries
60    pub time_to_idle: Duration,
61}
62
63impl Default for MokaCacheConfig {
64    fn default() -> Self {
65        Self {
66            max_capacity: 5000,
67            time_to_live: Duration::from_hours(1),
68            time_to_idle: Duration::from_mins(2),
69        }
70    }
71}
72
73/// Moka in-memory cache with per-key TTL support
74pub struct MokaCache {
75    /// Moka cache instance for raw bytes
76    cache: Cache<String, CacheEntry>,
77    /// Moka cache instance for typed objects (Zero-cost optimization)
78    typed_cache: Cache<String, TypedCacheEntry>,
79    /// Hit counter
80    hits: Arc<AtomicU64>,
81    /// Miss counter
82    misses: Arc<AtomicU64>,
83    /// Set counter
84    sets: Arc<AtomicU64>,
85    /// Coalesced requests counter
86    #[allow(dead_code)]
87    coalesced_requests: Arc<AtomicU64>,
88}
89
90impl MokaCache {
91    /// Create new Moka cache
92    ///
93    /// # Errors
94    ///
95    /// Returns an error if cache configuration is invalid.
96    pub fn new(config: MokaCacheConfig) -> CacheResult<Self> {
97        info!("Initializing Moka Cache");
98
99        let cache = Cache::builder()
100            .max_capacity(config.max_capacity)
101            .time_to_live(config.time_to_live)
102            .time_to_idle(config.time_to_idle)
103            .build();
104
105        let typed_cache = Cache::builder()
106            .max_capacity(config.max_capacity)
107            .time_to_live(config.time_to_live)
108            .time_to_idle(config.time_to_idle)
109            .build();
110
111        info!(
112            capacity = config.max_capacity,
113            "Moka Cache initialized with Byte and Typed storage"
114        );
115
116        Ok(Self {
117            cache,
118            typed_cache,
119            hits: Arc::new(AtomicU64::new(0)),
120            misses: Arc::new(AtomicU64::new(0)),
121            sets: Arc::new(AtomicU64::new(0)),
122            coalesced_requests: Arc::new(AtomicU64::new(0)),
123        })
124    }
125
126    /// Set a typed value in the L1 cache (zero-cost optimization)
127    ///
128    /// # Errors
129    ///
130    /// Returns an error if the key or value cannot be inserted into the cache.
131    pub async fn set_typed(
132        &self,
133        key: &str,
134        value: Arc<dyn Any + Send + Sync>,
135        ttl: Duration,
136    ) -> CacheResult<()> {
137        let entry = TypedCacheEntry::new(value, ttl);
138        self.typed_cache.insert(key.to_string(), entry).await;
139        self.sets.fetch_add(1, Ordering::Relaxed);
140        Ok(())
141    }
142
143    /// Get a typed value from the L1 cache
144    pub async fn get_typed(&self, key: &str) -> Option<Arc<dyn Any + Send + Sync>> {
145        match self.typed_cache.get(key).await {
146            Some(entry) => {
147                if entry.is_expired() {
148                    let _ = self.typed_cache.remove(key).await;
149                    self.misses.fetch_add(1, Ordering::Relaxed);
150                    None
151                } else {
152                    self.hits.fetch_add(1, Ordering::Relaxed);
153                    Some(entry.value)
154                }
155            }
156            _ => None,
157        }
158    }
159}
160
161// ===== Trait Implementations =====
162
163use crate::traits::{CacheBackend, L2CacheBackend};
164
165/// Implement `CacheBackend` trait for `MokaCache`
166impl CacheBackend for MokaCache {
167    fn get<'a>(&'a self, key: &'a str) -> BoxFuture<'a, Option<Bytes>> {
168        Box::pin(async move {
169            if let Some(entry) = self.cache.get(key).await {
170                if entry.is_expired() {
171                    let _ = self.cache.remove(key).await;
172                    self.misses.fetch_add(1, Ordering::Relaxed);
173                    None
174                } else {
175                    self.hits.fetch_add(1, Ordering::Relaxed);
176                    Some(entry.value)
177                }
178            } else {
179                self.misses.fetch_add(1, Ordering::Relaxed);
180                None
181            }
182        })
183    }
184
185    fn set_with_ttl<'a>(
186        &'a self,
187        key: &'a str,
188        value: Bytes,
189        ttl: Duration,
190    ) -> BoxFuture<'a, CacheResult<()>> {
191        Box::pin(async move {
192            let entry = CacheEntry::new(value, ttl);
193            self.cache.insert(key.to_string(), entry).await;
194            self.sets.fetch_add(1, Ordering::Relaxed);
195            debug!(key = %key, ttl_secs = %ttl.as_secs(), "[Moka] Cached key bytes with TTL");
196            Ok(())
197        })
198    }
199
200    fn remove<'a>(&'a self, key: &'a str) -> BoxFuture<'a, CacheResult<()>> {
201        Box::pin(async move {
202            self.cache.invalidate(key).await;
203            self.typed_cache.invalidate(key).await;
204            Ok(())
205        })
206    }
207
208    fn health_check(&self) -> BoxFuture<'_, bool> {
209        Box::pin(async move {
210            let test_key = "health_check_moka";
211            let test_value = Bytes::from("health_check");
212
213            match self
214                .set_with_ttl(test_key, test_value.clone(), Duration::from_mins(1))
215                .await
216            {
217                Ok(()) => match self.get(test_key).await {
218                    Some(retrieved) => {
219                        let _ = self.remove(test_key).await;
220                        retrieved == test_value
221                    }
222                    None => false,
223                },
224                Err(_) => false,
225            }
226        })
227    }
228
229    fn remove_pattern<'a>(&'a self, pattern: &'a str) -> BoxFuture<'a, CacheResult<()>> {
230        Box::pin(async move {
231            let mut keys_to_invalidate = Vec::new();
232            for (key, _) in &self.cache {
233                if crate::backends::matches_pattern(&key, pattern) {
234                    keys_to_invalidate.push((*key).clone());
235                }
236            }
237            for key in &keys_to_invalidate {
238                self.cache.invalidate(key).await;
239            }
240
241            let mut typed_keys_to_invalidate = Vec::new();
242            for (key, _) in &self.typed_cache {
243                if crate::backends::matches_pattern(&key, pattern) {
244                    typed_keys_to_invalidate.push((*key).clone());
245                }
246            }
247            for key in &typed_keys_to_invalidate {
248                self.typed_cache.invalidate(key).await;
249            }
250
251            // Ensure background invalidation tasks are processed
252            self.cache.run_pending_tasks().await;
253            self.typed_cache.run_pending_tasks().await;
254
255            debug!(pattern = %pattern, moka_count = %keys_to_invalidate.len(), typed_count = %typed_keys_to_invalidate.len(), "[Moka] Invalidated matching entries due to pattern '{}' request", pattern);
256            Ok(())
257        })
258    }
259
260    fn name(&self) -> &'static str {
261        "Moka"
262    }
263}
264
265impl L2CacheBackend for MokaCache {
266    fn get_with_ttl<'a>(
267        &'a self,
268        key: &'a str,
269    ) -> BoxFuture<'a, Option<(Bytes, Option<Duration>)>> {
270        Box::pin(async move {
271            // Moka doesn't easily expose remaining TTL for an entry
272            if let Some(entry) = self.cache.get(key).await {
273                if entry.is_expired() {
274                    let _ = self.cache.remove(key).await;
275                    self.misses.fetch_add(1, Ordering::Relaxed);
276                    None
277                } else {
278                    self.hits.fetch_add(1, Ordering::Relaxed);
279                    let now = Instant::now();
280                    let remaining = if entry.expires_at > now {
281                        Some(entry.expires_at.duration_since(now))
282                    } else {
283                        None
284                    };
285                    Some((entry.value, remaining))
286                }
287            } else {
288                self.misses.fetch_add(1, Ordering::Relaxed);
289                None
290            }
291        })
292    }
293}
294
295/// Cache statistics
296#[allow(dead_code)]
297#[derive(Debug, Clone)]
298pub struct CacheStats {
299    pub hits: u64,
300    pub misses: u64,
301    pub sets: u64,
302    pub coalesced_requests: u64,
303    pub size: u64,
304}