feat: page cache upper level prepopulation

nomt/src/lib.rs

@@ -276,6 +276,12 @@ impl<T: HashAlgorithm> Nomt<T> {
         let root_page = store.load_page(ROOT_PAGE_ID)?;
         let page_cache = PageCache::new(root_page, &o, metrics.clone());
         let root = compute_root_node::<T>(&page_cache, &store);
+
+        if o.prepopulate_page_cache {
+            let io_handle = store.io_pool().make_handle();
+            merkle::prepopulate_cache(io_handle, &page_cache, &store, 3)?;
+        }
+
         Ok(Self {
             merkle_update_pool: UpdatePool::new(o.commit_concurrency, o.warm_up),
             page_cache,

nomt/src/merkle/cache_prepopulate.rs

@@ -0,0 +1,93 @@
+//! Utility for prepopulating the first N layers of the cache.
+
+use crate::{
+    io::IoHandle,
+    page_cache::{PageCache, PageMut},
+    store::{PageLoad, PageLoader, Store},
+};
+
+use nomt_core::page_id::{ChildPageIndex, PageId, MAX_PAGE_DEPTH, NUM_CHILDREN, ROOT_PAGE_ID};
+
+/// Prepopulate the given number of levels of the page tree into the page cache.
+///
+/// This function blocks until the prepopulation has finished.
+pub fn prepopulate(
+    io_handle: IoHandle,
+    page_cache: &PageCache,
+    store: &Store,
+    levels: usize,
+) -> anyhow::Result<()> {
+    let page_loader = store.page_loader();
+    let mut loads = Vec::new();
+
+    let levels = std::cmp::min(levels, MAX_PAGE_DEPTH);
+
+    // dispatch all page loads recursively.
+    dispatch_recursive(ROOT_PAGE_ID, &page_loader, &io_handle, &mut loads, levels)?;
+
+    let mut completed = 0;
+
+    // wait on I/O results.
+    while completed < loads.len() {
+        let complete_io = io_handle.recv()?;
+        complete_io.result?;
+        let load_index = complete_io.command.user_data as usize;
+        let load = &mut loads[load_index];
+
+        // UNWRAP: all submitted requests are of kind Read(FatPage).
+        if let Some((page, bucket)) = load.try_complete(complete_io.command.kind.unwrap_buf()) {
+            completed += 1;
+            page_cache.insert(
+                load.page_id().clone(),
+                PageMut::pristine_with_data(page).freeze(),
+                bucket,
+            );
+        } else {
+            // misprobe. try again.
+            if !page_loader.probe(load, &io_handle, complete_io.command.user_data)? {
+                // guaranteed empty.
+                completed += 1;
+            }
+        }
+    }
+
+    Ok(())
+}
+
+// dispatch page loads for all the children of the given page.
+fn dispatch_recursive(
+    page_id: PageId,
+    page_loader: &PageLoader,
+    io_handle: &IoHandle,
+    loads: &mut Vec<PageLoad>,
+    levels_remaining: usize,
+) -> anyhow::Result<()> {
+    if levels_remaining == 0 {
+        return Ok(());
+    }
+
+    for child_index in 0..NUM_CHILDREN {
+        // UNWRAP: all indices up to NUM_CHILDREN are allowed.
+        let child_index = ChildPageIndex::new(child_index as u8).unwrap();
+
+        // UNWRAP: depth is not out of bounds and child index is valid.
+        let child_page_id = page_id.child_page_id(child_index).unwrap();
+
+        let mut page_load = page_loader.start_load(child_page_id.clone());
+
+        let next_index = loads.len() as u64;
+        if page_loader.probe(&mut page_load, io_handle, next_index)? {
+            // probe has been dispatched.
+            loads.push(page_load);
+            dispatch_recursive(
+                child_page_id,
+                page_loader,
+                io_handle,
+                loads,
+                levels_remaining - 1,
+            )?;
+        }
+    }
+
+    Ok(())
+}

nomt/src/merkle/mod.rs

@@ -27,11 +27,13 @@ use crate::{
 };
 use threadpool::ThreadPool;
 
+mod cache_prepopulate;
 mod page_set;
 mod page_walker;
 mod seek;
 mod worker;
 
+pub use cache_prepopulate::prepopulate as prepopulate_cache;
 pub use page_walker::UpdatedPage;
 
 /// Updated pages produced by update workers.

nomt/src/options.rs

@@ -25,6 +25,9 @@ pub struct Options {
     /// The maximum size of the leaf cache specified in MiB, rounded down
     /// to the nearest byte multiple of [`crate::io::PAGE_SIZE`].
     pub(crate) leaf_cache_size: usize,
+    /// Whether to prepopulate the upper layers of the page cache on startup.
+    /// This incurs some I/O on startup but leads to predictable worst-case performance.
+    pub(crate) prepopulate_page_cache: bool,
 }
 
 impl Options {
@@ -48,6 +51,7 @@ impl Options {
             preallocate_ht: true,
             page_cache_size: 256,
             leaf_cache_size: 256,
+            prepopulate_page_cache: false,
         }
     }
 
@@ -150,6 +154,14 @@ impl Options {
     pub fn leaf_cache_size(&mut self, leaf_cache_size: usize) {
         self.leaf_cache_size = leaf_cache_size;
     }
+
+    /// Sets whether to prepopulate the upper layers of the page cache on startup.
+    ///
+    /// This incurs some I/O on startup but leads to predictable worst-case performance.
+    /// Default: false
+    pub fn prepopulate_page_cache(&mut self, prepopulate: bool) {
+        self.prepopulate_page_cache = prepopulate;
+    }
 }
 
 #[test]