rabia_kvstore_example/
store.rs

1//! # KVStore Implementation
2//!
3//! Production-grade key-value store with consensus integration and change notifications.
4//! This is focused purely on the storage operations and data management.
5
6use crate::notifications::{ChangeNotification, ChangeType, NotificationBus};
7use crate::operations::{KVOperation, KVResult, StoreError};
8use dashmap::DashMap;
9use parking_lot::RwLock;
10use serde::{Deserialize, Serialize};
11use std::collections::HashMap;
12use std::sync::Arc;
13use std::time::{SystemTime, UNIX_EPOCH};
14use tokio::sync::watch;
15use tracing::{debug, info};
16
17/// Configuration for the KVStore
18#[derive(Debug, Clone)]
19pub struct KVStoreConfig {
20    /// Maximum number of keys to store
21    pub max_keys: usize,
22    /// Enable change notifications
23    pub enable_notifications: bool,
24    /// Snapshot frequency (number of operations)
25    pub snapshot_frequency: usize,
26    /// Enable compression for large values
27    pub enable_compression: bool,
28    /// Maximum value size in bytes
29    pub max_value_size: usize,
30}
31
32impl Default for KVStoreConfig {
33    fn default() -> Self {
34        Self {
35            max_keys: 1_000_000,
36            enable_notifications: true,
37            snapshot_frequency: 10_000,
38            enable_compression: false,
39            max_value_size: 1024 * 1024, // 1MB
40        }
41    }
42}
43
44/// Value entry in the store with metadata
45#[derive(Debug, Clone, Serialize, Deserialize)]
46pub struct ValueEntry {
47    pub value: String,
48    pub version: u64,
49    pub created_at: u64,
50    pub updated_at: u64,
51    pub size: usize,
52}
53
54impl ValueEntry {
55    pub fn new(value: String) -> Self {
56        let now = SystemTime::now()
57            .duration_since(UNIX_EPOCH)
58            .unwrap()
59            .as_millis() as u64;
60        let size = value.len();
61
62        Self {
63            value,
64            version: 1,
65            created_at: now,
66            updated_at: now,
67            size,
68        }
69    }
70
71    pub fn update(&mut self, new_value: String) {
72        self.value = new_value;
73        self.version += 1;
74        self.updated_at = SystemTime::now()
75            .duration_since(UNIX_EPOCH)
76            .unwrap()
77            .as_millis() as u64;
78        self.size = self.value.len();
79    }
80}
81
82/// Store statistics
83#[derive(Debug, Clone, Default)]
84pub struct StoreStats {
85    pub total_keys: usize,
86    pub total_operations: u64,
87    pub memory_usage_bytes: usize,
88    pub last_snapshot_at: u64,
89    pub operations_since_snapshot: usize,
90}
91
92/// Snapshot of the store state
93#[derive(Debug, Clone, Serialize, Deserialize)]
94pub struct StoreSnapshot {
95    pub data: HashMap<String, ValueEntry>,
96    pub version: u64,
97    pub created_at: u64,
98    pub checksum: u64,
99}
100
101/// Production-grade key-value store
102pub struct KVStore {
103    /// Configuration
104    pub(crate) config: KVStoreConfig,
105
106    /// Main data storage
107    data: Arc<DashMap<String, ValueEntry>>,
108
109    /// Store statistics
110    stats: Arc<RwLock<StoreStats>>,
111
112    /// Global version counter
113    version: Arc<std::sync::atomic::AtomicU64>,
114
115    /// Notification bus for change events
116    notification_bus: Arc<NotificationBus>,
117
118    /// Shutdown signal
119    shutdown_tx: watch::Sender<bool>,
120    #[allow(dead_code)]
121    shutdown_rx: watch::Receiver<bool>,
122}
123
124impl KVStore {
125    /// Create a new KVStore instance
126    pub async fn new(config: KVStoreConfig) -> Result<Self, StoreError> {
127        let (shutdown_tx, shutdown_rx) = watch::channel(false);
128
129        let store = Self {
130            config: config.clone(),
131            data: Arc::new(DashMap::new()),
132            stats: Arc::new(RwLock::new(StoreStats::default())),
133            version: Arc::new(std::sync::atomic::AtomicU64::new(0)),
134            notification_bus: Arc::new(NotificationBus::new()),
135            shutdown_tx,
136            shutdown_rx,
137        };
138
139        info!("KVStore initialized with config: {:?}", config);
140        Ok(store)
141    }
142
143    /// Set a key-value pair
144    pub async fn set(&self, key: &str, value: &str) -> Result<KVResult, StoreError> {
145        self.validate_key(key)?;
146        self.validate_value(value)?;
147
148        let old_value = if let Some(mut entry) = self.data.get_mut(key) {
149            let old = entry.value.clone();
150            entry.update(value.to_string());
151            Some(old)
152        } else {
153            if self.data.len() >= self.config.max_keys {
154                return Err(StoreError::StoreFull);
155            }
156            self.data
157                .insert(key.to_string(), ValueEntry::new(value.to_string()));
158            None
159        };
160
161        self.increment_operation_count();
162
163        // Send notification if enabled
164        if self.config.enable_notifications {
165            let change_type = if old_value.is_some() {
166                ChangeType::Updated
167            } else {
168                ChangeType::Created
169            };
170
171            let notification = ChangeNotification {
172                key: key.to_string(),
173                change_type,
174                old_value,
175                new_value: Some(value.to_string()),
176                version: self.get_version(),
177                timestamp: SystemTime::now()
178                    .duration_since(UNIX_EPOCH)
179                    .unwrap()
180                    .as_millis() as u64,
181            };
182
183            self.notification_bus.publish(notification).await;
184        }
185
186        debug!("SET operation: key={}, value_len={}", key, value.len());
187        Ok(KVResult::Success)
188    }
189
190    /// Get a value by key
191    pub async fn get(&self, key: &str) -> Result<Option<String>, StoreError> {
192        self.validate_key(key)?;
193
194        let result = self.data.get(key).map(|entry| entry.value.clone());
195        self.increment_operation_count();
196
197        debug!("GET operation: key={}, found={}", key, result.is_some());
198        Ok(result)
199    }
200
201    /// Get a value with metadata
202    pub async fn get_with_metadata(&self, key: &str) -> Result<Option<ValueEntry>, StoreError> {
203        self.validate_key(key)?;
204
205        let result = self.data.get(key).map(|entry| entry.clone());
206        self.increment_operation_count();
207
208        debug!(
209            "GET_META operation: key={}, found={}",
210            key,
211            result.is_some()
212        );
213        Ok(result)
214    }
215
216    /// Delete a key
217    pub async fn delete(&self, key: &str) -> Result<KVResult, StoreError> {
218        self.validate_key(key)?;
219
220        let old_value = self.data.remove(key).map(|(_, entry)| entry.value);
221        self.increment_operation_count();
222
223        // Send notification if enabled and key existed
224        if self.config.enable_notifications && old_value.is_some() {
225            let notification = ChangeNotification {
226                key: key.to_string(),
227                change_type: ChangeType::Deleted,
228                old_value: old_value.clone(),
229                new_value: None,
230                version: self.get_version(),
231                timestamp: SystemTime::now()
232                    .duration_since(UNIX_EPOCH)
233                    .unwrap()
234                    .as_millis() as u64,
235            };
236
237            self.notification_bus.publish(notification).await;
238        }
239
240        debug!(
241            "DELETE operation: key={}, existed={}",
242            key,
243            old_value.is_some()
244        );
245
246        if old_value.is_some() {
247            Ok(KVResult::Success)
248        } else {
249            Ok(KVResult::NotFound)
250        }
251    }
252
253    /// Check if a key exists
254    pub async fn exists(&self, key: &str) -> Result<bool, StoreError> {
255        self.validate_key(key)?;
256
257        let exists = self.data.contains_key(key);
258        self.increment_operation_count();
259
260        debug!("EXISTS operation: key={}, exists={}", key, exists);
261        Ok(exists)
262    }
263
264    /// List all keys (with optional prefix filter)
265    pub async fn keys(&self, prefix: Option<&str>) -> Result<Vec<String>, StoreError> {
266        let keys: Vec<String> = if let Some(prefix) = prefix {
267            self.data
268                .iter()
269                .filter(|entry| entry.key().starts_with(prefix))
270                .map(|entry| entry.key().clone())
271                .collect()
272        } else {
273            self.data.iter().map(|entry| entry.key().clone()).collect()
274        };
275
276        self.increment_operation_count();
277        debug!("KEYS operation: prefix={:?}, count={}", prefix, keys.len());
278        Ok(keys)
279    }
280
281    /// Get the number of keys in the store
282    pub async fn size(&self) -> usize {
283        self.data.len()
284    }
285
286    /// Clear all data
287    pub async fn clear(&self) -> Result<KVResult, StoreError> {
288        let old_size = self.data.len();
289        self.data.clear();
290        self.increment_operation_count();
291
292        // Send bulk notification if enabled
293        if self.config.enable_notifications && old_size > 0 {
294            let notification = ChangeNotification {
295                key: "*".to_string(),
296                change_type: ChangeType::Cleared,
297                old_value: Some(format!("{} keys", old_size)),
298                new_value: None,
299                version: self.get_version(),
300                timestamp: SystemTime::now()
301                    .duration_since(UNIX_EPOCH)
302                    .unwrap()
303                    .as_millis() as u64,
304            };
305
306            self.notification_bus.publish(notification).await;
307        }
308
309        info!("CLEAR operation: removed {} keys", old_size);
310        Ok(KVResult::Success)
311    }
312
313    /// Process a batch of operations atomically
314    pub async fn apply_batch(
315        &self,
316        operations: Vec<KVOperation>,
317    ) -> Result<Vec<KVResult>, StoreError> {
318        let mut results = Vec::with_capacity(operations.len());
319
320        // In a production implementation, this would use transactions
321        // For now, we apply operations sequentially
322        for operation in operations {
323            let result = match operation {
324                KVOperation::Set { key, value } => self.set(&key, &value).await?,
325                KVOperation::Get { key } => {
326                    let value = self.get(&key).await?;
327                    if value.is_some() {
328                        KVResult::Success
329                    } else {
330                        KVResult::NotFound
331                    }
332                }
333                KVOperation::Delete { key } => self.delete(&key).await?,
334                KVOperation::Exists { key } => {
335                    let exists = self.exists(&key).await?;
336                    if exists {
337                        KVResult::Success
338                    } else {
339                        KVResult::NotFound
340                    }
341                }
342            };
343            results.push(result);
344        }
345
346        debug!("BATCH operation: {} operations processed", results.len());
347        Ok(results)
348    }
349
350    /// Create a snapshot of the current state
351    pub async fn create_snapshot(&self) -> Result<StoreSnapshot, StoreError> {
352        let data: HashMap<String, ValueEntry> = self
353            .data
354            .iter()
355            .map(|entry| (entry.key().clone(), entry.value().clone()))
356            .collect();
357
358        let version = self.get_version();
359        let created_at = SystemTime::now()
360            .duration_since(UNIX_EPOCH)
361            .unwrap()
362            .as_millis() as u64;
363
364        // Simple checksum calculation
365        let checksum = self.calculate_checksum(&data);
366
367        let snapshot = StoreSnapshot {
368            data,
369            version,
370            created_at,
371            checksum,
372        };
373
374        // Update stats
375        {
376            let mut stats = self.stats.write();
377            stats.last_snapshot_at = created_at;
378            stats.operations_since_snapshot = 0;
379        }
380
381        info!(
382            "Snapshot created: version={}, keys={}",
383            version,
384            snapshot.data.len()
385        );
386        Ok(snapshot)
387    }
388
389    /// Restore from a snapshot
390    pub async fn restore_snapshot(&self, snapshot: StoreSnapshot) -> Result<(), StoreError> {
391        // Verify checksum
392        let calculated_checksum = self.calculate_checksum(&snapshot.data);
393        if calculated_checksum != snapshot.checksum {
394            return Err(StoreError::InvalidSnapshot);
395        }
396
397        // Clear and restore data
398        self.data.clear();
399        for (key, value) in snapshot.data {
400            self.data.insert(key, value);
401        }
402
403        self.version
404            .store(snapshot.version, std::sync::atomic::Ordering::Release);
405
406        info!(
407            "Snapshot restored: version={}, keys={}",
408            snapshot.version,
409            self.data.len()
410        );
411        Ok(())
412    }
413
414    /// Get store statistics
415    pub async fn get_stats(&self) -> StoreStats {
416        let mut stats = self.stats.read().clone();
417        stats.total_keys = self.data.len();
418        stats.memory_usage_bytes = self.estimate_memory_usage();
419        stats
420    }
421
422    /// Get notification bus for subscribing to changes
423    pub fn notification_bus(&self) -> Arc<NotificationBus> {
424        self.notification_bus.clone()
425    }
426
427    /// Shutdown the store
428    pub async fn shutdown(&self) -> Result<(), StoreError> {
429        info!("Shutting down KVStore");
430        let _ = self.shutdown_tx.send(true);
431        Ok(())
432    }
433
434    /// Get the current version
435    pub fn current_version(&self) -> u64 {
436        self.version.load(std::sync::atomic::Ordering::Acquire)
437    }
438
439    /// Get all data as a HashMap for state machine operations
440    pub fn get_all_data(&self) -> HashMap<String, ValueEntry> {
441        self.data
442            .iter()
443            .map(|entry| (entry.key().clone(), entry.value().clone()))
444            .collect()
445    }
446
447    /// Set the store version (for state restoration)
448    pub fn set_version(&self, version: u64) {
449        self.version
450            .store(version, std::sync::atomic::Ordering::Release);
451    }
452
453    /// Clear and set data from a HashMap (for state restoration)
454    pub fn set_all_data(&self, data: HashMap<String, ValueEntry>) {
455        self.data.clear();
456        for (key, value) in data {
457            self.data.insert(key, value);
458        }
459    }
460
461    // Private helper methods
462
463    fn validate_key(&self, key: &str) -> Result<(), StoreError> {
464        if key.is_empty() {
465            return Err(StoreError::InvalidKey("Key cannot be empty".to_string()));
466        }
467        if key.len() > 256 {
468            return Err(StoreError::InvalidKey("Key too long".to_string()));
469        }
470        Ok(())
471    }
472
473    fn validate_value(&self, value: &str) -> Result<(), StoreError> {
474        if value.len() > self.config.max_value_size {
475            return Err(StoreError::ValueTooLarge);
476        }
477        Ok(())
478    }
479
480    fn increment_operation_count(&self) {
481        let mut stats = self.stats.write();
482        stats.total_operations += 1;
483        stats.operations_since_snapshot += 1;
484    }
485
486    fn get_version(&self) -> u64 {
487        self.version
488            .fetch_add(1, std::sync::atomic::Ordering::AcqRel)
489    }
490
491    fn calculate_checksum(&self, data: &HashMap<String, ValueEntry>) -> u64 {
492        use std::collections::hash_map::DefaultHasher;
493        use std::hash::{Hash, Hasher};
494
495        let mut hasher = DefaultHasher::new();
496        for (key, value) in data {
497            key.hash(&mut hasher);
498            value.value.hash(&mut hasher);
499            value.version.hash(&mut hasher);
500        }
501        hasher.finish()
502    }
503
504    fn estimate_memory_usage(&self) -> usize {
505        let mut total = 0;
506        for entry in self.data.iter() {
507            total += entry.key().len();
508            total += entry.value().size;
509            total += std::mem::size_of::<ValueEntry>();
510        }
511        total
512    }
513}
514
515#[cfg(test)]
516mod tests {
517    use super::*;
518
519    #[tokio::test]
520    async fn test_basic_operations() {
521        let config = KVStoreConfig::default();
522        let store = KVStore::new(config).await.unwrap();
523
524        // Test SET
525        let result = store.set("key1", "value1").await.unwrap();
526        assert!(matches!(result, KVResult::Success));
527
528        // Test GET
529        let value = store.get("key1").await.unwrap();
530        assert_eq!(value.unwrap(), "value1");
531
532        // Test EXISTS
533        let exists = store.exists("key1").await.unwrap();
534        assert!(exists);
535
536        // Test DELETE
537        let result = store.delete("key1").await.unwrap();
538        assert!(matches!(result, KVResult::Success));
539
540        // Test GET after DELETE
541        let value = store.get("key1").await.unwrap();
542        assert!(value.is_none());
543    }
544
545    #[tokio::test]
546    async fn test_batch_operations() {
547        let config = KVStoreConfig::default();
548        let store = KVStore::new(config).await.unwrap();
549
550        let operations = vec![
551            KVOperation::Set {
552                key: "key1".to_string(),
553                value: "value1".to_string(),
554            },
555            KVOperation::Set {
556                key: "key2".to_string(),
557                value: "value2".to_string(),
558            },
559            KVOperation::Get {
560                key: "key1".to_string(),
561            },
562        ];
563
564        let results = store.apply_batch(operations).await.unwrap();
565        assert_eq!(results.len(), 3);
566        assert!(matches!(results[0], KVResult::Success));
567        assert!(matches!(results[1], KVResult::Success));
568        assert!(matches!(results[2], KVResult::Success));
569    }
570
571    #[tokio::test]
572    async fn test_snapshot_and_restore() {
573        let config = KVStoreConfig::default();
574        let store = KVStore::new(config).await.unwrap();
575
576        // Add some data
577        store.set("key1", "value1").await.unwrap();
578        store.set("key2", "value2").await.unwrap();
579
580        // Create snapshot
581        let snapshot = store.create_snapshot().await.unwrap();
582        assert_eq!(snapshot.data.len(), 2);
583
584        // Clear store
585        store.clear().await.unwrap();
586        assert_eq!(store.size().await, 0);
587
588        // Restore snapshot
589        store.restore_snapshot(snapshot).await.unwrap();
590        assert_eq!(store.size().await, 2);
591
592        let value = store.get("key1").await.unwrap();
593        assert_eq!(value.unwrap(), "value1");
594    }
595}
rabia_kvstore_example/store.rs

rabia_kvstore_example/
store.rs
