Repository: cberner/redb
Pull request #47 (2021-07-25 18:09:25)

Introduce API to retrieve btree stats


Patch:
diff --git a/src/btree.rs b/src/btree.rs
index b4d2245..2ff0062 100644
--- a/src/btree.rs
+++ b/src/btree.rs
@@ -3,7 +3,7 @@ use crate::btree::RangeIterState::{
 };
 use crate::page_manager::{Page, PageImpl, PageMut, PageNumber, TransactionalMemory};
 use crate::types::{RedbKey, RedbValue};
-use std::cmp::Ordering;
+use std::cmp::{max, Ordering};
 use std::convert::TryInto;
 use std::marker::PhantomData;
 use std::ops::{Bound, RangeBounds};
@@ -640,6 +640,24 @@ impl<'a: 'b, 'b> InternalBuilder<'a, 'b> {
     }
 }
 
+pub(in crate) fn tree_height<'a>(page: PageImpl<'a>, manager: &'a TransactionalMemory) -> usize {
+    let node_mem = page.memory();
+    match node_mem[0] {
+        LEAF => 1,
+        INTERNAL => {
+            let accessor = InternalAccessor::new(&page);
+            let left_page = accessor.lte_page();
+            let right_page = accessor.gt_page();
+            let left_height = tree_height(manager.get_page(left_page), manager);
+            let right_height = tree_height(manager.get_page(right_page), manager);
+            let child_height = max(left_height, right_height);
+
+            child_height + 1
+        }
+        _ => unreachable!(),
+    }
+}
+
 // Returns the page number of the sub-tree with this key deleted, or None if the sub-tree is empty.
 // If key is not found, guaranteed not to modify the tree
 #[allow(clippy::needless_return)]
diff --git a/src/db.rs b/src/db.rs
index 1c2449e..ee348dd 100644
--- a/src/db.rs
+++ b/src/db.rs
@@ -1,4 +1,4 @@
-use crate::storage::Storage;
+use crate::storage::{DbStats, Storage};
 use crate::table::Table;
 use crate::types::{RedbKey, RedbValue};
 use crate::Error;
@@ -49,4 +49,8 @@ impl Database {
         self.storage.commit(Some(root))?;
         Table::new(id, &self.storage)
     }
+
+    pub fn stats(&self) -> Result<DbStats, Error> {
+        self.storage.storage_stats()
+    }
 }
diff --git a/src/storage.rs b/src/storage.rs
index 9dac656..97fe7ae 100644
--- a/src/storage.rs
+++ b/src/storage.rs
@@ -1,5 +1,5 @@
 use crate::btree::{
-    lookup_in_raw, make_single_leaf, reserve_single_leaf, tree_delete, tree_insert,
+    lookup_in_raw, make_single_leaf, reserve_single_leaf, tree_delete, tree_height, tree_insert,
     tree_insert_reserve, AccessGuardMut, BtreeEntry, BtreeRangeIter,
 };
 use crate::page_manager::{Page, PageImpl, PageNumber, TransactionalMemory};
@@ -13,6 +13,21 @@ use std::ops::{RangeBounds, RangeFull};
 // The table of name -> table_id mappings
 const TABLE_TABLE_ID: u64 = 0;
 
+#[derive(Debug)]
+pub struct DbStats {
+    tree_height: usize,
+}
+
+impl DbStats {
+    fn new(tree_height: usize) -> Self {
+        DbStats { tree_height }
+    }
+
+    pub fn tree_height(&self) -> usize {
+        self.tree_height
+    }
+}
+
 pub(in crate) struct Storage {
     mem: TransactionalMemory,
 }
@@ -108,6 +123,14 @@ impl Storage {
         self.mem.rollback_uncommited_writes()
     }
 
+    pub(in crate) fn storage_stats(&self) -> Result<DbStats, Error> {
+        let tree_height = self
+            .get_root_page_number()
+            .map(|p| tree_height(self.mem.get_page(p), &self.mem))
+            .unwrap_or(0);
+        Ok(DbStats::new(tree_height))
+    }
+
     pub(in crate) fn get<K: RedbKey + ?Sized, V: RedbValue + ?Sized>(
         &self,
         table_id: u64,


Test patch:
diff --git a/tests/integration_tests.rs b/tests/integration_tests.rs
index 138aac7..9ef54d7 100644
--- a/tests/integration_tests.rs
+++ b/tests/integration_tests.rs
@@ -54,3 +54,29 @@ fn persistence() {
         }
     }
 }
+
+#[test]
+#[ignore]
+fn tree_balance() {
+    let tmpfile: NamedTempFile = NamedTempFile::new().unwrap();
+
+    let db = unsafe { Database::open(tmpfile.path()).unwrap() };
+    let mut table: Table<[u8], [u8]> = db.open_table(b"x").unwrap();
+
+    let elements = 129usize;
+
+    let mut txn = table.begin_write().unwrap();
+    for i in (0..elements).rev() {
+        txn.insert(&i.to_be_bytes(), b"").unwrap();
+    }
+    txn.commit().unwrap();
+
+    let expected_height = (elements as f32).log2() as usize + 1;
+    let height = db.stats().unwrap().tree_height();
+    assert!(
+        height <= expected_height,
+        "height={} expected={}",
+        height,
+        expected_height
+    );
+}