opt: preload needed beatree leaves before update

nomt/src/beatree/allocator/mod.rs

@@ -12,8 +12,8 @@ use free_list::FreeList;
 
 mod free_list;
 
-/// The number of a page.
-#[derive(Debug, Clone, Copy, PartialEq, Eq, PartialOrd, Ord)]
+/// The number of a page
+#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash, PartialOrd, Ord)]
 pub struct PageNumber(pub u32);
 
 impl PageNumber {
@@ -64,15 +64,8 @@ impl AllocatorReader {
 
     /// Returns the page with the specified page number. Blocks the current thread.
     pub fn query(&self, pn: PageNumber) -> Box<Page> {
-        let page = Box::new(Page::zeroed());
-
-        let command = IoCommand {
-            kind: IoKind::Read(self.store_file.as_raw_fd(), pn.0 as u64, page),
-            user_data: 0,
-        };
-
         self.io_handle
-            .send(command)
+            .send(self.io_command(pn, 0))
             .expect("I/O store worker dropped");
 
         // wait for completion
@@ -82,6 +75,21 @@ impl AllocatorReader {
 
         page
     }
+
+    /// Get a reference to the I/O handle.
+    pub fn io_handle(&self) -> &IoHandle {
+        &self.io_handle
+    }
+
+    /// Create an I/O command for querying a page by number.
+    pub fn io_command(&self, pn: PageNumber, user_data: u64) -> IoCommand {
+        let page = Box::new(Page::zeroed());
+
+        IoCommand {
+            kind: IoKind::Read(self.store_file.as_raw_fd(), pn.0 as u64, page),
+            user_data,
+        }
+    }
 }
 
 impl AllocatorWriter {

nomt/src/beatree/leaf/store.rs

@@ -3,7 +3,7 @@ use crate::{
         allocator::{AllocatorCommitOutput, AllocatorReader, AllocatorWriter, PageNumber},
         leaf::node::LeafNode,
     },
-    io::{IoPool, Page},
+    io::{IoPool, IoCommand, IoHandle, Page},
 };
 
 use std::fs::File;
@@ -58,6 +58,15 @@ impl LeafStoreReader {
     pub fn query(&self, pn: PageNumber) -> Box<Page> {
         self.allocator_reader.query(pn)
     }
+
+    pub fn io_handle(&self) -> &IoHandle {
+        self.allocator_reader.io_handle()
+    }
+
+    /// Create an I/O command for querying a page by number.
+    pub fn io_command(&self, pn: PageNumber, user_data: u64) -> IoCommand {
+        self.allocator_reader.io_command(pn, user_data)
+    }
 }
 
 impl LeafStoreWriter {

nomt/src/beatree/ops/update/mod.rs

@@ -1,7 +1,7 @@
 use anyhow::Result;
 use bitvec::prelude::*;
 
-use std::collections::BTreeMap;
+use std::collections::{BTreeMap, HashMap};
 
 use crate::beatree::{
     allocator::PageNumber,
@@ -49,6 +49,8 @@ pub fn update(
     leaf_writer: &mut LeafStoreWriter,
     bbn_writer: &mut bbn::BbnStoreWriter,
 ) -> Result<Vec<BranchId>> {
+    let leaf_pages = preload_leaves(leaf_reader, &bbn_index, &bnp, changeset.keys().cloned())?;
+
     let mut ctx = Ctx {
         bbn_index,
         bbn_writer,
@@ -57,7 +59,7 @@ pub fn update(
         leaf_writer,
     };
 
-    let mut updater = Updater::new(&ctx);
+    let mut updater = Updater::new(&ctx, leaf_pages);
     for (key, value_change) in changeset {
         updater.ingest(*key, value_change.clone(), &mut ctx);
     }
@@ -79,10 +81,11 @@ struct Updater {
     branch_updater: BranchUpdater,
     leaf_updater: LeafUpdater,
     obsolete_branches: Vec<BranchId>,
+    leaf_pages: HashMap<PageNumber, LeafNode>,
 }
 
 impl Updater {
-    fn new(ctx: &Ctx) -> Self {
+    fn new(ctx: &Ctx, mut leaf_pages: HashMap<PageNumber, LeafNode>) -> Self {
         let first = ctx.bbn_index.first();
 
         // UNWRAP: all nodes in index must exist.
@@ -115,9 +118,8 @@ impl Updater {
             })
             .map(|(id, separator)| BaseLeaf {
                 id,
-                node: LeafNode {
-                    inner: ctx.leaf_reader.query(id),
-                },
+                node: leaf_pages.remove(&id)
+                    .unwrap_or_else(|| LeafNode { inner: ctx.leaf_reader.query(id) }),
                 iter_pos: 0,
                 separator,
             });
@@ -140,6 +142,7 @@ impl Updater {
             branch_updater,
             leaf_updater,
             obsolete_branches: Vec::new(),
+            leaf_pages,
         }
     }
 
@@ -209,9 +212,8 @@ impl Updater {
     fn reset_leaf_base(&mut self, target: Key, ctx: &Ctx) -> Result<(), ()> {
         let branch = self.branch_updater.base().ok_or(())?;
         let (i, leaf_pn) = super::search_branch(&branch.node, target).ok_or(())?;
-        let leaf = LeafNode {
-            inner: ctx.leaf_reader.query(leaf_pn),
-        };
+        let leaf = self.leaf_pages.remove(&leaf_pn)
+            .unwrap_or_else(|| LeafNode { inner: ctx.leaf_reader.query(leaf_pn) });
 
         let separator = reconstruct_key(branch.node.prefix(), branch.node.separator(i));
 
@@ -258,3 +260,37 @@ pub fn reconstruct_key(prefix: &BitSlice<u8, Msb0>, separator: &BitSlice<u8, Msb
     key.view_bits_mut::<Msb0>()[prefix.len()..][..separator.len()].copy_from_bitslice(separator);
     key
 }
+
+fn preload_leaves(
+    leaf_reader: &LeafStoreReader,
+    bbn_index: &Index,
+    bnp: &BranchNodePool,
+    keys: impl IntoIterator<Item=Key>,
+) -> Result<HashMap<PageNumber, LeafNode>> {
+    let mut leaf_pages = HashMap::new();
+    let mut last_pn = None;
+
+    let mut submissions = 0;
+    for key in keys {
+        let Some(branch_id) = bbn_index.lookup(key) else { continue };
+        // UNWRAP: all branches in index exist.
+        let branch = bnp.checkout(branch_id).unwrap();
+        let Some((_, leaf_pn)) = super::search_branch(&branch, key) else { continue };
+        if last_pn == Some(leaf_pn) { continue }
+        last_pn = Some(leaf_pn);
+        leaf_reader.io_handle().send(leaf_reader.io_command(leaf_pn, leaf_pn.0 as u64))
+            .expect("I/O Pool Disconnected");
+
+        submissions += 1;
+    }
+
+    for _ in 0..submissions {
+        let completion = leaf_reader.io_handle().recv().expect("I/O Pool Disconnected");
+        completion.result?;
+        let pn = PageNumber(completion.command.user_data as u32);
+        let page = completion.command.kind.unwrap_buf();
+        leaf_pages.insert(pn, LeafNode { inner: page });
+    }
+
+    Ok(leaf_pages)
+}