threatflux_cache/backends/
memory.rs

1//! In-memory storage backend
2
3use async_trait::async_trait;
4use std::collections::HashMap;
5use std::sync::Arc;
6use tokio::sync::RwLock;
7
8use crate::backends::{StorageKey, StorageMeta, StorageValue};
9use crate::storage::EntryMap;
10use crate::{CacheEntry, Result, StorageBackend};
11
12/// In-memory storage backend
13#[allow(clippy::type_complexity)]
14pub struct MemoryBackend<K, V, M = ()>
15where
16    K: StorageKey,
17    V: StorageValue,
18    M: StorageMeta,
19{
20    data: Arc<RwLock<HashMap<K, Vec<CacheEntry<K, V, M>>>>>,
21}
22
23impl<K, V, M> MemoryBackend<K, V, M>
24where
25    K: StorageKey,
26    V: StorageValue,
27    M: StorageMeta,
28{
29    /// Create a new memory backend
30    pub fn new() -> Self {
31        Self {
32            data: Arc::new(RwLock::new(HashMap::new())),
33        }
34    }
35}
36
37impl<K, V, M> Default for MemoryBackend<K, V, M>
38where
39    K: StorageKey,
40    V: StorageValue,
41    M: StorageMeta,
42{
43    fn default() -> Self {
44        Self::new()
45    }
46}
47
48impl<K, V, M> Clone for MemoryBackend<K, V, M>
49where
50    K: StorageKey,
51    V: StorageValue,
52    M: StorageMeta,
53{
54    fn clone(&self) -> Self {
55        Self {
56            data: Arc::clone(&self.data),
57        }
58    }
59}
60
61#[async_trait]
62impl<K, V, M> StorageBackend for MemoryBackend<K, V, M>
63where
64    K: StorageKey,
65    V: StorageValue,
66    M: StorageMeta,
67{
68    type Key = K;
69    type Value = V;
70    type Metadata = M;
71
72    async fn save(&self, entries: &EntryMap<K, V, M>) -> Result<()> {
73        let mut data = self.data.write().await;
74        *data = entries.clone();
75        Ok(())
76    }
77
78    async fn load(&self) -> Result<EntryMap<K, V, M>> {
79        let data = self.data.read().await;
80        Ok(data.clone())
81    }
82
83    async fn remove(&self, key: &K) -> Result<()> {
84        let mut data = self.data.write().await;
85        data.remove(key);
86        Ok(())
87    }
88
89    async fn clear(&self) -> Result<()> {
90        let mut data = self.data.write().await;
91        data.clear();
92        Ok(())
93    }
94
95    async fn contains(&self, key: &K) -> Result<bool> {
96        let data = self.data.read().await;
97        Ok(data.contains_key(key))
98    }
99
100    async fn size_bytes(&self) -> Result<u64> {
101        let data = self.data.read().await;
102
103        // Estimate size based on number of entries
104        let total_entries = data
105            .values()
106            .fold(0usize, |total, values| total.saturating_add(values.len()));
107        let estimated_size =
108            total_entries.saturating_mul(std::mem::size_of::<CacheEntry<K, V, M>>());
109
110        Ok(u64::try_from(estimated_size).unwrap_or(u64::MAX))
111    }
112}
113
114#[cfg(test)]
115mod tests {
116    use super::*;
117
118    #[tokio::test]
119    async fn test_memory_backend_clone() {
120        let backend1: MemoryBackend<String, String> = MemoryBackend::new();
121        let backend2 = backend1.clone();
122
123        // Changes in one should be reflected in the other
124        let mut entries = HashMap::new();
125        let entry = CacheEntry::new("key1".to_string(), "value1".to_string());
126        entries.insert("key1".to_string(), vec![entry]);
127
128        backend1.save(&entries).await.unwrap();
129        assert!(backend2.contains(&"key1".to_string()).await.unwrap());
130    }
131}
threatflux_cache/backends/memory.rs

threatflux_cache/backends/
memory.rs
