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);