diff --git a/crates/store/src/combined.rs b/crates/store/src/combined.rs index 8f3ca7e..9b9e3ec 100644 --- a/crates/store/src/combined.rs +++ b/crates/store/src/combined.rs @@ -1,562 +1,578 @@ //! Generic transaction module for atomic operations across multiple stores. //! //! # Example //! //! This module provides a generic transaction mechanism for atomic operations across //! different types of stores. Each store implementation can plug into this system //! by implementing the `TransactionProvider` trait. //! //! ```no_run //! use store::{Transaction, TransactionContext}; //! //! // Assuming you have stores and trees //! # fn example_usage() -> store::Result<()> { //! # let temp_dir = tempfile::tempdir().unwrap(); //! # let db = sled::open(temp_dir.path())?; //! # let range_store = store::RangeStore::open(&db)?; //! # let namespace_store = store::NamespaceStore::open(&db)?; //! # let metadata_tree = db.open_tree("metadata")?; //! //! // Create a transaction with the stores you want to include //! let transaction = Transaction::new() //! .with_store("ranges", &range_store) //! .with_store("namespaces", &namespace_store) //! .with_tree(&metadata_tree); //! //! // Execute the transaction //! let result = transaction.execute(|ctx| { //! // Access stores by name //! let range_trees = ctx.store_trees("ranges")?; //! let namespace_trees = ctx.store_trees("namespaces")?; //! //! // Access additional trees by index //! let metadata = ctx.tree(0)?; //! //! metadata.insert("operation", "test")?; //! //! Ok(()) //! })?; //! # Ok(()) //! # } //! ``` use crate::{Result, Error}; use sled::Transactional; use std::collections::HashMap; /// Helper function to convert transaction errors fn convert_transaction_error(e: sled::transaction::ConflictableTransactionError, default_error: Error) -> Error { match e { sled::transaction::ConflictableTransactionError::Storage(storage_err) => Error::StoreError(storage_err), sled::transaction::ConflictableTransactionError::Abort(_) => default_error, _ => Error::StoreError(sled::Error::Unsupported("Unknown transaction error".to_string())), } } /// Trait for types that can provide trees to a transaction pub trait TransactionProvider { /// Return the trees that should be included in a transaction fn transaction_trees(&self) -> Vec<&sled::Tree>; } /// Implement TransactionProvider for individual trees impl TransactionProvider for sled::Tree { fn transaction_trees(&self) -> Vec<&sled::Tree> { vec![self] } } /// Implement TransactionProvider for RangeStore impl TransactionProvider for crate::RangeStore { fn transaction_trees(&self) -> Vec<&sled::Tree> { vec![&self.names, &self.map, &self.assign] } } /// Implement TransactionProvider for NamespaceStore impl TransactionProvider for crate::NamespaceStore { fn transaction_trees(&self) -> Vec<&sled::Tree> { vec![&self.names, &self.spaces] } } /// Implement TransactionProvider for NetworkStore impl TransactionProvider for crate::NetworkStore { fn transaction_trees(&self) -> Vec<&sled::Tree> { let mut trees = self.namespaces.transaction_trees(); trees.push(&self.networks); trees } } /// Generic transaction context provided to transaction operations pub struct TransactionContext<'ctx> { store_map: HashMap, // name -> (start_idx, end_idx) trees: Vec<&'ctx sled::transaction::TransactionalTree>, transactional_trees: &'ctx [sled::transaction::TransactionalTree], additional_trees_start: usize, } impl<'ctx> TransactionContext<'ctx> { /// Create a new transaction context fn new( store_map: HashMap, trees: Vec<&'ctx sled::transaction::TransactionalTree>, transactional_trees: &'ctx [sled::transaction::TransactionalTree], additional_trees_start: usize, ) -> Self { Self { store_map, trees, transactional_trees, additional_trees_start, } } /// Get trees for a store by name pub fn store_trees(&self, store_name: &str) -> Result<&[&sled::transaction::TransactionalTree]> { let (start_idx, end_idx) = self.store_map .get(store_name) .ok_or_else(|| Error::StoreError(sled::Error::Unsupported(format!("Store '{}' not found in transaction", store_name))))?; Ok(&self.trees[*start_idx..*end_idx]) } /// Access additional trees by index pub fn tree(&self, index: usize) -> Result<&sled::transaction::TransactionalTree> { self.trees.get(self.additional_trees_start + index) .copied() .ok_or_else(|| Error::StoreError(sled::Error::Unsupported(format!("Tree at index {} not found", index)))) } /// Access a raw transactional tree by its absolute index pub fn raw_tree(&self, index: usize) -> Result<&sled::transaction::TransactionalTree> { self.trees.get(index) .copied() .ok_or_else(|| Error::StoreError(sled::Error::Unsupported(format!("Raw tree at index {} not found", index)))) } /// Access the entire slice of transactional trees pub fn all_trees(&self) -> &[sled::transaction::TransactionalTree] { self.transactional_trees } /// Get store map for debugging or extension purposes pub fn store_map(&self) -> &HashMap { &self.store_map } } /// Generic transaction struct for atomic operations across multiple stores pub struct Transaction<'a> { stores: HashMap, additional_trees: Vec<&'a sled::Tree>, } impl<'a> Transaction<'a> { /// Create a new empty transaction pub fn new() -> Self { Self { stores: HashMap::new(), additional_trees: Vec::new(), } } /// Add a store with a name identifier pub fn with_store(mut self, name: &str, store: &'a T) -> Self { self.stores.insert(name.to_string(), store); self } /// Add a single tree to the transaction pub fn with_tree(mut self, tree: &'a sled::Tree) -> Self { self.additional_trees.push(tree); self } /// Execute a transaction with the configured stores pub fn execute(&self, operations: F) -> Result where F: Fn(&TransactionContext) -> Result, { // Collect all trees for the transaction let mut all_trees = Vec::new(); let mut store_map = HashMap::new(); // Add trees from stores for (name, store) in &self.stores { let start_idx = all_trees.len(); all_trees.extend(store.transaction_trees()); let end_idx = all_trees.len(); store_map.insert(name.clone(), (start_idx, end_idx)); } // Add additional trees let additional_trees_start = all_trees.len(); all_trees.extend(&self.additional_trees); // Execute the transaction let result = all_trees.transaction(|trees| { let context = TransactionContext::new( store_map.clone(), trees.into_iter().collect(), trees, additional_trees_start, ); operations(&context).map_err(|e| match e { Error::StoreError(store_err) => sled::transaction::ConflictableTransactionError::Storage(store_err), _ => sled::transaction::ConflictableTransactionError::Abort(()), }) }).map_err(|e| match e { sled::transaction::TransactionError::Abort(()) => Error::StoreError(sled::Error::Unsupported("Transaction aborted".to_string())), sled::transaction::TransactionError::Storage(storage_err) => Error::StoreError(storage_err), })?; Ok(result) } } impl<'a> Default for Transaction<'a> { fn default() -> Self { Self::new() } } /// Legacy alias for backward compatibility pub type CombinedTransaction<'a> = Transaction<'a>; /// Legacy alias for backward compatibility pub type CombinedTransactionContext<'ctx> = TransactionContext<'ctx>; #[cfg(test)] mod tests { use super::*; use crate::{RangeStore, NamespaceStore}; use tempfile::tempdir; fn create_test_stores() -> Result<(RangeStore, NamespaceStore, sled::Db)> { let temp_dir = tempdir().unwrap(); let db = sled::open(temp_dir.path())?; let range_store = RangeStore::open(&db)?; let namespace_store = NamespaceStore::open(&db)?; Ok((range_store, namespace_store, db)) } #[test] fn test_generic_transaction_basic() -> Result<()> { let (range_store, namespace_store, db) = create_test_stores()?; // Setup: define range and namespace range_store.define("test_range", 100)?; namespace_store.define("test_namespace")?; // Create additional tree for testing let extra_tree = db.open_tree("extra")?; let transaction = Transaction::new() .with_store("ranges", &range_store) .with_store("namespaces", &namespace_store) .with_tree(&extra_tree); // Execute transaction using generic interface transaction.execute(|ctx| { // Access range store trees let range_trees = ctx.store_trees("ranges")?; assert_eq!(range_trees.len(), 3); // names, map, assign // Access namespace store trees let namespace_trees = ctx.store_trees("namespaces")?; assert_eq!(namespace_trees.len(), 2); // names, spaces // Use additional tree let tree = ctx.tree(0)?; tree.insert("test_key", "test_value") .map_err(|e| Error::StoreError(sled::Error::Unsupported(format!("Insert failed: {}", e))))?; Ok(()) })?; // Verify the additional tree was modified let value = extra_tree.get("test_key")?; assert_eq!(value, Some("test_value".as_bytes().into())); Ok(()) } #[test] fn test_transaction_with_single_store() -> Result<()> { let (range_store, _, _) = create_test_stores()?; // Setup range_store.define("test_range", 50)?; let transaction = Transaction::new() .with_store("ranges", &range_store); // Execute transaction with just one store transaction.execute(|ctx| { let range_trees = ctx.store_trees("ranges")?; assert_eq!(range_trees.len(), 3); // Verify we can access the trees let _names_tree = range_trees[0]; let _map_tree = range_trees[1]; let _assign_tree = range_trees[2]; Ok(()) })?; Ok(()) } #[test] fn test_transaction_with_only_trees() -> Result<()> { let (_, _, db) = create_test_stores()?; let tree1 = db.open_tree("tree1")?; let tree2 = db.open_tree("tree2")?; let transaction = Transaction::new() .with_tree(&tree1) .with_tree(&tree2); transaction.execute(|ctx| { let t1 = ctx.tree(0)?; let t2 = ctx.tree(1)?; t1.insert("key1", "value1") .map_err(|e| Error::StoreError(sled::Error::Unsupported(format!("Insert failed: {}", e))))?; t2.insert("key2", "value2") .map_err(|e| Error::StoreError(sled::Error::Unsupported(format!("Insert failed: {}", e))))?; Ok(()) })?; // Verify the trees were modified let value1 = tree1.get("key1")?; assert_eq!(value1, Some("value1".as_bytes().into())); let value2 = tree2.get("key2")?; assert_eq!(value2, Some("value2".as_bytes().into())); Ok(()) } #[test] fn test_multiple_stores_same_type() -> Result<()> { let (range_store1, _, db) = create_test_stores()?; let range_store2 = RangeStore::open(&db)?; // Setup both stores range_store1.define("range1", 50)?; range_store2.define("range2", 100)?; let transaction = Transaction::new() .with_store("ranges1", &range_store1) .with_store("ranges2", &range_store2); transaction.execute(|ctx| { let trees1 = ctx.store_trees("ranges1")?; let trees2 = ctx.store_trees("ranges2")?; assert_eq!(trees1.len(), 3); assert_eq!(trees2.len(), 3); // Verify different stores have different trees assert_ne!(trees1[0] as *const _, trees2[0] as *const _); Ok(()) })?; Ok(()) } #[test] fn test_store_not_found_error() -> Result<()> { let (range_store, _, _) = create_test_stores()?; let transaction = Transaction::new() .with_store("ranges", &range_store); let result: Result<()> = transaction.execute(|ctx| { // Try to access a store that doesn't exist let _trees = ctx.store_trees("nonexistent")?; Ok(()) }); assert!(result.is_err()); Ok(()) } #[test] fn test_tree_index_out_of_bounds() -> Result<()> { let (_, _, db) = create_test_stores()?; let tree = db.open_tree("single_tree")?; let transaction = Transaction::new() .with_tree(&tree); let result: Result<()> = transaction.execute(|ctx| { // Try to access tree at index that doesn't exist let _tree = ctx.tree(5)?; Ok(()) }); assert!(result.is_err()); Ok(()) } #[test] fn test_transaction_rollback() -> Result<()> { let (_, _, db) = create_test_stores()?; let tree = db.open_tree("rollback_test")?; // First, insert some initial data tree.insert("initial", "data")?; let transaction = Transaction::new() .with_tree(&tree); // Execute a transaction that should fail let result: Result<()> = transaction.execute(|ctx| { let t = ctx.tree(0)?; // Insert some data t.insert("temp", "value") .map_err(|e| Error::StoreError(sled::Error::Unsupported(format!("Insert failed: {}", e))))?; // Force an error to trigger rollback Err(Error::StoreError(sled::Error::Unsupported("Forced error".to_string()))) }); assert!(result.is_err()); // Verify rollback - temp key should not exist let temp_value = tree.get("temp")?; assert_eq!(temp_value, None); // But initial data should still be there let initial_value = tree.get("initial")?; assert_eq!(initial_value, Some("data".as_bytes().into())); Ok(()) } #[test] fn test_complex_multi_store_transaction() -> Result<()> { let (range_store, namespace_store, db) = create_test_stores()?; // Setup range_store.define("ip_pool", 100)?; namespace_store.define("users")?; let metadata_tree = db.open_tree("metadata")?; let logs_tree = db.open_tree("logs")?; let transaction = Transaction::new() .with_store("ranges", &range_store) .with_store("namespaces", &namespace_store) .with_tree(&metadata_tree) .with_tree(&logs_tree); // Complex transaction transaction.execute(|ctx| { let range_trees = ctx.store_trees("ranges")?; let namespace_trees = ctx.store_trees("namespaces")?; let metadata = ctx.tree(0)?; let logs = ctx.tree(1)?; // Verify we have the right number of trees assert_eq!(range_trees.len(), 3); assert_eq!(namespace_trees.len(), 2); // Use metadata tree metadata.insert("operation", "complex_transaction") .map_err(|e| Error::StoreError(sled::Error::Unsupported(format!("Insert failed: {}", e))))?; // Use logs tree logs.insert("log_entry_1", "Started complex operation") .map_err(|e| Error::StoreError(sled::Error::Unsupported(format!("Insert failed: {}", e))))?; Ok(()) })?; // Verify all operations succeeded let op_value = metadata_tree.get("operation")?; assert_eq!(op_value, Some("complex_transaction".as_bytes().into())); let log_value = logs_tree.get("log_entry_1")?; assert_eq!(log_value, Some("Started complex operation".as_bytes().into())); Ok(()) } #[test] fn test_raw_tree_access() -> Result<()> { let (range_store, namespace_store, db) = create_test_stores()?; let extra_tree = db.open_tree("extra")?; let transaction = Transaction::new() .with_store("ranges", &range_store) .with_store("namespaces", &namespace_store) .with_tree(&extra_tree); transaction.execute(|ctx| { // Test raw tree access by absolute index let tree0 = ctx.raw_tree(0)?; // First range store tree let tree3 = ctx.raw_tree(3)?; // First namespace store tree let tree5 = ctx.raw_tree(5)?; // Extra tree // All should be valid trees tree0.insert("raw0", "value0") .map_err(|e| Error::StoreError(sled::Error::Unsupported(format!("Insert failed: {}", e))))?; tree3.insert("raw3", "value3") .map_err(|e| Error::StoreError(sled::Error::Unsupported(format!("Insert failed: {}", e))))?; tree5.insert("raw5", "value5") .map_err(|e| Error::StoreError(sled::Error::Unsupported(format!("Insert failed: {}", e))))?; // Test accessing out of bounds let invalid_result = ctx.raw_tree(10); assert!(invalid_result.is_err()); Ok(()) })?; Ok(()) } #[test] fn test_store_map_access() -> Result<()> { let (range_store, namespace_store, _) = create_test_stores()?; let transaction = Transaction::new() .with_store("my_ranges", &range_store) .with_store("my_namespaces", &namespace_store); transaction.execute(|ctx| { let store_map = ctx.store_map(); // Verify store map contains our stores assert!(store_map.contains_key("my_ranges")); assert!(store_map.contains_key("my_namespaces")); - // Verify ranges - let (start, end) = store_map.get("my_ranges").unwrap(); - assert_eq!(*start, 0); - assert_eq!(*end, 3); // RangeStore has 3 trees + // Get the ranges for verification + let (ranges_start, ranges_end) = store_map.get("my_ranges").unwrap(); + let (namespaces_start, namespaces_end) = store_map.get("my_namespaces").unwrap(); - // Verify namespaces - let (start, end) = store_map.get("my_namespaces").unwrap(); - assert_eq!(*start, 3); - assert_eq!(*end, 5); // NamespaceStore has 2 trees + // Verify that RangeStore has 3 trees and NamespaceStore has 2 trees + assert_eq!(*ranges_end - *ranges_start, 3); // RangeStore has 3 trees + assert_eq!(*namespaces_end - *namespaces_start, 2); // NamespaceStore has 2 trees + + // Verify that ranges are contiguous and cover all trees (order-independent) + let mut all_indices = vec![*ranges_start, *ranges_end, *namespaces_start, *namespaces_end]; + all_indices.sort(); + + // Total should be 5 trees (3 + 2), starting from 0 + assert_eq!(all_indices[0], 0); // First store starts at 0 + assert_eq!(all_indices[3], 5); // Last tree ends at 5 + + // Verify stores are contiguous (no gaps between them) + let ranges_contiguous = (*ranges_end - *ranges_start) == 3; + let namespaces_contiguous = (*namespaces_end - *namespaces_start) == 2; + assert!(ranges_contiguous && namespaces_contiguous); + + // Verify one store ends where the other begins (no overlap, no gaps) + let ranges_then_namespaces = *ranges_end == *namespaces_start; + let namespaces_then_ranges = *namespaces_end == *ranges_start; + assert!(ranges_then_namespaces || namespaces_then_ranges); Ok(()) })?; Ok(()) } } \ No newline at end of file