Skip to main content

multi_tier_cache/backends/
dashmap_cache.rs

1use crate::error::CacheResult;
2use crate::traits::{CacheBackend, L2CacheBackend};
3use bytes::Bytes;
4use dashmap::DashMap;
5use futures_util::future::BoxFuture;
6use std::sync::Arc;
7use std::sync::atomic::{AtomicU64, Ordering};
8use std::time::{Duration, Instant};
9use tracing::{debug, info};
10
11/// Cache entry with expiration tracking
12#[derive(Debug, Clone)]
13struct CacheEntry {
14    value: Bytes,
15    expires_at: Option<Instant>,
16}
17
18impl CacheEntry {
19    fn new(value: Bytes, ttl: Duration) -> Self {
20        Self {
21            value,
22            expires_at: Some(Instant::now() + ttl),
23        }
24    }
25
26    fn is_expired(&self) -> bool {
27        self.expires_at
28            .is_some_and(|expires_at| Instant::now() > expires_at)
29    }
30}
31
32pub struct DashMapCache {
33    /// Concurrent `HashMap`
34    map: Arc<DashMap<String, CacheEntry>>,
35    /// Optional maximum capacity limit
36    max_capacity: Option<usize>,
37    /// Hit counter
38    hits: Arc<AtomicU64>,
39    /// Miss counter
40    misses: Arc<AtomicU64>,
41    /// Set counter
42    sets: Arc<AtomicU64>,
43}
44
45impl DashMapCache {
46    /// Create new `DashMap` cache
47    pub fn new() -> Self {
48        info!("Initializing DashMap Cache (concurrent HashMap)");
49
50        Self {
51            map: Arc::new(DashMap::new()),
52            max_capacity: None,
53            hits: Arc::new(AtomicU64::new(0)),
54            misses: Arc::new(AtomicU64::new(0)),
55            sets: Arc::new(AtomicU64::new(0)),
56        }
57    }
58
59    /// Create new `DashMap` cache with capacity limit
60    #[must_use]
61    pub fn new_with_capacity(capacity: usize) -> Self {
62        info!(capacity = capacity, "Initializing DashMap Cache with capacity limit");
63
64        Self {
65            map: Arc::new(DashMap::with_capacity(capacity)),
66            max_capacity: Some(capacity),
67            hits: Arc::new(AtomicU64::new(0)),
68            misses: Arc::new(AtomicU64::new(0)),
69            sets: Arc::new(AtomicU64::new(0)),
70        }
71    }
72
73    /// Set maximum capacity limit
74    #[must_use]
75    pub fn with_max_capacity(mut self, capacity: usize) -> Self {
76        self.max_capacity = Some(capacity);
77        self
78    }
79
80    /// Cleanup expired entries
81    pub fn cleanup_expired(&self) -> usize {
82        let mut removed = 0;
83        self.map.retain(|_, entry| {
84            if entry.is_expired() {
85                removed += 1;
86                false
87            } else {
88                true
89            }
90        });
91        if removed > 0 {
92            debug!(count = removed, "[DashMap] Cleaned up expired entries");
93        }
94        removed
95    }
96
97    /// Get current cache size
98    #[must_use]
99    pub fn len(&self) -> usize {
100        self.map.len()
101    }
102
103    /// Check if cache is empty
104    #[must_use]
105    pub fn is_empty(&self) -> bool {
106        self.map.is_empty()
107    }
108}
109
110impl Default for DashMapCache {
111    fn default() -> Self {
112        Self::new()
113    }
114}
115
116// ===== Trait Implementations =====
117
118/// Implement `CacheBackend` trait for `DashMapCache`
119impl CacheBackend for DashMapCache {
120    fn get<'a>(&'a self, key: &'a str) -> BoxFuture<'a, Option<Bytes>> {
121        Box::pin(async move {
122            match self.map.get(key) {
123                Some(entry) => {
124                    if entry.is_expired() {
125                        drop(entry);
126                        self.map.remove(key);
127                        None
128                    } else {
129                        Some(entry.value.clone())
130                    }
131                }
132                None => None,
133            }
134        })
135    }
136
137    fn set_with_ttl<'a>(
138        &'a self,
139        key: &'a str,
140        value: Bytes,
141        ttl: Duration,
142    ) -> BoxFuture<'a, CacheResult<()>> {
143        Box::pin(async move {
144            // If capacity limit reached and not an existing key, evict expired entries
145            if let Some(cap) = self.max_capacity
146                && self.map.len() >= cap
147                && !self.map.contains_key(key)
148            {
149                self.cleanup_expired();
150            }
151
152            let entry = CacheEntry::new(value, ttl);
153            self.map.insert(key.to_string(), entry);
154            self.sets.fetch_add(1, Ordering::Relaxed);
155            debug!(key = %key, ttl_secs = %ttl.as_secs(), "[DashMap] Cached key bytes with TTL");
156            Ok(())
157        })
158    }
159
160    fn remove<'a>(&'a self, key: &'a str) -> BoxFuture<'a, CacheResult<()>> {
161        Box::pin(async move {
162            self.map.remove(key);
163            Ok(())
164        })
165    }
166
167    fn remove_pattern<'a>(&'a self, pattern: &'a str) -> BoxFuture<'a, CacheResult<()>> {
168        Box::pin(async move {
169            self.map
170                .retain(|key, _| !crate::backends::matches_pattern(key, pattern));
171            Ok(())
172        })
173    }
174
175    fn health_check(&self) -> BoxFuture<'_, bool> {
176        Box::pin(async move { true })
177    }
178
179    fn name(&self) -> &'static str {
180        "DashMap"
181    }
182}
183
184impl L2CacheBackend for DashMapCache {
185    fn get_with_ttl<'a>(
186        &'a self,
187        key: &'a str,
188    ) -> BoxFuture<'a, Option<(Bytes, Option<Duration>)>> {
189        Box::pin(async move {
190            if let Some(entry) = self.map.get(key) {
191                if entry.is_expired() {
192                    drop(entry);
193                    self.map.remove(key);
194                    self.misses.fetch_add(1, Ordering::Relaxed);
195                    None
196                } else {
197                    let now = Instant::now();
198                    if let Some(expires_at) = entry.expires_at {
199                        let ttl = expires_at.checked_duration_since(now);
200                        if ttl.is_none() {
201                            // Expired
202                            drop(entry);
203                            self.map.remove(key);
204                            self.misses.fetch_add(1, Ordering::Relaxed);
205                            return None;
206                        }
207                        self.hits.fetch_add(1, Ordering::Relaxed);
208                        Some((entry.value.clone(), ttl))
209                    } else {
210                        self.hits.fetch_add(1, Ordering::Relaxed);
211                        Some((entry.value.clone(), None))
212                    }
213                }
214            } else {
215                self.misses.fetch_add(1, Ordering::Relaxed);
216                None
217            }
218        })
219    }
220}
221
222#[cfg(test)]
223mod tests {
224    use super::*;
225
226    #[tokio::test]
227    async fn test_dashmap_cache_basic_and_capacity() {
228        let cache = DashMapCache::new_with_capacity(2);
229        assert_eq!(cache.len(), 0);
230
231        cache
232            .set_with_ttl("k1", Bytes::from("v1"), Duration::from_millis(50))
233            .await
234            .unwrap();
235        cache
236            .set_with_ttl("k2", Bytes::from("v2"), Duration::from_secs(10))
237            .await
238            .unwrap();
239        assert_eq!(cache.len(), 2);
240
241        // Wait for k1 to expire
242        tokio::time::sleep(Duration::from_millis(60)).await;
243
244        // Inserting k3 when len >= cap triggers cleanup_expired, evicting expired k1
245        cache
246            .set_with_ttl("k3", Bytes::from("v3"), Duration::from_secs(10))
247            .await
248            .unwrap();
249
250        assert_eq!(cache.get("k1").await, None);
251        assert_eq!(cache.get("k2").await, Some(Bytes::from("v2")));
252        assert_eq!(cache.get("k3").await, Some(Bytes::from("v3")));
253    }
254}