diff --git a/crates/store/src/combined.rs b/crates/store/src/combined.rs index 458c99d..3137051 100644 --- a/crates/store/src/combined.rs +++ b/crates/store/src/combined.rs @@ -1,782 +1,1183 @@ //! 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 range_store and namespace_store instances //! // and an additional tree for metadata //! # 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)?; //! # range_store.define("ip_addresses", 256)?; //! # namespace_store.define("users")?; //! # let metadata_tree = db.open_tree("metadata")?; //! //! // Create a transaction with the stores you want to include //! let transaction = Transaction::new() //! .with_store(&range_store) //! .with_store(&namespace_store) //! .with_tree(&metadata_tree); //! //! // Execute the transaction //! let (ip_bit, reserved) = transaction.execute(|ctx| { //! // Reserve a username using the namespace store's transaction methods //! let reserved = ctx.use_namespace().reserve("users", "alice", "user_data")?; //! //! // Assign an IP address using the range store's transaction methods //! let ip_bit = ctx.use_range().assign("ip_addresses", "192.168.1.100")?; //! //! // Store metadata in additional tree //! let tree = ctx.tree(0).ok_or_else(|| //! store::Error::StoreError(sled::Error::Unsupported("Tree not found".to_string())) //! )?; //! //! tree.insert("last_assignment", "alice") //! .map_err(|e| store::Error::StoreError(e))?; //! //! Ok((ip_bit, reserved)) //! })?; //! //! println!("Assigned IP bit position: {}, Reserved: {}", ip_bit, reserved); //! # Ok(()) //! # } //! ``` use crate::{Result, Error, RangeStore, NamespaceStore}; use sled::Transactional; /// 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 RangeStore { fn transaction_trees(&self) -> Vec<&sled::Tree> { vec![&self.names, &self.map, &self.assign] } } /// Implement TransactionProvider for NamespaceStore impl TransactionProvider for NamespaceStore { fn transaction_trees(&self) -> Vec<&sled::Tree> { vec![&self.names, &self.spaces] } } /// RangeTransactionContext provides range-specific transaction operations pub struct RangeTransactionContext<'a, 'ctx> { store: &'a RangeStore, ranges_names: &'ctx sled::transaction::TransactionalTree, ranges_map: &'ctx sled::transaction::TransactionalTree, ranges_assign: &'ctx sled::transaction::TransactionalTree, } impl<'a, 'ctx> RangeTransactionContext<'a, 'ctx> { /// Assign a value to a range within the transaction pub fn assign(&self, range_name: &str, value: &str) -> Result { self.store.assign_in_transaction( self.ranges_names, self.ranges_map, self.ranges_assign, range_name, value, ).map_err(|e| convert_transaction_error(e, Error::RangeFull(range_name.to_string()))) } /// Get range assignment details within the transaction pub fn get(&self, range_name: &str, bit_position: u64) -> Result> { self.store.get_in_transaction( self.ranges_names, self.ranges_assign, range_name, bit_position, ).map_err(|e| convert_transaction_error(e, Error::UndefinedRange(range_name.to_string()))) } + + /// Unassign a specific bit position within the transaction + pub fn unassign_bit(&self, range_name: &str, bit_position: u64) -> Result { + self.store.unassign_bit_in_transaction( + self.ranges_names, + self.ranges_map, + self.ranges_assign, + range_name, + bit_position, + ).map_err(|e| convert_transaction_error(e, Error::BitOutOfRange(range_name.to_string(), bit_position))) + } + + + + /// Check if a range exists within the transaction + pub fn exists(&self, range_name: &str) -> Result { + self.store.exists_in_transaction( + self.ranges_names, + range_name, + ).map_err(|e| convert_transaction_error(e, Error::UndefinedRange(range_name.to_string()))) + } + + /// Get range info (id and size) within the transaction + pub fn info(&self, range_name: &str) -> Result> { + self.store.info_in_transaction( + self.ranges_names, + range_name, + ).map_err(|e| convert_transaction_error(e, Error::UndefinedRange(range_name.to_string()))) + } } /// NamespaceTransactionContext provides namespace-specific transaction operations pub struct NamespaceTransactionContext<'a, 'ctx> { store: &'a NamespaceStore, namespace_names: &'ctx sled::transaction::TransactionalTree, namespace_spaces: &'ctx sled::transaction::TransactionalTree, } impl<'a, 'ctx> NamespaceTransactionContext<'a, 'ctx> { /// Reserve a key in a namespace within the transaction pub fn reserve(&self, namespace: &str, key: &str, value: &str) -> Result { self.store.reserve_in_transaction( self.namespace_names, self.namespace_spaces, namespace, key, value, ).map_err(|e| convert_transaction_error(e, Error::NamespaceKeyReserved(namespace.to_string(), key.to_string()))) } /// Get a value from a namespace within the transaction pub fn get(&self, namespace: &str, key: &str) -> Result> { self.store.get_in_transaction( self.namespace_names, self.namespace_spaces, namespace, key, ).map_err(|e| convert_transaction_error(e, Error::UndefinedNamespace(namespace.to_string()))) } /// Remove a key from a namespace within the transaction pub fn remove(&self, namespace: &str, key: &str) -> Result { self.store.remove_in_transaction( self.namespace_names, self.namespace_spaces, namespace, key, ).map_err(|e| convert_transaction_error(e, Error::UndefinedNamespace(namespace.to_string()))) } + + /// Update a key in a namespace within the transaction + pub fn update(&self, namespace: &str, key: &str, value: &str) -> Result { + self.store.update_in_transaction( + self.namespace_names, + self.namespace_spaces, + namespace, + key, + value, + ).map_err(|e| convert_transaction_error(e, Error::UndefinedNamespace(namespace.to_string()))) + } + + + + /// Check if a namespace exists within the transaction + pub fn namespace_exists(&self, namespace: &str) -> Result { + self.store.namespace_exists_in_transaction( + self.namespace_names, + namespace, + ).map_err(|e| convert_transaction_error(e, Error::UndefinedNamespace(namespace.to_string()))) + } + + /// Check if a key exists in a namespace within the transaction + pub fn key_exists(&self, namespace: &str, key: &str) -> Result { + self.store.key_exists_in_transaction( + self.namespace_names, + self.namespace_spaces, + namespace, + key, + ).map_err(|e| convert_transaction_error(e, Error::UndefinedNamespace(namespace.to_string()))) + } } /// Generic transaction context provided to transaction operations pub struct TransactionContext<'a, 'ctx> { range_store: Option<&'a RangeStore>, namespace_store: Option<&'a NamespaceStore>, trees: Vec<&'ctx sled::transaction::TransactionalTree>, transactional_trees: &'ctx [sled::transaction::TransactionalTree], } impl<'a, 'ctx> TransactionContext<'a, 'ctx> { /// Create a new transaction context fn new( range_store: Option<&'a RangeStore>, namespace_store: Option<&'a NamespaceStore>, trees: Vec<&'ctx sled::transaction::TransactionalTree>, transactional_trees: &'ctx [sled::transaction::TransactionalTree], ) -> Self { Self { range_store, namespace_store, trees, transactional_trees, } } /// Access range store operations pub fn use_range(&self) -> RangeTransactionContext<'a, 'ctx> { let store = self.range_store.expect("RangeStore not included in transaction"); // RangeStore requires 3 trees: names, map, assign RangeTransactionContext { store, ranges_names: self.trees[0], ranges_map: self.trees[1], ranges_assign: self.trees[2], } } /// Access namespace store operations pub fn use_namespace(&self) -> NamespaceTransactionContext<'a, 'ctx> { let store = self.namespace_store.expect("NamespaceStore not included in transaction"); // The index depends on whether RangeStore was included let base_index = if self.range_store.is_some() { 3 } else { 0 }; // NamespaceStore requires 2 trees: names, spaces NamespaceTransactionContext { store, namespace_names: self.trees[base_index], namespace_spaces: self.trees[base_index + 1], } } /// Access additional trees by index pub fn tree(&self, index: usize) -> Option<&sled::transaction::TransactionalTree> { let base_index = if self.range_store.is_some() { 3 } else { 0 }; let base_index = base_index + if self.namespace_store.is_some() { 2 } else { 0 }; self.trees.get(base_index + index).copied() } /// Access a raw transactional tree by its absolute index /// Used by extensions that might need direct access pub fn raw_tree(&self, index: usize) -> Option<&sled::transaction::TransactionalTree> { self.trees.get(index).copied() } /// Access the entire slice of transactional trees /// Used by extensions that might need direct access pub fn all_trees(&self) -> &[sled::transaction::TransactionalTree] { self.transactional_trees } } /// Generic transaction struct for atomic operations across multiple stores pub struct Transaction<'a> { range_store: Option<&'a RangeStore>, namespace_store: Option<&'a NamespaceStore>, additional_trees: Vec<&'a sled::Tree>, } impl<'a> Transaction<'a> { /// Create a new empty transaction pub fn new() -> Self { Self { range_store: None, namespace_store: None, additional_trees: Vec::new(), } } /// Add a RangeStore to the transaction pub fn with_range_store(mut self, store: &'a RangeStore) -> Self { self.range_store = Some(store); self } /// Add a NamespaceStore to the transaction pub fn with_namespace_store(mut self, store: &'a NamespaceStore) -> Self { self.namespace_store = Some(store); self } /// Add a generic store that implements TransactionProvider to the transaction pub fn with_store(mut self, store: &'a T) -> Self { // This is a convenience method for future store types let trees = store.transaction_trees(); self.additional_trees.extend(trees); 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(); // Add trees from RangeStore if present if let Some(range_store) = self.range_store { all_trees.extend(range_store.transaction_trees()); } // Add trees from NamespaceStore if present if let Some(namespace_store) = self.namespace_store { all_trees.extend(namespace_store.transaction_trees()); } // Add additional trees all_trees.extend(&self.additional_trees); // Execute the transaction let result = all_trees.transaction(|trees| { let context = TransactionContext::new( self.range_store, self.namespace_store, trees.into_iter().collect(), trees, ); 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<'a, 'ctx> = TransactionContext<'a, 'ctx>; /// Extension trait system for adding custom functionality to TransactionContext pub trait TransactionExtension<'a, 'ctx> { /// Create a new instance of the extension from a transaction context fn from_context(ctx: &'ctx TransactionContext<'a, 'ctx>) -> Self; } /// Registry for custom store types in transactions pub struct StoreRegistry<'a> { stores: Vec<(&'a dyn TransactionProvider, &'static str)>, } impl<'a> StoreRegistry<'a> { pub fn new() -> Self { Self { stores: Vec::new(), } } /// Register a store with a type identifier pub fn register(mut self, store: &'a T, type_name: &'static str) -> Self { self.stores.push((store, type_name)); self } /// Get all registered stores pub fn stores(&self) -> &[(&'a dyn TransactionProvider, &'static str)] { &self.stores } } /// Extended transaction builder that supports custom stores pub struct ExtendedTransaction<'a> { range_store: Option<&'a RangeStore>, namespace_store: Option<&'a NamespaceStore>, additional_trees: Vec<&'a sled::Tree>, custom_stores: StoreRegistry<'a>, } impl<'a> ExtendedTransaction<'a> { /// Create a new extended transaction pub fn new() -> Self { Self { range_store: None, namespace_store: None, additional_trees: Vec::new(), custom_stores: StoreRegistry::new(), } } /// Add a RangeStore to the transaction pub fn with_range_store(mut self, store: &'a RangeStore) -> Self { self.range_store = Some(store); self } /// Add a NamespaceStore to the transaction pub fn with_namespace_store(mut self, store: &'a NamespaceStore) -> Self { self.namespace_store = Some(store); self } /// Add a custom store with type information pub fn with_custom_store(mut self, store: &'a T, type_name: &'static str) -> Self { self.custom_stores = self.custom_stores.register(store, type_name); 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 the transaction with store type information pub fn execute(&self, operations: F) -> Result where F: Fn(&ExtendedTransactionContext) -> Result, { // Collect all trees for the transaction let mut all_trees = Vec::new(); let mut store_map = Vec::new(); // Add trees from RangeStore if present if let Some(range_store) = self.range_store { let start_idx = all_trees.len(); all_trees.extend(range_store.transaction_trees()); store_map.push(("range", start_idx, all_trees.len())); } // Add trees from NamespaceStore if present if let Some(namespace_store) = self.namespace_store { let start_idx = all_trees.len(); all_trees.extend(namespace_store.transaction_trees()); store_map.push(("namespace", start_idx, all_trees.len())); } // Add trees from custom stores for (store, type_name) in self.custom_stores.stores() { let start_idx = all_trees.len(); all_trees.extend(store.transaction_trees()); store_map.push((type_name, start_idx, all_trees.len())); } // Add additional trees let additional_start = all_trees.len(); all_trees.extend(&self.additional_trees); // Execute the transaction let result = all_trees.transaction(|trees| { let context = ExtendedTransactionContext::new( self.range_store, self.namespace_store, trees.into_iter().collect(), trees, store_map.clone(), additional_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) } } /// Extended transaction context with support for custom stores pub struct ExtendedTransactionContext<'a, 'ctx> { range_store: Option<&'a RangeStore>, namespace_store: Option<&'a NamespaceStore>, trees: Vec<&'ctx sled::transaction::TransactionalTree>, transactional_trees: &'ctx [sled::transaction::TransactionalTree], store_map: Vec<(&'static str, usize, usize)>, // (type_name, start_idx, end_idx) additional_trees_start: usize, } impl<'a, 'ctx> ExtendedTransactionContext<'a, 'ctx> { fn new( range_store: Option<&'a RangeStore>, namespace_store: Option<&'a NamespaceStore>, trees: Vec<&'ctx sled::transaction::TransactionalTree>, transactional_trees: &'ctx [sled::transaction::TransactionalTree], store_map: Vec<(&'static str, usize, usize)>, additional_trees_start: usize, ) -> Self { Self { range_store, namespace_store, trees, transactional_trees, store_map, additional_trees_start, } } /// Access range store operations pub fn use_range(&self) -> RangeTransactionContext<'a, 'ctx> { let store = self.range_store.expect("RangeStore not included in transaction"); // Find the range store in the store map let (_, start_idx, _) = self.store_map .iter() .find(|(name, _, _)| *name == "range") .expect("Range store not found in store map"); RangeTransactionContext { store, ranges_names: self.trees[*start_idx], ranges_map: self.trees[*start_idx + 1], ranges_assign: self.trees[*start_idx + 2], } } /// Access namespace store operations pub fn use_namespace(&self) -> NamespaceTransactionContext<'a, 'ctx> { let store = self.namespace_store.expect("NamespaceStore not included in transaction"); // Find the namespace store in the store map let (_, start_idx, _) = self.store_map .iter() .find(|(name, _, _)| *name == "namespace") .expect("Namespace store not found in store map"); NamespaceTransactionContext { store, namespace_names: self.trees[*start_idx], namespace_spaces: self.trees[*start_idx + 1], } } /// Access trees for a custom store by type name pub fn custom_store_trees(&self, type_name: &str) -> Option<&[&sled::transaction::TransactionalTree]> { self.store_map .iter() .find(|(name, _, _)| *name == type_name) .map(|(_, start_idx, end_idx)| &self.trees[*start_idx..*end_idx]) } /// Access additional trees by index pub fn tree(&self, index: usize) -> Option<&sled::transaction::TransactionalTree> { self.trees.get(self.additional_trees_start + index).copied() } /// Access a raw transactional tree by its absolute index pub fn raw_tree(&self, index: usize) -> Option<&sled::transaction::TransactionalTree> { self.trees.get(index).copied() } /// 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) -> &[(&'static str, usize, usize)] { &self.store_map } } impl<'a> Default for ExtendedTransaction<'a> { fn default() -> Self { Self::new() } } #[cfg(test)] mod tests { use super::*; 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_combined_range_and_namespace_assignment() -> 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_range_store(&range_store) .with_namespace_store(&namespace_store) .with_tree(&extra_tree); // Execute combined transaction let (bit_position, reserved) = transaction.execute(|ctx| { // Reserve namespace key let reserved = ctx.use_namespace().reserve("test_namespace", "my_key", "my_value")?; // Assign range value let bit_position = ctx.use_range().assign("test_range", "range_value")?; // Use additional tree if let Some(tree) = ctx.tree(0) { tree.insert("extra_key", "extra_value") .map_err(|e| Error::StoreError(sled::Error::Unsupported(format!("Tree insert failed: {}", e))))?; } Ok((bit_position, reserved)) })?; // Verify results assert!(bit_position < 100); // Should be within range size assert!(reserved); let namespace_value = namespace_store.get("test_namespace", "my_key")?; assert_eq!(namespace_value, Some("my_value".to_string())); let extra_value = extra_tree.get("extra_key")?; assert_eq!(extra_value, Some("extra_value".as_bytes().into())); Ok(()) } #[test] fn test_transaction_rollback_on_error() -> Result<()> { let (range_store, namespace_store, _) = create_test_stores()?; // Setup: define range and namespace range_store.define("test_range", 100)?; namespace_store.define("test_namespace")?; // Reserve a key first namespace_store.reserve("test_namespace", "existing_key", "existing_value")?; let transaction = Transaction::new() .with_range_store(&range_store) .with_namespace_store(&namespace_store); // Execute transaction that should fail let result = transaction.execute(|ctx| { // This should succeed let _bit_pos = ctx.use_range().assign("test_range", "range_value")?; // This should fail (key already exists) let _reserved = ctx.use_namespace().reserve("test_namespace", "existing_key", "new_value")?; Ok(()) }); // Transaction should have failed assert!(result.is_err()); // Verify that no range assignment was made (transaction rolled back) let ranges = range_store.list_range("test_range")?; assert!(ranges.is_empty()); Ok(()) } #[test] fn test_read_operations_in_transaction() -> Result<()> { let (range_store, namespace_store, _) = create_test_stores()?; // Setup range_store.define("test_range", 100)?; namespace_store.define("test_namespace")?; namespace_store.reserve("test_namespace", "existing_key", "existing_value")?; let transaction = Transaction::new() .with_namespace_store(&namespace_store) .with_range_store(&range_store); // Execute transaction with reads let (bit_position, existing_value) = transaction.execute(|ctx| { // Read existing value let existing_value = ctx.use_namespace().get("test_namespace", "existing_key")?; // Assign new range value let bit_position = ctx.use_range().assign("test_range", "new_range_value")?; Ok((bit_position, existing_value)) })?; assert!(bit_position < 100); // Should be within range size assert_eq!(existing_value, Some("existing_value".to_string())); Ok(()) } #[test] fn test_transaction_with_just_namespace_store() -> Result<()> { let (_, namespace_store, _) = create_test_stores()?; // Setup namespace_store.define("test_namespace")?; let transaction = Transaction::new() .with_namespace_store(&namespace_store); // Execute transaction with just namespace operations let success = transaction.execute(|ctx| { ctx.use_namespace().reserve("test_namespace", "new_key", "new_value") })?; assert!(success); // Verify the key was reserved let value = namespace_store.get("test_namespace", "new_key")?; assert_eq!(value, Some("new_value".to_string())); Ok(()) } #[test] fn test_transaction_with_just_range_store() -> Result<()> { let (range_store, _, _) = create_test_stores()?; // Setup range_store.define("test_range", 100)?; let transaction = Transaction::new() .with_range_store(&range_store); // Execute transaction with just range operations let bit_position = transaction.execute(|ctx| { ctx.use_range().assign("test_range", "test_value") })?; assert!(bit_position < 100); // Verify the assignment let ranges = range_store.list_range("test_range")?; assert_eq!(ranges.len(), 1); Ok(()) } #[test] fn test_extended_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 = ExtendedTransaction::new() .with_range_store(&range_store) .with_namespace_store(&namespace_store) .with_tree(&extra_tree); // Execute transaction let (bit_position, reserved) = transaction.execute(|ctx| { // Reserve namespace key let reserved = ctx.use_namespace().reserve("test_namespace", "my_key", "my_value")?; // Assign range value let bit_position = ctx.use_range().assign("test_range", "range_value")?; // Use additional tree if let Some(tree) = ctx.tree(0) { tree.insert("extra_key", "extra_value") .map_err(|e| Error::StoreError(sled::Error::Unsupported(format!("Tree insert failed: {}", e))))?; } Ok((bit_position, reserved)) })?; // Verify results assert!(bit_position < 100); assert!(reserved); let namespace_value = namespace_store.get("test_namespace", "my_key")?; assert_eq!(namespace_value, Some("my_value".to_string())); let extra_value = extra_tree.get("extra_key")?; assert_eq!(extra_value, Some("extra_value".as_bytes().into())); Ok(()) } + + #[test] + fn test_enhanced_range_transaction_methods() -> Result<()> { + let (range_store, namespace_store, _) = create_test_stores()?; + + // Setup + range_store.define("test_range", 50)?; + namespace_store.define("test_namespace")?; + + let transaction = Transaction::new() + .with_range_store(&range_store) + .with_namespace_store(&namespace_store); + + // Test enhanced range methods in transaction + let (_bit1, bit2, _success) = transaction.execute(|ctx| { + let range_ctx = ctx.use_range(); + + // Test range existence + assert!(range_ctx.exists("test_range")?); + assert!(!range_ctx.exists("nonexistent")?); + + // Test range info + let info = range_ctx.info("test_range")?; + assert!(info.is_some()); + let (_, size) = info.unwrap(); + assert_eq!(size, 50); + + // Assign some values + let bit1 = range_ctx.assign("test_range", "first_value")?; + let bit2 = range_ctx.assign("test_range", "second_value")?; + + // Test get + let value1 = range_ctx.get("test_range", bit1)?; + assert_eq!(value1, Some("first_value".to_string())); + + let value2 = range_ctx.get("test_range", bit2)?; + assert_eq!(value2, Some("second_value".to_string())); + + // Test unassign_bit + let unassigned = range_ctx.unassign_bit("test_range", bit1)?; + assert!(unassigned); + + // Verify it's gone + let value_after = range_ctx.get("test_range", bit1)?; + assert_eq!(value_after, None); + + Ok((bit1, bit2, true)) + })?; + + // Verify transaction committed properly + let final_assignments = range_store.list_range("test_range")?; + assert_eq!(final_assignments.len(), 1); // Only second value should remain + assert_eq!(final_assignments[0], (bit2, "second_value".to_string())); + + Ok(()) + } + + #[test] + fn test_enhanced_namespace_transaction_methods() -> Result<()> { + let (range_store, namespace_store, _) = create_test_stores()?; + + // Setup + range_store.define("test_range", 50)?; + namespace_store.define("test_namespace")?; + + let transaction = Transaction::new() + .with_range_store(&range_store) + .with_namespace_store(&namespace_store); + + // Test enhanced namespace methods in transaction + transaction.execute(|ctx| { + let ns_ctx = ctx.use_namespace(); + + // Test namespace existence + assert!(ns_ctx.namespace_exists("test_namespace")?); + assert!(!ns_ctx.namespace_exists("nonexistent")?); + + // Reserve some keys + assert!(ns_ctx.reserve("test_namespace", "key1", "value1")?); + assert!(ns_ctx.reserve("test_namespace", "key2", "value2")?); + assert!(ns_ctx.reserve("test_namespace", "key3", "value3")?); + + // Test key existence + assert!(ns_ctx.key_exists("test_namespace", "key1")?); + assert!(!ns_ctx.key_exists("test_namespace", "nonexistent")?); + + // Test update + assert!(ns_ctx.update("test_namespace", "key1", "updated_value")?); + assert!(!ns_ctx.update("test_namespace", "nonexistent", "value")?); + + // Test get after update + let value = ns_ctx.get("test_namespace", "key1")?; + assert_eq!(value, Some("updated_value".to_string())); + + // Test key existence for all keys + assert!(ns_ctx.key_exists("test_namespace", "key1")?); + assert!(ns_ctx.key_exists("test_namespace", "key2")?); + assert!(ns_ctx.key_exists("test_namespace", "key3")?); + + // Verify individual values + let value2 = ns_ctx.get("test_namespace", "key2")?; + assert_eq!(value2, Some("value2".to_string())); + let value3 = ns_ctx.get("test_namespace", "key3")?; + assert_eq!(value3, Some("value3".to_string())); + + // Test remove + assert!(ns_ctx.remove("test_namespace", "key2")?); + assert!(!ns_ctx.remove("test_namespace", "key2")?); // Already removed + + // Verify key is gone + assert!(!ns_ctx.key_exists("test_namespace", "key2")?); + + Ok(()) + })?; + + // Verify transaction committed properly + assert!(namespace_store.key_exists("test_namespace", "key1")?); + assert!(namespace_store.key_exists("test_namespace", "key3")?); + assert!(!namespace_store.key_exists("test_namespace", "key2")?); + + let value1 = namespace_store.get("test_namespace", "key1")?; + assert_eq!(value1, Some("updated_value".to_string())); + let value3 = namespace_store.get("test_namespace", "key3")?; + assert_eq!(value3, Some("value3".to_string())); + + Ok(()) + } + + #[test] + fn test_cross_store_operations() -> Result<()> { + let (range_store, namespace_store, db) = create_test_stores()?; + + // Setup + range_store.define("ip_pool", 100)?; + range_store.define("port_pool", 1000)?; + namespace_store.define("users")?; + namespace_store.define("sessions")?; + + let metadata_tree = db.open_tree("metadata")?; + + let transaction = Transaction::new() + .with_range_store(&range_store) + .with_namespace_store(&namespace_store) + .with_tree(&metadata_tree); + + // Complex cross-store operation + let (user_created, ip_assigned, port_assigned, session_id) = transaction.execute(|ctx| { + let ns_ctx = ctx.use_namespace(); + let range_ctx = ctx.use_range(); + + // Create a user + let user_created = ns_ctx.reserve("users", "alice", "Alice Smith")?; + + // Assign IP and port + let ip_bit = range_ctx.assign("ip_pool", "192.168.1.100")?; + let port_bit = range_ctx.assign("port_pool", "8080")?; + + // Create session with cross-references + let session_data = format!("user:alice,ip_bit:{},port_bit:{}", ip_bit, port_bit); + let session_created = ns_ctx.reserve("sessions", "sess_001", &session_data)?; + + // Store metadata + if let Some(tree) = ctx.tree(0) { + tree.insert("last_user", "alice") + .map_err(|e| Error::StoreError(sled::Error::Unsupported(format!("Failed: {}", e))))?; + tree.insert("assignments_count", &2u64.to_be_bytes()) + .map_err(|e| Error::StoreError(sled::Error::Unsupported(format!("Failed: {}", e))))?; + } + + Ok((user_created, ip_bit, port_bit, session_created)) + })?; + + // Verify all operations succeeded + assert!(user_created); + assert!(ip_assigned < 100); + assert!(port_assigned < 1000); + assert!(session_id); + + // Verify data consistency + let user_data = namespace_store.get("users", "alice")?; + assert_eq!(user_data, Some("Alice Smith".to_string())); + + let session_data = namespace_store.get("sessions", "sess_001")?; + assert!(session_data.is_some()); + assert!(session_data.unwrap().contains("user:alice")); + + let ip_value = range_store.get("ip_pool", ip_assigned)?; + assert_eq!(ip_value, Some("192.168.1.100".to_string())); + + let port_value = range_store.get("port_pool", port_assigned)?; + assert_eq!(port_value, Some("8080".to_string())); + + let last_user = metadata_tree.get("last_user")?; + assert_eq!(last_user, Some("alice".as_bytes().into())); + + Ok(()) + } + + #[test] + fn test_transaction_composition_flexibility() -> Result<()> { + let (range_store, namespace_store, db) = create_test_stores()?; + + // Setup + range_store.define("test_range", 50)?; + namespace_store.define("test_namespace")?; + + let tree1 = db.open_tree("tree1")?; + let tree2 = db.open_tree("tree2")?; + + // Test transaction with only range store + let transaction_range_only = Transaction::new() + .with_range_store(&range_store); + + let bit1 = transaction_range_only.execute(|ctx| { + ctx.use_range().assign("test_range", "range_only_value") + })?; + + // Test transaction with only namespace store + let transaction_ns_only = Transaction::new() + .with_namespace_store(&namespace_store); + + let reserved = transaction_ns_only.execute(|ctx| { + ctx.use_namespace().reserve("test_namespace", "ns_only_key", "ns_only_value") + })?; + + // Test transaction with only trees + let transaction_trees_only = Transaction::new() + .with_tree(&tree1) + .with_tree(&tree2); + + transaction_trees_only.execute(|ctx| { + if let Some(t1) = ctx.tree(0) { + t1.insert("tree1_key", "tree1_value") + .map_err(|e| Error::StoreError(sled::Error::Unsupported(format!("Failed: {}", e))))?; + } + if let Some(t2) = ctx.tree(1) { + t2.insert("tree2_key", "tree2_value") + .map_err(|e| Error::StoreError(sled::Error::Unsupported(format!("Failed: {}", e))))?; + } + Ok(()) + })?; + + // Test mixed composition + let transaction_mixed = Transaction::new() + .with_namespace_store(&namespace_store) + .with_tree(&tree1); + + transaction_mixed.execute(|ctx| { + ctx.use_namespace().reserve("test_namespace", "mixed_key", "mixed_value")?; + if let Some(tree) = ctx.tree(0) { + tree.insert("mixed_tree_key", "mixed_tree_value") + .map_err(|e| Error::StoreError(sled::Error::Unsupported(format!("Failed: {}", e))))?; + } + Ok(()) + })?; + + // Verify all operations + assert!(bit1 < 50); + assert!(reserved); + + let range_value = range_store.get("test_range", bit1)?; + assert_eq!(range_value, Some("range_only_value".to_string())); + + let ns_value = namespace_store.get("test_namespace", "ns_only_key")?; + assert_eq!(ns_value, Some("ns_only_value".to_string())); + + let mixed_ns_value = namespace_store.get("test_namespace", "mixed_key")?; + assert_eq!(mixed_ns_value, Some("mixed_value".to_string())); + + let tree1_value = tree1.get("tree1_key")?; + assert_eq!(tree1_value, Some("tree1_value".as_bytes().into())); + + let tree2_value = tree2.get("tree2_key")?; + assert_eq!(tree2_value, Some("tree2_value".as_bytes().into())); + + let mixed_tree_value = tree1.get("mixed_tree_key")?; + assert_eq!(mixed_tree_value, Some("mixed_tree_value".as_bytes().into())); + + Ok(()) + } + + #[test] + fn test_error_propagation_and_rollback() -> Result<()> { + let (range_store, namespace_store, _) = create_test_stores()?; + + // Setup + range_store.define("small_range", 2)?; // Very small range + namespace_store.define("test_namespace")?; + + // Fill up the range + range_store.assign("small_range", "value1")?; + range_store.assign("small_range", "value2")?; + + // Reserve a namespace key + namespace_store.reserve("test_namespace", "existing", "existing_value")?; + + let transaction = Transaction::new() + .with_range_store(&range_store) + .with_namespace_store(&namespace_store); + + // Test rollback on range full error + let result = transaction.execute(|ctx| { + // This should succeed + ctx.use_namespace().reserve("test_namespace", "temp_key", "temp_value")?; + + // This should fail (range is full) + ctx.use_range().assign("small_range", "overflow_value")?; + + Ok(()) + }); + + assert!(result.is_err()); + + // Verify rollback - temp_key should not exist + assert!(!namespace_store.key_exists("test_namespace", "temp_key")?); + + // Test rollback on namespace conflict + let result2 = transaction.execute(|ctx| { + // Free up a bit first + ctx.use_range().unassign_bit("small_range", 0)?; + + // This should succeed + ctx.use_range().assign("small_range", "new_value")?; + + // This should fail (key already exists) + ctx.use_namespace().reserve("test_namespace", "existing", "different_value")?; + + Ok(()) + }); + + assert!(result2.is_err()); + + // Verify rollback - range should still have original values + let range_assignments = range_store.list_range("small_range")?; + assert_eq!(range_assignments.len(), 2); + assert!(range_assignments.iter().any(|(_, v)| v == "value1")); + assert!(range_assignments.iter().any(|(_, v)| v == "value2")); + assert!(!range_assignments.iter().any(|(_, v)| v == "new_value")); + + Ok(()) + } } \ No newline at end of file diff --git a/crates/store/src/namespace.rs b/crates/store/src/namespace.rs index ac29ec0..149aafe 100644 --- a/crates/store/src/namespace.rs +++ b/crates/store/src/namespace.rs @@ -1,358 +1,645 @@ // This module implements "namespaces" which are buckets of unique values // that maps to keys elsewhere in storage. // // the `names` tree is a k/v of `name` -> `id` where `id` is a u64 from `generate_id`. // the `spaces` tree is a k/v of `id` -> `name` where `id` is a u64 the `names` tree id and `name` is a str. use crate::Result; use crate::Error; use sled::Transactional; #[derive(Debug, Clone)] pub struct NamespaceStore { pub(crate) names: sled::Tree, pub(crate) spaces: sled::Tree, } impl NamespaceStore { pub fn open(db: &sled::Db) -> Result { Ok(NamespaceStore { names: db.open_tree("namespaces/1/names")?, spaces: db.open_tree("namespaces/1/spaces")?, }) } /// Reserve a key in a namespace within an existing transaction context - pub(crate) fn reserve_in_transaction( + pub fn reserve_in_transaction( &self, names: &sled::transaction::TransactionalTree, spaces: &sled::transaction::TransactionalTree, namespace: &str, key: &str, value: &str, ) -> sled::transaction::ConflictableTransactionResult { // Get ID of the namespace from `names` let namespace_id = match names.get(namespace)? { Some(id_bytes) => { let id_array: [u8; 8] = id_bytes.as_ref().try_into() .map_err(|_| sled::transaction::ConflictableTransactionError::Abort(()))?; u64::from_be_bytes(id_array) } None => return Err(sled::transaction::ConflictableTransactionError::Abort(())), }; // Create composite key: namespace_id + key let mut composite_key = Vec::new(); composite_key.extend_from_slice(&namespace_id.to_be_bytes()); composite_key.extend_from_slice(key.as_bytes()); // Check if key already exists if spaces.get(&composite_key)?.is_some() { return Err(sled::transaction::ConflictableTransactionError::Abort(())); } // Insert the key-value pair spaces.insert(composite_key, value.as_bytes())?; Ok(true) } /// Get a value from a namespace within an existing transaction context - pub(crate) fn get_in_transaction( + pub fn get_in_transaction( &self, names: &sled::transaction::TransactionalTree, spaces: &sled::transaction::TransactionalTree, namespace: &str, key: &str, ) -> sled::transaction::ConflictableTransactionResult, ()> { // Get ID of the namespace from `names` let namespace_id = match names.get(namespace)? { Some(id_bytes) => { let id_array: [u8; 8] = id_bytes.as_ref().try_into() .map_err(|_| sled::transaction::ConflictableTransactionError::Abort(()))?; u64::from_be_bytes(id_array) } None => return Ok(None), // namespace doesn't exist }; // Create composite key: namespace_id + key let mut composite_key = Vec::new(); composite_key.extend_from_slice(&namespace_id.to_be_bytes()); composite_key.extend_from_slice(key.as_bytes()); // Get value from spaces match spaces.get(&composite_key)? { Some(value_bytes) => { let value = String::from_utf8(value_bytes.to_vec()) .map_err(|_| sled::transaction::ConflictableTransactionError::Abort(()))?; Ok(Some(value)) } None => Ok(None), } } /// Remove a key from a namespace within an existing transaction context - pub(crate) fn remove_in_transaction( + pub fn remove_in_transaction( &self, names: &sled::transaction::TransactionalTree, spaces: &sled::transaction::TransactionalTree, namespace: &str, key: &str, ) -> sled::transaction::ConflictableTransactionResult { // Get ID of the namespace from `names` let namespace_id = match names.get(namespace)? { Some(id_bytes) => { let id_array: [u8; 8] = id_bytes.as_ref().try_into() .map_err(|_| sled::transaction::ConflictableTransactionError::Abort(()))?; u64::from_be_bytes(id_array) } None => return Ok(false), // namespace doesn't exist }; // Create composite key: namespace_id + key let mut composite_key = Vec::new(); composite_key.extend_from_slice(&namespace_id.to_be_bytes()); composite_key.extend_from_slice(key.as_bytes()); // Remove the key-value pair Ok(spaces.remove(composite_key)?.is_some()) } + /// Update a key in a namespace within an existing transaction context + /// Returns true if the key existed and was updated, false if it didn't exist + pub fn update_in_transaction( + &self, + names: &sled::transaction::TransactionalTree, + spaces: &sled::transaction::TransactionalTree, + namespace: &str, + key: &str, + value: &str, + ) -> sled::transaction::ConflictableTransactionResult { + // Get ID of the namespace from `names` + let namespace_id = match names.get(namespace)? { + Some(id_bytes) => { + let id_array: [u8; 8] = id_bytes.as_ref().try_into() + .map_err(|_| sled::transaction::ConflictableTransactionError::Abort(()))?; + u64::from_be_bytes(id_array) + } + None => return Ok(false), // namespace doesn't exist + }; + + // Create composite key: namespace_id + key + let mut composite_key = Vec::new(); + composite_key.extend_from_slice(&namespace_id.to_be_bytes()); + composite_key.extend_from_slice(key.as_bytes()); + + // Check if key exists and update it + match spaces.get(&composite_key)? { + Some(_) => { + spaces.insert(composite_key, value.as_bytes())?; + Ok(true) + } + None => Ok(false), + } + } + + + + /// Check if a namespace exists within an existing transaction context + pub fn namespace_exists_in_transaction( + &self, + names: &sled::transaction::TransactionalTree, + namespace: &str, + ) -> sled::transaction::ConflictableTransactionResult { + Ok(names.get(namespace)?.is_some()) + } + + /// Check if a key exists in a namespace within an existing transaction context + pub fn key_exists_in_transaction( + &self, + names: &sled::transaction::TransactionalTree, + spaces: &sled::transaction::TransactionalTree, + namespace: &str, + key: &str, + ) -> sled::transaction::ConflictableTransactionResult { + // Get ID of the namespace from `names` + let namespace_id = match names.get(namespace)? { + Some(id_bytes) => { + let id_array: [u8; 8] = id_bytes.as_ref().try_into() + .map_err(|_| sled::transaction::ConflictableTransactionError::Abort(()))?; + u64::from_be_bytes(id_array) + } + None => return Ok(false), // namespace doesn't exist + }; + + // Create composite key: namespace_id + key + let mut composite_key = Vec::new(); + composite_key.extend_from_slice(&namespace_id.to_be_bytes()); + composite_key.extend_from_slice(key.as_bytes()); + + Ok(spaces.get(&composite_key)?.is_some()) + } + // define a namespace. // inserts a key `namespace` into `names`, where the value is a random u64. pub fn define(&self, namespace: &str) -> Result<()> { self.names.transaction(|db| { match db.get(namespace)? { Some(_) => Ok(()), None => { let id = db.generate_id()?; db.insert(namespace, &id.to_be_bytes())?; Ok(()) } } })?; Ok(()) } pub fn resolve(&self, namespace: &str, key: &str) -> Result> { let result = (&self.names, &self.spaces).transaction(|(names, spaces)| { - // Get ID of the namespace from `names` - let namespace_id = match names.get(namespace)? { - Some(id_bytes) => { - let id_array: [u8; 8] = id_bytes.as_ref().try_into() - .map_err(|_| sled::transaction::ConflictableTransactionError::Abort(()))?; - u64::from_be_bytes(id_array) - } - None => return Err(sled::transaction::ConflictableTransactionError::Abort(())), - }; - - // Create composite key: namespace_id + key - let mut composite_key = Vec::new(); - composite_key.extend_from_slice(&namespace_id.to_be_bytes()); - composite_key.extend_from_slice(key.as_bytes()); - - // Check if key exists in spaces - match spaces.get(&composite_key)? { - Some(value_bytes) => { - let value = String::from_utf8(value_bytes.to_vec()) - .map_err(|_| sled::transaction::ConflictableTransactionError::Abort(()))?; - Ok(Some(value)) - } - None => Ok(None), - } + self.get_in_transaction(names, spaces, namespace, key) }).map_err(|e| match e { sled::transaction::TransactionError::Abort(()) => Error::UndefinedNamespace(namespace.to_string()), sled::transaction::TransactionError::Storage(storage_err) => Error::StoreError(storage_err), })?; Ok(result) } pub fn reserve(&self, namespace: &str, key: &str, value: &str) -> Result { let result = (&self.names, &self.spaces).transaction(|(names, spaces)| { - // Get ID of the namespace from `names` - let namespace_id = match names.get(namespace)? { - Some(id_bytes) => { - let id_array: [u8; 8] = id_bytes.as_ref().try_into() - .map_err(|_| sled::transaction::ConflictableTransactionError::Abort(()))?; - u64::from_be_bytes(id_array) - } - None => return Err(sled::transaction::ConflictableTransactionError::Abort(())), - }; - - // Create composite key: namespace_id + key - let mut composite_key = Vec::new(); - composite_key.extend_from_slice(&namespace_id.to_be_bytes()); - composite_key.extend_from_slice(key.as_bytes()); - - // Check if key already exists - if spaces.get(&composite_key)?.is_some() { - return Err(sled::transaction::ConflictableTransactionError::Abort(())); - } - - // Insert the key-value pair - spaces.insert(composite_key, value.as_bytes())?; - Ok(true) + self.reserve_in_transaction(names, spaces, namespace, key, value) }).map_err(|e| match e { sled::transaction::TransactionError::Abort(()) => Error::NamespaceKeyReserved(namespace.to_string(), key.to_string()), sled::transaction::TransactionError::Storage(storage_err) => Error::StoreError(storage_err), })?; Ok(result) } pub fn get(&self, namespace: &str, key: &str) -> Result> { let result = (&self.names, &self.spaces).transaction(|(names, spaces)| { - // Get ID of the namespace from `names` - let namespace_id = match names.get(namespace)? { - Some(id_bytes) => { - let id_array: [u8; 8] = id_bytes.as_ref().try_into() - .map_err(|_| sled::transaction::ConflictableTransactionError::Abort(()))?; - u64::from_be_bytes(id_array) - } - None => return Ok(None), // namespace doesn't exist - }; - - // Create composite key: namespace_id + key - let mut composite_key = Vec::new(); - composite_key.extend_from_slice(&namespace_id.to_be_bytes()); - composite_key.extend_from_slice(key.as_bytes()); - - // Get value from spaces - match spaces.get(&composite_key)? { - Some(value_bytes) => { - let value = String::from_utf8(value_bytes.to_vec()) - .map_err(|_| sled::transaction::ConflictableTransactionError::Abort(()))?; - Ok(Some(value)) - } - None => Ok(None), - } + self.get_in_transaction(names, spaces, namespace, key) }).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) } pub fn remove(&self, namespace: &str, key: &str) -> Result { let result = (&self.names, &self.spaces).transaction(|(names, spaces)| { - // Get ID of the namespace from `names` - let namespace_id = match names.get(namespace)? { - Some(id_bytes) => { - let id_array: [u8; 8] = id_bytes.as_ref().try_into() - .map_err(|_| sled::transaction::ConflictableTransactionError::Abort(()))?; - u64::from_be_bytes(id_array) - } - None => return Ok(false), // namespace doesn't exist - }; + self.remove_in_transaction(names, spaces, namespace, key) + }).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) + } + + /// Update a key in a namespace + pub fn update(&self, namespace: &str, key: &str, value: &str) -> Result { + let result = (&self.names, &self.spaces).transaction(|(names, spaces)| { + self.update_in_transaction(names, spaces, namespace, key, value) + }).map_err(|e| match e { + sled::transaction::TransactionError::Abort(()) => Error::UndefinedNamespace(namespace.to_string()), + sled::transaction::TransactionError::Storage(storage_err) => Error::StoreError(storage_err), + })?; + Ok(result) + } - // Create composite key: namespace_id + key - let mut composite_key = Vec::new(); - composite_key.extend_from_slice(&namespace_id.to_be_bytes()); - composite_key.extend_from_slice(key.as_bytes()); - // Remove the key-value pair - Ok(spaces.remove(composite_key)?.is_some()) + + /// Check if a namespace exists + pub fn namespace_exists(&self, namespace: &str) -> Result { + Ok(self.names.get(namespace)?.is_some()) + } + + /// Check if a key exists in a namespace + pub fn key_exists(&self, namespace: &str, key: &str) -> Result { + let result = (&self.names, &self.spaces).transaction(|(names, spaces)| { + self.key_exists_in_transaction(names, spaces, namespace, key) }).map_err(|e| match e { - sled::transaction::TransactionError::Abort(()) => Error::StoreError(sled::Error::Unsupported("Transaction aborted".to_string())), + sled::transaction::TransactionError::Abort(()) => Error::UndefinedNamespace(namespace.to_string()), sled::transaction::TransactionError::Storage(storage_err) => Error::StoreError(storage_err), })?; Ok(result) } + + /// List all defined namespaces + pub fn list_namespaces(&self) -> Result> { + let mut namespaces = Vec::new(); + + for result in self.names.iter() { + let (key, _) = result?; + let name = String::from_utf8(key.to_vec())?; + namespaces.push(name); + } + + namespaces.sort(); + Ok(namespaces) + } + + /// List all keys in a namespace + pub fn list_keys(&self, namespace: &str) -> Result> { + // Get namespace ID first + let namespace_id = match self.names.get(namespace)? { + Some(id_bytes) => { + let id_array: [u8; 8] = id_bytes.as_ref().try_into() + .map_err(|_| Error::StoreError(sled::Error::Unsupported("Invalid namespace ID".to_string())))?; + u64::from_be_bytes(id_array) + } + None => return Ok(Vec::new()), + }; + + let mut keys = Vec::new(); + let prefix = namespace_id.to_be_bytes(); + + for result in self.spaces.scan_prefix(&prefix) { + let (key, _) = result?; + if key.len() > 8 { + let key_part = &key[8..]; + let key_str = String::from_utf8(key_part.to_vec())?; + keys.push(key_str); + } + } + + keys.sort(); + Ok(keys) + } + + /// List all key-value pairs in a namespace + pub fn list_all(&self, namespace: &str) -> Result> { + // Get namespace ID first + let namespace_id = match self.names.get(namespace)? { + Some(id_bytes) => { + let id_array: [u8; 8] = id_bytes.as_ref().try_into() + .map_err(|_| Error::StoreError(sled::Error::Unsupported("Invalid namespace ID".to_string())))?; + u64::from_be_bytes(id_array) + } + None => return Ok(Vec::new()), + }; + + let mut pairs = Vec::new(); + let prefix = namespace_id.to_be_bytes(); + + for result in self.spaces.scan_prefix(&prefix) { + let (key, value) = result?; + if key.len() > 8 { + let key_part = &key[8..]; + let key_str = String::from_utf8(key_part.to_vec())?; + let value_str = String::from_utf8(value.to_vec())?; + pairs.push((key_str, value_str)); + } + } + + pairs.sort_by(|a, b| a.0.cmp(&b.0)); + Ok(pairs) + } } #[cfg(test)] mod tests { use super::*; use tempfile::tempdir; fn create_test_store() -> Result { let temp_dir = tempdir().unwrap(); let db = sled::open(temp_dir.path())?; NamespaceStore::open(&db) } #[test] fn test_define_namespace() -> Result<()> { let store = create_test_store()?; // Define a namespace store.define("test_namespace")?; // Defining the same namespace again should not error store.define("test_namespace")?; + // Check that namespace exists + assert!(store.namespace_exists("test_namespace")?); + assert!(!store.namespace_exists("nonexistent")?); + Ok(()) } #[test] fn test_reserve_and_get() -> Result<()> { let store = create_test_store()?; // Define namespace first store.define("test_namespace")?; // Reserve a key let success = store.reserve("test_namespace", "test_key", "test_value")?; assert!(success); // Get the value let value = store.get("test_namespace", "test_key")?; assert_eq!(value, Some("test_value".to_string())); + // Check key exists + assert!(store.key_exists("test_namespace", "test_key")?); + assert!(!store.key_exists("test_namespace", "nonexistent")?); + Ok(()) } #[test] fn test_reserve_duplicate_key() -> Result<()> { let store = create_test_store()?; // Define namespace first store.define("test_namespace")?; // Reserve a key store.reserve("test_namespace", "test_key", "test_value")?; // Try to reserve the same key again - should fail let result = store.reserve("test_namespace", "test_key", "other_value"); assert!(result.is_err()); Ok(()) } + #[test] + fn test_update() -> Result<()> { + let store = create_test_store()?; + + // Define namespace first + store.define("test_namespace")?; + + // Reserve a key + store.reserve("test_namespace", "test_key", "original_value")?; + + // Update the key + let updated = store.update("test_namespace", "test_key", "new_value")?; + assert!(updated); + + // Verify the update + let value = store.get("test_namespace", "test_key")?; + assert_eq!(value, Some("new_value".to_string())); + + // Try to update a non-existent key + let not_updated = store.update("test_namespace", "nonexistent", "value")?; + assert!(!not_updated); + + Ok(()) + } + #[test] fn test_remove() -> Result<()> { let store = create_test_store()?; // Define namespace first store.define("test_namespace")?; // Reserve a key store.reserve("test_namespace", "test_key", "test_value")?; // Remove the key let removed = store.remove("test_namespace", "test_key")?; assert!(removed); // Try to get the removed key let value = store.get("test_namespace", "test_key")?; assert_eq!(value, None); + // Check key no longer exists + assert!(!store.key_exists("test_namespace", "test_key")?); + // Try to remove a non-existent key let removed_again = store.remove("test_namespace", "test_key")?; assert!(!removed_again); Ok(()) } + #[test] + fn test_list_operations() -> Result<()> { + let store = create_test_store()?; + + // Define namespaces + store.define("namespace1")?; + store.define("namespace2")?; + + // Add some data + store.reserve("namespace1", "key1", "value1")?; + store.reserve("namespace1", "key2", "value2")?; + store.reserve("namespace1", "key3", "value3")?; + store.reserve("namespace2", "key1", "different_value")?; + + // Test list_namespaces + let namespaces = store.list_namespaces()?; + assert!(namespaces.contains(&"namespace1".to_string())); + assert!(namespaces.contains(&"namespace2".to_string())); + assert_eq!(namespaces.len(), 2); + + // Test list_keys + let keys = store.list_keys("namespace1")?; + assert_eq!(keys, vec!["key1", "key2", "key3"]); + + let keys2 = store.list_keys("namespace2")?; + assert_eq!(keys2, vec!["key1"]); + + // Test list_all + let pairs = store.list_all("namespace1")?; + assert_eq!(pairs, vec![ + ("key1".to_string(), "value1".to_string()), + ("key2".to_string(), "value2".to_string()), + ("key3".to_string(), "value3".to_string()), + ]); + + Ok(()) + } + #[test] fn test_undefined_namespace() -> Result<()> { let store = create_test_store()?; // Try to get from undefined namespace let value = store.get("undefined_namespace", "test_key")?; assert_eq!(value, None); // Try to remove from undefined namespace let removed = store.remove("undefined_namespace", "test_key")?; assert!(!removed); + // Try to update in undefined namespace + let updated = store.update("undefined_namespace", "test_key", "value")?; + assert!(!updated); + + // Check key exists returns false for undefined namespace + assert!(!store.key_exists("undefined_namespace", "test_key")?); + + Ok(()) + } + + #[test] + fn test_transaction_methods() -> Result<()> { + let store = create_test_store()?; + + // Define namespace + store.define("test_namespace")?; + + // Test transaction methods + let result = (&store.names, &store.spaces).transaction(|(names, spaces)| { + // Reserve in transaction + let reserved = store.reserve_in_transaction(names, spaces, "test_namespace", "key1", "value1")?; + assert_eq!(reserved, true); + + // Get in transaction + let value = store.get_in_transaction(names, spaces, "test_namespace", "key1")?; + assert_eq!(value, Some("value1".to_string())); + + // Update in transaction + let updated = store.update_in_transaction(names, spaces, "test_namespace", "key1", "new_value")?; + assert_eq!(updated, true); + + // Check key exists in transaction + let exists = store.key_exists_in_transaction(names, spaces, "test_namespace", "key1")?; + assert_eq!(exists, true); + + // Check key exists after update + let exists_after_update = store.key_exists_in_transaction(names, spaces, "test_namespace", "key1")?; + assert_eq!(exists_after_update, true); + + // Check namespace exists in transaction + let ns_exists = store.namespace_exists_in_transaction(names, "test_namespace")?; + assert_eq!(ns_exists, true); + + // Remove in transaction + let removed = store.remove_in_transaction(names, spaces, "test_namespace", "key1")?; + assert_eq!(removed, true); + + Ok(()) + }).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) + } + + #[test] + fn test_transaction_rollback() -> Result<()> { + let store = create_test_store()?; + + // Define namespace and add initial data + store.define("test_namespace")?; + store.reserve("test_namespace", "existing_key", "existing_value")?; + + // Attempt a transaction that should fail + let result = (&store.names, &store.spaces).transaction(|(names, spaces)| { + // This should succeed + let _reserved = store.reserve_in_transaction(names, spaces, "test_namespace", "new_key", "new_value")?; + + // This should fail (key already exists) + let _duplicate = store.reserve_in_transaction(names, spaces, "test_namespace", "existing_key", "different_value")?; + + Ok(()) + }).map_err(|e| match e { + sled::transaction::TransactionError::Abort(()) => Error::NamespaceKeyReserved("test_namespace".to_string(), "existing_key".to_string()), + sled::transaction::TransactionError::Storage(storage_err) => Error::StoreError(storage_err), + }); + + // Transaction should have failed + assert!(result.is_err()); + + // Verify rollback - new_key should not exist + assert!(!store.key_exists("test_namespace", "new_key")?); + + // Existing key should still have original value + let value = store.get("test_namespace", "existing_key")?; + assert_eq!(value, Some("existing_value".to_string())); + + Ok(()) + } + + #[test] + fn test_multiple_namespaces() -> Result<()> { + let store = create_test_store()?; + + // Define multiple namespaces + store.define("users")?; + store.define("products")?; + store.define("orders")?; + + // Add data to each namespace + store.reserve("users", "alice", "Alice Smith")?; + store.reserve("users", "bob", "Bob Jones")?; + store.reserve("products", "laptop", "MacBook Pro")?; + store.reserve("products", "mouse", "Magic Mouse")?; + store.reserve("orders", "001", "Alice's laptop order")?; + + // Test isolation between namespaces + assert_eq!(store.get("users", "alice")?, Some("Alice Smith".to_string())); + assert_eq!(store.get("products", "laptop")?, Some("MacBook Pro".to_string())); + assert_eq!(store.get("orders", "001")?, Some("Alice's laptop order".to_string())); + + // Keys with same name in different namespaces should be separate + store.reserve("users", "test", "user_test")?; + store.reserve("products", "test", "product_test")?; + + assert_eq!(store.get("users", "test")?, Some("user_test".to_string())); + assert_eq!(store.get("products", "test")?, Some("product_test".to_string())); + + // Verify namespace listing + let namespaces = store.list_namespaces()?; + assert!(namespaces.contains(&"users".to_string())); + assert!(namespaces.contains(&"products".to_string())); + assert!(namespaces.contains(&"orders".to_string())); + Ok(()) } } diff --git a/crates/store/src/range.rs b/crates/store/src/range.rs index b79d8f0..6e8f114 100644 --- a/crates/store/src/range.rs +++ b/crates/store/src/range.rs @@ -1,614 +1,915 @@ //! # Range Store Module //! //! This module implements "ranges" which are buckets of unique range assignments //! that map values to specific bit positions within fixed-size ranges. This is useful //! for allocating unique identifiers, session tokens, or any scenario where you need //! to assign sequential positions to values. //! //! ## Storage Design //! //! The RangeStore uses three sled trees: //! //! - **`names` tree**: Maps range names to their metadata //! - Key: `name` (string) //! - Value: `id || size` where `id` is a u64 from `generate_id()` and `size` is a u64 //! //! - **`map` tree**: Stores bitmaps for each range to track allocated positions //! - Key: `id` (u64 from the `names` tree) //! - Value: `bitmap` (binary data representing which bits are allocated) //! //! - **`assign` tree**: Maps range positions to their assigned values //! - Key: `id || bit_position` where `id` is the range ID and `bit_position` is a u64 //! - Value: `value` (the string value assigned to this position) //! //! ## API Operations //! //! - `define(name, size)`: Create a new range with the given name and maximum size //! - `assign(name, value)`: Assign a value to the next available position in the range //! - `get(name, position)`: Retrieve the value at a specific position //! - `list_range(name)`: List all assignments in a range //! - `list_ranges()`: List all defined ranges //! - `unassign(name, value)`: Remove a value from the range and free its position //! //! ## Example Usage //! //! ```rust //! use store::{open, Result}; //! //! fn example() -> Result<()> { //! let store = open()?; //! let ranges = store.ranges(); //! //! // Define a range for user IDs //! ranges.define("user_ids", 1000)?; //! //! // Assign some users //! let alice_id = ranges.assign("user_ids", "alice@example.com")?; // Returns 0 //! let bob_id = ranges.assign("user_ids", "bob@example.com")?; // Returns 1 //! //! // Get a user by position //! let email = ranges.get("user_ids", alice_id)?; // Some("alice@example.com") //! //! // Remove a user (frees the position for reuse) //! ranges.unassign("user_ids", "bob@example.com")?; //! //! // Next assignment will reuse Bob's position //! let charlie_id = ranges.assign("user_ids", "charlie@example.com")?; // Returns 1 //! //! Ok(()) //! } //! ``` use crate::Result; use crate::Error; use sled::Transactional; #[derive(Debug, Clone)] pub struct RangeStore { pub(crate) names: sled::Tree, pub(crate) map: sled::Tree, pub(crate) assign: sled::Tree, } impl RangeStore { pub fn open(db: &sled::Db) -> Result { Ok(RangeStore { names: db.open_tree("ranges/1/names")?, map: db.open_tree("ranges/1/map")?, assign: db.open_tree("ranges/1/assign")?, }) } /// Assign a value to a range within an existing transaction context - pub(crate) fn assign_in_transaction( + pub fn assign_in_transaction( &self, names: &sled::transaction::TransactionalTree, map: &sled::transaction::TransactionalTree, assign: &sled::transaction::TransactionalTree, range_name: &str, value: &str, ) -> sled::transaction::ConflictableTransactionResult { // Get range info let (range_id, range_size) = match names.get(range_name)? { Some(data) => { if data.len() != 16 { return Err(sled::transaction::ConflictableTransactionError::Abort(())); } let id_bytes: [u8; 8] = data[0..8].try_into() .map_err(|_| sled::transaction::ConflictableTransactionError::Abort(()))?; let size_bytes: [u8; 8] = data[8..16].try_into() .map_err(|_| sled::transaction::ConflictableTransactionError::Abort(()))?; (u64::from_be_bytes(id_bytes), u64::from_be_bytes(size_bytes)) } None => return Err(sled::transaction::ConflictableTransactionError::Abort(())), }; // Get current bitmap let mut bitmap = match map.get(&range_id.to_be_bytes())? { Some(data) => data.to_vec(), None => return Err(sled::transaction::ConflictableTransactionError::Abort(())), }; // Find first free bit let mut bit_position = None; for bit in 0..range_size { let byte_index = (bit / 8) as usize; let bit_index = bit % 8; if byte_index < bitmap.len() { let mask = 1u8 << bit_index; if (bitmap[byte_index] & mask) == 0 { // This bit is free bitmap[byte_index] |= mask; bit_position = Some(bit); break; } } } let bit_pos = match bit_position { Some(pos) => pos, None => return Err(sled::transaction::ConflictableTransactionError::Abort(())), }; // Update bitmap map.insert(&range_id.to_be_bytes(), bitmap)?; // Store assignment let mut assign_key = Vec::new(); assign_key.extend_from_slice(&range_id.to_be_bytes()); assign_key.extend_from_slice(&bit_pos.to_be_bytes()); assign.insert(assign_key, value.as_bytes())?; Ok(bit_pos) } /// Get a value from a specific bit position within an existing transaction context - pub(crate) fn get_in_transaction( + /// Get range assignment details within an existing transaction context + pub fn get_in_transaction( &self, names: &sled::transaction::TransactionalTree, assign: &sled::transaction::TransactionalTree, range_name: &str, bit_position: u64, ) -> sled::transaction::ConflictableTransactionResult, ()> { // Get range info let (range_id, range_size) = match names.get(range_name)? { Some(data) => { if data.len() != 16 { return Ok(None); } let id_bytes: [u8; 8] = data[0..8].try_into() .map_err(|_| sled::transaction::ConflictableTransactionError::Abort(()))?; let size_bytes: [u8; 8] = data[8..16].try_into() .map_err(|_| sled::transaction::ConflictableTransactionError::Abort(()))?; (u64::from_be_bytes(id_bytes), u64::from_be_bytes(size_bytes)) } None => return Ok(None), }; if bit_position >= range_size { return Ok(None); } // Get assignment let mut assign_key = Vec::new(); assign_key.extend_from_slice(&range_id.to_be_bytes()); assign_key.extend_from_slice(&bit_position.to_be_bytes()); match assign.get(&assign_key)? { Some(value_bytes) => { let value = String::from_utf8(value_bytes.to_vec()) .map_err(|_| sled::transaction::ConflictableTransactionError::Abort(()))?; Ok(Some(value)) } None => Ok(None), } } + /// Unassign a specific bit position within an existing transaction context + pub fn unassign_bit_in_transaction( + &self, + names: &sled::transaction::TransactionalTree, + map: &sled::transaction::TransactionalTree, + assign: &sled::transaction::TransactionalTree, + range_name: &str, + bit_position: u64, + ) -> sled::transaction::ConflictableTransactionResult { + // Get range info + let (range_id, range_size) = match names.get(range_name)? { + Some(data) => { + if data.len() != 16 { + return Ok(false); + } + let id_bytes: [u8; 8] = data[0..8].try_into() + .map_err(|_| sled::transaction::ConflictableTransactionError::Abort(()))?; + let size_bytes: [u8; 8] = data[8..16].try_into() + .map_err(|_| sled::transaction::ConflictableTransactionError::Abort(()))?; + (u64::from_be_bytes(id_bytes), u64::from_be_bytes(size_bytes)) + } + None => return Ok(false), + }; + + if bit_position >= range_size { + return Ok(false); + } + + // Remove assignment + let mut assign_key = Vec::new(); + assign_key.extend_from_slice(&range_id.to_be_bytes()); + assign_key.extend_from_slice(&bit_position.to_be_bytes()); + let removed = assign.remove(assign_key)?.is_some(); + + if removed { + // Update bitmap + let mut bitmap = match map.get(&range_id.to_be_bytes())? { + Some(data) => data.to_vec(), + None => return Err(sled::transaction::ConflictableTransactionError::Abort(())), + }; + + let byte_index = (bit_position / 8) as usize; + let bit_index = bit_position % 8; + + if byte_index < bitmap.len() { + let mask = 1u8 << bit_index; + bitmap[byte_index] &= !mask; // Clear the bit + map.insert(&range_id.to_be_bytes(), bitmap)?; + } + } + + Ok(removed) + } + + + + /// Check if a range exists within an existing transaction context + pub fn exists_in_transaction( + &self, + names: &sled::transaction::TransactionalTree, + range_name: &str, + ) -> sled::transaction::ConflictableTransactionResult { + Ok(names.get(range_name)?.is_some()) + } + + /// Get range info (id and size) within an existing transaction context + pub fn info_in_transaction( + &self, + names: &sled::transaction::TransactionalTree, + range_name: &str, + ) -> sled::transaction::ConflictableTransactionResult, ()> { + match names.get(range_name)? { + Some(data) => { + if data.len() != 16 { + return Ok(None); + } + let id_bytes: [u8; 8] = data[0..8].try_into() + .map_err(|_| sled::transaction::ConflictableTransactionError::Abort(()))?; + let size_bytes: [u8; 8] = data[8..16].try_into() + .map_err(|_| sled::transaction::ConflictableTransactionError::Abort(()))?; + Ok(Some((u64::from_be_bytes(id_bytes), u64::from_be_bytes(size_bytes)))) + } + None => Ok(None), + } + } + // Define a new range with a given name and size pub fn define(&self, name: &str, size: u64) -> Result<()> { (&self.names, &self.map).transaction(|(names, map)| { match names.get(name)? { Some(_) => Ok(()), // Range already exists None => { let id = names.generate_id()?; // Store name -> (id, size) let mut value = Vec::new(); value.extend_from_slice(&id.to_be_bytes()); value.extend_from_slice(&size.to_be_bytes()); names.insert(name, value)?; // Initialize bitmap for the range (all zeros) let bitmap_size = (size + 7) / 8; // Round up to nearest byte let bitmap = vec![0u8; bitmap_size as usize]; map.insert(&id.to_be_bytes(), bitmap)?; Ok(()) } } })?; Ok(()) } // Assign the next free bit in the range to a value pub fn assign(&self, range_name: &str, value: &str) -> Result { let result = (&self.names, &self.map, &self.assign).transaction(|(names, map, assign)| { // Get range info let (range_id, range_size) = match names.get(range_name)? { Some(data) => { if data.len() != 16 { return Err(sled::transaction::ConflictableTransactionError::Abort(())); } let id_bytes: [u8; 8] = data[0..8].try_into() .map_err(|_| sled::transaction::ConflictableTransactionError::Abort(()))?; let size_bytes: [u8; 8] = data[8..16].try_into() .map_err(|_| sled::transaction::ConflictableTransactionError::Abort(()))?; (u64::from_be_bytes(id_bytes), u64::from_be_bytes(size_bytes)) } None => return Err(sled::transaction::ConflictableTransactionError::Abort(())), }; // Get current bitmap let mut bitmap = match map.get(&range_id.to_be_bytes())? { Some(data) => data.to_vec(), None => return Err(sled::transaction::ConflictableTransactionError::Abort(())), }; // Find first free bit let mut bit_position = None; for bit in 0..range_size { let byte_index = (bit / 8) as usize; let bit_index = bit % 8; if byte_index < bitmap.len() { let mask = 1u8 << bit_index; if (bitmap[byte_index] & mask) == 0 { // This bit is free bitmap[byte_index] |= mask; bit_position = Some(bit); break; } } } let bit_pos = match bit_position { Some(pos) => pos, None => return Err(sled::transaction::ConflictableTransactionError::Abort(())), }; // Update bitmap map.insert(&range_id.to_be_bytes(), bitmap)?; // Store assignment let mut assign_key = Vec::new(); assign_key.extend_from_slice(&range_id.to_be_bytes()); assign_key.extend_from_slice(&bit_pos.to_be_bytes()); assign.insert(assign_key, value.as_bytes())?; Ok(bit_pos) }).map_err(|e| match e { sled::transaction::TransactionError::Abort(()) => Error::RangeFull(range_name.to_string()), sled::transaction::TransactionError::Storage(storage_err) => Error::StoreError(storage_err), })?; Ok(result) } // Get a value from a specific bit position in the range pub fn get(&self, range_name: &str, bit_position: u64) -> Result> { let result = (&self.names, &self.assign).transaction(|(names, assign)| { // Get range info let (range_id, range_size) = match names.get(range_name)? { Some(data) => { if data.len() != 16 { return Ok(None); } let id_bytes: [u8; 8] = data[0..8].try_into() .map_err(|_| sled::transaction::ConflictableTransactionError::Abort(()))?; let size_bytes: [u8; 8] = data[8..16].try_into() .map_err(|_| sled::transaction::ConflictableTransactionError::Abort(()))?; (u64::from_be_bytes(id_bytes), u64::from_be_bytes(size_bytes)) } None => return Ok(None), }; if bit_position >= range_size { return Err(sled::transaction::ConflictableTransactionError::Abort(())); } // Get assignment let mut assign_key = Vec::new(); assign_key.extend_from_slice(&range_id.to_be_bytes()); assign_key.extend_from_slice(&bit_position.to_be_bytes()); match assign.get(assign_key)? { Some(value_bytes) => { let value = String::from_utf8(value_bytes.to_vec()) .map_err(|_| sled::transaction::ConflictableTransactionError::Abort(()))?; Ok(Some(value)) } None => Ok(None), } }).map_err(|e| match e { sled::transaction::TransactionError::Abort(()) => Error::BitOutOfRange(range_name.to_string(), bit_position), sled::transaction::TransactionError::Storage(storage_err) => Error::StoreError(storage_err), })?; Ok(result) } // List all assignments in a range pub fn list_range(&self, range_name: &str) -> Result> { // Get range info first let range_id = match self.names.get(range_name)? { Some(data) => { if data.len() != 16 { return Ok(Vec::new()); } let id_bytes: [u8; 8] = data[0..8].try_into() .map_err(|_| Error::StoreError(sled::Error::Unsupported("Invalid range data".to_string())))?; u64::from_be_bytes(id_bytes) } None => return Ok(Vec::new()), }; let mut assignments = Vec::new(); let prefix = range_id.to_be_bytes(); for result in self.assign.scan_prefix(&prefix) { let (key, value) = result?; if key.len() == 16 { // 8 bytes for range_id + 8 bytes for bit_position let bit_position_bytes: [u8; 8] = key[8..16].try_into() .map_err(|_| Error::StoreError(sled::Error::Unsupported("Invalid key format".to_string())))?; let bit_position = u64::from_be_bytes(bit_position_bytes); let value_str = String::from_utf8(value.to_vec())?; assignments.push((bit_position, value_str)); } } assignments.sort_by_key(|&(bit_pos, _)| bit_pos); Ok(assignments) } // List all ranges pub fn list_ranges(&self) -> Result> { let mut ranges = Vec::new(); for result in self.names.iter() { let (key, value) = result?; if value.len() == 16 { let name = String::from_utf8(key.to_vec())?; let size_bytes: [u8; 8] = value[8..16].try_into() .map_err(|_| Error::StoreError(sled::Error::Unsupported("Invalid size data".to_string())))?; let size = u64::from_be_bytes(size_bytes); ranges.push((name, size)); } } ranges.sort_by(|a, b| a.0.cmp(&b.0)); Ok(ranges) } + // Unassign a specific bit position from the range + pub fn unassign_bit(&self, range_name: &str, bit_position: u64) -> Result { + let result = (&self.names, &self.map, &self.assign).transaction(|(names, map, assign)| { + self.unassign_bit_in_transaction(names, map, assign, range_name, bit_position) + }).map_err(|e| match e { + sled::transaction::TransactionError::Abort(()) => Error::BitOutOfRange(range_name.to_string(), bit_position), + sled::transaction::TransactionError::Storage(storage_err) => Error::StoreError(storage_err), + })?; + Ok(result) + } + // Unassign a value from the range (find by value and remove) pub fn unassign(&self, range_name: &str, value: &str) -> Result { // First find the assignment outside of transaction let range_id = match self.names.get(range_name)? { Some(data) => { if data.len() != 16 { return Ok(false); } let id_bytes: [u8; 8] = data[0..8].try_into() .map_err(|_| Error::StoreError(sled::Error::Unsupported("Invalid range data".to_string())))?; u64::from_be_bytes(id_bytes) } None => return Ok(false), }; // Find the assignment with this value let prefix = range_id.to_be_bytes(); let mut found_bit = None; for result in self.assign.scan_prefix(&prefix) { let (key, stored_value) = result?; if key.len() == 16 { let stored_value_str = String::from_utf8(stored_value.to_vec())?; if stored_value_str == value { let bit_position_bytes: [u8; 8] = key[8..16].try_into() .map_err(|_| Error::StoreError(sled::Error::Unsupported("Invalid key format".to_string())))?; found_bit = Some(u64::from_be_bytes(bit_position_bytes)); break; } } } let bit_position = match found_bit { Some(pos) => pos, None => return Ok(false), }; // Now perform the removal in a transaction let result = (&self.map, &self.assign).transaction(|(map, assign)| { // Remove assignment let mut assign_key = Vec::new(); assign_key.extend_from_slice(&range_id.to_be_bytes()); assign_key.extend_from_slice(&bit_position.to_be_bytes()); assign.remove(assign_key)?; // Update bitmap let mut bitmap = match map.get(&range_id.to_be_bytes())? { Some(data) => data.to_vec(), None => return Err(sled::transaction::ConflictableTransactionError::Abort(())), }; let byte_index = (bit_position / 8) as usize; let bit_index = bit_position % 8; if byte_index < bitmap.len() { let mask = 1u8 << bit_index; bitmap[byte_index] &= !mask; // Clear the bit map.insert(&range_id.to_be_bytes(), bitmap)?; } Ok(true) }).map_err(|e| match e { sled::transaction::TransactionError::Abort(()) => Error::ValueNotInRange(range_name.to_string(), value.to_string()), sled::transaction::TransactionError::Storage(storage_err) => Error::StoreError(storage_err), })?; Ok(result) } + } #[cfg(test)] mod tests { use super::*; use tempfile::tempdir; fn create_test_store() -> Result { let temp_dir = tempdir().unwrap(); let db = sled::open(temp_dir.path())?; RangeStore::open(&db) } + #[test] + fn test_transaction_methods() -> Result<()> { + let store = create_test_store()?; + + // Define a range + store.define("test_range", 100)?; + + // Test transaction methods + let result = (&store.names, &store.map, &store.assign).transaction(|(names, map, assign)| { + // Assign in transaction + let bit_pos = store.assign_in_transaction(names, map, assign, "test_range", "tx_value")?; + assert!(bit_pos < 100); + + // Get in transaction + let value = store.get_in_transaction(names, assign, "test_range", bit_pos)?; + assert_eq!(value, Some("tx_value".to_string())); + + // Check if range exists in transaction + let exists = store.exists_in_transaction(names, "test_range")?; + assert_eq!(exists, true); + + // Get range info in transaction + let info = store.info_in_transaction(names, "test_range")?; + assert!(info.is_some()); + let (_, size) = info.unwrap(); + assert_eq!(size, 100); + + // Verify assignment exists + let value_check = store.get_in_transaction(names, assign, "test_range", bit_pos)?; + assert_eq!(value_check, Some("tx_value".to_string())); + + // Unassign bit in transaction + let unassigned = store.unassign_bit_in_transaction(names, map, assign, "test_range", bit_pos)?; + assert_eq!(unassigned, true); + + // Verify it's gone + let value_after = store.get_in_transaction(names, assign, "test_range", bit_pos)?; + assert_eq!(value_after, None); + + Ok(()) + }).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) + } + + #[test] + fn test_transaction_rollback() -> Result<()> { + let store = create_test_store()?; + + // Define a range and assign initial value + store.define("test_range", 5)?; // Small range to easily fill + let initial_bit = store.assign("test_range", "initial_value")?; + + // Fill up the range + for i in 1..5 { + store.assign("test_range", &format!("value_{}", i))?; + } + + // Attempt a transaction that should fail (range full) + let result = (&store.names, &store.map, &store.assign).transaction(|(names, map, assign)| { + // This should succeed + let _unassigned = store.unassign_bit_in_transaction(names, map, assign, "test_range", initial_bit)?; + + // This should fail (range is still full after we just freed one bit) + // Actually, this should succeed since we freed a bit, so let's try to assign to a full range differently + // Let's assign multiple values to ensure failure + store.assign_in_transaction(names, map, assign, "test_range", "temp1")?; + store.assign_in_transaction(names, map, assign, "test_range", "temp2")?; // This should fail + + Ok(()) + }).map_err(|e| match e { + sled::transaction::TransactionError::Abort(()) => Error::RangeFull("test_range".to_string()), + sled::transaction::TransactionError::Storage(storage_err) => Error::StoreError(storage_err), + }); + + // Transaction should have failed + assert!(result.is_err()); + + // Verify rollback - initial value should still be there + let value = store.get("test_range", initial_bit)?; + assert_eq!(value, Some("initial_value".to_string())); + + Ok(()) + } + + #[test] + fn test_bit_out_of_range() -> Result<()> { + let store = create_test_store()?; + + // Define a small range + store.define("small_range", 10)?; + + // Try to get bit beyond range + let result = store.get("small_range", 15); + assert!(result.is_err()); + + // Try to unassign bit beyond range + let result = store.unassign_bit("small_range", 15); + assert!(result.is_ok()); + assert_eq!(result?, false); + + Ok(()) + } + + #[test] + fn test_multiple_ranges() -> Result<()> { + let store = create_test_store()?; + + // Define multiple ranges + store.define("range_a", 50)?; + store.define("range_b", 100)?; + store.define("range_c", 25)?; + + // Assign values to each range + let bit_a = store.assign("range_a", "value_a")?; + let bit_b = store.assign("range_b", "value_b")?; + let bit_c = store.assign("range_c", "value_c")?; + + // Verify isolation between ranges + assert_eq!(store.get("range_a", bit_a)?, Some("value_a".to_string())); + assert_eq!(store.get("range_b", bit_b)?, Some("value_b".to_string())); + assert_eq!(store.get("range_c", bit_c)?, Some("value_c".to_string())); + + // Cross-range access: accessing bit positions that exist in one range but not another + // Since all ranges start at bit 0, we need to test with positions beyond smaller ranges + + // Test accessing valid bit positions across ranges + // All ranges have their first assignment at bit 0, so this tests basic isolation + assert_eq!(store.get("range_a", bit_a)?, Some("value_a".to_string())); + + // Test accessing bit positions that are out of range + // range_c has only 25 bits (0-24), so bit 25 and above should give an error + let result = store.get("range_c", 25); + assert!(result.is_err()); + + // range_a has 50 bits (0-49), so bit 50 and above should give an error + let result = store.get("range_a", 50); + assert!(result.is_err()); + + // List all ranges + let ranges = store.list_ranges()?; + assert_eq!(ranges.len(), 3); + assert!(ranges.iter().any(|(name, size)| name == "range_a" && *size == 50)); + assert!(ranges.iter().any(|(name, size)| name == "range_b" && *size == 100)); + assert!(ranges.iter().any(|(name, size)| name == "range_c" && *size == 25)); + + // List assignments in each range + let assignments_a = store.list_range("range_a")?; + assert_eq!(assignments_a.len(), 1); + assert_eq!(assignments_a[0], (bit_a, "value_a".to_string())); + + let assignments_b = store.list_range("range_b")?; + assert_eq!(assignments_b.len(), 1); + assert_eq!(assignments_b[0], (bit_b, "value_b".to_string())); + + let assignments_c = store.list_range("range_c")?; + assert_eq!(assignments_c.len(), 1); + assert_eq!(assignments_c[0], (bit_c, "value_c".to_string())); + + Ok(()) + } + + #[test] + fn test_sequential_bit_assignment() -> Result<()> { + let store = create_test_store()?; + + // Define a range + store.define("sequential", 10)?; + + // Assign values and verify they get sequential bits + let mut bits = Vec::new(); + for i in 0..5 { + let bit = store.assign("sequential", &format!("seq_{}", i))?; + bits.push(bit); + } + + // Should get bits 0, 1, 2, 3, 4 in that order + for (i, &bit) in bits.iter().enumerate() { + assert_eq!(bit, i as u64); + } + + // Unassign bit 2 + store.unassign("sequential", "seq_2")?; + + // Next assignment should reuse bit 2 + let next_bit = store.assign("sequential", "reused")?; + assert_eq!(next_bit, 2); + + Ok(()) + } + #[test] fn test_define_range() -> Result<()> { let store = create_test_store()?; // Define a range - store.define("test_range", 64)?; + store.define("test_range", 100)?; // Defining the same range again should not error - store.define("test_range", 64)?; + store.define("test_range", 100)?; + + // Check that ranges are listed + let ranges = store.list_ranges()?; + assert!(ranges.iter().any(|(name, size)| name == "test_range" && *size == 100)); Ok(()) } #[test] fn test_assign_and_get() -> Result<()> { let store = create_test_store()?; - // Define range first - store.define("test_range", 64)?; + // Define a range first + store.define("test_range", 100)?; // Assign a value let bit_pos = store.assign("test_range", "test_value")?; - assert_eq!(bit_pos, 0); // First assignment should be at position 0 + assert!(bit_pos < 100); // Get the value let value = store.get("test_range", bit_pos)?; assert_eq!(value, Some("test_value".to_string())); + // Test getting non-existent bit + let empty = store.get("test_range", 99)?; + assert_eq!(empty, None); + Ok(()) } #[test] fn test_multiple_assignments() -> Result<()> { let store = create_test_store()?; // Define range store.define("test_range", 64)?; // Assign multiple values let pos1 = store.assign("test_range", "value1")?; let pos2 = store.assign("test_range", "value2")?; let pos3 = store.assign("test_range", "value3")?; assert_eq!(pos1, 0); assert_eq!(pos2, 1); assert_eq!(pos3, 2); // Verify all values assert_eq!(store.get("test_range", 0)?, Some("value1".to_string())); assert_eq!(store.get("test_range", 1)?, Some("value2".to_string())); assert_eq!(store.get("test_range", 2)?, Some("value3".to_string())); Ok(()) } #[test] fn test_list_range() -> Result<()> { let store = create_test_store()?; // Define range store.define("test_range", 64)?; // Assign some values store.assign("test_range", "alpha")?; store.assign("test_range", "beta")?; store.assign("test_range", "gamma")?; // List assignments let assignments = store.list_range("test_range")?; assert_eq!(assignments.len(), 3); assert_eq!(assignments[0], (0, "alpha".to_string())); assert_eq!(assignments[1], (1, "beta".to_string())); assert_eq!(assignments[2], (2, "gamma".to_string())); Ok(()) } #[test] fn test_list_ranges() -> Result<()> { let store = create_test_store()?; // Define multiple ranges store.define("range_a", 32)?; store.define("range_b", 64)?; store.define("range_c", 128)?; // List all ranges let ranges = store.list_ranges()?; assert_eq!(ranges.len(), 3); assert_eq!(ranges[0], ("range_a".to_string(), 32)); assert_eq!(ranges[1], ("range_b".to_string(), 64)); assert_eq!(ranges[2], ("range_c".to_string(), 128)); Ok(()) } #[test] fn test_unassign() -> Result<()> { let store = create_test_store()?; // Define range store.define("test_range", 64)?; // Assign values store.assign("test_range", "value1")?; store.assign("test_range", "value2")?; store.assign("test_range", "value3")?; // Unassign middle value let removed = store.unassign("test_range", "value2")?; assert!(removed); // Verify it's gone assert_eq!(store.get("test_range", 1)?, None); // Other values should still be there assert_eq!(store.get("test_range", 0)?, Some("value1".to_string())); assert_eq!(store.get("test_range", 2)?, Some("value3".to_string())); // Assign a new value - should reuse the freed slot let new_pos = store.assign("test_range", "new_value")?; assert_eq!(new_pos, 1); // Should reuse position 1 Ok(()) } #[test] fn test_range_full() -> Result<()> { let store = create_test_store()?; // Define small range store.define("small_range", 2)?; // Fill it up store.assign("small_range", "value1")?; store.assign("small_range", "value2")?; // Try to assign one more - should fail let result = store.assign("small_range", "value3"); assert!(result.is_err()); Ok(()) } #[test] fn test_undefined_range() -> Result<()> { let store = create_test_store()?; // Try operations on undefined range let result = store.assign("undefined_range", "value"); assert!(result.is_err()); let value = store.get("undefined_range", 0)?; assert_eq!(value, None); let assignments = store.list_range("undefined_range")?; assert_eq!(assignments.len(), 0); let removed = store.unassign("undefined_range", "value")?; assert!(!removed); Ok(()) } }