Page MenuHomePhabricator

No OneTemporary

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<T>(e: sled::transaction::ConflictableTransactionError<T>, 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<u64> {
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<Option<String>> {
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<bool> {
+ 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<bool> {
+ 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<Option<(u64, u64)>> {
+ 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<bool> {
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<Option<String>> {
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<bool> {
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<bool> {
+ 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<bool> {
+ 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<bool> {
+ 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<T: TransactionProvider + 'a>(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<F, R>(&self, operations: F) -> Result<R>
where
F: Fn(&TransactionContext) -> Result<R>,
{
// 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<T: TransactionProvider>(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<T: TransactionProvider>(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<F, R>(&self, operations: F) -> Result<R>
where
F: Fn(&ExtendedTransactionContext) -> Result<R>,
{
// 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<Self> {
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<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(
+ pub 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(
+ pub 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())
}
+ /// 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<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());
+
+ // 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<bool, ()> {
+ 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<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());
+
+ 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<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),
- }
+ 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<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)
+ 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<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),
- }
+ 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<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
- };
+ 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<bool> {
+ 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<bool> {
+ 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<bool> {
+ 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<Vec<String>> {
+ 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<Vec<String>> {
+ // 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<Vec<(String, String)>> {
+ // 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<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")?;
+ // 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<Self> {
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<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(
+ /// 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<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),
}
}
+ /// 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<bool, ()> {
+ // 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<bool, ()> {
+ 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<Option<(u64, u64)>, ()> {
+ 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<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 specific bit position from the range
+ pub fn unassign_bit(&self, range_name: &str, bit_position: u64) -> Result<bool> {
+ 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<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_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(())
}
}

File Metadata

Mime Type
text/x-diff
Expires
Sun, Jun 8, 3:50 PM (1 d, 12 h)
Storage Engine
blob
Storage Format
Raw Data
Storage Handle
47547
Default Alt Text
(109 KB)

Event Timeline