diff --git a/crates/store/src/combined.rs b/crates/store/src/combined.rs new file mode 100644 index 0000000..ca13586 --- /dev/null +++ b/crates/store/src/combined.rs @@ -0,0 +1,324 @@ +//! Combined transaction module for atomic operations across range and namespace stores. +//! +//! # Example +//! +//! This module provides atomic transactions across range and namespace stores. +//! The stores keep their APIs in their respective modules, and this module +//! provides a way to combine their operations transactionally. +//! +//! ```no_run +//! use store::{CombinedTransaction}; +//! +//! // 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")?; +//! +//! let combined = CombinedTransaction::new( +//! &range_store, +//! &namespace_store, +//! vec![&metadata_tree], +//! ); +//! +//! let (ip_bit, reserved) = combined.execute(|ctx| { +//! // Reserve a username +//! let reserved = ctx.reserve_namespace("users", "alice", "user_data")?; +//! +//! // Assign an IP address from the range +//! let ip_bit = ctx.assign_range("ip_addresses", "192.168.1.100")?; +//! +//! // Store metadata in additional tree (error handling for transaction context) +//! if let Some(tree) = ctx.additional_tree(0) { +//! tree.insert("last_assignment", "alice") +//! .map_err(|_| store::Error::StoreError( +//! sled::Error::Unsupported("Metadata insert failed".to_string()) +//! ))?; +//! } +//! +//! 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())), + } +} + +/// Combined transaction operations for range and namespace stores +pub struct CombinedTransaction<'a> { + range_store: &'a RangeStore, + namespace_store: &'a NamespaceStore, + additional_trees: Vec<&'a sled::Tree>, +} + +impl<'a> CombinedTransaction<'a> { + pub fn new( + range_store: &'a RangeStore, + namespace_store: &'a NamespaceStore, + additional_trees: Vec<&'a sled::Tree>, + ) -> Self { + Self { + range_store, + namespace_store, + additional_trees, + } + } + + /// Execute a combined transaction with range and namespace operations + pub fn execute(&self, operations: F) -> Result + where + F: Fn(&CombinedTransactionContext) -> Result, + { + // Collect all trees for the transaction + let mut all_trees = vec![ + &self.range_store.names, + &self.range_store.map, + &self.range_store.assign, + &self.namespace_store.names, + &self.namespace_store.spaces, + ]; + all_trees.extend(&self.additional_trees); + + // Execute the transaction + let result = all_trees.transaction(|trees| { + let context = CombinedTransactionContext { + range_store: self.range_store, + namespace_store: self.namespace_store, + range_names: &trees[0], + range_map: &trees[1], + range_assign: &trees[2], + namespace_names: &trees[3], + namespace_spaces: &trees[4], + additional_trees: trees[5..].iter().collect(), + }; + + 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) + } +} + +/// Context provided to transaction operations +pub struct CombinedTransactionContext<'a> { + range_store: &'a RangeStore, + namespace_store: &'a NamespaceStore, + range_names: &'a sled::transaction::TransactionalTree, + range_map: &'a sled::transaction::TransactionalTree, + range_assign: &'a sled::transaction::TransactionalTree, + namespace_names: &'a sled::transaction::TransactionalTree, + namespace_spaces: &'a sled::transaction::TransactionalTree, + additional_trees: Vec<&'a sled::transaction::TransactionalTree>, +} + +impl<'a> CombinedTransactionContext<'a> { + /// Assign a value to a range within the transaction + pub fn assign_range(&self, range_name: &str, value: &str) -> Result { + self.range_store.assign_in_transaction( + self.range_names, + self.range_map, + self.range_assign, + range_name, + value, + ).map_err(|e| convert_transaction_error(e, Error::RangeFull(range_name.to_string()))) + } + + /// Reserve a key in a namespace within the transaction + pub fn reserve_namespace(&self, namespace: &str, key: &str, value: &str) -> Result { + self.namespace_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_namespace(&self, namespace: &str, key: &str) -> Result> { + self.namespace_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_namespace(&self, namespace: &str, key: &str) -> Result { + self.namespace_store.remove_in_transaction( + self.namespace_names, + self.namespace_spaces, + namespace, + key, + ).map_err(|e| convert_transaction_error(e, Error::UndefinedNamespace(namespace.to_string()))) + } + + /// Get range assignment details within the transaction + pub fn get_range(&self, range_name: &str, bit_position: u64) -> Result> { + self.range_store.get_in_transaction( + self.range_names, + self.range_assign, + range_name, + bit_position, + ).map_err(|e| convert_transaction_error(e, Error::UndefinedRange(range_name.to_string()))) + } + + /// Access additional trees by index + pub fn additional_tree(&self, index: usize) -> Option<&sled::transaction::TransactionalTree> { + self.additional_trees.get(index).copied() + } +} + +#[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 combined = CombinedTransaction::new( + &range_store, + &namespace_store, + vec![&extra_tree], + ); + + // Execute combined transaction + let (bit_position, reserved) = combined.execute(|ctx| { + // Reserve namespace key + let reserved = ctx.reserve_namespace("test_namespace", "my_key", "my_value")?; + + // Assign range value + let bit_position = ctx.assign_range("test_range", "range_value")?; + + // Use additional tree + if let Some(tree) = ctx.additional_tree(0) { + tree.insert("extra_key", "extra_value").map_err(|_| Error::StoreError(sled::Error::Unsupported("Tree insert failed".to_string())))?; + } + + 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 combined = CombinedTransaction::new( + &range_store, + &namespace_store, + vec![], + ); + + // Execute transaction that should fail + let result = combined.execute(|ctx| { + // This should succeed + let _bit_pos = ctx.assign_range("test_range", "range_value")?; + + // This should fail (key already exists) + let _reserved = ctx.reserve_namespace("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 combined = CombinedTransaction::new( + &range_store, + &namespace_store, + vec![], + ); + + // Execute transaction with reads + let (bit_position, existing_value) = combined.execute(|ctx| { + // Read existing value + let existing_value = ctx.get_namespace("test_namespace", "existing_key")?; + + // Assign new range value + let bit_position = ctx.assign_range("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(()) + } +} \ No newline at end of file diff --git a/crates/store/src/lib.rs b/crates/store/src/lib.rs index c39f7a0..7f60157 100644 --- a/crates/store/src/lib.rs +++ b/crates/store/src/lib.rs @@ -1,12 +1,15 @@ mod error; mod store; pub use error::{Result, Error}; pub use store::Store; pub use store::open; pub mod namespace; pub use namespace::NamespaceStore; pub mod range; pub use range::RangeStore; + +pub mod combined; +pub use combined::{CombinedTransaction, CombinedTransactionContext}; diff --git a/crates/store/src/namespace.rs b/crates/store/src/namespace.rs index bcfc10c..ac29ec0 100644 --- a/crates/store/src/namespace.rs +++ b/crates/store/src/namespace.rs @@ -1,263 +1,358 @@ // 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 { - names: sled::Tree, - spaces: sled::Tree, + pub(crate) names: sled::Tree, + pub(crate) spaces: sled::Tree, } impl NamespaceStore { - pub(crate) fn open(db: &sled::Db) -> Result { + pub fn open(db: &sled::Db) -> Result { Ok(NamespaceStore { - names: db.open_tree("namespaces/1/n")?, - spaces: db.open_tree("namespaces/1/s")?, + 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( + &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( + &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( + &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()) + } + // 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), } }).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) }).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), } }).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 }; // 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()) }).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) } } #[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")?; 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())); 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_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); // Try to remove a non-existent key let removed_again = store.remove("test_namespace", "test_key")?; assert!(!removed_again); 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); Ok(()) } } diff --git a/crates/store/src/range.rs b/crates/store/src/range.rs index d026e82..b79d8f0 100644 --- a/crates/store/src/range.rs +++ b/crates/store/src/range.rs @@ -1,508 +1,614 @@ //! # 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 { - names: sled::Tree, - map: sled::Tree, - assign: sled::Tree, + pub(crate) names: sled::Tree, + pub(crate) map: sled::Tree, + pub(crate) assign: sled::Tree, } impl RangeStore { - pub(crate) fn open(db: &sled::Db) -> Result { + pub fn open(db: &sled::Db) -> Result { Ok(RangeStore { - names: db.open_tree("ranges/1/n")?, - map: db.open_tree("ranges/1/m")?, - assign: db.open_tree("ranges/1/a")?, + 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( + &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( + &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), + } + } + // 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 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_define_range() -> Result<()> { let store = create_test_store()?; // Define a range store.define("test_range", 64)?; // Defining the same range again should not error store.define("test_range", 64)?; Ok(()) } #[test] fn test_assign_and_get() -> Result<()> { let store = create_test_store()?; // Define range first store.define("test_range", 64)?; // Assign a value let bit_pos = store.assign("test_range", "test_value")?; assert_eq!(bit_pos, 0); // First assignment should be at position 0 // Get the value let value = store.get("test_range", bit_pos)?; assert_eq!(value, Some("test_value".to_string())); 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(()) } }