diff --git a/src/btreemap.rs b/src/btreemap.rs index afdf65472..16c64d171 100644 --- a/src/btreemap.rs +++ b/src/btreemap.rs @@ -60,7 +60,7 @@ use crate::{ }; use allocator::Allocator; pub use iter::Iter; -use node::{DerivedPageSize, Entry, Node, NodeType, PageSize, Version}; +use node::{DerivedPageSize, Entry, Node, NodeType, PageSize, Version, MINIMUM_PAGE_SIZE}; use node_cache::NodeCache; pub use node_cache::NodeCacheMetrics; use std::borrow::Cow; @@ -236,6 +236,9 @@ const DEFAULT_NODE_CACHE_NUM_SLOTS: usize = 16; /// - Use when your type's serialized size can vary or has no fixed maximum /// - Recommended for most custom types, especially those containing Strings or Vecs /// - Example: `const BOUND: Bound = Bound::Unbounded;` +/// - The map stores its nodes in 1024-byte pages by default. If your entries are +/// typically much smaller, pick a smaller page size with +/// [`init_with_page_size`](BTreeMap::init_with_page_size) to save memory. /// /// - **Bounded (`Bound::Bounded{ max_size, is_fixed_size }`)**: /// - Use when you know the maximum serialized size of your type @@ -295,21 +298,54 @@ where /// If the memory provided already contains a `BTreeMap`, then that /// map is loaded. Otherwise, a new `BTreeMap` instance is created. pub fn init(memory: M) -> Self { + if Self::contains_map(&memory) { + BTreeMap::load(memory) + } else { + BTreeMap::new(memory) + } + } + + /// Initializes a `BTreeMap` that stores its nodes in pages of `page_size` bytes. + /// + /// If the memory provided already contains a `BTreeMap`, then that map is + /// loaded and `page_size` is ignored, as a map's page size is fixed when + /// it's created. Otherwise, a new `BTreeMap` is created with the given + /// page size. See [`new_with_page_size`](Self::new_with_page_size) for + /// guidance on choosing a page size. + /// + /// # Panics + /// + /// Panics if a new map is created and `page_size` is less than 128 bytes. + /// + /// # Examples + /// + /// ```rust + /// use ic_stable_structures::{BTreeMap, DefaultMemoryImpl}; + /// + /// // Entries are unbounded but typically ~20 bytes (key and value + /// // combined), so a page much smaller than the default 1024 bytes fits + /// // a typical node. + /// let map: BTreeMap = + /// BTreeMap::init_with_page_size(DefaultMemoryImpl::default(), 384); + /// ``` + pub fn init_with_page_size(memory: M, page_size: u32) -> Self { + if Self::contains_map(&memory) { + BTreeMap::load(memory) + } else { + BTreeMap::new_with_page_size(memory, page_size) + } + } + + // Returns true if the memory already contains a `BTreeMap`. + fn contains_map(memory: &M) -> bool { if memory.size() == 0 { - // Memory is empty. Create a new map. - return BTreeMap::new(memory); + return false; } // Check if the magic in the memory corresponds to a BTreeMap. - let mut dst = vec![0; 3]; + let mut dst = [0; 3]; memory.read(0, &mut dst); - if dst != MAGIC { - // No BTreeMap found. Create a new instance. - BTreeMap::new(memory) - } else { - // The memory already contains a BTreeMap. Load it. - BTreeMap::load(memory) - } + dst == *MAGIC } /// Configures the number of node-cache slots during construction. @@ -438,21 +474,11 @@ where /// This is exposed only in testing. #[cfg(test)] pub fn init_v1(memory: M) -> Self { - if memory.size() == 0 { - // Memory is empty. Create a new map. - return BTreeMap::new_v1(memory); - } - - // Check if the magic in the memory corresponds to a BTreeMap. - let mut dst = vec![0; 3]; - memory.read(0, &mut dst); - if dst != MAGIC { - // No BTreeMap found. Create a new instance. - BTreeMap::new_v1(memory) - } else { - // The memory already contains a BTreeMap. Load it, making sure - // we don't migrate the BTreeMap to v2. + if Self::contains_map(&memory) { + // Load the map, making sure we don't migrate the BTreeMap to v2. BTreeMap::load_helper(memory, false) + } else { + BTreeMap::new_v1(memory) } } @@ -493,6 +519,65 @@ where _ => PageSize::Value(DEFAULT_PAGE_SIZE), }; + Self::new_helper(memory, page_size) + } + + /// Creates a new `BTreeMap` that stores its nodes in pages of `page_size` bytes. + /// + /// Each node of the tree is allocated a page, and a node that doesn't fit + /// in its page continues into overflow pages, so keys and values of any + /// size can be stored regardless of the page size. The page size trades + /// off memory usage against performance: + /// + /// * A page that's too large wastes memory, as every node occupies a full + /// page no matter how few bytes it uses. + /// * A page that's too small makes nodes spill into overflow pages, each + /// of which adds a few bytes of overhead and extra reads and writes + /// whenever the node is loaded or saved. + /// + /// A good page size is one that fits a typical node. A node holds at most + /// 11 entries, and a 4-byte length is stored alongside each unbounded key + /// and value, so a full leaf node takes roughly + /// `15 + 11 * (4 + key_size) + 11 * (4 + value_size)` bytes, and an + /// internal node another 96 bytes for the addresses of its children. + /// + /// [`new`](Self::new) picks a page size automatically: if both keys and + /// values are bounded, it's derived from their maximum sizes, and + /// otherwise it's 1024 bytes. That can waste most of each page when + /// entries are unbounded but typically small, which is when a smaller + /// page size pays off. + /// + /// The page size is stored in memory, and can't be changed once the map + /// is created. + /// + /// # Panics + /// + /// Panics if `page_size` is less than 128 bytes. + /// + /// # Examples + /// + /// ```rust + /// use ic_stable_structures::{BTreeMap, DefaultMemoryImpl}; + /// + /// // Entries of ~20 bytes (key and value combined) need ~420 bytes per + /// // full internal node, and nodes are rarely full, so 384-byte pages + /// // fit nearly every node. + /// let mut map: BTreeMap, Vec, _> = + /// BTreeMap::new_with_page_size(DefaultMemoryImpl::default(), 384); + /// + /// // Entries larger than a page are still supported. + /// map.insert(vec![1], vec![0; 10_000]); + /// assert_eq!(map.get(&vec![1]), Some(vec![0; 10_000])); + /// ``` + pub fn new_with_page_size(memory: M, page_size: u32) -> Self { + assert!( + page_size >= MINIMUM_PAGE_SIZE, + "page_size must be at least {MINIMUM_PAGE_SIZE} bytes, got {page_size}", + ); + Self::new_helper(memory, PageSize::Value(page_size)) + } + + fn new_helper(memory: M, page_size: PageSize) -> Self { let btree = Self { root_addr: NULL, allocator: Allocator::new( diff --git a/src/btreemap/node.rs b/src/btreemap/node.rs index 2efb5cd6e..8b704240f 100644 --- a/src/btreemap/node.rs +++ b/src/btreemap/node.rs @@ -15,6 +15,7 @@ mod v1; mod v2; use io::NodeReader; +pub(crate) use v2::MINIMUM_PAGE_SIZE; // The minimum degree to use in the btree. // This constant is taken from Rust's std implementation of BTreeMap. diff --git a/src/btreemap/node/v2.rs b/src/btreemap/node/v2.rs index 2037c0a98..d5e402c90 100644 --- a/src/btreemap/node/v2.rs +++ b/src/btreemap/node/v2.rs @@ -87,7 +87,7 @@ pub(super) const PAGE_OVERFLOW_DATA_OFFSET: Bytes = Bytes::new(11); // The minimum size a page can have. // Rationale: a page size needs to at least store the header (15 bytes) + all the children // addresses (88 bytes). We round that up to 128 to get a nice binary number. -const MINIMUM_PAGE_SIZE: u32 = 128; +pub(crate) const MINIMUM_PAGE_SIZE: u32 = 128; impl Node { /// Creates a new v2 node at the given address. diff --git a/src/btreemap/proptests.rs b/src/btreemap/proptests.rs index 8e1938232..c0505d4a8 100644 --- a/src/btreemap/proptests.rs +++ b/src/btreemap/proptests.rs @@ -1,7 +1,7 @@ use crate::{ btreemap::{ tests::{b, make_memory, run_btree_test}, - BTreeMap, + BTreeMap, MINIMUM_PAGE_SIZE, }, storable::Blob, Memory, @@ -69,6 +69,21 @@ fn comprehensive(#[strategy(pvec(operation_strategy(), 100..5_000))] ops: Vec, +) { + let mem = make_memory(); + let mut btree = BTreeMap::new_with_page_size(mem, MINIMUM_PAGE_SIZE); + let mut std_btree = StdBTreeMap::new(); + + for op in ops.into_iter() { + execute_operation(&mut std_btree, &mut btree, op); + } +} + // Same as `comprehensive` but with the node cache enabled. #[proptest(cases = 10)] fn comprehensive_cached(#[strategy(pvec(operation_strategy(), 100..5_000))] ops: Vec) { @@ -239,9 +254,22 @@ fn iter_count_test(#[strategy(0..250u8)] start: u8, #[strategy(#start..255u8)] e #[proptest] fn no_memory_leaks(#[strategy(pvec(pvec(0..u8::MAX, 100..10_000), 100))] keys: Vec>) { - let mem = make_memory(); - let mut btree = BTreeMap::new(mem); + run_no_memory_leaks(keys, BTreeMap::new(make_memory())); +} + +// Same as `no_memory_leaks` but with the smallest page size allowed, where +// every node spans many overflow pages. +#[proptest(cases = 32)] +fn no_memory_leaks_min_page_size( + #[strategy(pvec(pvec(0..u8::MAX, 100..10_000), 100))] keys: Vec>, +) { + run_no_memory_leaks( + keys, + BTreeMap::new_with_page_size(make_memory(), MINIMUM_PAGE_SIZE), + ); +} +fn run_no_memory_leaks(keys: Vec>, mut btree: BTreeMap, (), M>) { // Insert entries. for k in keys.iter() { btree.insert(k.clone(), ()); diff --git a/src/btreemap/tests.rs b/src/btreemap/tests.rs index 0d7216fcf..ed81290aa 100644 --- a/src/btreemap/tests.rs +++ b/src/btreemap/tests.rs @@ -1586,6 +1586,57 @@ fn accepts_small_or_equal_value_sizes() { let _btree: BTreeMap, Blob<3>, _> = BTreeMap::init(btree.into_memory()); } +#[test] +fn new_with_page_size_persists_the_page_size() { + let mut btree: BTreeMap, Vec, _> = BTreeMap::new_with_page_size(make_memory(), 200); + btree.insert(vec![1], vec![2]); + + let btree: BTreeMap, Vec, _> = BTreeMap::load(btree.into_memory()); + assert_eq!(btree.version, Version::V2(PageSize::Value(200))); + assert_eq!(btree.get(&vec![1]), Some(vec![2])); +} + +#[test] +fn new_with_page_size_overrides_the_page_size_of_bounded_types() { + let btree: BTreeMap = BTreeMap::new_with_page_size(make_memory(), 4096); + assert_eq!(btree.version, Version::V2(PageSize::Value(4096))); +} + +#[test] +#[should_panic(expected = "page_size must be at least 128 bytes, got 127")] +fn new_with_page_size_rejects_pages_below_the_minimum() { + let _btree: BTreeMap, Vec, _> = BTreeMap::new_with_page_size(make_memory(), 127); +} + +#[test] +fn init_with_page_size_creates_a_map_with_the_page_size() { + let btree: BTreeMap, Vec, _> = BTreeMap::init_with_page_size(make_memory(), 256); + assert_eq!(btree.version, Version::V2(PageSize::Value(256))); +} + +#[test] +fn init_with_page_size_keeps_the_page_size_of_an_existing_map() { + let mut btree: BTreeMap, Vec, _> = + BTreeMap::init_with_page_size(make_memory(), 256); + btree.insert(vec![1], vec![2]); + + // The page size is fixed when the map is created, so a different page + // size passed when reinitializing is ignored. + let btree: BTreeMap, Vec, _> = + BTreeMap::init_with_page_size(btree.into_memory(), 1024); + assert_eq!(btree.version, Version::V2(PageSize::Value(256))); + assert_eq!(btree.get(&vec![1]), Some(vec![2])); + + // Likewise for a map created with the default page size. + let btree: BTreeMap, Vec, _> = BTreeMap::new(make_memory()); + let btree: BTreeMap, Vec, _> = + BTreeMap::init_with_page_size(btree.into_memory(), 256); + assert_eq!( + btree.version, + Version::V2(PageSize::Value(DEFAULT_PAGE_SIZE)) + ); +} + fn bruteforce_range_search() { let (key, value) = (K::build, V::build); run_btree_test(|mut stable_map| {