diff --git a/Cargo.toml b/Cargo.toml index 0a29fcb..2cb8e8f 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -40,6 +40,7 @@ uuid = { version = "1.0", optional = true } futures = { version = "0.3", optional = true } tokio = { version = "1.0", features = ["rt-multi-thread", "macros"], optional = true } pyo3 = { version = "0.22", features = ["extension-module"], optional = true } +rocksdb = { version = "0.22", optional = true } [dev-dependencies] bytes = "1.10.1" @@ -50,7 +51,7 @@ tempfile = "3.0" tokio = { version = "1.0", features = ["full"] } [features] -default = ["digest_base64", "prolly_balance_max_nodes", "git", "sql"] +default = ["digest_base64", "prolly_balance_max_nodes", "git", "sql", "rocksdb_storage"] tracing = ["dep:tracing"] digest_base64 = ["dep:base64"] prolly_balance_max_nodes = [] @@ -58,6 +59,7 @@ prolly_balance_rolling_hash = [] git = ["dep:gix", "dep:clap", "dep:lru", "dep:hex", "dep:chrono"] sql = ["dep:gluesql-core", "dep:async-trait", "dep:uuid", "dep:futures", "dep:tokio"] python = ["dep:pyo3"] +rocksdb_storage = ["dep:rocksdb", "dep:lru"] [[bin]] name = "git-prolly" @@ -78,6 +80,10 @@ name = "git_prolly_bench" harness = false required-features = ["git", "sql"] +[[bench]] +name = "storage_bench" +harness = false + [[example]] name = "proof_visualization" path = "examples/proof_visualization.rs" @@ -97,5 +103,10 @@ name = "git_merge" path = "examples/git_merge.rs" required-features = ["git"] +[[example]] +name = "rocksdb_storage" +path = "examples/rocksdb_storage.rs" +required-features = ["rocksdb_storage"] + [workspace] members = ["examples/rig_versioned_memory", "examples/financial_advisor"] diff --git a/benches/README.md b/benches/README.md index b3ae732..8b16efb 100644 --- a/benches/README.md +++ b/benches/README.md @@ -32,6 +32,19 @@ Git versioning and SQL integration: - **Time travel queries**: Historical data queries - **Concurrent operations**: Parallel table operations +### 4. Storage Backend Comparison (`storage_bench.rs`) +Compares different storage backend implementations: +- **Insert performance**: Sequential inserts across storage backends +- **Read performance**: Random key lookups +- **Batch operations**: Bulk insert performance +- **Direct node operations**: Low-level storage API performance + +Supported backends: +- InMemoryNodeStorage (default) +- FileNodeStorage (default) +- RocksDBNodeStorage (requires `rocksdb_storage` feature) +- GitNodeStorage (requires `git` feature) + ## Running Benchmarks ### Run all benchmarks: @@ -49,6 +62,12 @@ cargo bench --bench sql_bench --features sql # Git-Prolly benchmarks (requires both git and sql features) cargo bench --bench git_prolly_bench --features git,sql + +# Storage backend benchmarks +cargo bench --bench storage_bench + +# Storage benchmarks with RocksDB included +cargo bench --bench storage_bench --features rocksdb_storage ``` ### Run specific benchmark within a suite: diff --git a/benches/storage_bench.rs b/benches/storage_bench.rs new file mode 100644 index 0000000..c8b3491 --- /dev/null +++ b/benches/storage_bench.rs @@ -0,0 +1,317 @@ +/* +Licensed under the Apache License, Version 2.0 (the "License"); +you may not use this file except in compliance with the License. +You may obtain a copy of the License at + + http://www.apache.org/licenses/LICENSE-2.0 + +Unless required by applicable law or agreed to in writing, software +distributed under the License is distributed on an "AS IS" BASIS, +WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +See the License for the specific language governing permissions and +limitations under the License. +*/ + +use criterion::{black_box, criterion_group, criterion_main, BenchmarkId, Criterion}; +use prollytree::config::TreeConfig; +use prollytree::storage::{FileNodeStorage, InMemoryNodeStorage, NodeStorage}; +use prollytree::tree::{ProllyTree, Tree}; +use tempfile::TempDir; + +#[cfg(feature = "rocksdb_storage")] +use prollytree::storage::RocksDBNodeStorage; + +fn generate_test_data(size: usize) -> (Vec>, Vec>) { + let mut keys = Vec::new(); + let mut values = Vec::new(); + + for i in 0..size { + keys.push(format!("key_{:08}", i).into_bytes()); + values.push(format!("value_data_{:08}_padding_to_make_it_larger", i).into_bytes()); + } + + (keys, values) +} + +fn bench_storage_insert(c: &mut Criterion) { + let mut group = c.benchmark_group("storage_insert"); + group.sample_size(10); + + for &size in &[1000, 5000, 10000] { + let (keys, values) = generate_test_data(size); + + // InMemory Storage + group.bench_with_input( + BenchmarkId::new("InMemory", size), + &(&keys, &values), + |b, (keys, values)| { + b.iter(|| { + let storage = InMemoryNodeStorage::<32>::new(); + let config = TreeConfig::<32>::default(); + let mut tree = ProllyTree::new(storage, config); + + for (key, value) in keys.iter().zip(values.iter()) { + tree.insert(black_box(key.clone()), black_box(value.clone())); + } + }); + }, + ); + + // File Storage + group.bench_with_input( + BenchmarkId::new("File", size), + &(&keys, &values), + |b, (keys, values)| { + let temp_dir = TempDir::new().unwrap(); + b.iter(|| { + let storage = FileNodeStorage::<32>::new(temp_dir.path().to_path_buf()); + let config = TreeConfig::<32>::default(); + let mut tree = ProllyTree::new(storage, config); + + for (key, value) in keys.iter().zip(values.iter()) { + tree.insert(black_box(key.clone()), black_box(value.clone())); + } + }); + }, + ); + + // RocksDB Storage + #[cfg(feature = "rocksdb_storage")] + group.bench_with_input( + BenchmarkId::new("RocksDB", size), + &(&keys, &values), + |b, (keys, values)| { + let temp_dir = TempDir::new().unwrap(); + b.iter(|| { + let storage = RocksDBNodeStorage::<32>::new(temp_dir.path().join("rocksdb")) + .expect("Failed to create RocksDB storage"); + let config = TreeConfig::<32>::default(); + let mut tree = ProllyTree::new(storage, config); + + for (key, value) in keys.iter().zip(values.iter()) { + tree.insert(black_box(key.clone()), black_box(value.clone())); + } + }); + }, + ); + } + + group.finish(); +} + +fn bench_storage_read(c: &mut Criterion) { + let mut group = c.benchmark_group("storage_read"); + group.sample_size(10); + + for &size in &[1000, 5000, 10000] { + let (keys, values) = generate_test_data(size); + + // Prepare InMemory tree + let inmem_storage = InMemoryNodeStorage::<32>::new(); + let config = TreeConfig::<32>::default(); + let mut inmem_tree = ProllyTree::new(inmem_storage, config.clone()); + for (key, value) in keys.iter().zip(values.iter()) { + inmem_tree.insert(key.clone(), value.clone()); + } + + // Prepare File tree + let file_dir = TempDir::new().unwrap(); + let file_storage = FileNodeStorage::<32>::new(file_dir.path().to_path_buf()); + let mut file_tree = ProllyTree::new(file_storage, config.clone()); + for (key, value) in keys.iter().zip(values.iter()) { + file_tree.insert(key.clone(), value.clone()); + } + + // Prepare RocksDB tree + #[cfg(feature = "rocksdb_storage")] + let rocksdb_tree = { + let rocksdb_dir = TempDir::new().unwrap(); + let rocksdb_storage = + RocksDBNodeStorage::<32>::new(rocksdb_dir.path().join("rocksdb")).unwrap(); + let mut rocksdb_tree = ProllyTree::new(rocksdb_storage, config.clone()); + for (key, value) in keys.iter().zip(values.iter()) { + rocksdb_tree.insert(key.clone(), value.clone()); + } + (rocksdb_tree, rocksdb_dir) + }; + + // Benchmark InMemory reads + group.bench_with_input( + BenchmarkId::new("InMemory", size), + &(&keys, &inmem_tree), + |b, (keys, tree)| { + b.iter(|| { + for key in keys.iter() { + black_box(tree.find(key)); + } + }); + }, + ); + + // Benchmark File reads + group.bench_with_input( + BenchmarkId::new("File", size), + &(&keys, &file_tree), + |b, (keys, tree)| { + b.iter(|| { + for key in keys.iter() { + black_box(tree.find(key)); + } + }); + }, + ); + + // Benchmark RocksDB reads + #[cfg(feature = "rocksdb_storage")] + { + let (rocksdb_tree, _dir) = &rocksdb_tree; + group.bench_with_input( + BenchmarkId::new("RocksDB", size), + &(&keys, rocksdb_tree), + |b, (keys, tree)| { + b.iter(|| { + for key in keys.iter() { + black_box(tree.find(key)); + } + }); + }, + ); + } + } + + group.finish(); +} + +fn bench_storage_batch_operations(c: &mut Criterion) { + let mut group = c.benchmark_group("storage_batch_operations"); + group.sample_size(10); + + for &size in &[1000, 5000] { + let (keys, values) = generate_test_data(size); + + // Benchmark batch insert for different storage backends + group.bench_with_input( + BenchmarkId::new("InMemory_batch_insert", size), + &(&keys, &values), + |b, (keys, values)| { + b.iter(|| { + let storage = InMemoryNodeStorage::<32>::new(); + let config = TreeConfig::<32>::default(); + let mut tree = ProllyTree::new(storage, config); + tree.insert_batch(black_box(keys), black_box(values)); + }); + }, + ); + + #[cfg(feature = "rocksdb_storage")] + group.bench_with_input( + BenchmarkId::new("RocksDB_batch_insert", size), + &(&keys, &values), + |b, (keys, values)| { + let temp_dir = TempDir::new().unwrap(); + b.iter(|| { + let storage = RocksDBNodeStorage::<32>::new(temp_dir.path().join("rocksdb")) + .expect("Failed to create RocksDB storage"); + let config = TreeConfig::<32>::default(); + let mut tree = ProllyTree::new(storage, config); + tree.insert_batch(black_box(keys), black_box(values)); + }); + }, + ); + } + + group.finish(); +} + +fn bench_node_storage_direct(c: &mut Criterion) { + let mut group = c.benchmark_group("node_storage_direct"); + + // Create test nodes + let config = TreeConfig::<32>::default(); + let mut nodes = Vec::new(); + for i in 0..1000 { + let node = prollytree::node::ProllyNode { + keys: vec![format!("key_{:08}", i).into_bytes()], + key_schema: config.key_schema.clone(), + values: vec![format!("value_{:08}", i).into_bytes()], + value_schema: config.value_schema.clone(), + is_leaf: true, + level: 0, + base: config.base, + modulus: config.modulus, + min_chunk_size: config.min_chunk_size, + max_chunk_size: config.max_chunk_size, + pattern: config.pattern, + split: false, + merged: false, + encode_types: Vec::new(), + encode_values: Vec::new(), + }; + nodes.push((node.get_hash(), node)); + } + + // Benchmark direct node insertions + group.bench_function("InMemory_insert_nodes", |b| { + b.iter(|| { + let mut storage = InMemoryNodeStorage::<32>::new(); + for (hash, node) in &nodes { + storage.insert_node(black_box(hash.clone()), black_box(node.clone())); + } + }); + }); + + #[cfg(feature = "rocksdb_storage")] + group.bench_function("RocksDB_insert_nodes", |b| { + let temp_dir = TempDir::new().unwrap(); + b.iter(|| { + let mut storage = RocksDBNodeStorage::<32>::new(temp_dir.path().join("rocksdb")) + .expect("Failed to create RocksDB storage"); + for (hash, node) in &nodes { + storage.insert_node(black_box(hash.clone()), black_box(node.clone())); + } + }); + }); + + // Benchmark direct node reads + let mut inmem_storage = InMemoryNodeStorage::<32>::new(); + for (hash, node) in &nodes { + inmem_storage.insert_node(hash.clone(), node.clone()); + } + + group.bench_function("InMemory_read_nodes", |b| { + b.iter(|| { + for (hash, _) in &nodes { + black_box(inmem_storage.get_node_by_hash(black_box(hash))); + } + }); + }); + + #[cfg(feature = "rocksdb_storage")] + { + let temp_dir = TempDir::new().unwrap(); + let mut rocksdb_storage = + RocksDBNodeStorage::<32>::new(temp_dir.path().join("rocksdb")).unwrap(); + for (hash, node) in &nodes { + rocksdb_storage.insert_node(hash.clone(), node.clone()); + } + + group.bench_function("RocksDB_read_nodes", |b| { + b.iter(|| { + for (hash, _) in &nodes { + black_box(rocksdb_storage.get_node_by_hash(black_box(hash))); + } + }); + }); + } + + group.finish(); +} + +criterion_group!( + benches, + bench_storage_insert, + bench_storage_read, + bench_storage_batch_operations, + bench_node_storage_direct +); +criterion_main!(benches); diff --git a/examples/rocksdb_storage.rs b/examples/rocksdb_storage.rs new file mode 100644 index 0000000..fbcb678 --- /dev/null +++ b/examples/rocksdb_storage.rs @@ -0,0 +1,79 @@ +use prollytree::config::TreeConfig; +use prollytree::storage::RocksDBNodeStorage; +use prollytree::tree::{ProllyTree, Tree}; +use tempfile::TempDir; + +fn main() -> Result<(), Box> { + // Create a temporary directory for RocksDB + let temp_dir = TempDir::new()?; + let db_path = temp_dir.path().join("prolly_rocksdb"); + + println!("Creating RocksDB storage at: {:?}", db_path); + + // Create RocksDB storage + let storage = RocksDBNodeStorage::<32>::new(db_path)?; + + // Create a ProllyTree with RocksDB storage + let config = TreeConfig::<32>::default(); + let mut tree = ProllyTree::new(storage, config); + + // Insert some data + println!("Inserting data into ProllyTree with RocksDB storage..."); + for i in 0..1000 { + let key = format!("key_{:04}", i); + let value = format!("value_{:04}", i); + tree.insert(key.as_bytes().to_vec(), value.as_bytes().to_vec()); + } + + println!("Inserted 1000 key-value pairs"); + + // Retrieve and verify data + println!("Retrieving data..."); + for i in 0..10 { + let key = format!("key_{:04}", i); + let node = tree.find(key.as_bytes()); + if let Some(node) = node { + // Find the key in the node and get its corresponding value + if let Some(idx) = node.keys.iter().position(|k| k == key.as_bytes()) { + let value = &node.values[idx]; + println!(" {} -> {}", key, String::from_utf8_lossy(value)); + } + } + } + + // Get tree statistics + let root_hash = tree.get_root_hash(); + println!("\nTree root hash: {:?}", root_hash); + + // Demonstrate persistence by creating a new tree with the same storage + drop(tree); + println!("\nCreating new tree with same storage..."); + + // Re-open the storage + let storage2 = RocksDBNodeStorage::<32>::new(temp_dir.path().join("prolly_rocksdb"))?; + let config2 = TreeConfig::<32>::default(); + let tree2 = + ProllyTree::load_from_storage(storage2, config2).expect("Failed to load tree from storage"); + + // Verify data is still there + println!("Verifying persistence..."); + for i in 0..5 { + let key = format!("key_{:04}", i); + let node = tree2.find(key.as_bytes()); + if let Some(node) = node { + // Find the key in the node and get its corresponding value + if let Some(idx) = node.keys.iter().position(|k| k == key.as_bytes()) { + let value = &node.values[idx]; + println!( + " {} -> {} (persisted)", + key, + String::from_utf8_lossy(value) + ); + } + } + } + + println!("\nRocksDB storage example completed successfully!"); + + Ok(()) +} diff --git a/src/lib.rs b/src/lib.rs index b52e26b..422ac78 100644 --- a/src/lib.rs +++ b/src/lib.rs @@ -54,6 +54,8 @@ pub mod node; pub mod proof; #[cfg(feature = "python")] pub mod python; +#[cfg(feature = "rocksdb_storage")] +pub mod rocksdb; #[cfg(feature = "sql")] pub mod sql; pub mod storage; diff --git a/src/rocksdb/mod.rs b/src/rocksdb/mod.rs new file mode 100644 index 0000000..200b12b --- /dev/null +++ b/src/rocksdb/mod.rs @@ -0,0 +1,17 @@ +/* +Licensed under the Apache License, Version 2.0 (the "License"); +you may not use this file except in compliance with the License. +You may obtain a copy of the License at + + http://www.apache.org/licenses/LICENSE-2.0 + +Unless required by applicable law or agreed to in writing, software +distributed under the License is distributed on an "AS IS" BASIS, +WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +See the License for the specific language governing permissions and +limitations under the License. +*/ + +pub mod storage; + +pub use storage::RocksDBNodeStorage; diff --git a/src/rocksdb/storage.rs b/src/rocksdb/storage.rs new file mode 100644 index 0000000..7f5d3c1 --- /dev/null +++ b/src/rocksdb/storage.rs @@ -0,0 +1,368 @@ +/* +Licensed under the Apache License, Version 2.0 (the "License"); +you may not use this file except in compliance with the License. +You may obtain a copy of the License at + + http://www.apache.org/licenses/LICENSE-2.0 + +Unless required by applicable law or agreed to in writing, software +distributed under the License is distributed on an "AS IS" BASIS, +WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +See the License for the specific language governing permissions and +limitations under the License. +*/ + +use crate::digest::ValueDigest; +use crate::node::ProllyNode; +use crate::storage::NodeStorage; +use lru::LruCache; +use rocksdb::{ + BlockBasedOptions, Cache, DBCompressionType, Options, SliceTransform, WriteBatch, DB, +}; +use std::num::NonZeroUsize; +use std::path::PathBuf; +use std::sync::{Arc, Mutex}; + +const CONFIG_PREFIX: &[u8] = b"config:"; +const NODE_PREFIX: &[u8] = b"node:"; + +/// RocksDB-backed storage for ProllyTree nodes +/// +/// This storage implementation uses RocksDB as the persistent storage backend, +/// with an LRU cache for frequently accessed nodes to improve performance. +pub struct RocksDBNodeStorage { + db: Arc, + cache: Arc, ProllyNode>>>, +} + +impl Clone for RocksDBNodeStorage { + fn clone(&self) -> Self { + RocksDBNodeStorage { + db: self.db.clone(), + cache: Arc::new(Mutex::new(LruCache::new(NonZeroUsize::new(1000).unwrap()))), + } + } +} + +impl RocksDBNodeStorage { + /// Create a new RocksDBNodeStorage instance with default options + pub fn new(path: PathBuf) -> Result { + let opts = Self::default_options(); + let db = DB::open(&opts, path)?; + + Ok(RocksDBNodeStorage { + db: Arc::new(db), + cache: Arc::new(Mutex::new(LruCache::new(NonZeroUsize::new(1000).unwrap()))), + }) + } + + /// Create RocksDBNodeStorage with custom options + pub fn with_options(path: PathBuf, opts: Options) -> Result { + let db = DB::open(&opts, path)?; + + Ok(RocksDBNodeStorage { + db: Arc::new(db), + cache: Arc::new(Mutex::new(LruCache::new(NonZeroUsize::new(1000).unwrap()))), + }) + } + + /// Create RocksDBNodeStorage with custom cache size + pub fn with_cache_size(path: PathBuf, cache_size: usize) -> Result { + let opts = Self::default_options(); + let db = DB::open(&opts, path)?; + + Ok(RocksDBNodeStorage { + db: Arc::new(db), + cache: Arc::new(Mutex::new(LruCache::new( + NonZeroUsize::new(cache_size).unwrap_or(NonZeroUsize::new(1000).unwrap()), + ))), + }) + } + + /// Get default RocksDB options optimized for ProllyTree + pub fn default_options() -> Options { + let mut opts = Options::default(); + opts.create_if_missing(true); + opts.create_missing_column_families(true); + + // Optimize for ProllyTree's write-heavy workload + opts.set_write_buffer_size(128 * 1024 * 1024); // 128MB memtable + opts.set_max_write_buffer_number(4); + opts.set_min_write_buffer_number_to_merge(2); + + // Enable compression for storage efficiency + opts.set_compression_type(DBCompressionType::Lz4); + opts.set_bottommost_compression_type(DBCompressionType::Zstd); + + // Bloom filters for faster lookups + let mut block_opts = BlockBasedOptions::default(); + block_opts.set_bloom_filter(10.0, false); + + // Block cache for frequently accessed nodes + let cache = Cache::new_lru_cache(512 * 1024 * 1024); // 512MB block cache + block_opts.set_block_cache(&cache); + + opts.set_block_based_table_factory(&block_opts); + + // Use prefix extractor for efficient scans + let prefix_len = NODE_PREFIX.len() + N; + opts.set_prefix_extractor(SliceTransform::create_fixed_prefix(prefix_len)); + + opts + } + + /// Create a key for storing a node + fn node_key(hash: &ValueDigest) -> Vec { + let mut key = Vec::with_capacity(NODE_PREFIX.len() + N); + key.extend_from_slice(NODE_PREFIX); + key.extend_from_slice(&hash.0); + key + } + + /// Create a key for storing config + fn config_key(key: &str) -> Vec { + let mut result = Vec::with_capacity(CONFIG_PREFIX.len() + key.len()); + result.extend_from_slice(CONFIG_PREFIX); + result.extend_from_slice(key.as_bytes()); + result + } +} + +impl NodeStorage for RocksDBNodeStorage { + fn get_node_by_hash(&self, hash: &ValueDigest) -> Option> { + // Check cache first + if let Some(node) = self.cache.lock().unwrap().get(hash) { + return Some(node.clone()); + } + + // Fetch from RocksDB + let key = Self::node_key(hash); + match self.db.get(&key) { + Ok(Some(data)) => { + match bincode::deserialize::>(&data) { + Ok(mut node) => { + // Reset transient flags + node.split = false; + node.merged = false; + + // Update cache + self.cache.lock().unwrap().put(hash.clone(), node.clone()); + + Some(node) + } + Err(_) => None, + } + } + _ => None, + } + } + + fn insert_node(&mut self, hash: ValueDigest, node: ProllyNode) -> Option<()> { + // Update cache + self.cache.lock().unwrap().put(hash.clone(), node.clone()); + + // Serialize and store in RocksDB + match bincode::serialize(&node) { + Ok(data) => { + let key = Self::node_key(&hash); + match self.db.put(&key, data) { + Ok(_) => Some(()), + Err(_) => None, + } + } + Err(_) => None, + } + } + + fn delete_node(&mut self, hash: &ValueDigest) -> Option<()> { + // Remove from cache + self.cache.lock().unwrap().pop(hash); + + // Delete from RocksDB + let key = Self::node_key(hash); + match self.db.delete(&key) { + Ok(_) => Some(()), + Err(_) => None, + } + } + + fn save_config(&self, key: &str, config: &[u8]) { + let db_key = Self::config_key(key); + let _ = self.db.put(&db_key, config); + } + + fn get_config(&self, key: &str) -> Option> { + let db_key = Self::config_key(key); + self.db.get(&db_key).ok().flatten() + } +} + +/// Batch operations for RocksDBNodeStorage +impl RocksDBNodeStorage { + /// Insert multiple nodes in a single batch operation + pub fn batch_insert_nodes( + &mut self, + nodes: Vec<(ValueDigest, ProllyNode)>, + ) -> Result<(), rocksdb::Error> { + let mut batch = WriteBatch::default(); + let mut cache = self.cache.lock().unwrap(); + + for (hash, node) in nodes { + // Update cache + cache.put(hash.clone(), node.clone()); + + // Add to batch + match bincode::serialize(&node) { + Ok(data) => { + let key = Self::node_key(&hash); + batch.put(&key, data); + } + Err(_) => { + // Skip this entry if serialization fails + continue; + } + } + } + + self.db.write(batch) + } + + /// Delete multiple nodes in a single batch operation + pub fn batch_delete_nodes(&mut self, hashes: &[ValueDigest]) -> Result<(), rocksdb::Error> { + let mut batch = WriteBatch::default(); + let mut cache = self.cache.lock().unwrap(); + + for hash in hashes { + // Remove from cache + cache.pop(hash); + + // Add to batch + let key = Self::node_key(hash); + batch.delete(&key); + } + + self.db.write(batch) + } + + /// Flush all pending writes to disk + pub fn flush(&self) -> Result<(), rocksdb::Error> { + self.db.flush() + } +} + +#[cfg(test)] +mod tests { + use super::*; + use crate::config::TreeConfig; + use tempfile::TempDir; + + fn create_test_node() -> ProllyNode { + let config: TreeConfig = TreeConfig::default(); + ProllyNode { + keys: vec![b"key1".to_vec(), b"key2".to_vec()], + key_schema: config.key_schema.clone(), + values: vec![b"value1".to_vec(), b"value2".to_vec()], + value_schema: config.value_schema.clone(), + is_leaf: true, + level: 0, + base: config.base, + modulus: config.modulus, + min_chunk_size: config.min_chunk_size, + max_chunk_size: config.max_chunk_size, + pattern: config.pattern, + split: false, + merged: false, + encode_types: Vec::new(), + encode_values: Vec::new(), + } + } + + #[test] + fn test_rocksdb_basic_operations() { + let temp_dir = TempDir::new().unwrap(); + let mut storage = RocksDBNodeStorage::<32>::new(temp_dir.path().to_path_buf()).unwrap(); + + let node = create_test_node(); + let hash = node.get_hash(); + + // Test insert + assert!(storage.insert_node(hash.clone(), node.clone()).is_some()); + + // Test get + let retrieved = storage.get_node_by_hash(&hash); + assert!(retrieved.is_some()); + + let retrieved_node = retrieved.unwrap(); + assert_eq!(retrieved_node.keys, node.keys); + assert_eq!(retrieved_node.values, node.values); + assert_eq!(retrieved_node.is_leaf, node.is_leaf); + + // Test delete + assert!(storage.delete_node(&hash).is_some()); + assert!(storage.get_node_by_hash(&hash).is_none()); + } + + #[test] + fn test_config_operations() { + let temp_dir = TempDir::new().unwrap(); + let storage = RocksDBNodeStorage::<32>::new(temp_dir.path().to_path_buf()).unwrap(); + + let config_data = b"test config data"; + storage.save_config("test_key", config_data); + + let retrieved = storage.get_config("test_key"); + assert!(retrieved.is_some()); + assert_eq!(retrieved.unwrap(), config_data); + + // Test non-existent config + assert!(storage.get_config("non_existent").is_none()); + } + + #[test] + fn test_batch_operations() { + let temp_dir = TempDir::new().unwrap(); + let mut storage = RocksDBNodeStorage::<32>::new(temp_dir.path().to_path_buf()).unwrap(); + + // Create multiple nodes + let mut nodes = Vec::new(); + for i in 0..10 { + let mut node = create_test_node(); + node.keys[0] = format!("key{}", i).into_bytes(); + let hash = node.get_hash(); + nodes.push((hash, node)); + } + + // Batch insert + let hashes: Vec<_> = nodes.iter().map(|(h, _)| h.clone()).collect(); + assert!(storage.batch_insert_nodes(nodes.clone()).is_ok()); + + // Verify all inserted + for (hash, _) in &nodes { + assert!(storage.get_node_by_hash(hash).is_some()); + } + + // Batch delete + assert!(storage.batch_delete_nodes(&hashes).is_ok()); + + // Verify all deleted + for hash in &hashes { + assert!(storage.get_node_by_hash(hash).is_none()); + } + } + + #[test] + fn test_cache_functionality() { + let temp_dir = TempDir::new().unwrap(); + let mut storage = + RocksDBNodeStorage::<32>::with_cache_size(temp_dir.path().to_path_buf(), 2).unwrap(); + + let node1 = create_test_node(); + let hash1 = node1.get_hash(); + + // Insert and verify it's cached + storage.insert_node(hash1.clone(), node1.clone()); + + // Accessing should be from cache (we can't directly test this, but it should be fast) + assert!(storage.get_node_by_hash(&hash1).is_some()); + } +} diff --git a/src/storage.rs b/src/storage.rs index 5a6db82..5424f72 100644 --- a/src/storage.rs +++ b/src/storage.rs @@ -216,3 +216,6 @@ impl NodeStorage for FileNodeStorage { } } } + +#[cfg(feature = "rocksdb_storage")] +pub use crate::rocksdb::RocksDBNodeStorage;