//! RocksDB database wrapper for Synor. //! //! Provides a typed interface to the underlying RocksDB instance. use crate::cf; use rocksdb::{ColumnFamily, ColumnFamilyDescriptor, DBCompressionType, Options, WriteBatch, DB}; use std::path::Path; use std::sync::Arc; use thiserror::Error; /// Database configuration options. #[derive(Clone, Debug)] pub struct DatabaseConfig { /// Maximum number of open files. pub max_open_files: i32, /// Write buffer size in bytes. pub write_buffer_size: usize, /// Maximum number of write buffers. pub max_write_buffer_number: i32, /// Target file size for level-1 files. pub target_file_size_base: u64, /// Maximum total WAL size. pub max_total_wal_size: u64, /// Enable compression. pub enable_compression: bool, /// Block cache size in bytes. pub block_cache_size: usize, /// Enable statistics. pub enable_statistics: bool, /// Create database if missing. pub create_if_missing: bool, /// Number of background threads. pub parallelism: i32, } impl Default for DatabaseConfig { fn default() -> Self { DatabaseConfig { max_open_files: 512, write_buffer_size: 64 * 1024 * 1024, // 64 MB max_write_buffer_number: 3, target_file_size_base: 64 * 1024 * 1024, // 64 MB max_total_wal_size: 256 * 1024 * 1024, // 256 MB enable_compression: true, block_cache_size: 128 * 1024 * 1024, // 128 MB enable_statistics: false, create_if_missing: true, parallelism: 4, } } } impl DatabaseConfig { /// Creates a configuration optimized for SSDs. pub fn ssd() -> Self { DatabaseConfig { max_open_files: 1024, write_buffer_size: 128 * 1024 * 1024, max_write_buffer_number: 4, target_file_size_base: 128 * 1024 * 1024, max_total_wal_size: 512 * 1024 * 1024, enable_compression: true, block_cache_size: 512 * 1024 * 1024, enable_statistics: false, create_if_missing: true, parallelism: 8, } } /// Creates a configuration for low-memory environments. pub fn low_memory() -> Self { DatabaseConfig { max_open_files: 128, write_buffer_size: 16 * 1024 * 1024, max_write_buffer_number: 2, target_file_size_base: 32 * 1024 * 1024, max_total_wal_size: 64 * 1024 * 1024, enable_compression: true, block_cache_size: 32 * 1024 * 1024, enable_statistics: false, create_if_missing: true, parallelism: 2, } } /// Creates a configuration for testing. pub fn for_testing() -> Self { DatabaseConfig { max_open_files: 64, write_buffer_size: 4 * 1024 * 1024, max_write_buffer_number: 2, target_file_size_base: 4 * 1024 * 1024, max_total_wal_size: 16 * 1024 * 1024, enable_compression: false, block_cache_size: 8 * 1024 * 1024, enable_statistics: true, create_if_missing: true, parallelism: 1, } } /// Converts to RocksDB Options. fn to_options(&self) -> Options { let mut opts = Options::default(); opts.create_if_missing(self.create_if_missing); opts.create_missing_column_families(true); opts.set_max_open_files(self.max_open_files); opts.set_write_buffer_size(self.write_buffer_size); opts.set_max_write_buffer_number(self.max_write_buffer_number); opts.set_target_file_size_base(self.target_file_size_base); opts.set_max_total_wal_size(self.max_total_wal_size); opts.increase_parallelism(self.parallelism); if self.enable_compression { opts.set_compression_type(DBCompressionType::Lz4); } opts } } /// Database errors. #[derive(Debug, Error)] pub enum DbError { #[error("RocksDB error: {0}")] RocksDb(#[from] rocksdb::Error), #[error("Column family not found: {0}")] ColumnFamilyNotFound(String), #[error("Serialization error: {0}")] Serialization(String), #[error("Deserialization error: {0}")] Deserialization(String), #[error("Key not found")] NotFound, #[error("Database is closed")] Closed, #[error("Invalid key format")] InvalidKey, #[error("IO error: {0}")] Io(#[from] std::io::Error), } /// The main database handle. pub struct Database { db: Arc, path: String, } impl Database { /// Opens or creates a database at the given path. pub fn open>(path: P, config: &DatabaseConfig) -> Result { let path_str = path.as_ref().to_string_lossy().to_string(); let opts = config.to_options(); // Create column family descriptors let cf_descriptors: Vec = cf::all() .into_iter() .map(|name| { let mut cf_opts = Options::default(); if config.enable_compression { cf_opts.set_compression_type(DBCompressionType::Lz4); } ColumnFamilyDescriptor::new(name, cf_opts) }) .collect(); let db = DB::open_cf_descriptors(&opts, &path_str, cf_descriptors)?; Ok(Database { db: Arc::new(db), path: path_str, }) } /// Opens a database in read-only mode. pub fn open_read_only>( path: P, config: &DatabaseConfig, ) -> Result { let path_str = path.as_ref().to_string_lossy().to_string(); let opts = config.to_options(); let cf_names: Vec<&str> = cf::all(); let db = DB::open_cf_for_read_only(&opts, &path_str, cf_names, false)?; Ok(Database { db: Arc::new(db), path: path_str, }) } /// Returns the database path. pub fn path(&self) -> &str { &self.path } /// Gets a value by key from a column family. pub fn get(&self, cf_name: &str, key: &[u8]) -> Result>, DbError> { let cf = self.get_cf(cf_name)?; Ok(self.db.get_cf(cf, key)?) } /// Puts a key-value pair into a column family. pub fn put(&self, cf_name: &str, key: &[u8], value: &[u8]) -> Result<(), DbError> { let cf = self.get_cf(cf_name)?; self.db.put_cf(cf, key, value)?; Ok(()) } /// Deletes a key from a column family. pub fn delete(&self, cf_name: &str, key: &[u8]) -> Result<(), DbError> { let cf = self.get_cf(cf_name)?; self.db.delete_cf(cf, key)?; Ok(()) } /// Checks if a key exists in a column family. pub fn exists(&self, cf_name: &str, key: &[u8]) -> Result { let cf = self.get_cf(cf_name)?; Ok(self.db.get_cf(cf, key)?.is_some()) } /// Creates a new write batch. pub fn batch(&self) -> WriteBatchWrapper { WriteBatchWrapper { batch: WriteBatch::default(), db: Arc::clone(&self.db), } } /// Writes a batch atomically. pub fn write_batch(&self, batch: WriteBatchWrapper) -> Result<(), DbError> { self.db.write(batch.batch)?; Ok(()) } /// Iterates over all keys in a column family. pub fn iter( &self, cf_name: &str, ) -> Result, Box<[u8]>)> + '_, DbError> { let cf = self.get_cf(cf_name)?; let iter = self.db.iterator_cf(cf, rocksdb::IteratorMode::Start); Ok(iter.map(|r| r.unwrap())) } /// Iterates over keys with a given prefix. pub fn prefix_iter( &self, cf_name: &str, prefix: &[u8], ) -> Result, Box<[u8]>)> + '_, DbError> { let cf = self.get_cf(cf_name)?; let iter = self.db.prefix_iterator_cf(cf, prefix); let prefix_owned = prefix.to_vec(); Ok(iter .map(|r| r.unwrap()) .take_while(move |(k, _)| k.starts_with(&prefix_owned))) } /// Gets multiple values by keys. pub fn multi_get( &self, cf_name: &str, keys: &[&[u8]], ) -> Result>>, DbError> { let cf = self.get_cf(cf_name)?; let cf_keys: Vec<_> = keys.iter().map(|k| (cf, *k)).collect(); let results = self.db.multi_get_cf(cf_keys); results .into_iter() .map(|r| r.map_err(DbError::from)) .collect() } /// Flushes the WAL to disk. pub fn flush(&self) -> Result<(), DbError> { self.db.flush()?; Ok(()) } /// Compacts the entire database. pub fn compact(&self) -> Result<(), DbError> { for cf_name in cf::all() { if let Ok(cf) = self.get_cf(cf_name) { self.db.compact_range_cf(cf, None::<&[u8]>, None::<&[u8]>); } } Ok(()) } /// Returns database statistics if enabled. pub fn stats(&self) -> Option { self.db.property_value("rocksdb.stats").ok().flatten() } /// Returns the approximate size of a column family in bytes. pub fn cf_size(&self, cf_name: &str) -> Result { let cf = self.get_cf(cf_name)?; let size = self .db .property_int_value_cf(cf, "rocksdb.estimate-live-data-size") .map_err(|e| DbError::RocksDb(e))? .unwrap_or(0); Ok(size) } /// Gets a column family handle. fn get_cf(&self, name: &str) -> Result<&ColumnFamily, DbError> { self.db .cf_handle(name) .ok_or_else(|| DbError::ColumnFamilyNotFound(name.to_string())) } /// Returns an Arc to the underlying database. pub fn inner(&self) -> Arc { Arc::clone(&self.db) } } impl std::fmt::Debug for Database { fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { f.debug_struct("Database") .field("path", &self.path) .finish() } } /// A wrapper around RocksDB WriteBatch with column family support. pub struct WriteBatchWrapper { batch: WriteBatch, db: Arc, } impl WriteBatchWrapper { /// Adds a put operation to the batch. pub fn put(&mut self, cf_name: &str, key: &[u8], value: &[u8]) -> Result<(), DbError> { let cf = self .db .cf_handle(cf_name) .ok_or_else(|| DbError::ColumnFamilyNotFound(cf_name.to_string()))?; self.batch.put_cf(cf, key, value); Ok(()) } /// Adds a delete operation to the batch. pub fn delete(&mut self, cf_name: &str, key: &[u8]) -> Result<(), DbError> { let cf = self .db .cf_handle(cf_name) .ok_or_else(|| DbError::ColumnFamilyNotFound(cf_name.to_string()))?; self.batch.delete_cf(cf, key); Ok(()) } /// Returns the number of operations in the batch. pub fn len(&self) -> usize { self.batch.len() } /// Returns true if the batch is empty. pub fn is_empty(&self) -> bool { self.batch.is_empty() } /// Clears all operations from the batch. pub fn clear(&mut self) { self.batch.clear(); } } /// Utility functions for key encoding. pub mod keys { /// Encodes a u64 as big-endian bytes for lexicographic ordering. pub fn encode_u64(value: u64) -> [u8; 8] { value.to_be_bytes() } /// Decodes a u64 from big-endian bytes. pub fn decode_u64(bytes: &[u8]) -> Option { if bytes.len() < 8 { return None; } let arr: [u8; 8] = bytes[..8].try_into().ok()?; Some(u64::from_be_bytes(arr)) } /// Encodes a u32 as big-endian bytes. pub fn encode_u32(value: u32) -> [u8; 4] { value.to_be_bytes() } /// Decodes a u32 from big-endian bytes. pub fn decode_u32(bytes: &[u8]) -> Option { if bytes.len() < 4 { return None; } let arr: [u8; 4] = bytes[..4].try_into().ok()?; Some(u32::from_be_bytes(arr)) } /// Creates a prefixed key. pub fn prefixed_key(prefix: u8, key: &[u8]) -> Vec { let mut result = Vec::with_capacity(1 + key.len()); result.push(prefix); result.extend_from_slice(key); result } /// Creates a composite key from multiple parts. pub fn composite_key(parts: &[&[u8]]) -> Vec { let total_len: usize = parts.iter().map(|p| p.len()).sum(); let mut result = Vec::with_capacity(total_len); for part in parts { result.extend_from_slice(part); } result } } #[cfg(test)] mod tests { use super::*; use tempfile::tempdir; #[test] fn test_database_open_close() { let dir = tempdir().unwrap(); let config = DatabaseConfig::for_testing(); let db = Database::open(dir.path(), &config).unwrap(); assert!(db.path().contains(dir.path().to_str().unwrap())); } #[test] fn test_put_get_delete() { let dir = tempdir().unwrap(); let config = DatabaseConfig::for_testing(); let db = Database::open(dir.path(), &config).unwrap(); // Put db.put(cf::HEADERS, b"key1", b"value1").unwrap(); // Get let value = db.get(cf::HEADERS, b"key1").unwrap(); assert_eq!(value, Some(b"value1".to_vec())); // Delete db.delete(cf::HEADERS, b"key1").unwrap(); let value = db.get(cf::HEADERS, b"key1").unwrap(); assert_eq!(value, None); } #[test] fn test_batch_operations() { let dir = tempdir().unwrap(); let config = DatabaseConfig::for_testing(); let db = Database::open(dir.path(), &config).unwrap(); let mut batch = db.batch(); batch.put(cf::HEADERS, b"key1", b"value1").unwrap(); batch.put(cf::HEADERS, b"key2", b"value2").unwrap(); batch.put(cf::BLOCKS, b"block1", b"data1").unwrap(); assert_eq!(batch.len(), 3); db.write_batch(batch).unwrap(); assert_eq!( db.get(cf::HEADERS, b"key1").unwrap(), Some(b"value1".to_vec()) ); assert_eq!( db.get(cf::HEADERS, b"key2").unwrap(), Some(b"value2".to_vec()) ); assert_eq!( db.get(cf::BLOCKS, b"block1").unwrap(), Some(b"data1".to_vec()) ); } #[test] fn test_exists() { let dir = tempdir().unwrap(); let config = DatabaseConfig::for_testing(); let db = Database::open(dir.path(), &config).unwrap(); assert!(!db.exists(cf::HEADERS, b"key1").unwrap()); db.put(cf::HEADERS, b"key1", b"value1").unwrap(); assert!(db.exists(cf::HEADERS, b"key1").unwrap()); } #[test] fn test_key_encoding() { let value: u64 = 12345678901234567890; let encoded = keys::encode_u64(value); let decoded = keys::decode_u64(&encoded).unwrap(); assert_eq!(value, decoded); let value: u32 = 1234567890; let encoded = keys::encode_u32(value); let decoded = keys::decode_u32(&encoded).unwrap(); assert_eq!(value, decoded); } #[test] fn test_prefixed_key() { let key = keys::prefixed_key(0x01, b"test"); assert_eq!(key, vec![0x01, b't', b'e', b's', b't']); } }