From 87c466d3eb0072b73dfd879d137b46b0de648ce9 Mon Sep 17 00:00:00 2001 From: Hamish Peebles Date: Tue, 28 Jul 2026 12:36:42 +0100 Subject: [PATCH] perf(btreemap): batch the per-field node reads and writes MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Loading a node performed one memory read per field: every child address, key size, key, and value size separately — dozens of reads per node. A memory access costs a fixed amount per call plus a fee per byte, so the load now fetches the front of the initial page in one read and parses the fields out of the buffer, with any field beyond the read falling back to an individual read as before. The read is sized so that it only ever buys bytes the load will use. For key types bounded by at most 64 bytes the header pins down the extent of the children and key areas, and that estimate is all that is fetched. For larger or unbounded key bounds the interleaved layout offers no extent worth prefetching — the sizes the walk needs sit between payloads it skips — so those types read exactly as before; an earlier draft that read a fixed 1 KiB for them regressed the large-value benchmarks by up to 8%. Saving mirrors this: fields are serialized into a buffer flushed in as few writes as possible, except that slices of 48 bytes or more are written directly, uncopied, flushing whatever is buffered around them. The threshold is measured, and lower than intuition suggests: copying into the buffer costs a few instructions per byte while a separate write costs a fixed amount for the call and its page mapping, which puts the crossover around forty bytes — earlier drafts with a 128- and 512-byte threshold both measured thousands of instructions per save slower on nodes with ~128-byte keys and values. The `NodeWriter` keeps sole ownership of the overflow-address field at bytes [7, 15): it updates the chain pointer in place as overflow pages are allocated and released, so the save skips over it rather than risk clobbering a pointer the writer just installed. The writer's unused field helpers are gone. Full canbench sweep: 147 improved / 0 regressed / 105 unchanged of 252, median -3.9%, best -23.5%; the largest positive movement is +1.8% on a benchmark whose code path is untouched (noise). Memory-manager maps gain the most (-16 to -23%), followed by maps with small keys and values (-10 to -19%). Co-Authored-By: Claude Fable 5 --- src/btreemap/node/io.rs | 16 -- src/btreemap/node/tests.rs | 29 ++++ src/btreemap/node/v2.rs | 292 ++++++++++++++++++++++++++----------- 3 files changed, 233 insertions(+), 104 deletions(-) diff --git a/src/btreemap/node/io.rs b/src/btreemap/node/io.rs index 4a2c97e8..1d43fcf2 100644 --- a/src/btreemap/node/io.rs +++ b/src/btreemap/node/io.rs @@ -281,22 +281,6 @@ impl<'a, M: Memory> NodeWriter<'a, M> { } } - pub fn write_u32(&mut self, offset: Address, val: u32) { - self.write(offset, &val.to_le_bytes()); - } - - pub fn write_u64(&mut self, offset: Address, val: u64) { - self.write(offset, &val.to_le_bytes()); - } - - pub fn write_struct(&mut self, t: &T, addr: Address) { - let slice = unsafe { - core::slice::from_raw_parts(t as *const _ as *const u8, core::mem::size_of::()) - }; - - self.write(addr, slice) - } - // Allocates a new page and appends it to the node's overflows. fn allocate_new_page(&mut self) { let new_page = self.allocator.allocate(); diff --git a/src/btreemap/node/tests.rs b/src/btreemap/node/tests.rs index ff1158ca..ffa46d10 100644 --- a/src/btreemap/node/tests.rs +++ b/src/btreemap/node/tests.rs @@ -309,3 +309,32 @@ fn can_call_node_value_multiple_times_on_same_index() { let value2 = node.value(0, &mem); assert_eq!(value1, value2); } + +#[test] +fn loading_near_the_end_of_memory_clamps_the_batched_read() { + use crate::storable::Blob; + + // A bounded-key node placed so that the load's batched-read estimate extends past + // the last byte of memory: the read must clamp to the memory's end rather than + // trap. (The node's own bytes always fit — the save wrote them — it is only the + // estimate, an upper bound based on the key bound, that can overshoot.) + let mem = make_memory(); + let page_size = PageSize::Value(128); + // The allocator's header write grows the memory to exactly one Wasm page. + let mut allocator = Allocator::new(mem.clone(), Address::from(0), Bytes::from(128u64)); + + // The node occupies 27 bytes (header 15, key size 4, key 4, value size 4), and its + // batched-read estimate is 55 bytes (header 15, key bound 32 + size field 4, value + // size 4). Placing it 30 bytes from the end fits the node but not the estimate. + let node_addr = Address::from(crate::WASM_PAGE_SIZE - 30); + let mut node: Node> = Node::new_v2(node_addr, NodeType::Leaf, page_size); + node.push_entry((Blob::try_from(&[7u8; 4][..]).unwrap(), vec![])); + node.save_v2(&mut allocator); + assert_eq!(mem.size(), 1); + + let node = Node::>::load(node_addr, page_size, &mem); + assert_eq!( + node.entries(&mem), + vec![(Blob::try_from(&[7u8; 4][..]).unwrap(), vec![])] + ); +} diff --git a/src/btreemap/node/v2.rs b/src/btreemap/node/v2.rs index 2037c0a9..c2aea1f8 100644 --- a/src/btreemap/node/v2.rs +++ b/src/btreemap/node/v2.rs @@ -73,7 +73,7 @@ use super::*; use crate::btreemap::Allocator; -use crate::{btreemap::node::io::NodeWriter, types::NULL}; +use crate::{btreemap::node::io::NodeWriter, types::NULL, WASM_PAGE_SIZE}; // Initial page pub(super) const OVERFLOW_ADDRESS_OFFSET: Bytes = Bytes::new(7); @@ -89,6 +89,32 @@ pub(super) const PAGE_OVERFLOW_DATA_OFFSET: Bytes = Bytes::new(11); // addresses (88 bytes). We round that up to 128 to get a nice binary number. const MINIMUM_PAGE_SIZE: u32 = 128; +// How much of the initial page a load reads in one batched read. A memory access costs a +// fixed amount per call plus a fee per byte, so batching many small reads into one pays +// for itself only while the batch stays small; a derived page can be many megabytes, and +// reading all of it to serve a handful of fields is a large regression. 1 KiB covers the +// header, the children and the key area of typical nodes. +const INITIAL_PAGE_READ_CAP: u64 = 1024; + +// Keys are included in the batched read only while their bound keeps the read small. +// The walk below needs the size fields, which interleave with the key payloads, so +// covering the sizes means also buying the key bytes between them — bytes the lazy key +// loads may never look at. Up to this bound the size fields the read covers save more +// than the skipped key bytes cost; beyond it the read would mostly buy ignored bytes, +// which measured as a regression on the large-key benchmarks. +const BATCHED_READ_MAX_KEY_SIZE: u32 = 64; + +// Slices at least this large are written directly rather than copied into the save +// buffer. Copying a slice into the buffer costs a few instructions per byte, while a +// separate write costs a fixed amount for the call and its page mapping — roughly a +// forty-byte copy. Below the threshold coalescing neighboring fields into one write +// wins; above it the copy costs more than the write it saves. +const DIRECT_WRITE_THRESHOLD: usize = 48; + +// The save buffer's initial capacity. Only slices below `DIRECT_WRITE_THRESHOLD` land in +// the buffer, so it stays small; this covers the typical node without regrowing. +const WRITE_BUFFER_INITIAL_CAPACITY: usize = 1024; + impl Node { /// Creates a new v2 node at the given address. pub fn new_v2(address: Address, node_type: NodeType, page_size: PageSize) -> Node { @@ -114,18 +140,6 @@ impl Node { #[cfg(feature = "bench_scope")] let _p = canbench_rs::bench_scope("node_load_v2"); // May add significant overhead. - // Load the node, including any overflows, into a buffer. - let overflows = read_overflows(address, memory); - - let reader = NodeReader { - address, - overflows: &overflows, - page_size, - memory, - }; - - let mut offset = Address::from(0); - // Load the node type. let node_type = match header.node_type { LEAF_NODE_TYPE => NodeType::Leaf, @@ -136,27 +150,96 @@ impl Node { // Load the number of entries let num_entries = header.num_entries as usize; + // Read the front of the initial page in one go. Almost everything the load needs + // — the children, the key sizes and small keys, and the value sizes — lives + // there, so this replaces a separate memory read per field. A memory read costs + // a fixed amount per call plus a fee per byte, so the read is kept as small as + // its purpose allows: for key types bounded small enough (see + // `BATCHED_READ_MAX_KEY_SIZE`) the header pins down the extent of the children + // and key areas, and the value sizes follow immediately when the values are + // empty, so that estimate is all that is fetched (the cap cannot bind under the + // current `CAPACITY` and key-size gate; it guards changes to either). For + // large or unbounded key types nothing worth reading ahead bounds the entry + // area, so only the metadata ahead of it is batched and the entry walk reads + // individually, as it did before batching. Fields beyond the read fall back to + // individual reads. The read is also clamped to the end of the memory, since a + // node's tail may extend past the last byte ever written. + let page = { + let children_bytes = match node_type { + NodeType::Internal => (num_entries as u64 + 1) * Address::size().get(), + NodeType::Leaf => 0, + }; + let target = match K::BOUND { + crate::storable::Bound::Bounded { + max_size, + is_fixed_size, + } if max_size <= BATCHED_READ_MAX_KEY_SIZE => { + let key_size_field = if is_fixed_size { 0 } else { U32_SIZE.get() }; + let keys_bytes = num_entries as u64 * (max_size as u64 + key_size_field); + let value_sizes_bytes = num_entries as u64 * U32_SIZE.get(); + (ENTRIES_OFFSET.get() + children_bytes + keys_bytes + value_sizes_bytes) + .min(INITIAL_PAGE_READ_CAP) + } + // Nothing bounds the entry area for large or unbounded key types, so + // no read pays for itself up front; every field falls back to an + // individual read, exactly as before batching. + _ => 0, + }; + let memory_bytes = memory.size() * WASM_PAGE_SIZE; + let len = (page_size.get() as u64) + .min(target) + .min(memory_bytes.saturating_sub(address.get())); + let mut page = vec![]; + if len > 0 { + read_to_vec(memory, address, &mut page, len as usize); + } + page + }; + + // Reads a little-endian u32 from the page, falling back to a memory read for the + // rare node whose fields extend past the initial page. + let page_u32 = |reader: &NodeReader, offset: Address| -> u32 { + let start = offset.get() as usize; + match page.get(start..start + 4) { + Some(bytes) => u32::from_le_bytes(bytes.try_into().unwrap()), + None => read_u32(reader, offset), + } + }; + + // Load the addresses of the node's overflow pages, if any. + let first_overflow = { + let start = OVERFLOW_ADDRESS_OFFSET.get() as usize; + Address::from(match page.get(start..start + 8) { + Some(bytes) => u64::from_le_bytes(bytes.try_into().unwrap()), + None => read_u64(memory, address + OVERFLOW_ADDRESS_OFFSET), + }) + }; + let overflows = read_overflows(first_overflow, memory); + + let reader = NodeReader { + address, + overflows: &overflows, + page_size, + memory, + }; + + let mut offset = Address::from(0); + // Load children if this is an internal node. offset += ENTRIES_OFFSET; let children = if node_type == NodeType::Internal { - // Empirical constant from benchmarks: individual reads are faster for ≤4 children. - // From 5 up, batch reads consistently perform better. - const CHILDREN_BATCH_READ_THRESHOLD: usize = 4; let count = num_entries + 1; - if count <= CHILDREN_BATCH_READ_THRESHOLD { - // For very small counts, use individual reads to avoid allocation overhead - let mut children = Vec::with_capacity(count); - for _ in 0..count { - children.push(Address::from(read_u64(&reader, offset))); - offset += Address::size(); - } - children - } else { - // For larger counts, use the optimized batch read - let children = read_address_vec(&reader, offset, count); - offset += Bytes::from(count) * Address::size(); - children - } + let byte_len = count * Address::size().get() as usize; + let start = offset.get() as usize; + let children = match page.get(start..start + byte_len) { + Some(bytes) => bytes + .chunks_exact(Address::size().get() as usize) + .map(|chunk| Address::from(u64::from_le_bytes(chunk.try_into().unwrap()))) + .collect(), + None => read_address_vec(&reader, offset, count), + }; + offset += Bytes::from(byte_len as u64); + children } else { vec![] }; @@ -173,20 +256,27 @@ impl Node { let key_size = if K::BOUND.is_fixed_size() { K::BOUND.max_size() } else { - let size = read_u32(&reader, offset); + let size = page_u32(&reader, offset); offset += U32_SIZE; size }; // Eager-load small keys, defer large ones. let key = if key_size <= EAGER_LOAD_KEY_SIZE_THRESHOLD { - read_to_vec( - &reader, - Address::from(offset.get()), - &mut buf, - key_size as usize, - ); - LazyKey::by_value(K::from_bytes(Cow::Borrowed(&buf))) + let start = offset.get() as usize; + let key_bytes = match page.get(start..start + key_size as usize) { + Some(bytes) => bytes, + None => { + read_to_vec( + &reader, + Address::from(offset.get()), + &mut buf, + key_size as usize, + ); + &buf[..] + } + }; + LazyKey::by_value(K::from_bytes(Cow::Borrowed(key_bytes))) } else { LazyKey::by_ref(key_offset, key_size) }; @@ -198,7 +288,7 @@ impl Node { // Load the values for (_key, value) in entries.iter_mut() { // Load the values lazily. - let value_size = read_u32(&reader, offset); + let value_size = page_u32(&reader, offset); *value = LazyValue::by_ref(Bytes::from(offset.get()), value_size); offset += U32_SIZE + Bytes::from(value_size as u64); } @@ -227,8 +317,10 @@ impl Node { self.entry(i, allocator.memory()); } - // Initialize a NodeWriter. The NodeWriter takes care of allocating/deallocating - // overflow pages as needed. + // The writer takes care of allocating and deallocating overflow pages as + // needed; the buffer below batches the many small field writes into few memory + // writes while passing large slices (big keys and values) straight through, + // uncopied. let mut writer = NodeWriter::new( self.address, self.overflows.clone(), @@ -236,69 +328,93 @@ impl Node { allocator, ); - // Write the header of the initial page. - let mut offset = Address::from(0); - let header = NodeHeader { - magic: *MAGIC, - version: LAYOUT_VERSION_2, - node_type: match self.node_type { - NodeType::Leaf => LEAF_NODE_TYPE, - NodeType::Internal => INTERNAL_NODE_TYPE, - }, - num_entries: self.entries.len() as u16, - }; - - writer.write_struct(&header, offset); - offset += NodeHeader::size(); - - // Add a null overflow address. - // This might get overwritten later in case the node does overflow. - writer.write_u64(offset, self.overflows.first().unwrap_or(&NULL).get()); - offset += Bytes::from(8u64); - - // Write the children. - let count = self.children.len(); - let byte_len = count * Address::size().get() as usize; - let mut bytes = Vec::with_capacity(byte_len); + // The node is serialized into `buf` and written out in a handful of writes + // instead of one write per field. `buf` always holds the pending bytes for the + // contiguous range starting at `buf_start`; slices of `DIRECT_WRITE_THRESHOLD` + // bytes or more skip it — flushing what is pending first — so large keys and + // values go straight to memory without the copy. + let mut buf = Vec::with_capacity(WRITE_BUFFER_INITIAL_CAPACITY); + let mut buf_start = Address::from(0); + + // The header. + buf.extend_from_slice(MAGIC); + buf.push(LAYOUT_VERSION_2); + buf.push(match self.node_type { + NodeType::Leaf => LEAF_NODE_TYPE, + NodeType::Internal => INTERNAL_NODE_TYPE, + }); + buf.extend_from_slice(&(self.entries.len() as u16).to_le_bytes()); + + // Flush at the overflow address, which lives at [7, 15) and is skipped over: it + // is owned by the `NodeWriter`, which updates it in place as overflow pages are + // allocated and released, and a write from here could clobber a chain pointer + // the writer just installed. Never writing it from here is sufficient because + // `finish` always leaves the field consistent: it writes the `NULL` terminator + // into it whenever the node ends up with no overflow pages, so even a fresh + // node at a reused address cannot retain a stale pointer. + debug_assert_eq!(buf.len() as u64, OVERFLOW_ADDRESS_OFFSET.get()); + writer.write(buf_start, &buf); + buf.clear(); + buf_start = Address::from(ENTRIES_OFFSET.get()); + + // The children. for child in &self.children { - bytes.extend_from_slice(&child.get().to_le_bytes()); + buf.extend_from_slice(&child.get().to_le_bytes()); } - writer.write(offset, &bytes); - offset += Bytes::from(byte_len); - // Write the keys. + // The keys, sized unless fixed in size. for i in 0..self.entries.len() { - let key = self.key(i, writer.memory()); - let key_bytes = key.to_bytes_checked(); - - // Write the size of the key if it isn't fixed in size. + let key_bytes = self.key(i, writer.memory()).to_bytes_checked(); if !K::BOUND.is_fixed_size() { - writer.write_u32(offset, key_bytes.len() as u32); - offset += U32_SIZE; + buf.extend_from_slice(&(key_bytes.len() as u32).to_le_bytes()); + } + let key_bytes: &[u8] = key_bytes.borrow(); + if key_bytes.len() >= DIRECT_WRITE_THRESHOLD { + buf_start = write_direct(&mut writer, buf_start, &mut buf, key_bytes); + } else { + buf.extend_from_slice(key_bytes); } - - // Write the key. - writer.write(offset, key_bytes.borrow()); - offset += Bytes::from(key_bytes.len()); } - // Write the values. + // The values, sized. for i in 0..self.entries.len() { - // Write the size of the value. let value = self.value(i, writer.memory()); - writer.write_u32(offset, value.len() as u32); - offset += U32_SIZE; + buf.extend_from_slice(&(value.len() as u32).to_le_bytes()); + if value.len() >= DIRECT_WRITE_THRESHOLD { + buf_start = write_direct(&mut writer, buf_start, &mut buf, value); + } else { + buf.extend_from_slice(value); + } + } - // Write the value. - writer.write(offset, value); - offset += Bytes::from(value.len()); + if !buf.is_empty() { + writer.write(buf_start, &buf); } self.overflows = writer.finish(); } } -fn read_overflows(address: Address, memory: &M) -> Vec
{ +/// Flushes the save buffer if it holds anything, writes `piece` directly after it, and +/// returns the new buffer start: just past `piece`. +fn write_direct( + writer: &mut NodeWriter<'_, M>, + buf_start: Address, + buf: &mut Vec, + piece: &[u8], +) -> Address { + let piece_start = buf_start + Bytes::from(buf.len() as u64); + if !buf.is_empty() { + writer.write(buf_start, buf); + buf.clear(); + } + writer.write(piece_start, piece); + piece_start + Bytes::from(piece.len() as u64) +} + +/// Walks the overflow-page chain starting from `first_overflow`, which the caller has +/// already read out of the node's initial page. +fn read_overflows(first_overflow: Address, memory: &M) -> Vec
{ #[repr(C, packed)] struct OverflowPageHeader { magic: [u8; 3], @@ -306,7 +422,7 @@ fn read_overflows(address: Address, memory: &M) -> Vec
{ } let mut overflows = vec![]; - let mut overflow = Address::from(read_u64(memory, address + OVERFLOW_ADDRESS_OFFSET)); + let mut overflow = first_overflow; while overflow != NULL { overflows.push(overflow);