Repository: cberner/redb
Pull request #127 (2022-01-12 05:56:09)

Fix some TODOs


Patch:
diff --git a/src/db.rs b/src/db.rs
index f9817ba..0ab6fd2 100644
--- a/src/db.rs
+++ b/src/db.rs
@@ -31,12 +31,11 @@ impl Database {
         let file = if path.as_ref().exists() && File::open(path.as_ref())?.metadata()?.len() > 0 {
             let existing_size = get_db_size(path.as_ref())?;
             if existing_size != db_size {
-                return Err(Error::DbSizeMismatch(format!(
-                    "Database {} is of size {} bytes, but you requested {} bytes",
-                    path.as_ref().to_string_lossy(),
-                    existing_size,
-                    db_size
-                )));
+                return Err(Error::DbSizeMismatch {
+                    path: path.as_ref().to_string_lossy().to_string(),
+                    size: existing_size,
+                    requested_size: db_size,
+                });
             }
             OpenOptions::new().read(true).write(true).open(path)?
         } else {
@@ -241,28 +240,21 @@ impl<'a> DatabaseTransaction<'a> {
     }
 
     pub fn commit(self) -> Result<(), Error> {
-        match self.commit_helper(false) {
-            Ok(_) => Ok(()),
-            Err(err) => match err {
-                // Rollback the transaction if we ran out of space during commit
-                // TODO: maybe rollback on other errors too?
-                Error::OutOfSpace => {
-                    self.abort()?;
-                    Err(err)
-                }
-                err => Err(err),
-            },
-        }
+        self.commit_helper(false)
     }
 
     /// Note: pages are only freed during commit(). So exclusively using this function may result
-    /// in an out-of-memory error
+    /// in Error::OutOfSpace
     pub fn non_durable_commit(self) -> Result<(), Error> {
-        match self.commit_helper(true) {
+        self.commit_helper(true)
+    }
+
+    pub fn commit_helper(self, non_durable: bool) -> Result<(), Error> {
+        match self.commit_helper_inner(non_durable) {
             Ok(_) => Ok(()),
             Err(err) => match err {
-                // Rollback the transaction if we ran out of space during commit
-                // TODO: maybe rollback on other errors too?
+                // Rollback the transaction if we ran out of space during commit, so that user may
+                // continue with another transaction (like a delete)
                 Error::OutOfSpace => {
                     self.abort()?;
                     Err(err)
@@ -272,7 +264,7 @@ impl<'a> DatabaseTransaction<'a> {
         }
     }
 
-    pub fn commit_helper(&self, non_durable: bool) -> Result<(), Error> {
+    pub fn commit_helper_inner(&self, non_durable: bool) -> Result<(), Error> {
         // Update all the table roots in the master table, before committing
         for (name, update) in self.pending_table_root_changes.borrow_mut().drain() {
             let new_root = self.storage.update_table_root(
diff --git a/src/error.rs b/src/error.rs
index f2e213b..0139d66 100644
--- a/src/error.rs
+++ b/src/error.rs
@@ -6,9 +6,13 @@ use std::{io, panic};
 pub enum Error {
     Corrupted(String),
     TableTypeMismatch(String),
-    DbSizeMismatch(String),
+    DbSizeMismatch {
+        path: String,
+        size: usize,
+        requested_size: usize,
+    },
     DoesNotExist(String),
-    LeakedWriteTransaction(String),
+    LeakedWriteTransaction(&'static panic::Location<'static>),
     OutOfSpace,
     Io(io::Error),
     LockPoisoned(&'static panic::Location<'static>),
@@ -28,8 +32,40 @@ impl From<io::Error> for Error {
 
 impl Display for Error {
     fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result {
-        // TODO: better formatting
-        write!(f, "{:?}", self)
+        match self {
+            Error::Corrupted(msg) => {
+                write!(f, "DB corrupted: {}", msg)
+            }
+            Error::TableTypeMismatch(msg) => {
+                write!(f, "{}", msg)
+            }
+            Error::DbSizeMismatch {
+                path,
+                size,
+                requested_size,
+            } => {
+                write!(
+                    f,
+                    "Database {} is of size {} bytes, but you requested {} bytes",
+                    path, size, requested_size
+                )
+            }
+            Error::DoesNotExist(msg) => {
+                write!(f, "{}", msg)
+            }
+            Error::LeakedWriteTransaction(location) => {
+                write!(f, "Leaked write transaction: {}", location)
+            }
+            Error::OutOfSpace => {
+                write!(f, "Database is out of space")
+            }
+            Error::Io(err) => {
+                write!(f, "I/O error: {}", err)
+            }
+            Error::LockPoisoned(location) => {
+                write!(f, "Poisoned internal lock: {}", location)
+            }
+        }
     }
 }
 
diff --git a/src/tree_store/page_store/mmap.rs b/src/tree_store/page_store/mmap.rs
index 1c6d912..69def45 100644
--- a/src/tree_store/page_store/mmap.rs
+++ b/src/tree_store/page_store/mmap.rs
@@ -1,29 +1,11 @@
 use crate::Error;
-use std::ops::{Deref, DerefMut, Range};
+use std::ops::Range;
 use std::slice;
 
 pub(crate) struct Mmap {
     inner: memmap2::MmapRaw,
 }
 
-// TODO: remove these, they're dangerous mixed with the get_memory() methods
-impl Deref for Mmap {
-    type Target = [u8];
-
-    #[inline]
-    fn deref(&self) -> &[u8] {
-        unsafe { slice::from_raw_parts(self.inner.as_ptr(), self.inner.len()) }
-    }
-}
-
-// TODO: remove these, they're dangerous mixed with the get_memory() methods
-impl DerefMut for Mmap {
-    #[inline]
-    fn deref_mut(&mut self) -> &mut [u8] {
-        unsafe { slice::from_raw_parts_mut(self.inner.as_mut_ptr(), self.inner.len()) }
-    }
-}
-
 impl Mmap {
     pub(crate) fn new(mmap: memmap2::MmapRaw) -> Self {
         Self { inner: mmap }
diff --git a/src/tree_store/page_store/page_allocator.rs b/src/tree_store/page_store/page_allocator.rs
index 4f29a9b..81506e2 100644
--- a/src/tree_store/page_store/page_allocator.rs
+++ b/src/tree_store/page_store/page_allocator.rs
@@ -298,7 +298,7 @@ impl BuddyAllocator {
         let mut orders = vec![];
         for _ in 0..=max_order {
             orders.push(PageAllocator::new(num_pages));
-            num_pages = Self::higher_order(num_pages);
+            num_pages = Self::next_higher_order(num_pages as u64) as usize;
         }
 
         Self { orders }
@@ -322,7 +322,7 @@ impl BuddyAllocator {
                 Self::get_order_bytes(data, order),
                 num_pages,
             ));
-            num_pages = Self::higher_order(num_pages);
+            num_pages = Self::next_higher_order(num_pages as u64) as usize;
             metadata_offset += 2 * size_of::<u64>();
             data_offset += required;
         }
@@ -347,7 +347,7 @@ impl BuddyAllocator {
         let mut required = size_of::<u64>() + 2 * (max_order + 1) * size_of::<u64>();
         for _ in 0..=max_order {
             required += PageAllocator::required_space(num_pages);
-            num_pages = Self::higher_order(num_pages);
+            num_pages = Self::next_higher_order(num_pages as u64) as usize;
         }
 
         required
@@ -383,11 +383,11 @@ impl BuddyAllocator {
         &mut data[offset..(offset + length)]
     }
 
-    fn higher_order(page_number: usize) -> usize {
+    fn next_higher_order(page_number: u64) -> u64 {
         page_number / 2
     }
 
-    fn buddy_page(page_number: usize) -> usize {
+    fn buddy_page(page_number: u64) -> u64 {
         page_number ^ 1
     }
 
@@ -423,8 +423,7 @@ impl BuddyAllocator {
         assert!(order < self.orders.len());
         if self.orders[order].is_allocated(Self::get_order_bytes(data, order), page_number) {
             // Need to split parent page
-            // TODO: make the use of usize vs u64 consistent
-            let upper_page = Self::higher_order(page_number as usize) as u64;
+            let upper_page = Self::next_higher_order(page_number);
             self.record_alloc(data, upper_page, order + 1);
 
             let (free1, free2) = (upper_page * 2, upper_page * 2 + 1);
@@ -446,17 +445,13 @@ impl BuddyAllocator {
             return;
         }
 
-        let buddy = Self::buddy_page(page_number as usize) as u64;
+        let buddy = Self::buddy_page(page_number);
         if self.orders[order].is_allocated(Self::get_order_bytes(data, order), buddy) {
             self.orders[order].free(Self::get_order_bytes(data, order), page_number);
         } else {
             // Merge into higher order page
             self.orders[order].record_alloc(Self::get_order_bytes(data, order), buddy);
-            self.free(
-                data,
-                Self::higher_order(page_number as usize) as u64,
-                order + 1,
-            );
+            self.free(data, Self::next_higher_order(page_number), order + 1);
         }
     }
 }
diff --git a/src/tree_store/page_store/page_manager.rs b/src/tree_store/page_store/page_manager.rs
index 47e4e60..5d63ebe 100644
--- a/src/tree_store/page_store/page_manager.rs
+++ b/src/tree_store/page_store/page_manager.rs
@@ -437,9 +437,12 @@ impl TransactionalMemory {
     }
 
     pub(crate) fn new(mmap: MmapRaw, requested_page_size: Option<usize>) -> Result<Self, Error> {
-        let mut mmap = Mmap::new(mmap);
+        let mmap = Mmap::new(mmap);
+        // Safety: we have exclusive access to the mmap
+        let all_memory = unsafe { mmap.get_memory_mut(0..mmap.len()) };
+
         let mutex = Mutex::new(MetapageGuard {});
-        if mmap[0..MAGICNUMBER.len()] != MAGICNUMBER {
+        if all_memory[0..MAGICNUMBER.len()] != MAGICNUMBER {
             let page_size = requested_page_size.unwrap_or_else(get_page_size);
             // Ensure that the database metadata fits into the first page
             assert!(page_size >= DB_METAPAGE_SIZE);
@@ -449,86 +452,86 @@ impl TransactionalMemory {
             let usable_pages = Self::calculate_usable_pages(mmap.len(), page_size, max_order);
 
             // Explicitly zero the memory
-            mmap[0..DB_METAPAGE_SIZE].copy_from_slice(&[0; DB_METAPAGE_SIZE]);
-            for i in &mut mmap[(usable_pages * page_size)..] {
+            all_memory[0..DB_METAPAGE_SIZE].copy_from_slice(&[0; DB_METAPAGE_SIZE]);
+            for i in &mut all_memory[(usable_pages * page_size)..] {
                 *i = 0
             }
 
             let allocator_state_size = BuddyAllocator::required_space(usable_pages, max_order);
 
             // Store the page & db size. These are immutable
-            mmap[PAGE_SIZE_OFFSET] = page_size.trailing_zeros() as u8;
+            all_memory[PAGE_SIZE_OFFSET] = page_size.trailing_zeros() as u8;
             let length = mmap.len() as u64;
-            mmap[DB_SIZE_OFFSET..(DB_SIZE_OFFSET + size_of::<u64>())]
+            all_memory[DB_SIZE_OFFSET..(DB_SIZE_OFFSET + size_of::<u64>())]
                 .copy_from_slice(&length.to_be_bytes());
 
             // Set to 1, so that we can mutate the first transaction state
-            mmap[PRIMARY_BIT_OFFSET] = 1;
+            all_memory[PRIMARY_BIT_OFFSET] = 1;
             let start = mmap.len() - 2 * allocator_state_size;
             let mut mutator =
-                TransactionMutator::new(get_secondary(&mut mmap), mutex.lock().unwrap());
+                TransactionMutator::new(get_secondary(all_memory), mutex.lock().unwrap());
             mutator.set_root_page(PageNumber::new(0, 0), 0);
             mutator.set_last_committed_transaction_id(0);
             mutator.set_allocator_dirty(false);
             mutator.set_allocator_data(start, allocator_state_size);
             drop(mutator);
             let allocator = BuddyAllocator::init_new(
-                &mut mmap[start..(start + allocator_state_size)],
+                &mut all_memory[start..(start + allocator_state_size)],
                 usable_pages,
                 max_order,
             );
             allocator.record_alloc(
-                &mut mmap[start..(start + allocator_state_size)],
+                &mut all_memory[start..(start + allocator_state_size)],
                 DB_METADATA_PAGE,
                 0,
             );
             // Make the state we just wrote the primary
-            mmap[PRIMARY_BIT_OFFSET] = 0;
+            all_memory[PRIMARY_BIT_OFFSET] = 0;
 
             // Initialize the secondary allocator state
             let start = mmap.len() - allocator_state_size;
             let mut mutator =
-                TransactionMutator::new(get_secondary(&mut mmap), mutex.lock().unwrap());
+                TransactionMutator::new(get_secondary(all_memory), mutex.lock().unwrap());
             mutator.set_allocator_dirty(false);
             mutator.set_allocator_data(start, allocator_state_size);
             drop(mutator);
             let allocator = BuddyAllocator::init_new(
-                &mut mmap[start..(start + allocator_state_size)],
+                &mut all_memory[start..(start + allocator_state_size)],
                 usable_pages,
                 max_order,
             );
             allocator.record_alloc(
-                &mut mmap[start..(start + allocator_state_size)],
+                &mut all_memory[start..(start + allocator_state_size)],
                 DB_METADATA_PAGE,
                 0,
             );
 
-            mmap[VERSION_OFFSET] = 1;
+            all_memory[VERSION_OFFSET] = 1;
 
             mmap.flush()?;
             // Write the magic number only after the data structure is initialized and written to disk
             // to ensure that it's crash safe
-            mmap[0..MAGICNUMBER.len()].copy_from_slice(&MAGICNUMBER);
+            all_memory[0..MAGICNUMBER.len()].copy_from_slice(&MAGICNUMBER);
             mmap.flush()?;
         }
 
-        let page_size = (1 << mmap[PAGE_SIZE_OFFSET]) as usize;
+        let page_size = (1 << all_memory[PAGE_SIZE_OFFSET]) as usize;
         if let Some(size) = requested_page_size {
             assert_eq!(page_size, size);
         }
         assert_eq!(
             u64::from_be_bytes(
-                mmap[DB_SIZE_OFFSET..(DB_SIZE_OFFSET + size_of::<u64>())]
+                all_memory[DB_SIZE_OFFSET..(DB_SIZE_OFFSET + size_of::<u64>())]
                     .try_into()
                     .unwrap()
             ) as usize,
             mmap.len()
         );
 
-        let accessor = TransactionAccessor::new(get_primary(&mmap), mutex.lock().unwrap());
+        let accessor = TransactionAccessor::new(get_primary(all_memory), mutex.lock().unwrap());
         let mut allocator_dirty = accessor.get_allocator_dirty();
         drop(accessor);
-        let accessor = TransactionAccessor::new(get_secondary(&mut mmap), mutex.lock().unwrap());
+        let accessor = TransactionAccessor::new(get_secondary(all_memory), mutex.lock().unwrap());
         allocator_dirty |= accessor.get_allocator_dirty();
         drop(accessor);
 
diff --git a/src/tree_store/page_store/utils.rs b/src/tree_store/page_store/utils.rs
index 0c3f324..89b7d03 100644
--- a/src/tree_store/page_store/utils.rs
+++ b/src/tree_store/page_store/utils.rs
@@ -1,5 +1,3 @@
-// TODO: add Windows version of this function
-
 #[cfg(unix)]
 pub(crate) fn get_page_size() -> usize {
     use libc::{sysconf, _SC_PAGESIZE};
diff --git a/src/tree_store/storage.rs b/src/tree_store/storage.rs
index f08df38..1db59d5 100644
--- a/src/tree_store/storage.rs
+++ b/src/tree_store/storage.rs
@@ -192,7 +192,7 @@ impl Storage {
 
     pub(crate) fn allocate_write_transaction(&self) -> Result<u128, Error> {
         if let Some(leaked) = self.leaked_write_transaction.get() {
-            return Err(Error::LeakedWriteTransaction(leaked.to_string()));
+            return Err(Error::LeakedWriteTransaction(leaked));
         }
 
         assert!(self.live_write_transaction.get().is_none());
@@ -500,13 +500,16 @@ impl Storage {
             )?;
             freed_table.table_root = Some(r);
 
-            // TODO: only update the master root on the last iteration
-            master_root = Some(self.insert::<[u8]>(
-                FREED_TABLE,
-                &freed_table.to_bytes(),
-                transaction_id,
-                master_root,
-            )?);
+            if self.pending_freed_pages.borrow().len() <= chunk_size {
+                // Update the master root, only on the last loop iteration (this may cause another
+                // iteration, but that's ok since it would have very few pages to process)
+                master_root = Some(self.insert::<[u8]>(
+                    FREED_TABLE,
+                    &freed_table.to_bytes(),
+                    transaction_id,
+                    master_root,
+                )?);
+            }
 
             let len = self.pending_freed_pages.borrow().len();
             access_guard.as_mut()[..8]


Test patch:
diff --git a/tests/integration_tests.rs b/tests/integration_tests.rs
index b5d3b11..887b47b 100644
--- a/tests/integration_tests.rs
+++ b/tests/integration_tests.rs
@@ -104,7 +104,7 @@ fn change_db_size() {
     drop(db);
 
     let db = unsafe { Database::open(tmpfile.path(), db_size * 2) };
-    assert!(matches!(db.err().unwrap(), Error::DbSizeMismatch(_)));
+    assert!(matches!(db.err().unwrap(), Error::DbSizeMismatch { .. }));
 }
 
 #[test]