Repository: cberner/redb
Pull request #315 (2022-08-03 00:34:43)

Fixes to multimap tables


Patch:
diff --git a/fuzz/fuzz_targets/fuzz_redb.rs b/fuzz/fuzz_targets/fuzz_redb.rs
index 32cf73d..c8d3fd2 100644
--- a/fuzz/fuzz_targets/fuzz_redb.rs
+++ b/fuzz/fuzz_targets/fuzz_redb.rs
@@ -167,11 +167,9 @@ fn exec_multimap_table(db: Database, transactions: &[FuzzTransaction]) -> Result
                             table.range(&start..&end).unwrap()
                         };
                         while let Some((ref_key, ref_values)) = reference_iter.next() {
-                            for ref_value_len in ref_values.iter() {
-                                let (key, value) = iter.next().unwrap();
-                                assert_eq!(*ref_key, key);
-                                assert_eq!(*ref_value_len, value.len());
-                            }
+                            let (key, value_iter) = iter.next().unwrap();
+                            assert_eq!(*ref_key, key);
+                            assert_multimap_value_eq(value_iter, Some(ref_values));
                         }
                         assert!(iter.next().is_none());
                     },
diff --git a/src/multimap_table.rs b/src/multimap_table.rs
index 678425e..73d9164 100644
--- a/src/multimap_table.rs
+++ b/src/multimap_table.rs
@@ -221,59 +221,17 @@ impl<'a, V: RedbKey + ?Sized + 'a> MultimapValueIter<'a, V> {
     }
 }
 
-struct RawMultimapValueIter<'a, V: RedbKey + ?Sized + 'a> {
-    inner: ValueIterState<'a, V>,
-}
-
-impl<'a, V: RedbKey + ?Sized + 'a> RawMultimapValueIter<'a, V> {
-    // TODO: implement Iter when GATs are stable
-    #[allow(clippy::should_implement_trait)]
-    fn next(&mut self) -> Option<&[u8]> {
-        match self.inner {
-            ValueIterState::Subtree(ref mut iter) => iter.next().map(|e| e.key()),
-            ValueIterState::InlineLeaf(ref mut iter) => iter.next_key(),
-        }
-    }
-
-    fn rev(self) -> Self {
-        Self {
-            inner: self.inner.reverse(),
-        }
-    }
-}
-
 #[doc(hidden)]
 pub struct MultimapRangeIter<'a, K: RedbKey + ?Sized + 'a, V: RedbKey + ?Sized + 'a> {
     inner: BtreeRangeIter<'a, K, DynamicCollection>,
