Menu
Public repository / master
AkurAI-Framework / crates/storage/src/tree.rs
Latest change 35c1784a2189dce748f85504e491110cffe15ebf - OLI-13: clippy 1.98 fixes (as_chunks, mem::take) by Ólafur Búi Ólafsson
//! A copy-on-write B+tree over the [`Pager`].
//!
//! Every mutation rewrites the path from the touched leaf back to the root into
//! freshly allocated pages, leaving the old pages untouched. The old tree stays
//! the committed truth until [`BTree::commit`] swaps the root pointer in the
//! meta page (see [`crate::meta`]) and fsyncs. A crash before that swap simply
//! leaves the previous tree intact — no torn updates, no WAL.
//!
//! Pages the rewrite orphans are recycled. They are held in `pending_free`
//! until commit (they still belong to the *committed* tree until the swap
//! lands), then folded into the durable free list. Allocation draws from the
//! committed free list first and extends the file only when it is empty.
//!
//! Keys and values are arbitrary byte strings; keys order lexicographically.
//! Deletion removes the entry and copies the path but does not yet merge
//! underfull nodes — a documented, correctness-preserving simplification
//! (rebalancing is a later step on the roadmap).
use std::collections::HashSet;
use std::io;
use std::path::Path;
use crate::meta::{Meta, FIRST_TREE_PAGE, META_SLOT_0, META_SLOT_1};
use crate::node::{Node, Val};
use crate::pager::{PageId, Pager, PAGE_SIZE};
/// A persistent, crash-safe ordered map backed by a single file.
pub struct BTree {
pager: Pager,
root: PageId,
txn_id: u64,
/// Pages free in the committed tree, available to hand out now.
free: Vec<PageId>,
/// Pages orphaned by the in-flight transaction; freed at commit.
pending_free: Vec<PageId>,
}
/// Outcome of inserting into a subtree: either the subtree was rewritten in
/// place, or it split into two and a separator must rise to the parent.
enum Ins {
Updated(PageId),
Split {
left: PageId,
sep: Vec<u8>,
right: PageId,
},
}
/// Result of a bounded range scan: collected items and an optional key for
/// resumption. When `continuation` is `Some`, the caller should use it as the
/// next `start` (inclusive) to fetch the following batch without repeats.
pub struct RangeBounded {
pub items: Vec<(Vec<u8>, Vec<u8>)>,
/// When `Some`, more entries exist. Resume with this as the next `start`.
pub continuation: Option<Vec<u8>>,
}
/// Internal signal from the recursive range walk.
enum RangeOutcome {
/// The range was fully exhausted.
Done,
/// A budget was hit; the walker should stop. The payload is the key that
/// the next `range_bounded` call must use as its inclusive `start`.
LimitReached(Vec<u8>),
/// A budget was hit exactly at the end of a leaf. The parent interior-node
/// handler must resolve this by finding the first key of the next leaf and
/// turning it into a `LimitReached`.
#[allow(dead_code)]
NeedNext,
}
impl BTree {
/// Open the tree at `path`, creating an empty one if the file is new.
pub fn open(path: impl AsRef<Path>) -> io::Result<BTree> {
let mut pager = Pager::open(path)?;
if pager.page_count()? == 0 {
return Self::initialize(pager);
}
let meta = Self::load_meta(&mut pager)?;
Ok(BTree {
pager,
root: meta.root,
txn_id: meta.txn_id,
free: meta.free,
pending_free: Vec::new(),
})
}
/// Lay down a fresh database: reserve the two meta slots, write an empty
/// leaf root, and commit both meta slots so either survives a first crash.
fn initialize(mut pager: Pager) -> io::Result<BTree> {
// Pages must be created in order; reserve slot 0, slot 1, then root.
let s0 = pager.allocate()?;
let s1 = pager.allocate()?;
debug_assert_eq!((s0, s1), (META_SLOT_0, META_SLOT_1));
let root = pager.allocate()?;
debug_assert_eq!(root, FIRST_TREE_PAGE);
pager.write_page(root, &Node::Leaf(Vec::new()).encode())?;
// Seed both slots; the odd-txn slot wins, so it is the live one.
pager.write_page(
META_SLOT_0,
&Meta {
txn_id: 0,
root,
free: vec![],
}
.encode(),
)?;
pager.write_page(
META_SLOT_1,
&Meta {
txn_id: 1,
root,
free: vec![],
}
.encode(),
)?;
pager.sync()?;
Ok(BTree {
pager,
root,
txn_id: 1,
free: Vec::new(),
pending_free: Vec::new(),
})
}
/// Read both meta slots and choose the valid one with the highest txn id.
fn load_meta(pager: &mut Pager) -> io::Result<Meta> {
let m0 = Meta::decode(&pager.read_page(META_SLOT_0)?);
let m1 = Meta::decode(&pager.read_page(META_SLOT_1)?);
match (m0, m1) {
(Some(a), Some(b)) => Ok(if a.txn_id >= b.txn_id { a } else { b }),
(Some(a), None) => Ok(a),
(None, Some(b)) => Ok(b),
(None, None) => Err(io::Error::new(
io::ErrorKind::InvalidData,
"both meta slots are corrupt — database is unrecoverable",
)),
}
}
// ---- public map operations -------------------------------------------
/// Look up `key`, returning a copy of its value if present.
pub fn get(&mut self, key: &[u8]) -> io::Result<Option<Vec<u8>>> {
let mut id = self.root;
let mut visited = HashSet::new();
loop {
if !visited.insert(id) {
return Err(io::Error::new(
io::ErrorKind::InvalidData,
"cycle in BTree page graph",
));
}
match self.read_node(id)? {
Node::Leaf(entries) => {
return match entries.binary_search_by(|(k, _)| k.as_slice().cmp(key)) {
Ok(i) => Ok(Some(self.load_value(&entries[i].1)?)),
Err(_) => Ok(None),
};
}
Node::Interior { keys, children } => {
id = children[route(&keys, key)];
}
}
}
}
/// Insert or replace `key`'s value. Durable only after [`BTree::commit`].
pub fn insert(&mut self, key: &[u8], value: &[u8]) -> io::Result<()> {
match self.insert_rec(self.root, key, value)? {
Ins::Updated(new_root) => self.root = new_root,
Ins::Split { left, sep, right } => {
self.root = self.write_node(&Node::Interior {
keys: vec![sep],
children: vec![left, right],
})?;
}
}
Ok(())
}
/// Remove `key`. Returns whether it was present. Durable after `commit`.
pub fn delete(&mut self, key: &[u8]) -> io::Result<bool> {
let (new_root, removed) = self.delete_rec(self.root, key)?;
if removed {
self.root = new_root;
}
Ok(removed)
}
/// Collect every `(key, value)` with `start <= key < end`, in key order.
pub fn range(&mut self, start: &[u8], end: &[u8]) -> io::Result<Vec<(Vec<u8>, Vec<u8>)>> {
let mut out = Vec::new();
self.range_rec(self.root, start, end, &mut out, &mut HashSet::new(), false)?;
Ok(out)
}
/// Best-effort range scan for repairing a corrupt page graph.
///
/// Repeated pages and unreadable overflow values are skipped. Normal callers
/// should use [`BTree::range`], which reports corruption instead.
pub fn salvage_range(
&mut self,
start: &[u8],
end: &[u8],
) -> io::Result<Vec<(Vec<u8>, Vec<u8>)>> {
let mut out = Vec::new();
self.range_rec(self.root, start, end, &mut out, &mut HashSet::new(), true)?;
Ok(out)
}
/// Bounded range scan: collect keys in `[start, end)` in key order, stopping
/// early when the next item would exceed either budget. The returned
/// [`RangeBounded`] carries a continuation key suitable for resuming the
/// scan without repeating items.
///
/// * `item_limit` — maximum number of items to return. Zero means return
/// nothing and give back a continuation (if any in-range key exists).
/// * `byte_limit` — maximum total bytes across all keys and values combined.
///
/// Both limits are checked _before_ each item is collected (pre-flight), so
/// the result never exceeds either limit and a zero limit returns zero items.
pub fn range_bounded(
&mut self,
start: &[u8],
end: &[u8],
item_limit: usize,
byte_limit: usize,
) -> io::Result<RangeBounded> {
let mut items = Vec::new();
let mut byte_count: usize = 0;
let outcome = self.range_bounded_rec(
self.root,
start,
end,
&mut items,
&mut byte_count,
item_limit,
byte_limit,
&mut HashSet::new(),
)?;
let continuation = match outcome {
RangeOutcome::Done => None,
RangeOutcome::LimitReached(k) if (k.as_slice() < end) => Some(k),
RangeOutcome::LimitReached(_) | RangeOutcome::NeedNext => None,
};
Ok(RangeBounded {
items,
continuation,
})
}
/// Make every mutation since the last commit durable, then atomically swap
/// in the new root by writing the next meta slot.
pub fn commit(&mut self) -> io::Result<()> {
// 1. The new tree's data pages must hit disk before the meta names them.
self.pager.sync()?;
// 2. Now the orphaned pages are safe to recycle; fold them in (deduped,
// so a page can never be handed out twice).
for id in std::mem::take(&mut self.pending_free) {
if !self.free.contains(&id) {
self.free.push(id);
}
}
// 3. Swap the root by writing the slot for the next transaction, then
// fsync so the swap itself is durable.
let meta = Meta {
txn_id: self.txn_id + 1,
root: self.root,
free: self.free.clone(),
};
self.pager.write_page(meta.slot(), &meta.encode())?;
self.pager.sync()?;
self.txn_id = meta.txn_id;
Ok(())
}
// ---- recursion --------------------------------------------------------
fn insert_rec(&mut self, id: PageId, key: &[u8], value: &[u8]) -> io::Result<Ins> {
match self.read_node(id)? {
Node::Leaf(mut entries) => {
let val = self.store_value(value)?;
match entries.binary_search_by(|(k, _)| k.as_slice().cmp(key)) {
Ok(i) => entries[i].1 = val,
Err(i) => entries.insert(i, (key.to_vec(), val)),
}
self.free_page(id);
self.write_leaf_split(entries)
}
Node::Interior { keys, mut children } => {
let idx = route(&keys, key);
let outcome = self.insert_rec(children[idx], key, value)?;
self.free_page(id);
match outcome {
Ins::Updated(child) => {
children[idx] = child;
Ok(Ins::Updated(
self.write_node(&Node::Interior { keys, children })?,
))
}
Ins::Split { left, sep, right } => {
let mut keys = keys;
children[idx] = left;
children.insert(idx + 1, right);
keys.insert(idx, sep);
self.write_interior_split(keys, children)
}
}
}
}
}
/// Write a leaf, splitting in half if it overflows a page.
fn write_leaf_split(&mut self, entries: Vec<(Vec<u8>, Val)>) -> io::Result<Ins> {
let node = Node::Leaf(entries);
if node.fits() {
return Ok(Ins::Updated(self.write_node(&node)?));
}
let Node::Leaf(entries) = node else {
unreachable!()
};
let mid = (1..entries.len())
.filter(|&mid| {
Node::Leaf(entries[..mid].to_vec()).fits()
&& Node::Leaf(entries[mid..].to_vec()).fits()
})
.min_by_key(|&mid| {
Node::Leaf(entries[..mid].to_vec())
.encoded_len()
.max(Node::Leaf(entries[mid..].to_vec()).encoded_len())
})
.ok_or_else(|| {
io::Error::new(
io::ErrorKind::InvalidInput,
"leaf entry exceeds page capacity",
)
})?;
let sep = entries[mid].0.clone();
let left = self.write_node(&Node::Leaf(entries[..mid].to_vec()))?;
let right = self.write_node(&Node::Leaf(entries[mid..].to_vec()))?;
Ok(Ins::Split { left, sep, right })
}
/// Write an interior node, splitting and lifting a separator if it overflows.
fn write_interior_split(
&mut self,
keys: Vec<Vec<u8>>,
children: Vec<PageId>,
) -> io::Result<Ins> {
let node = Node::Interior { keys, children };
if node.fits() {
return Ok(Ins::Updated(self.write_node(&node)?));
}
let Node::Interior { keys, children } = node else {
unreachable!()
};
let mid = (0..keys.len())
.filter(|&mid| {
Node::Interior {
keys: keys[..mid].to_vec(),
children: children[..=mid].to_vec(),
}
.fits()
&& Node::Interior {
keys: keys[mid + 1..].to_vec(),
children: children[mid + 1..].to_vec(),
}
.fits()
})
.min_by_key(|&mid| {
Node::Interior {
keys: keys[..mid].to_vec(),
children: children[..=mid].to_vec(),
}
.encoded_len()
.max(
Node::Interior {
keys: keys[mid + 1..].to_vec(),
children: children[mid + 1..].to_vec(),
}
.encoded_len(),
)
})
.ok_or_else(|| {
io::Error::new(
io::ErrorKind::InvalidInput,
"interior key exceeds page capacity",
)
})?;
let sep = keys[mid].clone();
let left = self.write_node(&Node::Interior {
keys: keys[..mid].to_vec(),
children: children[..=mid].to_vec(),
})?;
let right = self.write_node(&Node::Interior {
keys: keys[mid + 1..].to_vec(),
children: children[mid + 1..].to_vec(),
})?;
Ok(Ins::Split { left, sep, right })
}
fn delete_rec(&mut self, id: PageId, key: &[u8]) -> io::Result<(PageId, bool)> {
match self.read_node(id)? {
Node::Leaf(mut entries) => {
match entries.binary_search_by(|(k, _)| k.as_slice().cmp(key)) {
Ok(i) => {
entries.remove(i);
self.free_page(id);
Ok((self.write_node(&Node::Leaf(entries))?, true))
}
Err(_) => Ok((id, false)), // unchanged: no CoW, no orphan
}
}
Node::Interior { keys, mut children } => {
let idx = route(&keys, key);
let (child, removed) = self.delete_rec(children[idx], key)?;
if !removed {
return Ok((id, false));
}
children[idx] = child;
self.free_page(id);
Ok((self.write_node(&Node::Interior { keys, children })?, true))
}
}
}
fn range_rec(
&mut self,
id: PageId,
start: &[u8],
end: &[u8],
out: &mut Vec<(Vec<u8>, Vec<u8>)>,
visited: &mut HashSet<PageId>,
salvage: bool,
) -> io::Result<()> {
if !visited.insert(id) {
if salvage {
return Ok(());
}
return Err(io::Error::new(
io::ErrorKind::InvalidData,
"cycle in BTree page graph",
));
}
match self.read_node(id)? {
Node::Leaf(entries) => {
for (k, v) in entries {
if k.as_slice() >= start && k.as_slice() < end {
match self.load_value(&v) {
Ok(value) => out.push((k, value)),
Err(_) if salvage => {}
Err(error) => return Err(error),
}
}
}
Ok(())
}
Node::Interior { keys, children } => {
for i in 0..children.len() {
let lo = if i == 0 {
None
} else {
Some(keys[i - 1].as_slice())
};
let hi = keys.get(i).map(|k| k.as_slice());
if hi.is_some_and(|h| h <= start) {
continue;
}
if lo.is_some_and(|l| l >= end) {
break;
}
self.range_rec(children[i], start, end, out, visited, salvage)?;
}
Ok(())
}
}
}
/// Recursive core of [`BTree::range_bounded`].
///
/// Stops traversal as soon as either budget is exhausted, returning a
/// [`RangeOutcome`] that either names the continuation key or signals that
/// the parent must resolve the next leaf.
#[allow(clippy::too_many_arguments)]
fn range_bounded_rec(
&mut self,
id: PageId,
start: &[u8],
end: &[u8],
items: &mut Vec<(Vec<u8>, Vec<u8>)>,
byte_count: &mut usize,
item_limit: usize,
byte_limit: usize,
visited: &mut HashSet<PageId>,
) -> io::Result<RangeOutcome> {
if !visited.insert(id) {
return Err(io::Error::new(
io::ErrorKind::InvalidData,
"cycle in BTree page graph",
));
}
match self.read_node(id)? {
Node::Leaf(entries) => {
for (k, v) in entries.iter() {
if k.as_slice() < start {
continue;
}
if k.as_slice() >= end {
break;
}
// Pre-flight: stop before collecting the next item if it
// would exhaust either budget. This comes before load_value
// so exhausted budgets avoid any I/O (important for zero
// limits and large overflow values). The byte estimate uses
// Val metadata and is always exact.
if items.len() >= item_limit {
return Ok(RangeOutcome::LimitReached(k.clone()));
}
let cost = k.len()
+ match v {
Val::Inline(b) => b.len(),
Val::Overflow { len, .. } => *len as usize,
};
if *byte_count + cost > byte_limit {
return Ok(RangeOutcome::LimitReached(k.clone()));
}
let value = self.load_value(v)?;
items.push((k.clone(), value));
*byte_count += cost;
}
Ok(RangeOutcome::Done)
}
Node::Interior { keys, children } => {
for i in 0..children.len() {
let lo = if i == 0 {
None
} else {
Some(keys[i - 1].as_slice())
};
let hi = keys.get(i).map(|k| k.as_slice());
if hi.is_some_and(|h| h <= start) {
continue;
}
if lo.is_some_and(|l| l >= end) {
break;
}
match self.range_bounded_rec(
children[i],
start,
end,
items,
byte_count,
item_limit,
byte_limit,
visited,
)? {
RangeOutcome::Done => {} // keep going to next child
RangeOutcome::LimitReached(k) => return Ok(RangeOutcome::LimitReached(k)),
RangeOutcome::NeedNext => {
// Budget was exhausted before collecting any item
// from child `i` (e.g. the first candidate already
// hit a zero/byte limit). Walk forward through
// remaining children to find an in-range key,
// skipping empty subtrees.
for j in (i + 1)..children.len() {
let child_lo = keys[j - 1].as_slice();
if child_lo >= end {
break;
}
if let Some(first) = self.leaf_first_key(children[j])? {
if first.as_slice() < end {
return Ok(RangeOutcome::LimitReached(first));
}
}
}
// No remaining child yields an in-range key.
return Ok(RangeOutcome::Done);
}
}
}
Ok(RangeOutcome::Done)
}
}
}
/// Return the first (smallest) key stored anywhere under `id`, or `None`
/// if the subtree is empty (every leaf visited is empty).
fn leaf_first_key(&mut self, id: PageId) -> io::Result<Option<Vec<u8>>> {
let mut cur = id;
let mut visited = HashSet::new();
loop {
if !visited.insert(cur) {
return Err(io::Error::new(
io::ErrorKind::InvalidData,
"cycle in BTree page graph",
));
}
match self.read_node(cur)? {
Node::Leaf(entries) => {
return Ok(entries.first().map(|(k, _)| k.clone()));
}
Node::Interior { children, .. } => {
if children.is_empty() {
return Err(io::Error::new(
io::ErrorKind::InvalidData,
"empty interior node in BTree",
));
}
cur = children[0];
}
}
}
}
// ---- page plumbing ----------------------------------------------------
fn read_node(&mut self, id: PageId) -> io::Result<Node> {
Ok(Node::decode(&self.pager.read_page(id)?))
}
/// Allocate a page (recycling a free one first) and write `node` to it.
fn write_node(&mut self, node: &Node) -> io::Result<PageId> {
if !node.fits() {
return Err(io::Error::new(
io::ErrorKind::InvalidInput,
format!(
"node requires {} bytes but page capacity is {PAGE_SIZE}",
node.encoded_len()
),
));
}
let id = match self.free.pop() {
Some(id) => id,
None => self.pager.allocate()?,
};
self.pager.write_page(id, &node.encode())?;
Ok(id)
}
/// Mark a page orphaned by this transaction. It stays referenced by the
/// committed tree until [`BTree::commit`] folds it into the free list.
fn free_page(&mut self, id: PageId) {
self.pending_free.push(id);
}
// ---- large-value overflow ---------------------------------------------
/// Inline cutoff: values up to this stay in the leaf; larger ones spill to
/// overflow pages. Kept well under a page so any single entry still fits a
/// leaf — the pre-overflow single-large-value panic is exactly what this
/// prevents (a one-entry leaf can't be split).
const INLINE_MAX: usize = 1024;
/// Overflow page header: `next: u64` (0 ends the chain) + `chunk_len: u32`.
const OVF_HDR: usize = 12;
/// Turn a caller value into its on-leaf form, spilling a large value to a
/// freshly allocated chain of immutable overflow pages and keeping only a
/// stub on the leaf. Overflow pages live outside the copy-on-write tree; an
/// overwrite or delete simply orphans the old chain (leaked, never reused —
/// safe, the same stance the meta free list already takes). Each page is
/// `next: u64 | chunk_len: u32 | bytes`, with `next == 0` ending the chain.
fn store_value(&mut self, value: &[u8]) -> io::Result<Val> {
if value.len() <= Self::INLINE_MAX {
return Ok(Val::Inline(value.to_vec()));
}
let chunk_cap = PAGE_SIZE - Self::OVF_HDR;
// Write tail-first so each page can name its already-written successor.
let chunks: Vec<&[u8]> = value.chunks(chunk_cap).collect();
let mut next: PageId = 0;
let mut head: PageId = 0;
for chunk in chunks.into_iter().rev() {
let pid = self.pager.allocate()?;
let mut page = [0u8; PAGE_SIZE];
page[0..8].copy_from_slice(&next.to_le_bytes());
page[8..12].copy_from_slice(&(chunk.len() as u32).to_le_bytes());
page[Self::OVF_HDR..Self::OVF_HDR + chunk.len()].copy_from_slice(chunk);
self.pager.write_page(pid, &page)?;
next = pid;
head = pid;
}
Ok(Val::Overflow {
head,
len: value.len() as u64,
})
}
/// Materialize a leaf value, walking the overflow chain when present.
fn load_value(&mut self, v: &Val) -> io::Result<Vec<u8>> {
match v {
Val::Inline(b) => Ok(b.clone()),
Val::Overflow { head, len } => {
let mut out = Vec::with_capacity(*len as usize);
let mut pid = *head;
let mut visited = HashSet::new();
while pid != 0 {
if !visited.insert(pid) {
return Err(io::Error::new(
io::ErrorKind::InvalidData,
"cycle in BTree overflow chain",
));
}
let page = self.pager.read_page(pid)?;
let next = u64::from_le_bytes(page[0..8].try_into().unwrap());
let clen = u32::from_le_bytes(page[8..12].try_into().unwrap()) as usize;
if clen > PAGE_SIZE - Self::OVF_HDR || out.len() + clen > *len as usize {
return Err(io::Error::new(
io::ErrorKind::InvalidData,
"invalid BTree overflow length",
));
}
out.extend_from_slice(&page[Self::OVF_HDR..Self::OVF_HDR + clen]);
pid = next;
}
if out.len() != *len as usize {
return Err(io::Error::new(
io::ErrorKind::InvalidData,
"truncated BTree overflow value",
));
}
Ok(out)
}
}
}
}
/// Pick the child index to descend for `key` in an interior node: the first
/// `i` with `key < keys[i]`, else the last child. Mirrors [`Node::Interior`].
fn route(keys: &[Vec<u8>], key: &[u8]) -> usize {
keys.partition_point(|k| k.as_slice() <= key)
}
#[cfg(test)]
mod tests {
use super::*;
fn tmp(name: &str) -> std::path::PathBuf {
let mut p = std::env::temp_dir();
p.push(format!("akurai-tree-{}-{}.db", name, std::process::id()));
let _ = std::fs::remove_file(&p);
p
}
#[test]
fn get_on_empty_tree_returns_none() {
let path = tmp("empty");
let mut t = BTree::open(&path).unwrap();
assert_eq!(t.get(b"missing").unwrap(), None);
std::fs::remove_file(&path).unwrap();
}
#[test]
fn insert_then_get_round_trips() {
let path = tmp("rt");
let mut t = BTree::open(&path).unwrap();
t.insert(b"key", b"value").unwrap();
assert_eq!(t.get(b"key").unwrap(), Some(b"value".to_vec()));
assert_eq!(t.get(b"nope").unwrap(), None);
std::fs::remove_file(&path).unwrap();
}
#[test]
fn insert_replaces_existing_value() {
let path = tmp("replace");
let mut t = BTree::open(&path).unwrap();
t.insert(b"k", b"one").unwrap();
t.insert(b"k", b"two").unwrap();
assert_eq!(t.get(b"k").unwrap(), Some(b"two".to_vec()));
std::fs::remove_file(&path).unwrap();
}
#[test]
fn delete_removes_only_the_target() {
let path = tmp("del");
let mut t = BTree::open(&path).unwrap();
t.insert(b"a", b"1").unwrap();
t.insert(b"b", b"2").unwrap();
assert!(t.delete(b"a").unwrap());
assert!(!t.delete(b"a").unwrap()); // already gone
assert_eq!(t.get(b"a").unwrap(), None);
assert_eq!(t.get(b"b").unwrap(), Some(b"2".to_vec()));
std::fs::remove_file(&path).unwrap();
}
#[test]
fn many_keys_force_splits_and_stay_searchable() {
let path = tmp("split");
let mut t = BTree::open(&path).unwrap();
// Big values make pages overflow fast, forcing many splits.
let val = vec![b'x'; 200];
for i in 0..500u32 {
t.insert(&i.to_be_bytes(), &val).unwrap();
}
for i in 0..500u32 {
assert_eq!(
t.get(&i.to_be_bytes()).unwrap(),
Some(val.clone()),
"key {i}"
);
}
assert_eq!(t.get(&999u32.to_be_bytes()).unwrap(), None);
std::fs::remove_file(&path).unwrap();
}
#[test]
fn range_returns_sorted_window() {
let path = tmp("range");
let mut t = BTree::open(&path).unwrap();
for i in 0..300u32 {
t.insert(&i.to_be_bytes(), &[i as u8]).unwrap();
}
let got = t.range(&10u32.to_be_bytes(), &20u32.to_be_bytes()).unwrap();
let keys: Vec<u32> = got
.iter()
.map(|(k, _)| u32::from_be_bytes(k.as_slice().try_into().unwrap()))
.collect();
assert_eq!(keys, (10..20).collect::<Vec<_>>());
std::fs::remove_file(&path).unwrap();
}
#[test]
fn commit_then_reopen_sees_the_data() {
let path = tmp("durable");
{
let mut t = BTree::open(&path).unwrap();
for i in 0..200u32 {
t.insert(&i.to_be_bytes(), b"v").unwrap();
}
t.commit().unwrap();
}
let mut t = BTree::open(&path).unwrap();
for i in 0..200u32 {
assert_eq!(t.get(&i.to_be_bytes()).unwrap(), Some(b"v".to_vec()));
}
std::fs::remove_file(&path).unwrap();
}
#[test]
fn uncommitted_changes_vanish_on_reopen() {
let path = tmp("rollback");
{
let mut t = BTree::open(&path).unwrap();
t.insert(b"committed", b"yes").unwrap();
t.commit().unwrap();
t.insert(b"dirty", b"no").unwrap(); // never committed
}
let mut t = BTree::open(&path).unwrap();
assert_eq!(t.get(b"committed").unwrap(), Some(b"yes".to_vec()));
assert_eq!(t.get(b"dirty").unwrap(), None);
std::fs::remove_file(&path).unwrap();
}
#[test]
fn free_list_recycles_pages_across_commits() {
let path = tmp("recycle");
let mut t = BTree::open(&path).unwrap();
t.insert(b"k", &[b'a'; 100]).unwrap();
t.commit().unwrap();
let pages_after_first = t.pager.page_count().unwrap();
// Rewrites of the same key orphan a leaf each time; commits recycle it,
// so the file must not grow unboundedly.
for _ in 0..50 {
t.insert(b"k", &[b'b'; 100]).unwrap();
t.commit().unwrap();
}
let pages_after_many = t.pager.page_count().unwrap();
assert!(
pages_after_many <= pages_after_first + 2,
"expected page reuse, grew {pages_after_first} -> {pages_after_many}"
);
std::fs::remove_file(&path).unwrap();
}
#[test]
fn split_balances_skewed_leaf_by_bytes() {
let path = tmp("skewed-leaf-split");
let mut t = BTree::open(&path).unwrap();
for i in 0..10u32 {
t.insert(&i.to_be_bytes(), b"small").unwrap();
}
for i in 0..5u32 {
t.insert(&i.to_be_bytes(), &[i as u8; 900]).unwrap();
}
for i in 0..5u32 {
assert_eq!(t.get(&i.to_be_bytes()).unwrap(), Some(vec![i as u8; 900]));
}
std::fs::remove_file(&path).unwrap();
}
#[test]
fn large_value_spills_and_round_trips() {
let path = tmp("overflow");
let mut t = BTree::open(&path).unwrap();
// Several pages worth — the exact case that used to panic on encode.
let big = vec![b'Z'; 20_000];
t.insert(b"doc", &big).unwrap();
assert_eq!(t.get(b"doc").unwrap(), Some(big.clone()));
// Survives commit + reopen (overflow pages are synced before the swap).
t.commit().unwrap();
drop(t);
let mut t = BTree::open(&path).unwrap();
assert_eq!(t.get(b"doc").unwrap(), Some(big));
std::fs::remove_file(&path).unwrap();
}
#[test]
fn value_straddling_the_inline_cutoff_round_trips() {
let path = tmp("overflow-cutoff");
let mut t = BTree::open(&path).unwrap();
for n in [1024usize, 1025, 4083, 4084, 4085, 8200] {
let v = vec![b'c'; n];
t.insert(b"k", &v).unwrap();
assert_eq!(t.get(b"k").unwrap(), Some(v), "len {n}");
}
std::fs::remove_file(&path).unwrap();
}
#[test]
fn large_value_can_be_overwritten_by_small_and_back() {
let path = tmp("overflow-overwrite");
let mut t = BTree::open(&path).unwrap();
let big = vec![b'A'; 10_000];
t.insert(b"k", &big).unwrap();
t.insert(b"k", b"tiny").unwrap();
assert_eq!(t.get(b"k").unwrap(), Some(b"tiny".to_vec()));
let big2 = vec![b'B'; 9_999];
t.insert(b"k", &big2).unwrap();
assert_eq!(t.get(b"k").unwrap(), Some(big2));
std::fs::remove_file(&path).unwrap();
}
#[test]
fn large_and_small_values_coexist_in_range() {
let path = tmp("overflow-range");
let mut t = BTree::open(&path).unwrap();
let big = vec![b'q'; 8_000];
t.insert(&1u32.to_be_bytes(), b"small").unwrap();
t.insert(&2u32.to_be_bytes(), &big).unwrap();
t.insert(&3u32.to_be_bytes(), b"also").unwrap();
let got = t.range(&0u32.to_be_bytes(), &9u32.to_be_bytes()).unwrap();
assert_eq!(got.len(), 3);
assert_eq!(got[1].1, big);
std::fs::remove_file(&path).unwrap();
}
#[test]
fn many_large_values_force_overflow_and_stay_searchable() {
let path = tmp("overflow-many");
let mut t = BTree::open(&path).unwrap();
for i in 0..40u32 {
// Each value spans multiple overflow pages and varies by key.
let v = vec![i as u8; 6_000 + i as usize];
t.insert(&i.to_be_bytes(), &v).unwrap();
}
for i in 0..40u32 {
assert_eq!(
t.get(&i.to_be_bytes()).unwrap(),
Some(vec![i as u8; 6_000 + i as usize]),
"key {i}"
);
}
std::fs::remove_file(&path).unwrap();
}
#[test]
fn cycle_is_reported_and_salvage_keeps_reachable_leaf() {
let path = tmp("cycle");
let mut t = BTree::open(&path).unwrap();
let leaf = t
.pager
.allocate()
.and_then(|id| {
t.pager.write_page(
id,
&Node::Leaf(vec![(b"key".to_vec(), Val::Inline(b"value".to_vec()))]).encode(),
)?;
Ok(id)
})
.unwrap();
t.pager
.write_page(
t.root,
&Node::Interior {
keys: vec![b"z".to_vec()],
children: vec![leaf, t.root],
}
.encode(),
)
.unwrap();
assert_eq!(
t.range(b"a", b"zz").unwrap_err().kind(),
io::ErrorKind::InvalidData
);
assert_eq!(
t.salvage_range(b"a", b"zz").unwrap(),
vec![(b"key".to_vec(), b"value".to_vec())]
);
std::fs::remove_file(&path).unwrap();
}
// ---- bounded range ----------------------------------------------------
#[test]
fn range_bounded_item_limit_stops_early() {
let path = tmp("bounded-items");
let mut t = BTree::open(&path).unwrap();
for i in 0..50u32 {
t.insert(&i.to_be_bytes(), &[i as u8]).unwrap();
}
let rb = t
.range_bounded(&0u32.to_be_bytes(), &99u32.to_be_bytes(), 10, usize::MAX)
.unwrap();
assert_eq!(rb.items.len(), 10);
assert!(rb.continuation.is_some());
// Continuation should be > the last returned key.
let last_key = &rb.items.last().unwrap().0;
let cont = rb.continuation.as_ref().unwrap();
assert!(cont.as_slice() > last_key.as_slice());
std::fs::remove_file(&path).unwrap();
}
#[test]
fn range_bounded_byte_limit_stops_early() {
let path = tmp("bounded-bytes");
let mut t = BTree::open(&path).unwrap();
for i in 0..50u32 {
t.insert(&i.to_be_bytes(), &[i as u8; 20]).unwrap();
}
// Each item is 4 (key) + 20 (value) = 24 bytes. 5 items = 120 bytes.
let rb = t
.range_bounded(&0u32.to_be_bytes(), &99u32.to_be_bytes(), usize::MAX, 100)
.unwrap();
assert!(rb.items.len() < 50, "should stop before all items");
// Pre-flight check ensures the total never exceeds the byte limit.
let total_bytes: usize = rb.items.iter().map(|(k, v)| k.len() + v.len()).sum();
assert!(
total_bytes <= 100,
"total {total_bytes} must not exceed byte_limit 100"
);
assert!(!rb.items.is_empty());
assert!(rb.continuation.is_some());
std::fs::remove_file(&path).unwrap();
}
#[test]
fn range_bounded_resume_covers_all_items() {
let path = tmp("bounded-resume");
let mut t = BTree::open(&path).unwrap();
for i in 0..100u32 {
t.insert(&i.to_be_bytes(), &[i as u8]).unwrap();
}
let mut all: Vec<u32> = Vec::new();
let mut cursor: Vec<u8> = 0u32.to_be_bytes().to_vec();
let end = 100u32.to_be_bytes();
loop {
let rb = t.range_bounded(&cursor, &end, 7, usize::MAX).unwrap();
for (k, _) in &rb.items {
all.push(u32::from_be_bytes(k.as_slice().try_into().unwrap()));
}
match rb.continuation {
Some(c) => cursor = c,
None => break,
}
}
assert_eq!(all, (0..100).collect::<Vec<_>>());
std::fs::remove_file(&path).unwrap();
}
#[test]
fn range_bounded_exhaust_returns_no_continuation() {
let path = tmp("bounded-exhaust");
let mut t = BTree::open(&path).unwrap();
for i in 0..5u32 {
t.insert(&i.to_be_bytes(), &[i as u8]).unwrap();
}
let rb = t
.range_bounded(&0u32.to_be_bytes(), &99u32.to_be_bytes(), 100, usize::MAX)
.unwrap();
assert_eq!(rb.items.len(), 5);
assert!(rb.continuation.is_none());
std::fs::remove_file(&path).unwrap();
}
#[test]
fn range_bounded_empty_range_returns_nothing() {
let path = tmp("bounded-empty");
let mut t = BTree::open(&path).unwrap();
for i in 0..10u32 {
t.insert(&i.to_be_bytes(), &[i as u8]).unwrap();
}
// start > all keys
let rb = t
.range_bounded(&50u32.to_be_bytes(), &99u32.to_be_bytes(), 10, usize::MAX)
.unwrap();
assert!(rb.items.is_empty());
assert!(rb.continuation.is_none());
std::fs::remove_file(&path).unwrap();
}
// ---- bounded-range regressions -----------------------------------------
/// Zero item_limit must return zero items, not one.
#[test]
fn range_bounded_zero_item_limit_returns_nothing() {
let path = tmp("bounded-zero-items");
let mut t = BTree::open(&path).unwrap();
for i in 0..10u32 {
t.insert(&i.to_be_bytes(), &[i as u8]).unwrap();
}
let rb = t
.range_bounded(&0u32.to_be_bytes(), &99u32.to_be_bytes(), 0, usize::MAX)
.unwrap();
assert!(rb.items.is_empty(), "zero item_limit must collect nothing");
// A continuation must still be returned so the caller can resume.
assert!(
rb.continuation.is_some(),
"continuation must name the first in-range key"
);
std::fs::remove_file(&path).unwrap();
}
/// Zero byte_limit must return zero items, not one.
#[test]
fn range_bounded_zero_byte_limit_returns_nothing() {
let path = tmp("bounded-zero-bytes");
let mut t = BTree::open(&path).unwrap();
for i in 0..10u32 {
t.insert(&i.to_be_bytes(), &[i as u8]).unwrap();
}
let rb = t
.range_bounded(&0u32.to_be_bytes(), &99u32.to_be_bytes(), usize::MAX, 0)
.unwrap();
assert!(rb.items.is_empty(), "zero byte_limit must collect nothing");
assert!(
rb.continuation.is_some(),
"continuation must name the first in-range key"
);
std::fs::remove_file(&path).unwrap();
}
/// A zero byte_limit alongside a large overflow value must not trigger a
/// load/corruption error — the pre-flight check must short-circuit before
/// load_value allocates or reads any overflow data.
#[test]
fn range_bounded_zero_byte_limit_skips_large_overflow_load() {
let path = tmp("bounded-zero-overflow");
let mut t = BTree::open(&path).unwrap();
// Small inline values first, then a large overflow value.
for i in 0..5u32 {
t.insert(&i.to_be_bytes(), &[i as u8; 32]).unwrap();
}
let big = vec![b'Z'; 30_000];
t.insert(&5u32.to_be_bytes(), &big).unwrap();
// Zero byte_limit must return nothing without attempting to load the
// 30 kB overflow value.
let rb = t
.range_bounded(&0u32.to_be_bytes(), &99u32.to_be_bytes(), usize::MAX, 0)
.unwrap();
assert!(rb.items.is_empty());
assert!(rb.continuation.is_some());
std::fs::remove_file(&path).unwrap();
}
/// Continuation keys >= end are suppressed.
#[test]
fn range_bounded_continuation_never_exceeds_end() {
let path = tmp("bounded-cont-end");
let mut t = BTree::open(&path).unwrap();
for i in 0..10u32 {
t.insert(&i.to_be_bytes(), &[i as u8]).unwrap();
}
// Set item_limit so we stop partway, at a key inside [0, 5).
let rb = t
.range_bounded(&0u32.to_be_bytes(), &5u32.to_be_bytes(), 3, usize::MAX)
.unwrap();
assert_eq!(rb.items.len(), 3);
// The continuation must be strictly less than the exclusive end.
if let Some(c) = &rb.continuation {
assert!(
c.as_slice() < &5u32.to_be_bytes(),
"continuation {c:?} must be < end"
);
}
std::fs::remove_file(&path).unwrap();
}
/// When a budget is exhausted exactly at a leaf boundary, and the next
/// child is empty or entirely beyond the range, the scan must finish
/// cleanly without an in-range continuation.
#[test]
fn range_bounded_empty_subtree_beyond_range_is_skipped() {
let path = tmp("bounded-empty-subtree");
let mut t = BTree::open(&path).unwrap();
// Insert keys so that the tree has multiple leaves. Use values big
// enough to force splits and exercise interior-node routing.
for i in 0..200u32 {
t.insert(&i.to_be_bytes(), &[i as u8; 80]).unwrap();
}
// Scan the first few keys with a tight item_limit, then resume.
let end = 10u32.to_be_bytes();
let mut all: Vec<u32> = Vec::new();
let mut cursor: Vec<u8> = 0u32.to_be_bytes().to_vec();
loop {
let rb = t.range_bounded(&cursor, &end, 3, usize::MAX).unwrap();
for (k, _) in &rb.items {
all.push(u32::from_be_bytes(k.as_slice().try_into().unwrap()));
}
match rb.continuation {
Some(c) if c.as_slice() < &end => cursor = c,
Some(_) | None => break,
}
}
assert_eq!(all, (0..10).collect::<Vec<_>>());
std::fs::remove_file(&path).unwrap();
}
// ---- fault injection ---------------------------------------------------
#[test]
fn fault_write_page_propagates_error() {
let path = tmp("fault-write");
let mut t = BTree::open(&path).unwrap();
// First write succeeds (it's a fresh file with a few pages).
t.pager.inject_write_faults(1);
let err = t.insert(b"k", b"v").unwrap_err();
assert_eq!(err.kind(), io::ErrorKind::Other);
assert!(
err.to_string().contains("fault"),
"expected fault message, got: {err}"
);
std::fs::remove_file(&path).unwrap();
}
#[test]
fn fault_sync1_leaves_old_state_reopenable() {
let path = tmp("fault-sync1");
// Commit with data, then mutate, then fail the first sync.
{
let mut t = BTree::open(&path).unwrap();
t.insert(b"survivor", b"yes").unwrap();
t.commit().unwrap();
}
{
let mut t = BTree::open(&path).unwrap();
t.insert(b"casualty", b"no").unwrap();
// Fail the first sync (flush of data pages).
t.pager.inject_sync_fault(0);
let err = t.commit().unwrap_err();
assert_eq!(err.kind(), io::ErrorKind::Other);
assert!(
err.to_string().contains("fault"),
"expected fault message, got: {err}"
);
}
// Reopen — the failed commit must have left the old state intact.
let mut t = BTree::open(&path).unwrap();
assert_eq!(t.get(b"survivor").unwrap(), Some(b"yes".to_vec()));
assert_eq!(t.get(b"casualty").unwrap(), None);
std::fs::remove_file(&path).unwrap();
}
/// Reject the meta-page write during commit so the new transaction is never
/// recorded. The old meta slot must still validate and the uncommitted
/// mutation must be invisible on reopen — a fault-injection proof that a
/// rejected meta write preserves the previous committed tree.
#[test]
fn fault_meta_write_rejection_preserves_old_commit() {
let path = tmp("fault-sync2");
// --- commit a baseline ---
{
let mut t = BTree::open(&path).unwrap();
t.insert(b"survivor", b"yes").unwrap();
t.commit().unwrap(); // txn_id = 2 → slot 0
}
// --- mutate, then fail the meta-page write ---
{
let mut t = BTree::open(&path).unwrap();
// load_meta picks slot 0 (txn_id = 2), root points to survivor.
t.insert(b"casualty", b"no").unwrap();
// The insert rewrites the leaf (1 write_page). Fail the very next
// write_page, which is the meta write inside commit.
t.pager.inject_write_faults(1);
let err = t.commit().unwrap_err();
assert_eq!(err.kind(), io::ErrorKind::Other);
assert!(err.to_string().contains("fault"));
}
// --- reopen: old slot (0, txn_id = 2) wins; new tx was never stored ---
let mut t = BTree::open(&path).unwrap();
assert_eq!(t.get(b"survivor").unwrap(), Some(b"yes".to_vec()));
assert_eq!(t.get(b"casualty").unwrap(), None);
std::fs::remove_file(&path).unwrap();
}
#[test]
fn fault_committed_state_survives_reopen() {
// After a fault, a fresh commit should work and persist.
let path = tmp("fault-recover");
{
let mut t = BTree::open(&path).unwrap();
t.insert(b"a", b"1").unwrap();
t.commit().unwrap();
}
{
let mut t = BTree::open(&path).unwrap();
t.insert(b"b", b"2").unwrap();
t.pager.inject_sync_fault(0); // fail first sync
let _ = t.commit().unwrap_err();
}
// Now commit again without faults — should work.
{
let mut t = BTree::open(&path).unwrap();
t.insert(b"b", b"2").unwrap();
t.commit().unwrap();
}
let mut t = BTree::open(&path).unwrap();
assert_eq!(t.get(b"a").unwrap(), Some(b"1".to_vec()));
assert_eq!(t.get(b"b").unwrap(), Some(b"2".to_vec()));
std::fs::remove_file(&path).unwrap();
}
}