eshanized_polaris_core/
storage.rs

1//! Storage backends for persisting cluster state.
2
3use crate::errors::PolarisResult;
4use async_trait::async_trait;
5use bytes::Bytes;
6use std::sync::Arc;
7
8/// Storage backend trait
9#[async_trait]
10pub trait Storage: Send + Sync {
11    /// Get a value by key
12    async fn get(&self, key: &str) -> PolarisResult<Option<Bytes>>;
13
14    /// Set a value for a key
15    async fn set(&self, key: &str, value: Bytes) -> PolarisResult<()>;
16
17    /// Delete a key
18    async fn delete(&self, key: &str) -> PolarisResult<()>;
19
20    /// Check if a key exists
21    async fn exists(&self, key: &str) -> PolarisResult<bool>;
22
23    /// List all keys with a prefix
24    async fn list(&self, prefix: &str) -> PolarisResult<Vec<String>>;
25}
26
27/// In-memory storage backend (default)
28#[derive(Debug, Clone)]
29pub struct InMemoryStorage {
30    data: Arc<dashmap::DashMap<String, Bytes>>,
31}
32
33impl InMemoryStorage {
34    /// Create a new in-memory storage
35    pub fn new() -> Self {
36        Self {
37            data: Arc::new(dashmap::DashMap::new()),
38        }
39    }
40}
41
42impl Default for InMemoryStorage {
43    fn default() -> Self {
44        Self::new()
45    }
46}
47
48#[async_trait]
49impl Storage for InMemoryStorage {
50    async fn get(&self, key: &str) -> PolarisResult<Option<Bytes>> {
51        Ok(self.data.get(key).map(|v| v.clone()))
52    }
53
54    async fn set(&self, key: &str, value: Bytes) -> PolarisResult<()> {
55        self.data.insert(key.to_string(), value);
56        Ok(())
57    }
58
59    async fn delete(&self, key: &str) -> PolarisResult<()> {
60        self.data.remove(key);
61        Ok(())
62    }
63
64    async fn exists(&self, key: &str) -> PolarisResult<bool> {
65        Ok(self.data.contains_key(key))
66    }
67
68    async fn list(&self, prefix: &str) -> PolarisResult<Vec<String>> {
69        Ok(self
70            .data
71            .iter()
72            .filter(|entry| entry.key().starts_with(prefix))
73            .map(|entry| entry.key().clone())
74            .collect())
75    }
76}
77
78// RocksDB storage backend (commented out due to C dependencies)
79// Enable by uncommenting rocksdb in workspace Cargo.toml and rebuilding
80// 
81// #[cfg(feature = "rocksdb")]
82// pub struct RocksDbStorage {
83//     db: Arc<rocksdb::DB>,
84// }
85
86#[cfg(test)]
87mod tests {
88    use super::*;
89
90    #[tokio::test]
91    async fn test_in_memory_storage() {
92        let storage = InMemoryStorage::new();
93
94        // Test set and get
95        storage
96            .set("key1", Bytes::from("value1"))
97            .await
98            .unwrap();
99        let value = storage.get("key1").await.unwrap();
100        assert_eq!(value, Some(Bytes::from("value1")));
101
102        // Test exists
103        assert!(storage.exists("key1").await.unwrap());
104        assert!(!storage.exists("key2").await.unwrap());
105
106        // Test delete
107        storage.delete("key1").await.unwrap();
108        assert!(!storage.exists("key1").await.unwrap());
109    }
110
111    #[tokio::test]
112    async fn test_storage_list() {
113        let storage = InMemoryStorage::new();
114
115        storage.set("prefix:key1", Bytes::from("v1")).await.unwrap();
116        storage.set("prefix:key2", Bytes::from("v2")).await.unwrap();
117        storage.set("other:key3", Bytes::from("v3")).await.unwrap();
118
119        let keys = storage.list("prefix:").await.unwrap();
120        assert_eq!(keys.len(), 2);
121        assert!(keys.contains(&"prefix:key1".to_string()));
122        assert!(keys.contains(&"prefix:key2".to_string()));
123    }
124}