-    subtree_key: Option<Vec<u8>>,
-    subtree_value: Vec<u8>,
-    subtree_iter: RawMultimapValueIter<'a, V>,
     mem: &'a TransactionalMemory,
     _value_type: PhantomData<V>,
 }
 
 impl<'a, K: RedbKey + ?Sized + 'a, V: RedbKey + ?Sized + 'a> MultimapRangeIter<'a, K, V> {
-    fn new(
-        mut inner: BtreeRangeIter<'a, K, DynamicCollection>,
-        mem: &'a TransactionalMemory,
-    ) -> Self {
-        let (subtree_key, subtree_iter) = if let Some(entry) = inner.next() {
-            let collection = DynamicCollection::from_bytes(entry.value());
-            // TODO: optimize out this copy
-            (Some(entry.key().to_vec()), collection.iter(mem))
-        } else {
-            let iter =
-                MultimapValueIter::new_subtree(BtreeRangeIter::new::<RangeFull, &V>(.., None, mem));
-            (None, iter)
-        };
-
+    fn new(inner: BtreeRangeIter<'a, K, DynamicCollection>, mem: &'a TransactionalMemory) -> Self {
         Self {
             inner,
-            subtree_key,
-            subtree_value: vec![],
-            subtree_iter: RawMultimapValueIter {
-                inner: subtree_iter.inner,
-            },
             mem,
             _value_type: Default::default(),
         }
@@ -287,44 +245,19 @@ impl<'a, K: RedbKey + ?Sized + 'a, V: RedbKey + ?Sized + 'a> MultimapRangeIter<'
         &mut self,
     ) -> Option<(
         <<K as RedbValue>::View as WithLifetime>::Out,
-        <<V as RedbValue>::View as WithLifetime>::Out,
+        MultimapValueIter<V>,
     )> {
-        self.subtree_key.as_ref()?;
-
-        if let Some(value_bytes) = self.subtree_iter.next() {
-            let key = K::from_bytes(self.subtree_key.as_ref().unwrap());
-            // TODO: optimize out this copy
-            self.subtree_value = value_bytes.to_vec();
-            let value = V::from_bytes(&self.subtree_value);
-            return Some((key, value));
-        }
-
-        let (subtree_key, subtree_iter) = if let Some(entry) = self.inner.next() {
-            let collection = DynamicCollection::from_bytes(entry.value());
-            // TODO: optimize out this copy
-            (Some(entry.key().to_vec()), collection.iter(self.mem))
-        } else {
-            self.subtree_key = None;
-            return None;
-        };
-        self.subtree_key = subtree_key;
-        self.subtree_iter = RawMultimapValueIter {
-            inner: subtree_iter.inner,
-        };
+        let entry = self.inner.next()?;
+        let key = K::from_bytes(entry.key());
+        let collection = DynamicCollection::from_bytes(entry.value());
+        let iter = collection.iter(self.mem);
 
-        let value_bytes = self.subtree_iter.next().unwrap();
-        let key = K::from_bytes(self.subtree_key.as_ref().unwrap());
-        self.subtree_value = value_bytes.to_vec();
-        let value = V::from_bytes(&self.subtree_value);
-        Some((key, value))
+        Some((key, iter))
     }
 
     pub fn rev(self) -> Self {
         Self {
             inner: self.inner.reverse(),
-            subtree_key: self.subtree_key,
-            subtree_value: vec![],
-            subtree_iter: self.subtree_iter.rev(),
             mem: self.mem,
             _value_type: Default::default(),
         }
@@ -403,9 +336,6 @@ impl<'db, 'txn, K: RedbKey + ?Sized, V: RedbKey + ?Sized> MultimapTable<'db, 'tx
                             <() as RedbValue>::fixed_width(),
                             new_key_bytes,
                         );
-                        if position == 0 {
-                            builder.append(value.as_bytes().as_ref(), ().as_bytes().as_ref());
-                        }
                         for i in 0..accessor.num_pairs() {
                             if i == position {
                                 builder.append(value.as_bytes().as_ref(), ().as_bytes().as_ref());
@@ -644,8 +574,10 @@ impl<'db, 'txn, K: RedbKey + ?Sized, V: RedbKey + ?Sized> ReadableMultimapTable<
     fn len(&self) -> Result<usize> {
         let mut iter: MultimapRangeIter<K, V> = self.range(..)?;
         let mut count = 0;
-        while iter.next().is_some() {
-            count += 1;
+        while let Some((_, mut values)) = iter.next() {
+            while values.next().is_some() {
+                count += 1;
+            }
         }
         Ok(count)
     }
@@ -720,8 +652,10 @@ impl<'txn, K: RedbKey + ?Sized, V: RedbKey + ?Sized> ReadableMultimapTable<K, V>
     fn len(&self) -> Result<usize> {
         let mut iter: MultimapRangeIter<K, V> = self.range(..)?;
         let mut count = 0;
-        while iter.next().is_some() {
-            count += 1;
+        while let Some((_, mut values)) = iter.next() {
+            while values.next().is_some() {
+                count += 1;
+            }
         }
         Ok(count)
     }
diff --git a/src/tree_store/btree_base.rs b/src/tree_store/btree_base.rs
index 6c3416d..273e843 100644
--- a/src/tree_store/btree_base.rs
+++ b/src/tree_store/btree_base.rs
@@ -5,6 +5,7 @@ use crate::Result;
 use std::cmp::Ordering;
 use std::marker::PhantomData;
 use std::mem::size_of;
+use std::thread;
 
 pub(crate) const LEAF: u8 = 1;
 pub(crate) const BRANCH: u8 = 2;
@@ -197,7 +198,7 @@ impl<'a, V: RedbValue + ?Sized> Drop for AccessGuard<'a, V> {
                 if let EitherPage::Mutable(ref mut mut_page) = self.page {
                     let mut mutator = LeafMutator::new(mut_page, fixed_key_size, V::fixed_width());
                     mutator.remove(position);
-                } else {
+                } else if !thread::panicking() {
                     unreachable!();
                 }
             }
@@ -768,11 +769,13 @@ impl<'a> RawLeafBuilder<'a> {
 
 impl<'a> Drop for RawLeafBuilder<'a> {
     fn drop(&mut self) {
-        assert_eq!(self.pairs_written, self.num_pairs);
-        assert_eq!(
-            self.key_section_start() + self.provisioned_key_bytes,
-            self.key_end(self.num_pairs - 1)
-        );
+        if !thread::panicking() {
+            assert_eq!(self.pairs_written, self.num_pairs);
+            assert_eq!(
+                self.key_section_start() + self.provisioned_key_bytes,
+                self.key_end(self.num_pairs - 1)
+            );
+        }
     }
 }
 
@@ -1502,7 +1505,9 @@ impl<'a: 'b, 'b> RawBranchBuilder<'a, 'b> {
 
 impl<'a: 'b, 'b> Drop for RawBranchBuilder<'a, 'b> {
     fn drop(&mut self) {
-        assert_eq!(self.keys_written, self.num_keys);
+        if !thread::panicking() {
+            assert_eq!(self.keys_written, self.num_keys);
+        }
     }
 }
 


Test patch:
diff --git a/tests/integration_tests.rs b/tests/integration_tests.rs
index 43bd9fe..05305e4 100644
--- a/tests/integration_tests.rs
+++ b/tests/integration_tests.rs
@@ -5,8 +5,8 @@ use tempfile::NamedTempFile;
 use rand::prelude::SliceRandom;
 use rand::Rng;
 use redb::{
-    Database, DatabaseBuilder, Durability, Error, MultimapTableDefinition, ReadableTable,
-    TableDefinition,
+    Database, DatabaseBuilder, Durability, Error, MultimapTableDefinition, ReadableMultimapTable,
+    ReadableTable, TableDefinition,
 };
 
 const ELEMENTS: usize = 100;
@@ -736,6 +736,63 @@ fn regression12() {
     tx.commit().unwrap();
 }
 
+#[test]
+fn regression13() {
+    let tmpfile: NamedTempFile = NamedTempFile::new().unwrap();
+
+    let db_size = 1024 * 1024;
+    let db = unsafe { Database::create(tmpfile.path(), db_size).unwrap() };
+
+    let table_def: MultimapTableDefinition<u64, [u8]> = MultimapTableDefinition::new("x");
+
+    let mut tx = db.begin_write().unwrap();
+    tx.set_durability(Durability::None);
+    {
+        let mut t = tx.open_multimap_table(table_def).unwrap();
+        let value = vec![0; 1026];
+        t.insert(&539717, &value).unwrap();
+        let value = vec![0; 530];
+        t.insert(&539717, &value).unwrap();
+    }
+    tx.abort().unwrap();
+}
+
+#[test]
+fn regression14() {
+    let tmpfile: NamedTempFile = NamedTempFile::new().unwrap();
+
+    let db_size = 150256483323;
+    let db = unsafe { Database::create(tmpfile.path(), db_size).unwrap() };
+
+    let table_def: MultimapTableDefinition<u64, [u8]> = MultimapTableDefinition::new("x");
+
+    let mut tx = db.begin_write().unwrap();
+    tx.set_durability(Durability::None);
+    {
+        let mut t = tx.open_multimap_table(table_def).unwrap();
+        let value = vec![0; 1424];
+        t.insert(&539749, &value).unwrap();
+    }
+    tx.commit().unwrap();
+
+    let mut tx = db.begin_write().unwrap();
+    tx.set_durability(Durability::None);
+    {
+        let mut t = tx.open_multimap_table(table_def).unwrap();
+        let value = vec![0; 2230];
+        t.insert(&776971, &value).unwrap();
+
+        let mut iter = t.range(&514043..&(514043 + 514043)).unwrap().rev();
+        let (key, mut value_iter) = iter.next().unwrap();
+        assert_eq!(key, 776971);
+        assert_eq!(value_iter.next().unwrap(), &[0; 2230]);
+        let (key, mut value_iter) = iter.next().unwrap();
+        assert_eq!(key, 539749);
+        assert_eq!(value_iter.next().unwrap(), &[0; 1424]);
+    }
+    tx.abort().unwrap();
+}
+
 #[test]
 fn non_durable_read_isolation() {
     let tmpfile: NamedTempFile = NamedTempFile::new().unwrap();
diff --git a/tests/multimap_tests.rs b/tests/multimap_tests.rs
index 89db9ea..ec65ccd 100644
--- a/tests/multimap_tests.rs
+++ b/tests/multimap_tests.rs
@@ -101,15 +101,21 @@ fn range_query() {
     let start = b"0".as_ref();
     let end = b"1".as_ref();
     let mut iter = table.range(start..=end).unwrap();
-    for i in 0..10u8 {
-        let (key, value) = iter.next().unwrap();
-        if i < 5 {
-            assert_eq!(b"0", key);
-        } else {
-            assert_eq!(b"1", key);
-        }
+
+    let (key, mut values) = iter.next().unwrap();
+    for i in 0..5u8 {
+        assert_eq!(b"0", key);
+        let value = values.next().unwrap();
         assert_eq!(&[i], value);
     }
+
+    let (key, mut values) = iter.next().unwrap();
+    for i in 5..10u8 {
+        assert_eq!(b"1", key);
+        let value = values.next().unwrap();
+        assert_eq!(&[i], value);
+    }
+
     assert!(iter.next().is_none());
 }