Page MenuHomePhabricator
Files F73636

No OneTemporary
Actions

View Options

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<F, R>(&self, operations: F) -> Result<R>
+ where
+ F: Fn(&CombinedTransactionContext) -> Result<R>,
+ {
+ // 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<u64> {
+ 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<bool> {
+ 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<Option<String>> {
+ 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<bool> {
+ 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<Option<String>> {
+ 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<Self> {
+ pub fn open(db: &sled::Db) -> Result<Self> {
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<bool, ()> {
+ // 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<Option<String>, ()> {
+ // 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<bool, ()> {
+ // 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<Option<String>> {
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<bool> {
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<Option<String>> {
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<bool> {
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<NamespaceStore> {
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<Self> {
+ pub fn open(db: &sled::Db) -> Result<Self> {
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<u64, ()> {
+ // 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<Option<String>, ()> {
+ // 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<u64> {
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<Option<String>> {
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<Vec<(u64, String)>> {
// 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<Vec<(String, u64)>> {
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<bool> {
// 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<RangeStore> {
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(())
}
}

File Metadata

Mime Type
text/x-diff
Expires
Sun, Jun 8, 5:18 AM (1 d, 1 h)
Storage Engine
blob
Storage Format
Raw Data
Storage Handle
47543
Default Alt Text
(50 KB)

Event Timeline