From 867743842c17ab6987b18835d38b8dfe8c0eab69 Mon Sep 17 00:00:00 2001 From: Christopher Berner Date: Thu, 21 Dec 2023 09:18:33 -0800 Subject: [PATCH] Add range_arc() and get_arc() These methods return reference counted iterators from read only transactions --- src/lib.rs | 6 ++-- src/multimap_table.rs | 70 +++++++++++++++++++++++++++++++++++++++++ src/table.rs | 30 ++++++++++++++++++ tests/basic_tests.rs | 24 ++++++++++++++ tests/multimap_tests.rs | 58 ++++++++++++++++++++++++++++++++++ 5 files changed, 185 insertions(+), 3 deletions(-) diff --git a/src/lib.rs b/src/lib.rs index 219dcb8e..fa8ef5f9 100644 --- a/src/lib.rs +++ b/src/lib.rs @@ -62,11 +62,11 @@ pub use error::{ TransactionError, }; pub use multimap_table::{ - MultimapRange, MultimapTable, MultimapValue, ReadOnlyMultimapTable, - ReadOnlyUntypedMultimapTable, ReadableMultimapTable, + ArcMultimapRange, ArcMultimapValue, MultimapRange, MultimapTable, MultimapValue, + ReadOnlyMultimapTable, ReadOnlyUntypedMultimapTable, ReadableMultimapTable, }; pub use table::{ - Drain, DrainFilter, Range, ReadOnlyTable, ReadOnlyUntypedTable, ReadableTable, Table, + ArcRange, Drain, DrainFilter, Range, ReadOnlyTable, ReadOnlyUntypedTable, ReadableTable, Table, TableStats, }; pub use transactions::{DatabaseStats, Durability, ReadTransaction, WriteTransaction}; diff --git a/src/multimap_table.rs b/src/multimap_table.rs index 5eb6cc5c..35cb9809 100644 --- a/src/multimap_table.rs +++ b/src/multimap_table.rs @@ -566,6 +566,24 @@ enum ValueIterState<'a, V: RedbKey + 'static> { InlineLeaf(LeafKeyIter<'a, V>), } +pub struct ArcMultimapValue { + inner: MultimapValue<'static, V>, +} + +impl Iterator for ArcMultimapValue { + type Item = as Iterator>::Item; + + fn next(&mut self) -> Option { + self.inner.next() + } +} + +impl DoubleEndedIterator for ArcMultimapValue { + fn next_back(&mut self) -> Option { + self.inner.next_back() + } +} + pub struct MultimapValue<'a, V: RedbKey + 'static> { inner: Option>, freed_pages: Option>>>, @@ -665,6 +683,24 @@ impl<'a, V: RedbKey + 'static> Drop for MultimapValue<'a, V> { } } +pub struct ArcMultimapRange { + inner: MultimapRange<'static, K, V>, +} + +impl Iterator for ArcMultimapRange { + type Item = as Iterator>::Item; + + fn next(&mut self) -> Option { + self.inner.next() + } +} + +impl DoubleEndedIterator for ArcMultimapRange { + fn next_back(&mut self) -> Option { + self.inner.next_back() + } +} + pub struct MultimapRange<'a, K: RedbKey + 'static, V: RedbKey + 'static> { inner: BtreeRangeIter>, mem: Arc, @@ -1296,6 +1332,40 @@ impl<'txn, K: RedbKey + 'static, V: RedbKey + 'static> ReadOnlyMultimapTable<'tx _lifetime: Default::default(), }) } + + /// This method is like `get()`, but the iterator is reference counted and keeps the transaction + /// alive until it is dropped. + pub fn get_arc<'a>(&self, key: impl Borrow>) -> Result> + where + K: 'a, + { + let iter = if let Some(collection) = self.tree.get(key.borrow())? { + DynamicCollection::iter(collection, self.transaction_guard.clone(), self.mem.clone())? + } else { + MultimapValue::new_subtree( + BtreeRangeIter::new::>(&(..), None, self.mem.clone())?, + self.transaction_guard.clone(), + ) + }; + + Ok(ArcMultimapValue { inner: iter }) + } + + /// This method is like `range()`, but the iterator is reference counted and keeps the transaction + /// alive until it is dropped. + pub fn range_arc<'a, KR>( + &self, + range: impl RangeBounds + 'a, + ) -> Result> + where + K: 'a, + KR: Borrow> + 'a, + { + let inner = self.tree.range(&range)?; + let range = MultimapRange::new(inner, self.transaction_guard.clone(), self.mem.clone()); + + Ok(ArcMultimapRange { inner: range }) + } } impl<'txn, K: RedbKey + 'static, V: RedbKey + 'static> ReadableMultimapTable diff --git a/src/table.rs b/src/table.rs index 2541df77..e1d67227 100644 --- a/src/table.rs +++ b/src/table.rs @@ -439,6 +439,18 @@ impl<'txn, K: RedbKey + 'static, V: RedbValue + 'static> ReadOnlyTable<'txn, K, _lifetime: Default::default(), }) } + + /// This method is like `range()`, but the iterator is reference counted and keeps the transaction + /// alive until it is dropped. + pub fn range_arc<'a, KR>(&self, range: impl RangeBounds + 'a) -> Result> + where + K: 'a, + KR: Borrow> + 'a, + { + self.tree.range(&range).map(|x| ArcRange { + inner: Range::new(x, self.transaction_guard.clone()), + }) + } } impl<'txn, K: RedbKey + 'static, V: RedbValue + 'static> ReadableTable @@ -594,6 +606,24 @@ impl< } } +pub struct ArcRange { + inner: Range<'static, K, V>, +} + +impl Iterator for ArcRange { + type Item = Result<(AccessGuard<'static, K>, AccessGuard<'static, V>)>; + + fn next(&mut self) -> Option { + self.inner.next() + } +} + +impl DoubleEndedIterator for ArcRange { + fn next_back(&mut self) -> Option { + self.inner.next_back() + } +} + pub struct Range<'a, K: RedbKey + 'static, V: RedbValue + 'static> { inner: BtreeRangeIter, _transaction_guard: Arc, diff --git a/tests/basic_tests.rs b/tests/basic_tests.rs index 4138c832..3a3fb4c0 100644 --- a/tests/basic_tests.rs +++ b/tests/basic_tests.rs @@ -1246,6 +1246,30 @@ fn range_lifetime() { assert!(iter.next().is_none()); } +#[test] +fn range_arc() { + let tmpfile = create_tempfile(); + let db = Database::create(tmpfile.path()).unwrap(); + + let definition: TableDefinition<&str, &str> = TableDefinition::new("x"); + + let write_txn = db.begin_write().unwrap(); + { + let mut table = write_txn.open_table(definition).unwrap(); + table.insert("hello", "world").unwrap(); + } + write_txn.commit().unwrap(); + + let mut iter = { + let read_txn = db.begin_read().unwrap(); + let table = read_txn.open_table(definition).unwrap(); + let start = "hello".to_string(); + table.range_arc::<&str>(start.as_str()..).unwrap() + }; + assert_eq!(iter.next().unwrap().unwrap().1.value(), "world"); + assert!(iter.next().is_none()); +} + #[test] fn drain_lifetime() { let tmpfile = create_tempfile(); diff --git a/tests/multimap_tests.rs b/tests/multimap_tests.rs index 317096b8..72891900 100644 --- a/tests/multimap_tests.rs +++ b/tests/multimap_tests.rs @@ -172,6 +172,64 @@ fn range_lifetime() { assert!(iter.next().is_none()); } +#[test] +fn range_arc_lifetime() { + let tmpfile = create_tempfile(); + let db = Database::create(tmpfile.path()).unwrap(); + + let definition: MultimapTableDefinition<&str, &str> = MultimapTableDefinition::new("x"); + + let write_txn = db.begin_write().unwrap(); + { + let mut table = write_txn.open_multimap_table(definition).unwrap(); + table.insert("hello", "world").unwrap(); + } + write_txn.commit().unwrap(); + + let mut iter = { + let read_txn = db.begin_read().unwrap(); + let table = read_txn.open_multimap_table(definition).unwrap(); + let start = "hello".to_string(); + table.range_arc::<&str>(start.as_str()..).unwrap() + }; + assert_eq!( + iter.next() + .unwrap() + .unwrap() + .1 + .next() + .unwrap() + .unwrap() + .value(), + "world" + ); + assert!(iter.next().is_none()); +} + +#[test] +fn get_arc_lifetime() { + let tmpfile = create_tempfile(); + let db = Database::create(tmpfile.path()).unwrap(); + + let definition: MultimapTableDefinition<&str, &str> = MultimapTableDefinition::new("x"); + + let write_txn = db.begin_write().unwrap(); + { + let mut table = write_txn.open_multimap_table(definition).unwrap(); + table.insert("hello", "world").unwrap(); + } + write_txn.commit().unwrap(); + + let mut iter = { + let read_txn = db.begin_read().unwrap(); + let table = read_txn.open_multimap_table(definition).unwrap(); + let start = "hello".to_string(); + table.get_arc(start.as_str()).unwrap() + }; + assert_eq!(iter.next().unwrap().unwrap().value(), "world"); + assert!(iter.next().is_none()); +} + #[test] fn delete() { let tmpfile = create_tempfile();