aurora_db/storage/
cold.rs

1use crate::error::{AuroraError, Result};
2use sled::Db;
3use crate::types::{AuroraConfig, ColdStoreMode};
4
5pub struct ColdStore {
6    db: Db,
7    #[allow(dead_code)]
8    db_path: String,
9}
10
11impl ColdStore {
12    /// Creates a new ColdStore instance with default settings
13    pub fn new(path: &str) -> Result<Self> {
14        // Use the default configuration values
15        let config = AuroraConfig::default();
16        Self::with_config(
17            path, 
18            config.cold_cache_capacity_mb,
19            config.cold_flush_interval_ms,
20            config.cold_mode
21        )
22    }
23    
24    /// Creates a new ColdStore instance with custom configuration
25    pub fn with_config(
26        path: &str, 
27        cache_capacity_mb: usize,
28        flush_interval_ms: Option<u64>,
29        mode: ColdStoreMode,
30    ) -> Result<Self> {
31        let db_path = if !path.ends_with(".db") {
32            format!("{}.db", path)
33        } else {
34            path.to_string()
35        };
36
37        // Configure sled based on the provided config
38        let mut sled_config = sled::Config::new()
39            .path(&db_path)
40            .cache_capacity((cache_capacity_mb * 1024 * 1024) as u64)
41            .flush_every_ms(flush_interval_ms); // Don't need unwrap_or(0) - passing Option directly
42            
43        // Set the mode based on config
44        sled_config = match mode {
45            ColdStoreMode::HighThroughput => sled_config.mode(sled::Mode::HighThroughput),
46            ColdStoreMode::LowSpace => sled_config.mode(sled::Mode::LowSpace),
47        };
48
49        let db = sled_config.open()?;
50
51        Ok(Self { 
52            db,
53            db_path,
54        })
55    }
56
57
58    /// Retrieves a value from cold storage
59    pub fn get(&self, key: &str) -> Result<Option<Vec<u8>>> {
60        Ok(self.db.get(key.as_bytes())?.map(|ivec| ivec.to_vec()))
61    }
62
63    /// Stores a value in cold storage
64    pub fn set(&self, key: String, value: Vec<u8>) -> Result<()> {
65        self.db.insert(key.as_bytes(), value)?;
66        // Don't flush immediately - let the background flusher handle it
67        Ok(())
68    }
69
70    /// Removes a value from cold storage
71    pub fn delete(&self, key: &str) -> Result<()> {
72        self.db.remove(key.as_bytes())?;
73        // Don't flush immediately
74        Ok(())
75    }
76
77    /// Returns an iterator over all key-value pairs
78    pub fn scan(&self) -> impl Iterator<Item = Result<(String, Vec<u8>)>> + '_ {
79        self.db.iter().map(|result| {
80            result.map_err(|e| AuroraError::Storage(e)).and_then(|(key, value)| {
81                Ok((
82                    String::from_utf8(key.to_vec())
83                        .map_err(|_| AuroraError::Protocol("Invalid UTF-8 in key".to_string()))?,
84                    value.to_vec(),
85                ))
86            })
87        })
88    }
89
90    /// Returns an iterator over keys with a specific prefix
91    pub fn scan_prefix(&self, prefix: &str) -> impl Iterator<Item = Result<(String, Vec<u8>)>> + '_ {
92        self.db.scan_prefix(prefix.as_bytes()).map(|result| {
93            result.map_err(|e| AuroraError::Storage(e)).and_then(|(key, value)| {
94                Ok((
95                    String::from_utf8(key.to_vec())
96                        .map_err(|_| AuroraError::Protocol("Invalid UTF-8 in key".to_string()))?,
97                    value.to_vec(),
98                ))
99            })
100        })
101    }
102
103    /// Batch operations for better performance
104    pub fn batch_set(&self, pairs: Vec<(String, Vec<u8>)>) -> Result<()> {
105        let batch = pairs
106            .into_iter()
107            .fold(sled::Batch::default(), |mut batch, (key, value)| {
108                batch.insert(key.as_bytes(), value);
109                batch
110            });
111
112        self.db.apply_batch(batch)?;
113        self.db.flush()?;
114        Ok(())
115    }
116
117    /// Compacts the database to reclaim space
118    pub fn compact(&self) -> Result<()> {
119        self.db.flush()?;
120        // Note: In newer versions of sled, we would use compact_range
121        // For now, just flush and let sled handle internal compaction
122        Ok(())
123    }
124
125    /// Get database statistics
126    pub fn get_stats(&self) -> Result<ColdStoreStats> {
127        Ok(ColdStoreStats {
128            size_on_disk: self.estimated_size(),
129            tree_count: self.db.tree_names().len() as u64,
130        })
131    }
132
133    /// Get estimated size on disk
134    pub fn estimated_size(&self) -> u64 {
135        self.db.size_on_disk().unwrap_or(0)
136    }
137}
138
139impl Drop for ColdStore {
140    fn drop(&mut self) {
141        if let Err(e) = self.db.flush() {
142            eprintln!("Error flushing database: {}", e);
143        }
144    }
145}
146
147/// Custom statistics struct since sled doesn't expose one
148#[derive(Debug)]
149pub struct ColdStoreStats {
150    pub size_on_disk: u64,
151    pub tree_count: u64,
152}
aurora_db/storage/cold.rs

aurora_db/storage/
cold.rs
