Repository: cberner/redb
Pull request #78 (2021-10-04 03:06:09)

Various cleanups


Patch:
diff --git a/benches/lmdb_benchmark.rs b/benches/lmdb_benchmark.rs
index 4a4f058..b8132c2 100644
--- a/benches/lmdb_benchmark.rs
+++ b/benches/lmdb_benchmark.rs
@@ -109,7 +109,7 @@ fn main() {
     }
     {
         let tmpfile: NamedTempFile = NamedTempFile::new().unwrap();
-        let db = unsafe { redb::Database::open(tmpfile.path(), 16 * 1024 * 1024 * 1024).unwrap() };
+        let db = unsafe { redb::Database::open(tmpfile.path(), 4096 * 1024 * 1024).unwrap() };
         let table = RedbBenchTable::new(&db);
         benchmark(table);
     }
diff --git a/src/btree.rs b/src/btree.rs
index ed4b4bc..cf9911c 100644
--- a/src/btree.rs
+++ b/src/btree.rs
@@ -1339,7 +1339,9 @@ pub(in crate) fn tree_insert<'a, K: RedbKey + ?Sized>(
     value: &[u8],
     manager: &'a TransactionalMemory,
 ) -> (PageNumber, AccessGuardMut<'a>, Vec<PageNumber>) {
-    let (page1, more, guard, freed) = tree_insert_helper::<K>(page, table, key, value, manager);
+    let mut freed = vec![];
+    let (page1, more, guard) =
+        tree_insert_helper::<K>(page, table, key, value, &mut freed, manager);
 
     if let Some((table, key, page2)) = more {
         let index_page = make_index(table, &key, page1, page2, manager);
@@ -1355,12 +1357,12 @@ fn tree_insert_helper<'a, K: RedbKey + ?Sized>(
     table: u64,
     key: &[u8],
     value: &[u8],
+    freed: &mut Vec<PageNumber>,
     manager: &'a TransactionalMemory,
 ) -> (
     PageNumber,
     Option<(u64, Vec<u8>, PageNumber)>,
     AccessGuardMut<'a>,
-    Vec<PageNumber>,
 ) {
     let node_mem = page.memory();
     match node_mem[0] {
@@ -1378,7 +1380,6 @@ fn tree_insert_helper<'a, K: RedbKey + ?Sized>(
                             left_page,
                             Some((entry.table_id(), entry.key().to_vec(), right_page)),
                             guard,
-                            vec![],
                         )
                     }
                     Ordering::Equal => {
@@ -1392,7 +1393,9 @@ fn tree_insert_helper<'a, K: RedbKey + ?Sized>(
                             manager,
                         );
 
-                        (new_page, None, guard, vec![page.get_page_number()])
+                        freed.push(page.get_page_number());
+
+                        (new_page, None, guard)
                     }
                     Ordering::Greater => {
                         let right_table = entry.table_id();
@@ -1411,12 +1414,9 @@ fn tree_insert_helper<'a, K: RedbKey + ?Sized>(
                                 let right =
                                     make_single_leaf(right_table, right_key, right_value, manager);
 
-                                (
-                                    left,
-                                    Some((table, key.to_vec(), right)),
-                                    guard,
-                                    vec![page.get_page_number()],
-                                )
+                                freed.push(page.get_page_number());
+
+                                (left, Some((table, key.to_vec(), right)), guard)
                             }
                             Ordering::Equal => {
                                 let (new_page, guard) = make_mut_double_leaf_left::<K>(
@@ -1429,7 +1429,9 @@ fn tree_insert_helper<'a, K: RedbKey + ?Sized>(
                                     manager,
                                 );
 
-                                (new_page, None, guard, vec![page.get_page_number()])
+                                freed.push(page.get_page_number());
+
+                                (new_page, None, guard)
                             }
                             Ordering::Greater => {
                                 let (left, guard) = make_mut_double_leaf_left::<K>(
@@ -1438,12 +1440,9 @@ fn tree_insert_helper<'a, K: RedbKey + ?Sized>(
                                 let right =
                                     make_single_leaf(right_table, right_key, right_value, manager);
 
-                                (
-                                    left,
-                                    Some((left_table, left_key.to_vec(), right)),
-                                    guard,
-                                    vec![page.get_page_number()],
-                                )
+                                freed.push(page.get_page_number());
+
+                                (left, Some((left_table, left_key.to_vec(), right)), guard)
                             }
                         }
                     }
