Repository: cberner/redb
Pull request #226 (2022-03-19 03:56:42)

Various minor cleanups


Patch:
diff --git a/src/db.rs b/src/db.rs
index ad6c1f7..219df37 100644
--- a/src/db.rs
+++ b/src/db.rs
@@ -166,8 +166,8 @@ impl DatabaseBuilder {
     /// Whether to grow the database file dynamically.
     /// When set to true, the database file will start at a small size and grow as insertions are made
     /// When set to false, the database file will be statically sized
-    pub fn set_grow(&mut self, dynamic: bool) -> &mut Self {
-        self.dynamic_growth = dynamic;
+    pub fn set_dynamic_growth(&mut self, enabled: bool) -> &mut Self {
+        self.dynamic_growth = enabled;
         self
     }
 
diff --git a/src/tree_store/btree_utils.rs b/src/tree_store/btree_utils.rs
index 900cdb6..8d33780 100644
--- a/src/tree_store/btree_utils.rs
+++ b/src/tree_store/btree_utils.rs
@@ -818,7 +818,7 @@ unsafe fn tree_insert_helper<'a, K: RedbKey + ?Sized>(
                 let total_size =
                     LeafBuilder::required_bytes(accessor.num_pairs() + 1, total_pairs_size);
                 // Split the leaf if it's larger than a native page
-                let too_large = total_size > manager.native_page_size() && accessor.num_pairs() > 1;
+                let too_large = total_size > manager.get_page_size() && accessor.num_pairs() > 1;
                 if accessor.num_pairs() >= BTREE_ORDER || too_large {
                     // split
                     let division = accessor.num_pairs() / 2;
diff --git a/src/tree_store/page_store/page_manager.rs b/src/tree_store/page_store/page_manager.rs
index 91709bd..6cc8abc 100644
--- a/src/tree_store/page_store/page_manager.rs
+++ b/src/tree_store/page_store/page_manager.rs
@@ -47,6 +47,9 @@ const MAX_PAGE_ORDER: usize = 20;
 const MIN_USABLE_PAGES: usize = 10;
 const MIN_DESIRED_USABLE_BYTES: usize = 1024 * 1024;
 
+// TODO: set to 1, when version 1.0 is released
+const FILE_FORMAT_VERSION: u8 = 100;
+
 // Inspired by PNG's magic number
 const MAGICNUMBER: [u8; 9] = [b'r', b'e', b'd', b'b', 0x1A, 0x0A, 0xA9, 0x0D, 0x0A];
 const VERSION_OFFSET: usize = MAGICNUMBER.len();
@@ -55,8 +58,6 @@ const GOD_BYTE_OFFSET: usize = PAGE_SIZE_OFFSET + size_of::<u8>();
 const RESERVED: usize = 4;
 const REGION_SIZE_OFFSET: usize = GOD_BYTE_OFFSET + size_of::<u8>() + RESERVED;
 const DB_SIZE_OFFSET: usize = REGION_SIZE_OFFSET + size_of::<u64>();
-const UPGRADE_LOG_OFFSET: usize = DB_SIZE_OFFSET + size_of::<u64>();
-const UPGRADE_LOG_LEN_OFFSET: usize = UPGRADE_LOG_OFFSET + size_of::<u64>();
 const TRANSACTION_SIZE: usize = 128;
 const TRANSACTION_0_OFFSET: usize = 128;
 const TRANSACTION_1_OFFSET: usize = TRANSACTION_0_OFFSET + TRANSACTION_SIZE;
@@ -65,7 +66,6 @@ const DB_HEADER_SIZE: usize = TRANSACTION_1_OFFSET + TRANSACTION_SIZE;
 // God byte flags
 const PRIMARY_BIT: u8 = 1;
 const ALLOCATOR_STATE_DIRTY: u8 = 2;
-const UPGRADE_IN_PROGRESS: u8 = 4;
 
 // Structure of each commit slot
 const ROOT_PAGE_OFFSET: usize = 0;
@@ -83,23 +83,8 @@ fn ceil_log2(x: usize) -> usize {
 pub(crate) fn get_db_size(path: impl AsRef<Path>) -> Result<usize, io::Error> {
     let mut db_size = [0u8; size_of::<u64>()];
     let mut file = File::open(path)?;
-
-    file.seek(SeekFrom::Start(GOD_BYTE_OFFSET as u64))?;
-    let mut buffer = [0u8; 1];
-    file.read_exact(&mut buffer)?;
-    let god_byte = buffer[0];
-    if god_byte & UPGRADE_IN_PROGRESS == 0 {
-        file.seek(SeekFrom::Start(DB_SIZE_OFFSET as u64))?;
-    } else {
-        let mut log_size = [0u8; size_of::<u64>()];
-        file.seek(SeekFrom::Start(UPGRADE_LOG_LEN_OFFSET as u64))?;
-        file.read_exact(&mut log_size)?;
-        assert_eq!(u64::from_be_bytes(log_size) as usize, size_of::<u64>());
-        file.seek(SeekFrom::Start(UPGRADE_LOG_OFFSET as u64))?;
-    }
-
+    file.seek(SeekFrom::Start(DB_SIZE_OFFSET as u64))?;
     file.read_exact(&mut db_size)?;
-
     Ok(u64::from_be_bytes(db_size) as usize)
 }
 
@@ -177,6 +162,35 @@ impl<'a> MetadataAccessor<'a> {
         ) as usize
     }
 
+    fn set_max_capacity(&mut self, max_size: usize) {
+        self.header[DB_SIZE_OFFSET..DB_SIZE_OFFSET + size_of::<u64>()]
+            .copy_from_slice(&(max_size as u64).to_be_bytes());
+    }
+
+    fn get_magic_number(&self) -> [u8; MAGICNUMBER.len()] {
+        self.header[..MAGICNUMBER.len()].try_into().unwrap()
+    }
+
+    fn set_magic_number(&mut self) {
+        self.header[..MAGICNUMBER.len()].copy_from_slice(&MAGICNUMBER);
+    }
+
+    fn get_page_size(&self) -> usize {
+        1usize << self.header[PAGE_SIZE_OFFSET]
+    }
+
+    fn set_page_size(&mut self, page_size: usize) {
+        self.header[PAGE_SIZE_OFFSET] = page_size.trailing_zeros() as u8;
+    }
+
+    fn get_version(&self) -> u8 {
+        self.header[VERSION_OFFSET]
+    }
+
+    fn set_version(&mut self, version: u8) {
+        self.header[VERSION_OFFSET] = version;
+    }
+
     fn get_allocator_dirty(&self) -> bool {
         self.header[GOD_BYTE_OFFSET] & ALLOCATOR_STATE_DIRTY != 0
     }
@@ -189,7 +203,7 @@ impl<'a> MetadataAccessor<'a> {
         }
     }
 
-    fn get_regional_allocator(&mut self, region: usize, layout: &DataSectionLayout) -> &[u8] {
+    fn get_regional_allocator(&mut self, region: usize, layout: &DatabaseLayout) -> &[u8] {
         let base = layout.region_base_address(region);
         let relative = layout.region_layout(region).allocator_state;
         let absolute = (base + relative.start)..(base + relative.end);
@@ -204,7 +218,7 @@ impl<'a> MetadataAccessor<'a> {
     // as self, since self holds the metadata lock
     fn allocators_mut(
         &mut self,
-        layout: &DataSectionLayout,
+        layout: &DatabaseLayout,
     ) -> Result<(U64GroupedBitMapMut, RegionsAccessor)> {
         if !self.get_allocator_dirty() {
             self.set_allocator_dirty(true);
@@ -232,7 +246,7 @@ impl<'a> MetadataAccessor<'a> {
 // Safety: RegionAccessor may only access regional metadata, and no other references to it may exist
 struct RegionsAccessor<'a> {
     mmap: &'a Mmap,
-    layout: DataSectionLayout,
+    layout: DatabaseLayout,
 }
 
 impl<'a> RegionsAccessor<'a> {
@@ -247,17 +261,6 @@ impl<'a> RegionsAccessor<'a> {
     }
 }
 
-fn get_secondary(header: &mut [u8]) -> &mut [u8] {
-    let start = if header[GOD_BYTE_OFFSET] & PRIMARY_BIT == 0 {
-        TRANSACTION_1_OFFSET
-    } else {
-        TRANSACTION_0_OFFSET
-    };
-    let end = start + TRANSACTION_SIZE;
-
-    &mut header[start..end]
-}
-
 // On-disk format is:
 // lowest 20bits: page index within the region
 // second 20bits: region number
@@ -357,10 +360,9 @@ impl<'a> TransactionAccessor<'a> {
         )
     }
 
-    fn get_data_section_layout(&self) -> DataSectionLayout {
-        DataSectionLayout::from_be_bytes(
-            self.mem
-                [DATA_LAYOUT_OFFSET..(DATA_LAYOUT_OFFSET + DataSectionLayout::serialized_size())]
+    fn get_data_section_layout(&self) -> DatabaseLayout {
+        DatabaseLayout::from_be_bytes(
+            self.mem[DATA_LAYOUT_OFFSET..(DATA_LAYOUT_OFFSET + DatabaseLayout::serialized_size())]
                 .try_into()
                 .unwrap(),
         )
@@ -386,8 +388,8 @@ impl<'a> TransactionMutator<'a> {
             .copy_from_slice(&transaction_id.to_be_bytes());
     }
 
-    fn set_data_section_layout(&mut self, layout: &DataSectionLayout) {
-        self.mem[DATA_LAYOUT_OFFSET..(DATA_LAYOUT_OFFSET + DataSectionLayout::serialized_size())]
+    fn set_data_section_layout(&mut self, layout: &DatabaseLayout) {
+        self.mem[DATA_LAYOUT_OFFSET..(DATA_LAYOUT_OFFSET + DatabaseLayout::serialized_size())]
             .copy_from_slice(&layout.to_be_bytes());
     }
 }
@@ -568,14 +570,14 @@ impl RegionLayout {
 }
 
 #[derive(Clone)]
-struct DataSectionLayout {
+struct DatabaseLayout {
     db_header_bytes: usize,
     full_region_layout: RegionLayout,
     num_full_regions: usize,
     trailing_partial_region: Option<RegionLayout>,
 }
 
-impl DataSectionLayout {
+impl DatabaseLayout {
     fn calculate(
         db_capacity: usize,
         mut desired_usable_bytes: usize,
@@ -592,7 +594,7 @@ impl DataSectionLayout {
                 page_size,
             )
             .ok_or(Error::OutOfSpace)?;
-            Ok(DataSectionLayout {
+            Ok(DatabaseLayout {
                 db_header_bytes: min_header_size,
                 full_region_layout,
                 num_full_regions: 0,
@@ -610,7 +612,7 @@ impl DataSectionLayout {
                 db_capacity - db_header_bytes - num_full_regions * full_region_layout.len();
             let remaining_desired =
                 desired_usable_bytes - num_full_regions * full_region_layout.len();
-            Ok(DataSectionLayout {
+            Ok(DatabaseLayout {
                 db_header_bytes,
                 full_region_layout,
                 num_full_regions,
@@ -645,7 +647,6 @@ impl DataSectionLayout {
         self.num_full_regions * self.full_region_layout.usable_bytes() + trailing
     }
 
-    // TODO: move this to a HeaderLayout structure
     fn region_allocator_address_range(&self) -> Range<usize> {
         DB_HEADER_SIZE..self.db_header_bytes
     }
@@ -737,7 +738,6 @@ enum AllocationOp {
     FreeUncommitted(PageNumber),
 }
 
-// TODO: maybe the TransactionalMemory should be held in a Mutex, instead of all the mutexes inside it?
 pub(crate) struct TransactionalMemory {
     // Pages allocated since the last commit
     allocated_since_commit: Mutex<HashSet<PageNumber>>,
@@ -750,16 +750,16 @@ pub(crate) struct TransactionalMemory {
     // It would be nice if this was a RefCell<&[u8]> on the metadata. However, that would be
     // self-referential, since we also hold the mmap object
     metadata_guard: Mutex<MetadataGuard>,
-    layout: Mutex<DataSectionLayout>,
+    layout: Mutex<DatabaseLayout>,
     // The number of PageMut which are outstanding
     #[cfg(debug_assertions)]
     open_dirty_pages: Mutex<HashSet<PageNumber>>,
     // Indicates that a non-durable commit has been made, so reads should be served from the secondary meta page
     read_from_secondary: AtomicBool,
     page_size: usize,
-    // TODO: remove this? we have the layout
+    // We store these separately from the layout because they're static, and accessed on the get_page()
+    // code path where there is no locking
     region_size: usize,
-    // TODO: remove this? we have the layout
     header_size: usize,
 }
 
@@ -783,13 +783,9 @@ impl TransactionalMemory {
         }
 
         let mutex = Mutex::new(MetadataGuard {});
-        let magic_number: [u8; 9] = {
-            // Safety: we have exclusive access to the mmap
-            let header = unsafe { mmap.get_memory_mut(0..DB_HEADER_SIZE) };
-            header[..MAGICNUMBER.len()].try_into().unwrap()
-        };
+        let mut metadata = unsafe { MetadataAccessor::new(&mmap, mutex.lock().unwrap()) };
 
-        if magic_number != MAGICNUMBER {
+        if metadata.get_magic_number() != MAGICNUMBER {
             let page_size = requested_page_size.unwrap_or_else(get_page_size);
             assert!(page_size.is_power_of_two());
 
@@ -798,7 +794,7 @@ impl TransactionalMemory {
             } else {
                 max_capacity
             };
-            let layout = DataSectionLayout::calculate(max_capacity, starting_size, page_size)?;
+            let layout = DatabaseLayout::calculate(max_capacity, starting_size, page_size)?;
 
             if mmap.len() < layout.len() {
                 // Safety: We're growing the mmap
@@ -807,96 +803,69 @@ impl TransactionalMemory {
                 }
             }
 
-            let all_memory = unsafe { mmap.get_memory_mut(0..layout.len()) };
-
             // Explicitly zero the header
-            all_memory[0..layout.db_header_bytes].fill(0);
+            metadata.header.fill(0);
+
+            let (mut region_allocator, mut regions) = metadata.allocators_mut(&layout)?;
 
             // Initialize the region allocator
             let num_regions = layout.num_regions();
-            let mut region_allocator =
-                U64GroupedBitMapMut::new(&mut all_memory[layout.region_allocator_address_range()]);
+            for i in 0..num_regions {
+                region_allocator.clear(i);
+            }
             for i in num_regions..region_allocator.len() {
                 region_allocator.set(i);
             }
 
             // Initialize all the regional allocators
             for i in 0..num_regions {
-                let base = layout.region_base_address(i);
+                let mem = regions.get_regional_allocator_mut(i);
                 let region_layout = layout.region_layout(i);
-                let relative = &region_layout.allocator_state;
-                let absolute = (base + relative.start)..(base + relative.end);
                 BuddyAllocator::init_new(
-                    &mut all_memory[absolute],
+                    mem,
                     region_layout.num_pages,
                     layout.full_region_layout.num_pages,
                     region_layout.max_order(),
                 );
             }
+            // Set the allocator to not dirty, because the allocator initialization above will have
+            // dirtied it
+            metadata.set_allocator_dirty(false);
 
             // Store the page & db size. These are immutable
-            all_memory[PAGE_SIZE_OFFSET] = page_size.trailing_zeros() as u8;
-            all_memory[DB_SIZE_OFFSET..(DB_SIZE_OFFSET + size_of::<u64>())]
-                .copy_from_slice(&(max_capacity as u64).to_be_bytes());
+            metadata.set_page_size(page_size);
+            metadata.set_max_capacity(max_capacity);
+            metadata.set_version(FILE_FORMAT_VERSION);
 
-            // Set to 1, so that we can mutate the first transaction state
-            all_memory[GOD_BYTE_OFFSET] = PRIMARY_BIT;
-            let mut mutator = TransactionMutator::new(get_secondary(all_memory));
+            let mut mutator = metadata.secondary_slot_mut();
             mutator.set_root_page(PageNumber::null());
             mutator.set_last_committed_transaction_id(0);
             mutator.set_data_section_layout(&layout);
             drop(mutator);
             // Make the state we just wrote the primary
-            all_memory[GOD_BYTE_OFFSET] &= !PRIMARY_BIT;
+            metadata.swap_primary();
 
             // Initialize the secondary allocator state
-            let mut mutator = TransactionMutator::new(get_secondary(all_memory));
+            let mut mutator = metadata.secondary_slot_mut();
             mutator.set_data_section_layout(&layout);
             drop(mutator);
 
-            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
-            all_memory[0..MAGICNUMBER.len()].copy_from_slice(&MAGICNUMBER);
-            mmap.flush()?;
-        }
-
-        let all_memory = unsafe { mmap.get_memory_mut(0..mmap.len()) };
-        if all_memory[GOD_BYTE_OFFSET] & UPGRADE_IN_PROGRESS != 0 {
-            let upgrade_log_len = u64::from_be_bytes(
-                all_memory[UPGRADE_LOG_LEN_OFFSET..(UPGRADE_LOG_LEN_OFFSET + size_of::<u64>())]
-                    .try_into()
-                    .unwrap(),
-            ) as usize;
-            assert_eq!(upgrade_log_len, size_of::<u64>());
-            let rollback =
-                all_memory[UPGRADE_LOG_OFFSET..(UPGRADE_LOG_OFFSET + upgrade_log_len)].to_vec();
-            all_memory[DB_SIZE_OFFSET..(DB_SIZE_OFFSET + size_of::<u64>())]
-                .copy_from_slice(&rollback);
-            mmap.flush()?;
-            all_memory[GOD_BYTE_OFFSET] &= !UPGRADE_IN_PROGRESS;
+            metadata.set_magic_number();
             mmap.flush()?;
         }
 
-        let page_size = (1 << all_memory[PAGE_SIZE_OFFSET]) as usize;
+        let page_size = metadata.get_page_size();
         if let Some(size) = requested_page_size {
             assert_eq!(page_size, size);
         }
-
-        let allocator_dirty = all_memory[GOD_BYTE_OFFSET] & (ALLOCATOR_STATE_DIRTY) != 0;
-
-        // Ensure that all_memory reference does not exist, because it must not for the safety guarantees
-        // of MetadataAccessor to hold
-        #[allow(clippy::drop_ref)]
-        drop(all_memory);
-        // Safety: we just dropped our reference to all the mmap memory
-        let metadata = unsafe { MetadataAccessor::new(&mmap, mutex.lock().unwrap()) };
+        assert_eq!(metadata.get_version(), FILE_FORMAT_VERSION);
         let layout = metadata.primary_slot().get_data_section_layout();
         let region_size = layout.full_region_layout.len();
 
-        let regional_allocators = if allocator_dirty {
+        let regional_allocators = if metadata.get_allocator_dirty() {
             None
         } else {
             let full_regional_allocator = BuddyAllocator::new(
@@ -938,11 +907,10 @@ impl TransactionalMemory {
         Ok(self.lock_metadata().get_allocator_dirty())
     }
 
-    // Returns true if the repair is complete. If false, this method must be called again
     pub(crate) fn repair_allocator(
         &self,
         allocated_pages: impl Iterator<Item = PageNumber>,
-    ) -> Result<bool> {
+    ) -> Result<()> {
         let mut metadata = self.lock_metadata();
         let layout = self.layout.lock().unwrap();
         let (mut region_allocator, mut regions) = metadata.allocators_mut(&layout)?;
@@ -981,17 +949,7 @@ impl TransactionalMemory {
 
         metadata.set_allocator_dirty(false);
         self.mmap.flush()?;
-        // Retain the metadata lock until flush is done
-        drop(metadata);
 
-        Ok(true)
-    }
-
-    // TODO: remove this method
-    pub(crate) fn finalize_repair_allocator(&mut self) -> Result {
-        assert!(!self.needs_repair()?);
-        let metadata = self.lock_metadata();
-        let layout = self.layout.lock().unwrap();
         let full_regional_allocator = BuddyAllocator::new(
             layout.full_region_layout.num_pages,
             layout.full_region_layout.num_pages,
@@ -1041,7 +999,11 @@ impl TransactionalMemory {
         }
 
         metadata.swap_primary();
-        self.mmap.flush()?;
+        if eventual {
+            self.mmap.eventual_flush()?;
+        } else {
+            self.mmap.flush()?;
+        }
         drop(metadata);
 
         self.log_since_commit.lock().unwrap().clear();
@@ -1236,14 +1198,10 @@ impl TransactionalMemory {
         self.allocated_since_commit.lock().unwrap().contains(&page)
     }
 
-    pub(crate) fn native_page_size(&self) -> usize {
-        self.page_size
-    }
-
     fn allocate_helper(
         &self,
         metadata: &mut MetadataAccessor,
-        layout: &DataSectionLayout,
+        layout: &DatabaseLayout,
         required_order: usize,
     ) -> Result<PageNumber> {
         let regional_guard = self.regional_allocators.lock().unwrap();
@@ -1281,7 +1239,7 @@ impl TransactionalMemory {
         ))
     }
 
-    fn grow(&self, metadata: &mut MetadataAccessor, layout: &mut DataSectionLayout) -> Result<()> {
+    fn grow(&self, metadata: &mut MetadataAccessor, layout: &mut DatabaseLayout) -> Result<()> {
         let next_desired_size = if layout.num_full_regions > 0 {
             // Grow by 1 region
             // TODO: prune the trailing partial region, if it exists
@@ -1289,7 +1247,7 @@ impl TransactionalMemory {
         } else {
             layout.usable_bytes() * 2
         };
-        let new_layout = DataSectionLayout::calculate(
+        let new_layout = DatabaseLayout::calculate(
             metadata.get_max_capacity(),
             next_desired_size,
             self.page_size,
@@ -1461,61 +1419,17 @@ impl Drop for TransactionalMemory {
 mod test {
     use crate::db::TableDefinition;
     use crate::tree_store::page_store::page_manager::{
-        ALLOCATOR_STATE_DIRTY, DB_SIZE_OFFSET, GOD_BYTE_OFFSET, MAGICNUMBER, MIN_USABLE_PAGES,
-        UPGRADE_IN_PROGRESS, UPGRADE_LOG_LEN_OFFSET, UPGRADE_LOG_OFFSET,
+        ALLOCATOR_STATE_DIRTY, GOD_BYTE_OFFSET, MAGICNUMBER, MIN_USABLE_PAGES,
     };
     use crate::tree_store::page_store::utils::get_page_size;
     use crate::tree_store::page_store::TransactionalMemory;
-    use crate::{Database, Error, ReadableTable};
+    use crate::{Database, Error};
     use memmap2::MmapMut;
     use std::fs::OpenOptions;
-    use std::mem::size_of;
     use tempfile::NamedTempFile;
 
     const X: TableDefinition<[u8], [u8]> = TableDefinition::new("x");
 
-    #[test]
-    fn recover_upgrade() {
-        let tmpfile: NamedTempFile = NamedTempFile::new().unwrap();
-        let db = unsafe { Database::create(tmpfile.path(), 1024 * 1024).unwrap() };
-        let write_txn = db.begin_write().unwrap();
-        {
-            let mut table = write_txn.open_table(X).unwrap();
-            table.insert(b"hello", b"world").unwrap();
-        }
-        write_txn.commit().unwrap();
-        drop(db);
-
-        let file = OpenOptions::new()
-            .read(true)
-            .write(true)
-            .open(tmpfile.path())
-            .unwrap();
-
-        let mut mmap = unsafe { MmapMut::map_mut(&file) }.unwrap();
-        let god_byte = mmap[GOD_BYTE_OFFSET];
-        mmap[GOD_BYTE_OFFSET] = god_byte | UPGRADE_IN_PROGRESS;
-        let old_size = u64::from_be_bytes(
-            mmap[DB_SIZE_OFFSET..(DB_SIZE_OFFSET + size_of::<u64>())]
-                .try_into()
-                .unwrap(),
-        );
-        mmap[UPGRADE_LOG_OFFSET..(UPGRADE_LOG_OFFSET + size_of::<u64>())]
-            .copy_from_slice(&(old_size as u64).to_be_bytes());
-        mmap[UPGRADE_LOG_LEN_OFFSET..(UPGRADE_LOG_LEN_OFFSET + size_of::<u64>())]
-            .copy_from_slice(&8u64.to_be_bytes());
-        // Perform a partial increase of the db size. Intentionally don't resize the mmap file, so that it's invalid
-        mmap[DB_SIZE_OFFSET..(DB_SIZE_OFFSET + size_of::<u64>())]
-            .copy_from_slice(&(2 * old_size).to_be_bytes());
-        mmap.flush().unwrap();
-        drop(mmap);
-
-        let db2 = unsafe { Database::create(tmpfile.path(), 1024 * 1024).unwrap() };
-        let read_txn = db2.begin_read().unwrap();
-        let table = read_txn.open_table(X).unwrap();
-        assert_eq!(b"world", table.get(b"hello").unwrap().unwrap());
-    }
-
     #[test]
     fn repair_allocator() {
         let tmpfile: NamedTempFile = NamedTempFile::new().unwrap();
diff --git a/src/tree_store/storage.rs b/src/tree_store/storage.rs
index 9fb71c3..dd1bc27 100644
--- a/src/tree_store/storage.rs
+++ b/src/tree_store/storage.rs
@@ -216,27 +216,27 @@ impl Storage {
         page_size: Option<usize>,
         dynamic_growth: bool,
     ) -> Result<Storage> {
-        let mut mem = TransactionalMemory::new(file, max_capacity, page_size, dynamic_growth)?;
+        let mem = TransactionalMemory::new(file, max_capacity, page_size, dynamic_growth)?;
         if mem.needs_repair()? {
             let root = mem
                 .get_primary_root_page()
                 .expect("Tried to repair an empty database");
             // Clear the freed table. We're about to rebuild the allocator state by walking
             // all the reachable pages, which will implicitly free them
+            {
+                // Safety: we own the mem object, and just created it, so there can't be other references.
+                let mut bytes = unsafe {
+                    get_mut_value::<str>(mem.get_page_mut(root), FREED_TABLE.as_bytes(), &mem)
+                        .unwrap()
+                };
+                // TODO: this needs to be done in a transaction. Otherwise a torn write could happen
+                let mut header = TableHeader::from_bytes(bytes.as_mut());
+                header.table_root = None;
+                assert_eq!(bytes.as_mut().len(), header.to_bytes().len());
+                bytes.as_mut().copy_from_slice(&header.to_bytes());
+            }
 
-            // Safety: we own the mem object, and just created it, so there can't be other references.
-            let mut bytes = unsafe {
-                get_mut_value::<str>(mem.get_page_mut(root), FREED_TABLE.as_bytes(), &mem).unwrap()
-            };
-            let mut header = TableHeader::from_bytes(bytes.as_mut());
-            header.table_root = None;
-            assert_eq!(bytes.as_mut().len(), header.to_bytes().len());
-            bytes.as_mut().copy_from_slice(&header.to_bytes());
-        }
-        while mem.needs_repair()? {
-            let root = mem
-                .get_primary_root_page()
-                .expect("Tried to repair an empty database");
+            // Repair the allocator state
             let root_page = mem.get_page(root);
             let start = page_numbers_iter_start_state(root_page);
 
@@ -261,7 +261,6 @@ impl Storage {
 
             mem.repair_allocator(all_pages_iter)?;
         }
-        mem.finalize_repair_allocator()?;
 
         let mut next_transaction_id = mem.get_last_committed_transaction_id()? + 1;
         if mem.get_primary_root_page().is_none() {


Test patch:
diff --git a/tests/integration_tests.rs b/tests/integration_tests.rs
index 014b348..6dc8010 100644
--- a/tests/integration_tests.rs
+++ b/tests/integration_tests.rs
@@ -145,7 +145,7 @@ fn free() {
     let db_size = 8 * 1024 * 1024;
     let db = unsafe {
         Database::builder()
-            .set_grow(false)
+            .set_dynamic_growth(false)
             .create(tmpfile.path(), db_size)
             .unwrap()
     };