Menu
Public repository / 51aec02c44fe7eb9438a160c80c03b0b12141f7a
AkurAI-Framework / crates/storage/src/tree.rs
Latest change 5458532ad0ba451030d0554d03be8ff383ccddb8 - feat(ws,http): pure-std WebSocket (RFC 6455) + Server-Sent Events 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::io;
use std::path::Path;
use crate::meta::{Meta, FIRST_TREE_PAGE, META_SLOT_0, META_SLOT_1};
use crate::node::Node;
use crate::pager::{PageId, Pager};
/// 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,
},
}
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;
loop {
match self.read_node(id)? {
Node::Leaf(entries) => {
return Ok(entries
.binary_search_by(|(k, _)| k.as_slice().cmp(key))
.ok()
.map(|i| entries[i].1.clone()));
}
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)?;
Ok(out)
}
/// 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 self.pending_free.drain(..).collect::<Vec<_>>() {
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) => {
match entries.binary_search_by(|(k, _)| k.as_slice().cmp(key)) {
Ok(i) => entries[i].1 = value.to_vec(),
Err(i) => entries.insert(i, (key.to_vec(), value.to_vec())),
}
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>, Vec<u8>)>) -> 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 = entries.len() / 2;
let sep = entries[mid].0.clone();
let right_entries = entries[mid..].to_vec();
let left_entries = entries[..mid].to_vec();
let left = self.write_node(&Node::Leaf(left_entries))?;
let right = self.write_node(&Node::Leaf(right_entries))?;
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 = keys.len() / 2; // this key rises to the parent
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>)>,
) -> io::Result<()> {
match self.read_node(id)? {
Node::Leaf(entries) => {
for (k, v) in entries {
if k.as_slice() >= start && k.as_slice() < end {
out.push((k, v));
}
}
Ok(())
}
Node::Interior { keys, children } => {
for i in 0..children.len() {
// Subtree i covers keys in [lo, hi); skip if disjoint from
// the query window.
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; // entirely below the window
}
if lo.is_some_and(|l| l >= end) {
break; // this and all later subtrees are above the window
}
self.range_rec(children[i], start, end, out)?;
}
Ok(())
}
}
}
// ---- 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> {
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);
}
}
/// 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();
}
}