Page MenuHomePhabricator
Files F73692

No OneTemporary
Actions

View Options

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

File Metadata

Mime Type
text/x-diff
Expires
Sun, Jun 8, 10:06 AM (1 d, 6 h)
Storage Engine
blob
Storage Format
Raw Data
Storage Handle
47597
Default Alt Text
(21 KB)

Event Timeline