kv_core/storage/
memory.rs

1//! In-memory storage implementation
2
3use std::collections::HashMap;
4use std::sync::Arc;
5use tokio::sync::RwLock;
6use tracing::debug;
7
8use crate::{
9    KVResult, Key, Entry, DatabaseId, Storage, StorageStats,
10};
11
12/// In-memory storage using `HashMap`
13pub struct MemoryStorage {
14    /// Map of `database_id` -> (key -> entry)
15    databases: Arc<RwLock<HashMap<DatabaseId, HashMap<Key, Entry>>>>,
16}
17
18impl Default for MemoryStorage {
19    fn default() -> Self {
20        Self::new()
21    }
22}
23
24impl MemoryStorage {
25    /// Create a new in-memory storage
26    #[must_use] 
27    pub fn new() -> Self {
28        Self {
29            databases: Arc::new(RwLock::new(HashMap::new())),
30        }
31    }
32
33    /// Get or create database
34    async fn _get_or_create_database(&self, database_id: DatabaseId) -> HashMap<Key, Entry> {
35        let mut databases = self.databases.write().await;
36        databases.entry(database_id).or_insert_with(HashMap::new).clone()
37    }
38}
39
40#[async_trait::async_trait]
41impl Storage for MemoryStorage {
42    #[allow(clippy::significant_drop_tightening, clippy::option_if_let_else)]
43    async fn get(&self, database_id: DatabaseId, key: &Key) -> KVResult<Option<Entry>> {
44        let databases = self.databases.read().await;
45        if let Some(db) = databases.get(&database_id) {
46            Ok(db.get(key).cloned())
47        } else {
48            Ok(None)
49        }
50    }
51
52    #[allow(clippy::significant_drop_tightening)]
53    async fn set(&self, database_id: DatabaseId, key: Key, entry: Entry) -> KVResult<()> {
54        let mut databases = self.databases.write().await;
55        let db = databases.entry(database_id).or_insert_with(HashMap::new);
56        db.insert(key, entry);
57        Ok(())
58    }
59
60    #[allow(clippy::significant_drop_tightening, clippy::option_if_let_else)]
61    async fn delete(&self, database_id: DatabaseId, key: &Key) -> KVResult<bool> {
62        let mut databases = self.databases.write().await;
63        if let Some(db) = databases.get_mut(&database_id) {
64            Ok(db.remove(key).is_some())
65        } else {
66            Ok(false)
67        }
68    }
69
70    #[allow(clippy::significant_drop_tightening, clippy::option_if_let_else)]
71    async fn exists(&self, database_id: DatabaseId, key: &Key) -> KVResult<bool> {
72        let databases = self.databases.read().await;
73        if let Some(db) = databases.get(&database_id) {
74            Ok(db.contains_key(key))
75        } else {
76            Ok(false)
77        }
78    }
79
80    #[allow(clippy::significant_drop_tightening, clippy::option_if_let_else)]
81    async fn keys(&self, database_id: DatabaseId) -> KVResult<Vec<Key>> {
82        let databases = self.databases.read().await;
83        if let Some(db) = databases.get(&database_id) {
84            Ok(db.keys().cloned().collect())
85        } else {
86            Ok(Vec::new())
87        }
88    }
89
90    #[allow(clippy::significant_drop_tightening, clippy::option_if_let_else)]
91    async fn keys_pattern(&self, database_id: DatabaseId, pattern: &str) -> KVResult<Vec<Key>> {
92        let databases = self.databases.read().await;
93        if let Some(db) = databases.get(&database_id) {
94            let keys: Vec<Key> = db.keys()
95                .filter(|key| matches_pattern(key, pattern))
96                .cloned()
97                .collect();
98            Ok(keys)
99        } else {
100            Ok(Vec::new())
101        }
102    }
103
104    #[allow(clippy::significant_drop_tightening)]
105    async fn clear_database(&self, database_id: DatabaseId) -> KVResult<()> {
106        let mut databases = self.databases.write().await;
107        if let Some(db) = databases.get_mut(&database_id) {
108            db.clear();
109        }
110        Ok(())
111    }
112
113    #[allow(clippy::significant_drop_tightening, clippy::option_if_let_else)]
114    async fn get_stats(&self, database_id: DatabaseId) -> KVResult<StorageStats> {
115        let databases = self.databases.read().await;
116        if let Some(db) = databases.get(&database_id) {
117            let total_keys = db.len() as u64;
118            let memory_usage = std::mem::size_of_val(db) as u64;
119            
120            Ok(StorageStats {
121                total_keys,
122                memory_usage,
123                disk_usage: None,
124                last_flush: None,
125            })
126        } else {
127            Ok(StorageStats {
128                total_keys: 0,
129                memory_usage: 0,
130                disk_usage: None,
131                last_flush: None,
132            })
133        }
134    }
135
136    async fn flush(&self) -> KVResult<()> {
137        // No-op for in-memory storage
138        debug!("Memory storage flush (no-op)");
139        Ok(())
140    }
141
142    async fn close(&self) -> KVResult<()> {
143        debug!("Closing memory storage");
144        Ok(())
145    }
146}
147
148/// Simple pattern matching (supports * wildcard)
149fn matches_pattern(key: &str, pattern: &str) -> bool {
150    if pattern == "*" {
151        return true;
152    }
153    
154    if !pattern.contains('*') {
155        return key == pattern;
156    }
157    
158    // Simple wildcard matching
159    let pattern_parts: Vec<&str> = pattern.split('*').collect();
160    if pattern_parts.len() == 2 {
161        let prefix = pattern_parts[0];
162        let suffix = pattern_parts[1];
163        
164        if prefix.is_empty() {
165            key.ends_with(suffix)
166        } else if suffix.is_empty() {
167            key.starts_with(prefix)
168        } else {
169            key.starts_with(prefix) && key.ends_with(suffix)
170        }
171    } else {
172        // More complex pattern - for now, just do simple contains
173        key.contains(pattern.trim_matches('*'))
174    }
175}
176
177#[cfg(test)]
178mod tests {
179    use super::*;
180    use crate::{Value, Entry};
181
182    #[tokio::test]
183    async fn test_memory_storage_basic_operations() {
184        let storage = MemoryStorage::new();
185        let database_id = 0;
186        
187        // Test set and get
188        let entry = Entry::new(Value::String("test_value".to_string()), None);
189        storage.set(database_id, "test_key".to_string(), entry.clone()).await.unwrap();
190        
191        let retrieved = storage.get(database_id, &"test_key".to_string()).await.unwrap();
192        assert!(retrieved.is_some());
193        assert_eq!(retrieved.unwrap().value.as_string().unwrap(), "test_value");
194        
195        // Test exists
196        let exists = storage.exists(database_id, &"test_key".to_string()).await.unwrap();
197        assert!(exists);
198        
199        // Test delete
200        let deleted = storage.delete(database_id, &"test_key".to_string()).await.unwrap();
201        assert!(deleted);
202        
203        let exists_after = storage.exists(database_id, &"test_key".to_string()).await.unwrap();
204        assert!(!exists_after);
205    }
206
207    #[tokio::test]
208    async fn test_memory_storage_keys() {
209        let storage = MemoryStorage::new();
210        let database_id = 0;
211        
212        // Add some keys
213        let entry = Entry::new(Value::String("value".to_string()), None);
214        storage.set(database_id, "key1".to_string(), entry.clone()).await.unwrap();
215        storage.set(database_id, "key2".to_string(), entry.clone()).await.unwrap();
216        storage.set(database_id, "test_key".to_string(), entry.clone()).await.unwrap();
217        
218        // Test get all keys
219        let keys = storage.keys(database_id).await.unwrap();
220        assert_eq!(keys.len(), 3);
221        assert!(keys.contains(&"key1".to_string()));
222        assert!(keys.contains(&"key2".to_string()));
223        assert!(keys.contains(&"test_key".to_string()));
224        
225        // Test pattern matching
226        let test_keys = storage.keys_pattern(database_id, "key*").await.unwrap();
227        assert_eq!(test_keys.len(), 2);
228        assert!(test_keys.contains(&"key1".to_string()));
229        assert!(test_keys.contains(&"key2".to_string()));
230    }
231
232    #[tokio::test]
233    async fn test_memory_storage_clear() {
234        let storage = MemoryStorage::new();
235        let database_id = 0;
236        
237        // Add some keys
238        let entry = Entry::new(Value::String("value".to_string()), None);
239        storage.set(database_id, "key1".to_string(), entry.clone()).await.unwrap();
240        storage.set(database_id, "key2".to_string(), entry.clone()).await.unwrap();
241        
242        // Clear database
243        storage.clear_database(database_id).await.unwrap();
244        
245        // Check keys are gone
246        let keys = storage.keys(database_id).await.unwrap();
247        assert!(keys.is_empty());
248    }
249
250    #[tokio::test]
251    async fn test_memory_storage_stats() {
252        let storage = MemoryStorage::new();
253        let database_id = 0;
254        
255        // Add some keys
256        let entry = Entry::new(Value::String("value".to_string()), None);
257        storage.set(database_id, "key1".to_string(), entry.clone()).await.unwrap();
258        storage.set(database_id, "key2".to_string(), entry.clone()).await.unwrap();
259        
260        let stats = storage.get_stats(database_id).await.unwrap();
261        assert_eq!(stats.total_keys, 2);
262        assert!(stats.memory_usage > 0);
263        assert!(stats.disk_usage.is_none());
264    }
265}
kv_core/storage/memory.rs

kv_core/storage/
memory.rs
