Skip to content

Commit

Permalink
Pass TransactionGuard into Range
Browse files Browse the repository at this point in the history
  • Loading branch information
@cberner
cberner committed Jan 16, 2024
1 parent 680ac34 commit a159a84
Show file tree
Hide file tree
Showing 4 changed files with 120 additions and 42 deletions.
25 changes: 22 additions & 3 deletions src/db.rs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -215,7 +215,7 @@ impl<'a, K: RedbKey + 'static, V: RedbKey + 'static> Display for MultimapTableDe
}

pub(crate) struct TransactionGuard {
transaction_tracker: Arc<TransactionTracker>,
transaction_tracker: Option<Arc<TransactionTracker>>,
transaction_id: Option<TransactionId>,
write_transaction: bool,
}
Expand All @@ -226,7 +226,7 @@ impl TransactionGuard {
tracker: Arc<TransactionTracker>,
) -> Self {
Self {
transaction_tracker: tracker,
transaction_tracker: Some(tracker),
transaction_id: Some(transaction_id),
write_transaction: false,
}
Expand All @@ -237,12 +237,21 @@ impl TransactionGuard {
tracker: Arc<TransactionTracker>,
) -> Self {
Self {
transaction_tracker: tracker,
transaction_tracker: Some(tracker),
transaction_id: Some(transaction_id),
write_transaction: true,
}
}

// TODO: remove this hack
pub(crate) fn fake() -> Self {
Self {
transaction_tracker: None,
transaction_id: None,
write_transaction: false,
}
}

pub(crate) fn id(&self) -> TransactionId {
self.transaction_id.unwrap()
}
Expand All @@ -254,12 +263,19 @@ impl TransactionGuard {

impl Drop for TransactionGuard {
fn drop(&mut self) {
if self.transaction_tracker.is_none() {
return;
}
if let Some(transaction_id) = self.transaction_id {
if self.write_transaction {
self.transaction_tracker
.as_ref()
.unwrap()
.end_write_transaction(transaction_id);
} else {
self.transaction_tracker
.as_ref()
.unwrap()
.deallocate_read_transaction(transaction_id);
}
}
Expand Down Expand Up @@ -444,6 +460,7 @@ impl Database {
"internal savepoint table".to_string(),
savepoint_table_def.get_root(),
PageHint::None,
Arc::new(TransactionGuard::fake()),
mem.clone(),
)?;
for result in savepoint_table.range::<SavepointId>(..)? {
Expand Down Expand Up @@ -484,6 +501,7 @@ impl Database {
"internal freed table".to_string(),
freed_root,
PageHint::None,
Arc::new(TransactionGuard::fake()),
mem.clone(),
)?;
let lookup_key = FreedTableKey {
Expand Down Expand Up @@ -618,6 +636,7 @@ impl Database {
"internal freed table".to_string(),
freed_root,
PageHint::None,
Arc::new(TransactionGuard::fake()),
mem.clone(),
)?;
// The persistent savepoints might hold references to older freed trees that are partially processed.
Expand Down
96 changes: 65 additions & 31 deletions src/multimap_table.rs
View file
Original file line number Diff line number Diff line change
@@ -1,3 +1,4 @@
use crate::db::TransactionGuard;
use crate::multimap_table::DynamicCollectionType::{Inline, Subtree};
use crate::sealed::Sealed;
use crate::table::TableStats;
Expand Down Expand Up @@ -477,6 +478,7 @@ impl<V> DynamicCollection<V> {
impl<V: RedbKey> DynamicCollection<V> {
fn iter<'a>(
collection: AccessGuard<'a, &'static DynamicCollection<V>>,
guard: Arc<TransactionGuard>,
mem: Arc<TransactionalMemory>,
) -> Result<MultimapValue<'a, V>> {
Ok(match collection.value().collection_type() {
Expand All @@ -486,15 +488,14 @@ impl<V: RedbKey> DynamicCollection<V> {
V::fixed_width(),
<() as RedbValue>::fixed_width(),
);
MultimapValue::new_inline(leaf_iter)
MultimapValue::new_inline(leaf_iter, guard)
}
Subtree => {
let root = collection.value().as_subtree().0;
MultimapValue::new_subtree(BtreeRangeIter::new::<RangeFull, &V::SelfType<'_>>(
&(..),
Some(root),
mem,
)?)
MultimapValue::new_subtree(
BtreeRangeIter::new::<RangeFull, &V::SelfType<'_>>(&(..), Some(root), mem)?,
guard,
)
}
})
}
Expand All @@ -503,6 +504,7 @@ impl<V: RedbKey> DynamicCollection<V> {
collection: AccessGuard<'a, &'static DynamicCollection<V>>,
pages: Vec<PageNumber>,
freed_pages: Arc<Mutex<Vec<PageNumber>>>,
guard: Arc<TransactionGuard>,
mem: Arc<TransactionalMemory>,
) -> Result<MultimapValue<'a, V>> {
Ok(match collection.value().collection_type() {
Expand All @@ -512,7 +514,7 @@ impl<V: RedbKey> DynamicCollection<V> {
V::fixed_width(),
<() as RedbValue>::fixed_width(),
);
MultimapValue::new_inline(leaf_iter)
MultimapValue::new_inline(leaf_iter, guard)
}
Subtree => {
let root = collection.value().as_subtree().0;
Expand All @@ -521,7 +523,7 @@ impl<V: RedbKey> DynamicCollection<V> {
Some(root),
mem.clone(),
)?;
MultimapValue::new_subtree_free_on_drop(inner, freed_pages, pages, mem)
MultimapValue::new_subtree_free_on_drop(inner, freed_pages, pages, guard, mem)
}
})
}
Expand Down Expand Up @@ -568,16 +570,18 @@ pub struct MultimapValue<'a, V: RedbKey + 'static> {
inner: Option<ValueIterState<'a, V>>,
freed_pages: Option<Arc<Mutex<Vec<PageNumber>>>>,
free_on_drop: Vec<PageNumber>,
_transaction_guard: Arc<TransactionGuard>,
mem: Option<Arc<TransactionalMemory>>,
_value_type: PhantomData<V>,
}

impl<'a, V: RedbKey + 'static> MultimapValue<'a, V> {
fn new_subtree(inner: BtreeRangeIter<V, ()>) -> Self {
fn new_subtree(inner: BtreeRangeIter<V, ()>, guard: Arc<TransactionGuard>) -> Self {
Self {
inner: Some(ValueIterState::Subtree(inner)),
freed_pages: None,
free_on_drop: vec![],
_transaction_guard: guard,
mem: None,
_value_type: Default::default(),
}
Expand All @@ -587,22 +591,25 @@ impl<'a, V: RedbKey + 'static> MultimapValue<'a, V> {
inner: BtreeRangeIter<V, ()>,
freed_pages: Arc<Mutex<Vec<PageNumber>>>,
pages: Vec<PageNumber>,
guard: Arc<TransactionGuard>,
mem: Arc<TransactionalMemory>,
) -> Self {
Self {
inner: Some(ValueIterState::Subtree(inner)),
freed_pages: Some(freed_pages),
free_on_drop: pages,
_transaction_guard: guard,
mem: Some(mem),
_value_type: Default::default(),
}
}

fn new_inline(inner: LeafKeyIter<'a, V>) -> Self {
fn new_inline(inner: LeafKeyIter<'a, V>, guard: Arc<TransactionGuard>) -> Self {
Self {
inner: Some(ValueIterState::InlineLeaf(inner)),
freed_pages: None,
free_on_drop: vec![],
_transaction_guard: guard,
mem: None,
_value_type: Default::default(),
}
Expand Down Expand Up @@ -661,6 +668,7 @@ impl<'a, V: RedbKey + 'static> Drop for MultimapValue<'a, V> {
pub struct MultimapRange<'a, K: RedbKey + 'static, V: RedbKey + 'static> {
inner: BtreeRangeIter<K, &'static DynamicCollection<V>>,
mem: Arc<TransactionalMemory>,
transaction_guard: Arc<TransactionGuard>,
_key_type: PhantomData<K>,
_value_type: PhantomData<V>,
_lifetime: PhantomData<&'a ()>,
Expand All @@ -669,11 +677,13 @@ pub struct MultimapRange<'a, K: RedbKey + 'static, V: RedbKey + 'static> {
impl<'a, K: RedbKey + 'static, V: RedbKey + 'static> MultimapRange<'a, K, V> {
fn new(
inner: BtreeRangeIter<K, &'static DynamicCollection<V>>,
guard: Arc<TransactionGuard>,
mem: Arc<TransactionalMemory>,
) -> Self {
Self {
inner,
mem,
transaction_guard: guard,
_key_type: Default::default(),
_value_type: Default::default(),
_lifetime: Default::default(),
Expand All @@ -690,7 +700,14 @@ impl<'a, K: RedbKey + 'static, V: RedbKey + 'static> Iterator for MultimapRange<
let key = AccessGuard::with_owned_value(entry.key_data());
let (page, _, value_range) = entry.into_raw();
let collection = AccessGuard::with_page(page, value_range);
Some(DynamicCollection::iter(collection, self.mem.clone()).map(|iter| (key, iter)))
Some(
DynamicCollection::iter(
collection,
self.transaction_guard.clone(),
self.mem.clone(),
)
.map(|iter| (key, iter)),
)
}
Err(err) => Some(Err(err)),
}
Expand All @@ -706,7 +723,14 @@ impl<'a, K: RedbKey + 'static, V: RedbKey + 'static> DoubleEndedIterator
let key = AccessGuard::with_owned_value(entry.key_data());
let (page, _, value_range) = entry.into_raw();
let collection = AccessGuard::with_page(page, value_range);
Some(DynamicCollection::iter(collection, self.mem.clone()).map(|iter| (key, iter)))
Some(
DynamicCollection::iter(
collection,
self.transaction_guard.clone(),
self.mem.clone(),
)
.map(|iter| (key, iter)),
)
}
Err(err) => Some(Err(err)),
}
Expand Down Expand Up @@ -1065,14 +1089,14 @@ impl<'db, 'txn, K: RedbKey + 'static, V: RedbKey + 'static> MultimapTable<'db, '
collection,
pages,
self.freed_pages.clone(),
self.transaction.transaction_guard(),
self.mem.clone(),
)?
} else {
MultimapValue::new_subtree(BtreeRangeIter::new::<RangeFull, &V::SelfType<'_>>(
&(..),
None,
self.mem.clone(),
)?)
MultimapValue::new_subtree(
BtreeRangeIter::new::<RangeFull, &V::SelfType<'_>>(&(..), None, self.mem.clone())?,
self.transaction.transaction_guard(),
)
};

Ok(iter)
Expand All @@ -1087,14 +1111,14 @@ impl<'db, 'txn, K: RedbKey + 'static, V: RedbKey + 'static> ReadableMultimapTabl
where
K: 'a,
{
let guard = self.transaction.transaction_guard();
let iter = if let Some(collection) = self.tree.get(key.borrow())? {
DynamicCollection::iter(collection, self.mem.clone())?
DynamicCollection::iter(collection, guard, self.mem.clone())?
} else {
MultimapValue::new_subtree(BtreeRangeIter::new::<RangeFull, &V::SelfType<'_>>(
&(..),
None,
self.mem.clone(),
)?)
MultimapValue::new_subtree(
BtreeRangeIter::new::<RangeFull, &V::SelfType<'_>>(&(..), None, self.mem.clone())?,
guard,
)
};

Ok(iter)
Expand All @@ -1107,7 +1131,11 @@ impl<'db, 'txn, K: RedbKey + 'static, V: RedbKey + 'static> ReadableMultimapTabl
KR: Borrow<K::SelfType<'a>> + 'a,
{
let inner = self.tree.range(&range)?;
Ok(MultimapRange::new(inner, self.mem.clone()))
Ok(MultimapRange::new(
inner,
self.transaction.transaction_guard(),
self.mem.clone(),
))
}

fn stats(&self) -> Result<TableStats> {
Expand Down Expand Up @@ -1236,18 +1264,21 @@ impl<'txn> ReadOnlyUntypedMultimapTable<'txn> {
pub struct ReadOnlyMultimapTable<'txn, K: RedbKey + 'static, V: RedbKey + 'static> {
tree: Btree<'txn, K, &'static DynamicCollection<V>>,
mem: Arc<TransactionalMemory>,
transaction_guard: Arc<TransactionGuard>,
_value_type: PhantomData<V>,
}

impl<'txn, K: RedbKey + 'static, V: RedbKey + 'static> ReadOnlyMultimapTable<'txn, K, V> {
pub(crate) fn new(
root_page: Option<(PageNumber, Checksum)>,
hint: PageHint,
guard: Arc<TransactionGuard>,
mem: Arc<TransactionalMemory>,
) -> Result<ReadOnlyMultimapTable<'txn, K, V>> {
Ok(ReadOnlyMultimapTable {
tree: Btree::new(root_page, hint, mem.clone())?,
mem,
transaction_guard: guard,
_value_type: Default::default(),
})
}
Expand All @@ -1262,13 +1293,12 @@ impl<'txn, K: RedbKey + 'static, V: RedbKey + 'static> ReadableMultimapTable<K,
K: 'a,
{
let iter = if let Some(collection) = self.tree.get(key.borrow())? {
DynamicCollection::iter(collection, self.mem.clone())?
DynamicCollection::iter(collection, self.transaction_guard.clone(), self.mem.clone())?
} else {
MultimapValue::new_subtree(BtreeRangeIter::new::<RangeFull, &V::SelfType<'_>>(
&(..),
None,
self.mem.clone(),
)?)
MultimapValue::new_subtree(
BtreeRangeIter::new::<RangeFull, &V::SelfType<'_>>(&(..), None, self.mem.clone())?,
self.transaction_guard.clone(),
)
};

Ok(iter)
Expand All @@ -1280,7 +1310,11 @@ impl<'txn, K: RedbKey + 'static, V: RedbKey + 'static> ReadableMultimapTable<K,
KR: Borrow<K::SelfType<'a>> + 'a,
{
let inner = self.tree.range(&range)?;
Ok(MultimapRange::new(inner, self.mem.clone()))
Ok(MultimapRange::new(
inner,
self.transaction_guard.clone(),
self.mem.clone(),
))
}

fn stats(&self) -> Result<TableStats> {
Expand Down
Loading

0 comments on commit a159a84

Please sign in to comment.