@@ -1476,15 +1475,16 @@ fn tree_insert_helper<'a, K: RedbKey + ?Sized>(
                     ),
                 };
 
-                // TODO: should refactor to remove all these vec! allocations
-                (new_page, None, guard, vec![page.get_page_number()])
+                freed.push(page.get_page_number());
+
+                (new_page, None, guard)
             }
         }
         INTERNAL => {
             let accessor = InternalAccessor::new(&page);
             // Inserting into an internal page will always free it
             // TODO: unless the (key, value) pair is the same? That could result in no change
-            let mut freed = vec![page.get_page_number()];
+            freed.push(page.get_page_number());
             let mut children = vec![];
             let mut index_table_keys = vec![];
             let mut guard = None;
@@ -1495,14 +1495,14 @@ fn tree_insert_helper<'a, K: RedbKey + ?Sized>(
                     if cmp_keys::<K>(table, key, index_table, index_key).is_le() && guard.is_none()
                     {
                         let lte_page = accessor.child_page(i).unwrap();
-                        let (page1, more, guard2, child_freed) = tree_insert_helper::<K>(
+                        let (page1, more, guard2) = tree_insert_helper::<K>(
                             manager.get_page(lte_page),
                             table,
                             key,
                             value,
+                            freed,
                             manager,
                         );
-                        freed.extend_from_slice(&child_freed);
                         children.push(page1);
                         if let Some((index_table, index_key, page2)) = more {
                             index_table_keys.push((index_table, index_key));
@@ -1525,14 +1525,14 @@ fn tree_insert_helper<'a, K: RedbKey + ?Sized>(
                 // Already found the insertion place, so just copy
                 children.push(last_page);
             } else {
-                let (page1, more, guard2, child_freed) = tree_insert_helper::<K>(
+                let (page1, more, guard2) = tree_insert_helper::<K>(
                     manager.get_page(last_page),
                     table,
                     key,
                     value,
+                    freed,
                     manager,
                 );
-                freed.extend_from_slice(&child_freed);
                 children.push(page1);
                 if let Some((index_table, index_key, page2)) = more {
                     index_table_keys.push((index_table, index_key));
@@ -1557,7 +1557,7 @@ fn tree_insert_helper<'a, K: RedbKey + ?Sized>(
                 {
                     builder.write_nth_key(*table, key, *page_number, i);
                 }
-                (page.get_page_number(), None, guard, freed)
+                (page.get_page_number(), None, guard)
             } else {
                 let division = if total_key_len < MAX_KEY_SPACE_PER_PAGE {
                     // Use tree order if we did not run out of space
@@ -1595,7 +1595,6 @@ fn tree_insert_helper<'a, K: RedbKey + ?Sized>(
                     page.get_page_number(),
                     Some((*index_table, index_key.to_vec(), page2.get_page_number())),
                     guard,
-                    freed,
                 )
             }
         }
diff --git a/src/page_allocator.rs b/src/page_allocator.rs
index d6f31e2..12bb556 100644
--- a/src/page_allocator.rs
+++ b/src/page_allocator.rs
@@ -46,6 +46,10 @@ impl<'a> U64GroupedBitMap<'a> {
             x => Some(start_bit + x as usize),
         }
     }
+
+    fn count_unset(&self) -> usize {
+        self.data.iter().map(|x| x.count_zeros() as usize).sum()
+    }
 }
 
 pub(crate) struct PageAllocator {
@@ -177,6 +181,10 @@ impl PageAllocator {
         height
     }
 
+    pub(crate) fn count_free_pages(&self, data: &mut [u8]) -> usize {
+        self.get_level(data, self.get_height() - 1).count_unset()
+    }
+
     fn get_level<'a>(&self, data: &'a mut [u8], i: usize) -> U64GroupedBitMap<'a> {
         let (start, end) = self.tree_level_offsets[i];
         U64GroupedBitMap {
diff --git a/src/page_manager.rs b/src/page_manager.rs
index 27667aa..9fd5380 100644
--- a/src/page_manager.rs
+++ b/src/page_manager.rs
@@ -458,6 +458,17 @@ impl TransactionalMemory {
             open_pages: &self.open_dirty_pages,
         }
     }
+
+    pub(crate) fn count_free_pages(&self) -> usize {
+        let (mmap, guard) = self.acquire_mutable_metapage();
+        let accessor = TransactionAccessor::new(get_secondary(mmap), guard);
+        let (mem, guard) = self.acquire_mutable_page_allocator(accessor);
+        let count = self.page_allocator.count_free_pages(mem);
+        // Drop guard only after page_allocator.count_free() is completed
+        drop(guard);
+
+        count
+    }
 }
 
 impl Drop for TransactionalMemory {
diff --git a/src/storage.rs b/src/storage.rs
index fa3ff49..7dc0fd1 100644
--- a/src/storage.rs
+++ b/src/storage.rs
@@ -18,16 +18,24 @@ const TABLE_TABLE_ID: u64 = 0;
 #[derive(Debug)]
 pub struct DbStats {
     tree_height: usize,
+    free_pages: usize,
 }
 
 impl DbStats {
-    fn new(tree_height: usize) -> Self {
-        DbStats { tree_height }
+    fn new(tree_height: usize, free_pages: usize) -> Self {
+        DbStats {
+            tree_height,
+            free_pages,
+        }
     }
 
     pub fn tree_height(&self) -> usize {
         self.tree_height
     }
+
+    pub fn free_pages(&self) -> usize {
+        self.free_pages
+    }
 }
 
 #[derive(Copy, Clone, Eq, PartialEq, Debug)]
@@ -277,7 +285,7 @@ impl Storage {
             .get_root_page_number()
             .map(|p| tree_height(self.mem.get_page(p), &self.mem))
             .unwrap_or(0);
-        Ok(DbStats::new(tree_height))
+        Ok(DbStats::new(tree_height, self.mem.count_free_pages()))
     }
 
     #[allow(dead_code)]


Test patch:
diff --git a/tests/integration_tests.rs b/tests/integration_tests.rs
index 6d6e64b..6c9e918 100644
--- a/tests/integration_tests.rs
+++ b/tests/integration_tests.rs
@@ -60,29 +60,36 @@ fn persistence() {
 fn free() {
     let tmpfile: NamedTempFile = NamedTempFile::new().unwrap();
 
-    let db_size = 1024_1024;
+    let db_size = 512 * 1024;
     let db = unsafe { Database::open(tmpfile.path(), db_size).unwrap() };
     let mut table: Table<[u8], [u8]> = db.open_table(b"x").unwrap();
 
+    let free_pages = db.stats().unwrap().free_pages();
+
     let key = vec![0; 100];
     let value = vec![0; 1024];
-    // Write 10% of db space each iteration
-    let num_writes = db_size / 10 / (key.len() + value.len());
-    assert!(num_writes > 0);
-
-    // Fill the database 10 times, to be sure that memory is getting freed
-    for _ in 0..100 {
-        let mut txn = table.begin_write().unwrap();
-        {
-            for _ in 0..num_writes {
-                txn.insert(&key, &value).unwrap();
-                txn.remove(&key).unwrap();
-            }
+    // Write 25% of db space each iteration
+    let num_writes = db_size / 4 / (key.len() + value.len());
+    // Make sure an internal index page is required
+    assert!(num_writes > 64);
+
+    let mut txn = table.begin_write().unwrap();
+    {
+        for _ in 0..num_writes {
+            txn.insert(&key, &value).unwrap();
+        }
+    }
+    txn.commit().unwrap();
+
+    let mut txn = table.begin_write().unwrap();
+    {
+        for _ in 0..num_writes {
+            txn.remove(&key).unwrap();
         }
-        txn.commit().unwrap();
     }
+    txn.commit().unwrap();
 
-    // TODO: assert that no pages were leaked
+    assert_eq!(free_pages, db.stats().unwrap().free_pages());
 }
 
 #[test]