Repository: cberner/redb
Pull request #289 (2022-07-04 04:44:01)

Fix corruption when calling abort()


Patch:
diff --git a/fuzz/fuzz_targets/common.rs b/fuzz/fuzz_targets/common.rs
index 3e87eff..03490a0 100644
--- a/fuzz/fuzz_targets/common.rs
+++ b/fuzz/fuzz_targets/common.rs
@@ -96,6 +96,7 @@ pub(crate) enum FuzzOperation {
 pub(crate) struct FuzzTransaction {
     pub ops: Vec<FuzzOperation>,
     pub durable: bool,
+    pub commit: bool,
 }
 
 #[derive(Arbitrary, Debug)]
@@ -111,6 +112,9 @@ impl FuzzConfig {
         let mut consecutive_nondurable_commits = 0;
         let mut max_consecutive_nondurable_commits = 0;
         for transaction in self.transactions.iter() {
+            if !transaction.commit {
+                continue;
+            }
             for write in transaction.ops.iter() {
                 total_entries += 1;
                 total_value_bytes += match write {
diff --git a/fuzz/fuzz_targets/fuzz_redb.rs b/fuzz/fuzz_targets/fuzz_redb.rs
index a690c86..21982d4 100644
--- a/fuzz/fuzz_targets/fuzz_redb.rs
+++ b/fuzz/fuzz_targets/fuzz_redb.rs
@@ -22,6 +22,7 @@ fuzz_target!(|config: FuzzConfig| {
     let mut reference = BTreeMap::new();
 
     for transaction in config.transactions.iter() {
+        let reference_backup = reference.clone();
         let mut txn = db.begin_write().unwrap();
         // We're not trying to test crash safety, so don't bother with durability
         if !transaction.durable {
@@ -98,9 +99,14 @@ fuzz_target!(|config: FuzzConfig| {
                 }
             }
         }
-        if let Err(err) = txn.commit() {
-            assert!(matches!(err, Error::OutOfSpace) && config.oom_plausible());
-            return;
+        if transaction.commit {
+            if let Err(err) = txn.commit() {
+                assert!(matches!(err, Error::OutOfSpace) && config.oom_plausible());
+                return;
+            }
+        } else {
+            txn.abort().unwrap();
+            reference = reference_backup;
         }
     }
 });
diff --git a/src/tree_store/page_store/layout.rs b/src/tree_store/page_store/layout.rs
index 61d036a..8782ce6 100644
--- a/src/tree_store/page_store/layout.rs
+++ b/src/tree_store/page_store/layout.rs
@@ -225,7 +225,7 @@ impl DatabaseLayout {
         if db_capacity < min_header_size + MIN_USABLE_PAGES * page_size {
             return Err(Error::OutOfSpace);
         }
-        let result = if db_capacity - min_header_size <= full_region_layout.len() {
+        let result = if desired_usable_bytes <= full_region_layout.usable_bytes() {
             // Single region layout
             let region_layout = RegionLayout::calculate(
                 db_capacity - min_header_size,
@@ -256,6 +256,7 @@ impl DatabaseLayout {
                 db_capacity - db_header_bytes - num_full_regions * full_region_layout.len();
             let remaining_desired =
                 desired_usable_bytes - num_full_regions * max_usable_region_bytes;
+            assert!(num_full_regions > 0);
             DatabaseLayout {
                 db_header_bytes,
                 region_allocator_range,
@@ -274,6 +275,25 @@ impl DatabaseLayout {
         Ok(result)
     }
 
+    pub(super) fn create_allocators(&self) -> Vec<BuddyAllocator> {
+        let full_regional_allocator = BuddyAllocator::new(
+            self.full_region_layout().num_pages(),
+            self.full_region_layout().num_pages(),
+            self.full_region_layout().max_order(),
+        );
+        let mut allocators = vec![full_regional_allocator; self.num_full_regions()];
+        if let Some(region_layout) = self.trailing_region_layout() {
+            let trailing = BuddyAllocator::new(
+                region_layout.num_pages(),
+                self.full_region_layout().num_pages(),
+                region_layout.max_order(),
+            );
+            allocators.push(trailing);
+        }
+
+        allocators
+    }
+
     pub(super) fn full_region_layout(&self) -> &RegionLayout {
         &self.full_region_layout
     }
diff --git a/src/tree_store/page_store/page_manager.rs b/src/tree_store/page_store/page_manager.rs
index e5cc50b..52e7da5 100644
--- a/src/tree_store/page_store/page_manager.rs
+++ b/src/tree_store/page_store/page_manager.rs
@@ -552,22 +552,7 @@ impl TransactionalMemory {
         let regional_allocators = if metadata.get_allocator_dirty() {
             None
         } else {
-            let full_regional_allocator = BuddyAllocator::new(
-                layout.full_region_layout().num_pages(),
-                layout.full_region_layout().num_pages(),
-                layout.full_region_layout().max_order(),
-            );
-            let mut allocators = vec![full_regional_allocator; layout.num_full_regions()];
-            if let Some(region_layout) = layout.trailing_region_layout() {
-                let trailing = BuddyAllocator::new(
-                    region_layout.num_pages(),
-                    layout.full_region_layout().num_pages(),
-                    region_layout.max_order(),
-                );
-                allocators.push(trailing);
-            }
-
-            Some(allocators)
+            Some(layout.create_allocators())
         };
         drop(metadata);
 
@@ -635,25 +620,10 @@ impl TransactionalMemory {
 
         metadata.set_allocator_dirty(false);
         self.mmap.flush()?;
-
-        let full_regional_allocator = BuddyAllocator::new(
-            layout.full_region_layout().num_pages(),
-            layout.full_region_layout().num_pages(),
-            layout.full_region_layout().max_order(),
-        );
-        let mut allocators = vec![full_regional_allocator; layout.num_full_regions()];
-        if let Some(region_layout) = layout.trailing_region_layout() {
-            let trailing = BuddyAllocator::new(
-                region_layout.num_pages(),
-                layout.full_region_layout().num_pages(),
-                region_layout.max_order(),
-            );
-            allocators.push(trailing);
-        }
         drop(metadata);
 
         let mut guard = self.regional_allocators.lock().unwrap();
-        *guard = Some(allocators);
+        *guard = Some(layout.create_allocators());
 
         Ok(())
     }
@@ -682,9 +652,10 @@ impl TransactionalMemory {
         let mut layout = self.layout.lock().unwrap();
 
         // Trim surplus file space, before finalizing the commit
+        let mut shrunk = false;
         if self.dynamic_growth {
-            self.try_shrink(&mut metadata, &mut layout)?;
-        }
+            shrunk = self.try_shrink(&mut metadata, &mut layout)?;
+        };
 
         let mut secondary = metadata.secondary_slot_mut();
         secondary.set_last_committed_transaction_id(transaction_id);
@@ -706,6 +677,15 @@ impl TransactionalMemory {
         }
         drop(metadata);
 
+        // Safety: try_shrink() only removes unallocated free pages at the end of the database file
+        // references to unallocated pages are not allowed to exist, and we've now promoted the
+        // shrunked layout to the primary
+        if shrunk {
+            unsafe {
+                self.mmap.resize(layout.len())?;
+            }
+        }
+
         self.log_since_commit.lock().unwrap().clear();
         self.allocated_since_commit.lock().unwrap().clear();
         self.read_from_secondary.store(false, Ordering::Release);
@@ -746,7 +726,14 @@ impl TransactionalMemory {
         #[cfg(debug_assertions)]
         debug_assert!(self.open_dirty_pages.lock().unwrap().is_empty());
         let mut metadata = self.lock_metadata();
-        let regional_guard = self.regional_allocators.lock().unwrap();
+        // The layout to restore
+        let restore = if self.read_from_secondary.load(Ordering::Acquire) {
+            metadata.secondary_slot().get_data_section_layout()
+        } else {
+            metadata.primary_slot().get_data_section_layout()
+        };
+
+        let mut regional_guard = self.regional_allocators.lock().unwrap();
         let mut layout = self.layout.lock().unwrap();
         let (mut region_allocator, mut regions) = metadata.allocators_mut(&layout)?;
         for op in self.log_since_commit.lock().unwrap().drain(..).rev() {
@@ -773,11 +760,29 @@ impl TransactionalMemory {
             }
         }
         self.allocated_since_commit.lock().unwrap().clear();
+
+        // Shrinking only happens during commit
+        assert!(restore.len() <= layout.len());
         // Reset the layout, in case it changed during the writes
-        if self.read_from_secondary.load(Ordering::Acquire) {
-            *layout = metadata.secondary_slot().get_data_section_layout();
-        } else {
-            *layout = metadata.primary_slot().get_data_section_layout();
+        if restore.len() < layout.len() {
+            // Drop all the allocators for regions that are being deleted
+            regional_guard
+                .as_mut()
+                .unwrap()
+                .drain(restore.num_regions()..);
+            // Restore the size of the last region's allocator
+            let last_region_index = restore.num_regions() - 1;
+            let last_region = restore.region_layout(last_region_index);
+            let allocator_data = regions.get_regional_allocator_mut(last_region_index);
+            let last_allocator = &mut regional_guard.as_mut().unwrap()[last_region_index];
+            last_allocator.resize(allocator_data, last_region.num_pages());
+
+            *layout = restore;
+            // Safety: we've rollbacked the transaction, so any data in that was written into
+            // space that was grown during this transaction no longer exists
+            unsafe {
+                self.mmap.resize(layout.len())?;
+            }
         }
 
         Ok(())
@@ -885,11 +890,10 @@ impl TransactionalMemory {
     // Frees the page if it was allocated since the last commit. Returns true, if the page was freed
     // Safety: the caller must ensure that no references to the memory in `page` exist
     pub(crate) unsafe fn free_if_uncommitted(&self, page: PageNumber) -> Result<bool> {
-        // Zero fill the page to ensure that deleted data is not stored in the file
-        let mut mut_page = self.get_page_mut(page);
-        mut_page.memory_mut().fill(0);
-
         if self.allocated_since_commit.lock().unwrap().remove(&page) {
+            // Zero fill the page to ensure that deleted data is not stored in the file
+            let mut mut_page = self.get_page_mut(page);
+            mut_page.memory_mut().fill(0);
             let mut metadata = self.lock_metadata();
             let layout = self.layout.lock().unwrap();
             let (mut region_allocator, mut regions) = metadata.allocators_mut(&layout)?;
@@ -965,7 +969,7 @@ impl TransactionalMemory {
         &self,
         metadata: &mut MetadataAccessor,
         layout: &mut DatabaseLayout,
-    ) -> Result<()> {
+    ) -> Result<bool> {
         let mut allocators = self.regional_allocators.lock().unwrap();
         let last_region_index = layout.num_regions() - 1;
         let last_allocator = allocators.as_ref().unwrap()[last_region_index].clone();
@@ -974,7 +978,7 @@ impl TransactionalMemory {
         let trailing_free = last_allocator.trailing_free_pages(allocator_data);
         // TODO: is this the right shrinkage heuristic?
         if trailing_free < last_allocator.len() / 2 {
-            return Ok(());
+            return Ok(false);
         }
         let reduce_to_pages = if layout.num_regions() > 1 && trailing_free == last_allocator.len() {
             0
@@ -1011,13 +1015,7 @@ impl TransactionalMemory {
 
         *layout = new_layout;
 
-        // Safety: we only removed unallocated free pages at the end of the database file
-        // references to unallocated pages are not allowed to exist
-        unsafe {
-            self.mmap.resize(layout.len())?;
-        }
-
-        Ok(())
+        Ok(true)
     }
 
     fn grow(


Test patch:
diff --git a/tests/integration_tests.rs b/tests/integration_tests.rs
index 686a4fc..8763056 100644
--- a/tests/integration_tests.rs
+++ b/tests/integration_tests.rs
@@ -643,6 +643,76 @@ fn regression9() {
     tx.commit().unwrap();
 }
 
+#[test]
+fn regression10() {
+    let tmpfile: NamedTempFile = NamedTempFile::new().unwrap();
+
+    let db_size = 10 * 1024 * 1024;
+    let db = unsafe { Database::create(tmpfile.path(), db_size).unwrap() };
+
+    let table_def: TableDefinition<u64, [u8]> = TableDefinition::new("x");
+
+    let tx = db.begin_write().unwrap();
+    {
+        let mut t = tx.open_table(table_def).unwrap();
+        let v = vec![0u8; 1043];
+        t.insert(&118749, &v).unwrap();
+    }
+    tx.commit().unwrap();
+
+    let tx = db.begin_write().unwrap();
+    {
+        let mut t = tx.open_table(table_def).unwrap();
+        let v = vec![0u8; 952];
+        t.insert(&118757, &v).unwrap();
+    }
+    tx.abort().unwrap();
+
+    let tx = db.begin_write().unwrap();
+    {
+        let t = tx.open_table(table_def).unwrap();
+        t.get(&829513).unwrap();
+    }
+    tx.abort().unwrap();
+}
+
+#[test]
+fn regression11() {
+    let tmpfile: NamedTempFile = NamedTempFile::new().unwrap();
+
+    let db_size = 10 * 1024 * 1024;
+    let db = unsafe { Database::create(tmpfile.path(), db_size).unwrap() };
+
+    let table_def: TableDefinition<u64, [u8]> = TableDefinition::new("x");
+
+    let tx = db.begin_write().unwrap();
+    {
+        let mut t = tx.open_table(table_def).unwrap();
+        let v = vec![0u8; 1204];
+        t.insert(&118749, &v).unwrap();
+        let v = vec![0u8; 2062];
+        t.insert(&153697, &v).unwrap();
+        let v = vec![0u8; 2980];
+        t.insert(&110557, &v).unwrap();
+        let v = vec![0u8; 1999];
+        t.insert(&677853, &v).unwrap();
+    }
+    tx.commit().unwrap();
+
+    let tx = db.begin_write().unwrap();
+    {
+        let mut t = tx.open_table(table_def).unwrap();
+        let v = vec![0u8; 691];
+        t.insert(&103591, &v).unwrap();
+        let v = vec![0u8; 952];
+        t.insert(&118757, &v).unwrap();
+    }
+    tx.abort().unwrap();
+
+    let tx = db.begin_write().unwrap();
+    tx.commit().unwrap();
+}
+
 #[test]
 fn non_durable_read_isolation() {
     let tmpfile: NamedTempFile = NamedTempFile::new().unwrap();