Repository: cberner/redb
Pull request #286 (2022-07-03 15:24:39)

Make region size configurable


Patch:
diff --git a/src/db.rs b/src/db.rs
index 5df2a52..2d0859e 100644
--- a/src/db.rs
+++ b/src/db.rs
@@ -182,7 +182,7 @@ impl Database {
                 .open(path)?
         };
 
-        Database::new(file, db_size, None, true)
+        Database::new(file, db_size, None, None, true)
     }
 
     /// Opens an existing redb database.
@@ -194,7 +194,7 @@ impl Database {
         if File::open(path.as_ref())?.metadata()?.len() > 0 {
             let existing_size = get_db_size(path.as_ref())?;
             let file = OpenOptions::new().read(true).write(true).open(path)?;
-            Database::new(file, existing_size, None, true)
+            Database::new(file, existing_size, None, None, true)
         } else {
             Err(Error::Io(io::Error::from(ErrorKind::InvalidData)))
         }
@@ -208,6 +208,7 @@ impl Database {
         file: File,
         max_capacity: usize,
         page_size: Option<usize>,
+        region_size: Option<usize>,
         dynamic_growth: bool,
     ) -> Result<Self> {
         #[cfg(feature = "logging")]
@@ -215,7 +216,8 @@ impl Database {
             "Opening database {:?} with max size {}",
             &file, max_capacity
         );
-        let mem = TransactionalMemory::new(file, max_capacity, page_size, dynamic_growth)?;
+        let mem =
+            TransactionalMemory::new(file, max_capacity, page_size, region_size, dynamic_growth)?;
         if mem.needs_repair()? {
             let root = mem
                 .get_data_root()
@@ -338,6 +340,7 @@ impl Database {
 
 pub struct DatabaseBuilder {
     page_size: Option<usize>,
+    region_size: Option<usize>,
     dynamic_growth: bool,
 }
 
@@ -346,6 +349,7 @@ impl DatabaseBuilder {
     pub fn new() -> Self {
         Self {
             page_size: None,
+            region_size: None,
             dynamic_growth: true,
         }
     }
@@ -359,6 +363,16 @@ impl DatabaseBuilder {
         self
     }
 
+    /// Set the internal region size of the database
+    /// Smaller regions may allow the database to compact more effectively, but will limit the maximum
+    /// size of values that can be stored
+    /// Defaults to 4GiB
+    pub fn set_region_size(&mut self, size: usize) -> &mut Self {
+        assert!(size.is_power_of_two());
+        self.region_size = Some(size);
+        self
+    }
+
     /// 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
@@ -384,6 +398,12 @@ impl DatabaseBuilder {
             .create(true)
             .open(path)?;
 
-        Database::new(file, db_size, self.page_size, self.dynamic_growth)
+        Database::new(
+            file,
+            db_size,
+            self.page_size,
+            self.region_size,
+            self.dynamic_growth,
+        )
     }
 }
diff --git a/src/lib.rs b/src/lib.rs
index 7c732c1..3ed8c49 100644
--- a/src/lib.rs
+++ b/src/lib.rs
@@ -1,3 +1,5 @@
+#![allow(clippy::drop_non_drop)]
+
 pub use db::{Database, DatabaseBuilder, MultimapTableDefinition, TableDefinition};
 pub use error::Error;
 pub use multimap_table::{
diff --git a/src/tree_store/page_store/layout.rs b/src/tree_store/page_store/layout.rs
index ef3ffe0..61d036a 100644
--- a/src/tree_store/page_store/layout.rs
+++ b/src/tree_store/page_store/layout.rs
@@ -18,7 +18,7 @@ fn round_up_to_multiple_of(value: usize, multiple: usize) -> usize {
 
 // Regions are laid out starting with the allocator state header, followed by the pages aligned
 // to the next page
-#[derive(Clone)]
+#[derive(Clone, Debug)]
 pub(super) struct RegionLayout {
     num_pages: usize,
     // Offset where pages start
@@ -96,9 +96,8 @@ impl RegionLayout {
     }
 
     fn full_region_layout(max_usable_region_bytes: usize, page_size: usize) -> RegionLayout {
-        let max_order = Self::calculate_usable_order(max_usable_region_bytes, page_size).unwrap();
         let max_region_size = max_usable_region_bytes
-            + BuddyAllocator::required_space(max_usable_region_bytes / page_size, max_order);
+            + Self::header_with_padding(max_usable_region_bytes, page_size).unwrap();
 
         Self::calculate(
             max_region_size,
@@ -415,3 +414,16 @@ impl DatabaseLayout {
         }
     }
 }
+
+#[cfg(test)]
+mod test {
+    use crate::tree_store::page_store::layout::RegionLayout;
+
+    #[test]
+    fn full_layout() {
+        let layout = RegionLayout::full_region_layout(512 * 4096, 4096);
+        assert_eq!(layout.num_pages, 512);
+        assert_eq!(layout.page_size, 4096);
+        assert_eq!(layout.max_order, 9);
+    }
+}
diff --git a/src/tree_store/page_store/page_manager.rs b/src/tree_store/page_store/page_manager.rs
index 82749be..72aa3f8 100644
--- a/src/tree_store/page_store/page_manager.rs
+++ b/src/tree_store/page_store/page_manager.rs
@@ -436,6 +436,7 @@ impl TransactionalMemory {
         file: File,
         max_capacity: usize,
         requested_page_size: Option<usize>,
+        requested_region_size: Option<usize>,
         dynamic_growth: bool,
     ) -> Result<Self> {
         let page_size = requested_page_size.unwrap_or_else(get_page_size);
@@ -456,8 +457,9 @@ impl TransactionalMemory {
         let mut metadata = unsafe { MetadataAccessor::new(&mmap, mutex.lock().unwrap()) };
 
         if metadata.get_magic_number() != MAGICNUMBER {
-            let max_usable_region_bytes =
-                min(MAX_USABLE_REGION_SPACE, max_capacity.next_power_of_two());
+            let region_size = requested_region_size.unwrap_or(MAX_USABLE_REGION_SPACE);
+            assert!(region_size.is_power_of_two());
+            let max_usable_region_bytes = min(region_size, max_capacity.next_power_of_two());
 
             let starting_size = if dynamic_growth {
                 MIN_DESIRED_USABLE_BYTES
@@ -1192,7 +1194,7 @@ mod test {
         buffer[0] |= ALLOCATOR_STATE_DIRTY;
         file.write_all(&buffer).unwrap();
 
-        assert!(TransactionalMemory::new(file, max_size, None, true)
+        assert!(TransactionalMemory::new(file, max_size, None, None, true)
             .unwrap()
             .needs_repair()
             .unwrap());


Test patch:
diff --git a/tests/basic_tests.rs b/tests/basic_tests.rs
index 5b806d0..6425b63 100644
--- a/tests/basic_tests.rs
+++ b/tests/basic_tests.rs
@@ -5,6 +5,31 @@ use tempfile::NamedTempFile;
 const SLICE_TABLE: TableDefinition<[u8], [u8]> = TableDefinition::new("x");
 const U64_TABLE: TableDefinition<u64, u64> = TableDefinition::new("u64");
 
+#[test]
+fn custom_region_size() {
+    let tmpfile: NamedTempFile = NamedTempFile::new().unwrap();
+    let db = unsafe {
+        Database::builder()
+            .set_region_size(512 * 1024)
+            .create(tmpfile.path(), 20 * 1024 * 1024)
+            .unwrap()
+    };
+    let rows = 4 * 1024;
+    let write_txn = db.begin_write().unwrap();
+    {
+        let mut table = write_txn.open_table(SLICE_TABLE).unwrap();
+        let big_value = vec![0u8; 1024];
+        for i in 0..rows {
+            table.insert(&(i as u32).to_le_bytes(), &big_value).unwrap();
+        }
+    }
+    write_txn.commit().unwrap();
+
+    let read_txn = db.begin_read().unwrap();
+    let table = read_txn.open_table(SLICE_TABLE).unwrap();
+    assert_eq!(table.len().unwrap(), rows);
+}
+
 #[test]
 fn len() {
     let tmpfile: NamedTempFile = NamedTempFile::new().unwrap();