use crate::page::{PageKind, PageMut, HEADER_LEN, PAGE_SIZE};
use crate::pool::BufferPool;
use crate::io::AlignedRegion;
use crate::{Error, Result};
use std::collections::BinaryHeap;
use std::fs::File;
use std::io::{BufReader, BufWriter, Read, Write};
use std::path::{Path, PathBuf};
const MARKER_BIT: u32 = 1 << 31;
const SCRATCH_HEADER_LEN: usize = 12;
fn invalid_scratch(why: &'static str) -> std::io::Error {
std::io::Error::new(std::io::ErrorKind::InvalidData, why)
}
fn write_item(w: &mut impl Write, k: &[u8], v: &[u8], marker: bool) -> std::io::Result<()> {
let kl = u32::try_from(k.len()).map_err(|_| invalid_scratch("scratch key is too long"))?;
let vl = u32::try_from(v.len()).map_err(|_| invalid_scratch("scratch value is too long"))?;
if vl & MARKER_BIT != 0 {
return Err(invalid_scratch("scratch value is too long for the marker bit"));
}
let raw_vl = vl | if marker { MARKER_BIT } else { 0 };
let mut header = [0u8; SCRATCH_HEADER_LEN];
header[0..4].copy_from_slice(&kl.to_le_bytes());
header[4..8].copy_from_slice(&raw_vl.to_le_bytes());
let mut crc = crc32c::crc32c(&header[..8]);
crc = crc32c::crc32c_append(crc, k);
crc = crc32c::crc32c_append(crc, v);
header[8..12].copy_from_slice(&crc.to_le_bytes());
w.write_all(&header)?;
w.write_all(k)?;
w.write_all(v)
}
fn read_item(r: &mut impl Read, max_key_len: usize, max_val_len: usize)
-> std::io::Result<Option<(Vec<u8>, Vec<u8>, bool)>>
{
let mut h = [0u8; SCRATCH_HEADER_LEN];
match r.read_exact(&mut h[..1]) {
Ok(()) => {}
Err(e) if e.kind() == std::io::ErrorKind::UnexpectedEof => return Ok(None),
Err(e) => return Err(e),
}
r.read_exact(&mut h[1..])?;
let kl = u32::from_le_bytes(h[0..4].try_into().unwrap()) as usize;
let raw = u32::from_le_bytes(h[4..8].try_into().unwrap());
let marker = raw & MARKER_BIT != 0;
let vl = (raw & !MARKER_BIT) as usize;
if kl > max_key_len || vl > max_val_len {
return Err(invalid_scratch("scratch record length exceeds its writer's bound"));
}
let mut k = Vec::new();
k.try_reserve_exact(kl)
.map_err(|_| invalid_scratch("scratch key allocation exceeds address space"))?;
k.resize(kl, 0);
r.read_exact(&mut k)?;
let mut v = Vec::new();
v.try_reserve_exact(vl)
.map_err(|_| invalid_scratch("scratch value allocation exceeds address space"))?;
v.resize(vl, 0);
r.read_exact(&mut v)?;
let want_crc = u32::from_le_bytes(h[8..12].try_into().unwrap());
let mut crc = crc32c::crc32c(&h[..8]);
crc = crc32c::crc32c_append(crc, &k);
crc = crc32c::crc32c_append(crc, &v);
if crc != want_crc {
return Err(invalid_scratch("scratch record checksum mismatch"));
}
Ok(Some((k, v, marker)))
}
const SORT_MANIFEST_MAGIC: &[u8; 8] = b"KSRUN01\0";
const SORT_MANIFEST_MAX: u64 = 16 << 20;
struct SortManifest {
watermark: u64,
records: u64,
framed_bytes: u64,
max_key_len: usize,
max_val_len: usize,
runs: Vec<PathBuf>,
run_counts: Vec<u64>,
}
fn manifest_u64(bytes: &[u8], pos: &mut usize) -> std::io::Result<u64> {
let end = pos.checked_add(8).ok_or_else(|| invalid_scratch("sort manifest offset overflow"))?;
let raw = bytes.get(*pos..end).ok_or_else(|| invalid_scratch("sort manifest is truncated"))?;
*pos = end;
Ok(u64::from_le_bytes(raw.try_into().unwrap()))
}
fn write_sort_manifest(sort: &ExternalSort, generation: u64, watermark: u64) -> Result<()> {
let mut body = Vec::new();
body.extend_from_slice(SORT_MANIFEST_MAGIC);
for n in [generation, watermark, sort.records, sort.framed_bytes,
sort.max_key_len as u64, sort.max_val_len as u64, sort.runs.len() as u64] {
body.extend_from_slice(&n.to_le_bytes());
}
for (path, &count) in sort.runs.iter().zip(&sort.run_counts) {
let name = path.file_name().and_then(|s| s.to_str())
.ok_or_else(|| Error::Io(invalid_scratch("durable run has no UTF-8 file name")))?;
let name = name.as_bytes();
let len = u16::try_from(name.len())
.map_err(|_| Error::Io(invalid_scratch("durable run file name is too long")))?;
body.extend_from_slice(&len.to_le_bytes());
body.extend_from_slice(name);
body.extend_from_slice(&count.to_le_bytes());
}
let crc = crc32c::crc32c(&body);
body.extend_from_slice(&crc.to_le_bytes());
let tmp = sort.dir.join("manifest.tmp");
let published = sort.dir.join("manifest");
{
let mut file = File::create(&tmp)?;
file.write_all(&body)?;
crate::write_stats::add(
crate::write_stats::Phase::Manifest,
body.len() as u64,
);
file.sync_all()?;
}
std::fs::rename(&tmp, &published)?;
File::open(&sort.dir)?.sync_all()?;
Ok(())
}
fn read_sort_manifest(dir: &Path, expected_generation: u64) -> Result<SortManifest> {
let path = dir.join("manifest");
let len = std::fs::metadata(&path)?.len();
if len < (SORT_MANIFEST_MAGIC.len() + 7 * 8 + 4) as u64 || len > SORT_MANIFEST_MAX {
return Err(Error::Io(invalid_scratch("sort manifest length is invalid")));
}
let bytes = std::fs::read(&path)?;
let split = bytes.len() - 4;
let want = u32::from_le_bytes(bytes[split..].try_into().unwrap());
if crc32c::crc32c(&bytes[..split]) != want {
return Err(Error::Io(invalid_scratch("sort manifest checksum mismatch")));
}
if bytes.get(..8) != Some(SORT_MANIFEST_MAGIC) {
return Err(Error::Io(invalid_scratch("sort manifest magic mismatch")));
}
let mut pos = 8;
let generation = manifest_u64(&bytes[..split], &mut pos)?;
if generation != expected_generation {
return Err(Error::Io(invalid_scratch("sort manifest generation mismatch")));
}
let watermark = manifest_u64(&bytes[..split], &mut pos)?;
let records = manifest_u64(&bytes[..split], &mut pos)?;
let framed_bytes = manifest_u64(&bytes[..split], &mut pos)?;
let max_key_len = usize::try_from(manifest_u64(&bytes[..split], &mut pos)?)
.map_err(|_| Error::TooLarge)?;
let max_val_len = usize::try_from(manifest_u64(&bytes[..split], &mut pos)?)
.map_err(|_| Error::TooLarge)?;
let run_count = usize::try_from(manifest_u64(&bytes[..split], &mut pos)?)
.map_err(|_| Error::TooLarge)?;
if run_count > (split.saturating_sub(pos)) / 10 {
return Err(Error::Io(invalid_scratch("sort manifest run count exceeds its bytes")));
}
let mut runs = Vec::with_capacity(run_count);
let mut run_counts = Vec::with_capacity(run_count);
for _ in 0..run_count {
let end = pos.checked_add(2).ok_or(Error::TooLarge)?;
let raw = bytes.get(pos..end)
.ok_or_else(|| Error::Io(invalid_scratch("sort manifest run name is truncated")))?;
pos = end;
let name_len = u16::from_le_bytes(raw.try_into().unwrap()) as usize;
let end = pos.checked_add(name_len).ok_or(Error::TooLarge)?;
let raw_name = bytes.get(pos..end)
.ok_or_else(|| Error::Io(invalid_scratch("sort manifest run name is truncated")))?;
pos = end;
let name = std::str::from_utf8(raw_name)
.map_err(|_| Error::Io(invalid_scratch("sort manifest run name is not UTF-8")))?;
if name.contains('/') || name.contains('\\') || name == "." || name == ".." {
return Err(Error::Io(invalid_scratch("sort manifest run name escapes its directory")));
}
let path = dir.join(name);
if !path.is_file() {
return Err(Error::Io(invalid_scratch("sort manifest names a missing run")));
}
runs.push(path);
run_counts.push(manifest_u64(&bytes[..split], &mut pos)?);
}
if pos != split {
return Err(Error::Io(invalid_scratch("sort manifest has trailing bytes")));
}
if run_counts.iter().try_fold(0u64, |sum, &n| sum.checked_add(n)) != Some(records) {
return Err(Error::Io(invalid_scratch("sort manifest row counts disagree")));
}
Ok(SortManifest { watermark, records, framed_bytes, max_key_len, max_val_len,
runs, run_counts })
}
pub struct ExternalSort {
finished: bool,
dir: PathBuf,
arena: Vec<(Vec<u8>, Vec<u8>, bool)>,
arena_bytes: usize,
used: usize,
runs: Vec<PathBuf>,
run_counts: Vec<u64>,
max_key_len: usize,
max_val_len: usize,
records: u64,
framed_bytes: u64,
durable_generation: Option<u64>,
}
impl ExternalSort {
pub fn new(dir: &Path, arena_bytes: usize) -> Result<Self> {
Self::new_inner(dir, arena_bytes, None)
}
pub fn new_durable(dir: &Path, arena_bytes: usize, generation: u64) -> Result<Self> {
Self::new_inner(dir, arena_bytes, Some(generation))
}
fn new_inner(dir: &Path, arena_bytes: usize, durable_generation: Option<u64>) -> Result<Self> {
std::fs::create_dir_all(dir)?;
Ok(ExternalSort {
finished: false, dir: dir.to_path_buf(), arena: Vec::new(), arena_bytes, used: 0,
runs: Vec::new(), run_counts: Vec::new(), max_key_len: 0, max_val_len: 0,
records: 0, framed_bytes: 0, durable_generation,
})
}
pub fn push_flagged(&mut self, key: Vec<u8>, val: Vec<u8>, marker: bool) -> Result<()> {
self.push_inner(key, val, marker)
}
pub fn push(&mut self, key: Vec<u8>, val: Vec<u8>) -> Result<()> {
self.push_inner(key, val, false)
}
fn push_inner(&mut self, key: Vec<u8>, val: Vec<u8>, marker: bool) -> Result<()> {
self.records = self.records.checked_add(1).ok_or(Error::TooLarge)?;
self.framed_bytes = self.framed_bytes
.checked_add((SCRATCH_HEADER_LEN + key.len() + val.len()) as u64)
.ok_or(Error::TooLarge)?;
self.max_key_len = self.max_key_len.max(key.len());
self.max_val_len = self.max_val_len.max(val.len());
self.used += key.len() + val.len() + 48; self.arena.push((key, val, marker));
if self.used >= self.arena_bytes { self.spill()?; }
Ok(())
}
fn spill(&mut self) -> Result<()> {
if self.arena.is_empty() { return Ok(()); }
self.arena.sort_by(|a, b| a.0.cmp(&b.0));
let path = match self.durable_generation {
Some(generation) => self.dir.join(format!("g{generation:016x}-run-{:05}.dat", self.runs.len())),
None => self.dir.join(format!("run-{:05}.tmp", self.runs.len())),
};
let mut w = BufWriter::new(File::create(&path)?);
let count = u64::try_from(self.arena.len()).map_err(|_| Error::TooLarge)?;
let mut written = 0u64;
for (k, v, m) in self.arena.drain(..) {
written = written.checked_add(
(SCRATCH_HEADER_LEN + k.len() + v.len()) as u64,
).ok_or(Error::TooLarge)?;
write_item(&mut w, &k, &v, m)?;
}
w.flush()?;
crate::write_stats::add(crate::write_stats::Phase::SortScratch, written);
if self.durable_generation.is_some() { w.get_ref().sync_all()?; }
self.runs.push(path);
self.run_counts.push(count);
self.used = 0;
Ok(())
}
pub fn flush_run(&mut self) -> Result<()> { self.spill() }
pub fn profile(&self) -> (u64, u64, usize) {
(self.records, self.framed_bytes,
self.runs.len() + usize::from(!self.arena.is_empty()))
}
pub fn checkpoint(&mut self, watermark: u64) -> Result<std::time::Duration> {
let started = std::time::Instant::now();
let generation = self.durable_generation.ok_or_else(|| Error::Io(
invalid_scratch("checkpoint requested for a non-durable sorter")))?;
self.spill()?;
write_sort_manifest(self, generation, watermark)?;
Ok(started.elapsed())
}
pub fn reopen_durable(dir: &Path, arena_bytes: usize, expected_generation: u64)
-> Result<(Self, u64)>
{
let manifest = read_sort_manifest(dir, expected_generation)?;
let mut sort = Self::new_inner(dir, arena_bytes, Some(expected_generation))?;
sort.runs = manifest.runs;
sort.run_counts = manifest.run_counts;
sort.max_key_len = manifest.max_key_len;
sort.max_val_len = manifest.max_val_len;
sort.records = manifest.records;
sort.framed_bytes = manifest.framed_bytes;
Ok((sort, manifest.watermark))
}
pub fn discard_durable(mut self) -> Result<()> {
self.finished = true;
std::fs::remove_dir_all(&self.dir)?;
Ok(())
}
pub fn finish(mut self) -> Result<SortedRuns> {
self.spill()?;
self.finished = true;
Ok(SortedRuns {
dir: std::mem::take(&mut self.dir), runs: std::mem::take(&mut self.runs),
owns_dir: self.durable_generation.is_none(),
run_counts: std::mem::take(&mut self.run_counts),
max_key_len: self.max_key_len, max_val_len: self.max_val_len,
})
}
}
impl Drop for ExternalSort {
fn drop(&mut self) {
if !self.finished && self.durable_generation.is_none() {
let _ = std::fs::remove_dir_all(&self.dir);
}
}
}
pub struct SortedRuns {
dir: PathBuf,
runs: Vec<PathBuf>,
run_counts: Vec<u64>,
owns_dir: bool,
max_key_len: usize,
max_val_len: usize,
}
impl SortedRuns {
pub fn run_count(&self) -> usize { self.runs.len() }
pub fn discard(mut self) -> Result<()> {
self.owns_dir = false;
if self.dir.exists() { std::fs::remove_dir_all(&self.dir)?; }
Ok(())
}
pub const MAX_FANOUT: usize = 64;
fn merge_down(&mut self) -> Result<()> {
let mut pass_no = 0usize;
while self.runs.len() > Self::MAX_FANOUT {
let mut next: Vec<PathBuf> = Vec::new();
let mut next_counts: Vec<u64> = Vec::new();
for (gi, start) in (0..self.runs.len()).step_by(Self::MAX_FANOUT).enumerate() {
let end = (start + Self::MAX_FANOUT).min(self.runs.len());
let group = &self.runs[start..end];
let group_count = self.run_counts[start..end].iter().try_fold(0u64, |total, &n| {
total.checked_add(n).ok_or(Error::TooLarge)
})?;
if group.len() == 1 {
next.push(group[0].clone());
next_counts.push(group_count);
continue;
}
let out = self.dir.join(format!("pass-{}-{}-{:05}.tmp", pass_no, next.len(), gi));
let mut w = BufWriter::new(File::create(&out)?);
let mut part = SortedRuns {
dir: self.dir.clone(), runs: group.to_vec(), owns_dir: false,
run_counts: self.run_counts[start..end].to_vec(),
max_key_len: self.max_key_len, max_val_len: self.max_val_len,
};
let mut written = 0u64;
for item in part.iter_unbounded()? {
let (k, v, m) = item?;
write_item(&mut w, &k, &v, m)?;
written = written.checked_add(
(SCRATCH_HEADER_LEN + k.len() + v.len()) as u64,
).ok_or(Error::TooLarge)?;
}
w.flush()?;
crate::write_stats::add(crate::write_stats::Phase::SortScratch, written);
for p in group { let _ = std::fs::remove_file(p); }
next.push(out);
next_counts.push(group_count);
}
self.runs = next;
self.run_counts = next_counts;
pass_no += 1;
}
Ok(())
}
pub fn iter(&mut self) -> Result<MergeIter> {
self.merge_down()?;
self.iter_unbounded()
}
fn iter_unbounded(&mut self) -> Result<MergeIter> {
let mut readers = Vec::with_capacity(self.runs.len());
let per_reader = (256 * 1024 / self.runs.len().max(1)).max(4096);
for p in &self.runs {
readers.push(BufReader::with_capacity(per_reader, File::open(p)?));
}
let mut m = MergeIter {
readers, heap: BinaryHeap::new(), pending_err: None, done: false,
max_key_len: self.max_key_len, max_val_len: self.max_val_len,
expected: self.run_counts.iter().try_fold(0u64, |total, &n| {
total.checked_add(n).ok_or(Error::TooLarge)
})?,
yielded: 0,
};
for i in 0..m.readers.len() {
if let Err(e) = m.pull(i) {
if m.pending_err.is_none() { m.pending_err = Some(e); }
}
}
Ok(m)
}
}
impl Drop for SortedRuns {
fn drop(&mut self) {
if self.owns_dir { let _ = std::fs::remove_dir_all(&self.dir); }
}
}
struct Head { key: Vec<u8>, val: Vec<u8>, marker: bool, from: usize }
impl PartialEq for Head { fn eq(&self, o: &Self) -> bool { self.key == o.key } }
impl Eq for Head {}
impl Ord for Head {
fn cmp(&self, o: &Self) -> std::cmp::Ordering { o.key.cmp(&self.key).then(o.from.cmp(&self.from)) }
}
impl PartialOrd for Head { fn partial_cmp(&self, o: &Self) -> Option<std::cmp::Ordering> { Some(self.cmp(o)) } }
pub struct MergeIter {
readers: Vec<BufReader<File>>,
heap: BinaryHeap<Head>,
pending_err: Option<Error>,
done: bool,
max_key_len: usize,
max_val_len: usize,
expected: u64,
yielded: u64,
}
impl MergeIter {
fn pull(&mut self, i: usize) -> Result<()> {
match read_item(&mut self.readers[i], self.max_key_len, self.max_val_len) {
Ok(Some((k, v, m))) => { self.heap.push(Head { key: k, val: v, marker: m, from: i }); Ok(()) }
Ok(None) => Ok(()),
Err(e) => Err(e.into()),
}
}
}
impl Iterator for MergeIter {
type Item = Result<(Vec<u8>, Vec<u8>, bool)>;
fn next(&mut self) -> Option<Self::Item> {
if self.done { return None; }
if let Some(e) = self.pending_err.take() {
self.done = true;
return Some(Err(e));
}
let Some(h) = self.heap.pop() else {
self.done = true;
if self.yielded == self.expected { return None; }
return Some(Err(Error::Io(invalid_scratch(
"scratch run ended before its recorded item count",
))));
};
if self.yielded >= self.expected {
self.done = true;
return Some(Err(Error::Io(invalid_scratch(
"scratch run exceeded its recorded item count",
))));
}
self.yielded += 1;
if let Err(e) = self.pull(h.from) { self.pending_err = Some(e); }
Some(Ok((h.key, h.val, h.marker)))
}
}
fn enc_leaf(key: &[u8], val: &[u8], compact: bool) -> Vec<u8> {
crate::btree::enc_leaf(key,val,compact)
}
fn enc_interior(key: &[u8], child: u32) -> Vec<u8> {
let mut r = Vec::with_capacity(6 + key.len());
r.extend_from_slice(&(key.len() as u16).to_le_bytes());
r.extend_from_slice(key);
r.extend_from_slice(&child.to_le_bytes());
r
}
fn build_scratch(kind: PageKind, tree_id: u16, page_no: u32, recs: &[Vec<u8>]) -> Result<Vec<u8>> {
let mut scratch = vec![0u8; PAGE_SIZE];
{
let mut p = PageMut::init(&mut scratch, kind, tree_id, page_no);
for r in recs {
let at = p.nentries_pub();
p.insert_slot(at, r)?;
}
p.finalise(0);
}
Ok(scratch)
}
struct PackedPageWriter<'a> {
pool: &'a BufferPool,
region: AlignedRegion,
first: Option<u32>,
pages: usize,
}
impl<'a> PackedPageWriter<'a> {
const CAPACITY: usize = 1;
fn new(pool: &'a BufferPool) -> Result<Self> {
Ok(Self {
pool,
region: AlignedRegion::new(Self::CAPACITY * PAGE_SIZE)?,
first: None,
pages: 0,
})
}
fn push(&mut self, page_no: u32, scratch: &[u8]) -> Result<()> {
if scratch.len() != PAGE_SIZE { return Err(Error::TooLarge); }
if self.pages == Self::CAPACITY
|| self.first.is_some_and(|first| first + self.pages as u32 != page_no)
{
self.flush()?;
}
if self.first.is_none() { self.first = Some(page_no); }
let page = unsafe { self.region.page_mut(self.pages) };
page.copy_from_slice(scratch);
crate::page::seal(page, self.pool.stamp_generation());
self.pages += 1;
Ok(())
}
fn flush(&mut self) -> Result<()> {
if self.pages == 0 { return Ok(()); }
let len = self.pages * PAGE_SIZE;
let bytes = unsafe { self.region.prefix(len) };
self.pool.write_unpooled_run(self.first.unwrap(), bytes)?;
self.first = None;
self.pages = 0;
Ok(())
}
}
pub(crate) enum PageSink<'a> {
Direct(PackedPageWriter<'a>),
Pooled { pool: &'a BufferPool, held: Vec<crate::pool::PinnedWrite<'a>> },
}
impl<'a> PageSink<'a> {
pub(crate) fn direct(pool: &'a BufferPool) -> Result<Self> {
Ok(PageSink::Direct(PackedPageWriter::new(pool)?))
}
pub(crate) fn pooled(pool: &'a BufferPool) -> Self {
PageSink::Pooled { pool, held: Vec::new() }
}
fn reserve(&mut self) -> Result<u32> {
match self {
PageSink::Direct(writer) => writer.pool.allocate_unpooled(),
PageSink::Pooled { pool, held } => {
let guard = pool.allocate()?;
let page_no = guard.page_no();
held.push(guard);
Ok(page_no)
}
}
}
fn push(&mut self, page_no: u32, scratch: &[u8]) -> Result<()> {
match self {
PageSink::Direct(writer) => writer.push(page_no, scratch),
PageSink::Pooled { held, .. } => {
if scratch.len() != PAGE_SIZE { return Err(Error::TooLarge); }
let at = held.iter().position(|guard| guard.page_no() == page_no)
.ok_or(Error::Corrupt { page_no, why: "packed page was never reserved" })?;
let mut guard = held.remove(at);
guard.bytes_mut().copy_from_slice(scratch);
Ok(())
}
}
}
fn finish(&mut self) -> Result<()> {
match self {
PageSink::Direct(writer) => writer.flush(),
PageSink::Pooled { held, .. } => { held.clear(); Ok(()) }
}
}
}
struct Separators {
path: PathBuf,
w: Option<BufWriter<File>>,
count: u64,
max_key_len: usize,
framed_bytes: u64,
}
impl Separators {
fn create(dir: &Path, pack: u64, seq: &mut u64) -> Result<Self> {
let n = *seq;
*seq += 1;
let path = dir.join(format!("sep-{}-{}-{:05}.tmp", std::process::id(), pack, n));
let w = BufWriter::with_capacity(256 * 1024, File::create(&path)?);
Ok(Separators {
path,
w: Some(w),
count: 0,
max_key_len: 0,
framed_bytes: 0,
})
}
fn push(&mut self, key: &[u8], page_no: u32) -> Result<()> {
let w = self.w.as_mut().expect("separators pushed after seal");
write_item(w, key, &page_no.to_le_bytes(), false)?;
self.max_key_len = self.max_key_len.max(key.len());
self.count += 1;
self.framed_bytes = self.framed_bytes
.checked_add((SCRATCH_HEADER_LEN + key.len() + 4) as u64)
.ok_or(Error::TooLarge)?;
Ok(())
}
fn seal(&mut self) -> Result<()> {
if let Some(mut w) = self.w.take() {
w.flush()?;
crate::write_stats::add(
crate::write_stats::Phase::PackScratch,
self.framed_bytes,
);
}
Ok(())
}
fn reader(&self) -> Result<BufReader<File>> {
debug_assert!(self.w.is_none(), "separators read before seal");
Ok(BufReader::with_capacity(256 * 1024, File::open(&self.path)?))
}
}
impl Drop for Separators {
fn drop(&mut self) {
self.w = None;
let _ = std::fs::remove_file(&self.path);
}
}
fn short_level(seen: u64, want: u64) -> Error {
Error::Io(std::io::Error::new(
std::io::ErrorKind::InvalidData,
format!("a separator level came back with {seen} of {want} entries"),
))
}
fn read_sep(r: &mut BufReader<File>, max_key_len: usize) -> Result<Option<(Vec<u8>, u32)>> {
match read_item(r, max_key_len, 4)? {
None => Ok(None),
Some((k, v, _)) => {
if v.len() != 4 {
return Err(Error::Io(std::io::Error::new(
std::io::ErrorKind::InvalidData,
"separator spill record is not a 4-byte page number",
)));
}
Ok(Some((k, u32::from_le_bytes(v[..4].try_into().unwrap()))))
}
}
}
#[derive(Debug, Clone)]
pub(crate) struct PackedRange {
pub(crate) root: u32,
pub(crate) first_leaf: u32,
pub(crate) last_leaf: u32,
pub(crate) rows: u64,
pub(crate) min: Option<Vec<u8>>,
pub(crate) max: Option<Vec<u8>>,
}
pub(crate) fn pack_range<I>(
pool: &BufferPool,
tree_id: u16,
sorted: I,
fill: f32,
scratch_dir: &Path,
last_next: u32,
) -> Result<PackedRange>
where I: Iterator<Item = Result<(Vec<u8>, Vec<u8>, bool)>> {
let mut sink = PageSink::direct(pool)?;
pack_range_into(pool, tree_id, sorted, fill, scratch_dir, last_next, &mut sink)
}
pub(crate) fn pack_range_pooled<I>(
pool: &BufferPool,
tree_id: u16,
sorted: I,
fill: f32,
scratch_dir: &Path,
last_next: u32,
) -> Result<PackedRange>
where I: Iterator<Item = Result<(Vec<u8>, Vec<u8>, bool)>> {
let mut sink = PageSink::pooled(pool);
let packed = pack_range_into(pool, tree_id, sorted, fill, scratch_dir, last_next, &mut sink);
sink.finish()?;
packed
}
fn pack_range_into<I>(
pool: &BufferPool,
tree_id: u16,
sorted: I,
fill: f32,
scratch_dir: &Path,
last_next: u32,
sink: &mut PageSink<'_>,
) -> Result<PackedRange>
where I: Iterator<Item = Result<(Vec<u8>, Vec<u8>, bool)>> {
let usable = ((PAGE_SIZE - HEADER_LEN) as f32 * fill) as usize;
let capacity = PAGE_SIZE - HEADER_LEN;
std::fs::create_dir_all(scratch_dir)?;
static PACK_SEQ: std::sync::atomic::AtomicU64 = std::sync::atomic::AtomicU64::new(0);
let pack = PACK_SEQ.fetch_add(1, std::sync::atomic::Ordering::Relaxed);
let mut seq = 0u64;
let mut level = Separators::create(scratch_dir, pack, &mut seq)?;
let mut cur: Option<(u32, Vec<Vec<u8>>, Vec<u8>, usize)> = None; let mut first_leaf: Option<u32> = None;
let mut last_leaf: Option<u32> = None;
let mut rows = 0u64;
let mut range_min: Option<Vec<u8>> = None;
let mut range_max: Option<Vec<u8>> = None;
let flush_leaf = |cur: &mut Option<(u32, Vec<Vec<u8>>, Vec<u8>, usize)>,
level: &mut Separators,
next_leaf: u32,
sink: &mut PageSink<'_>| -> Result<()> {
if let Some((no, recs, first, _)) = cur.take() {
let mut scratch = build_scratch(PageKind::Leaf, tree_id, no, &recs)?;
{
let mut page = PageMut::reopen(&mut scratch);
page.set_next_leaf(next_leaf);
page.finalise(0);
}
sink.push(no, &scratch)?;
level.push(&first, no)?;
}
Ok(())
};
let mut prev_key: Option<Vec<u8>> = None;
for item in sorted {
let (k, v, is_marker) = item?;
if prev_key.as_deref() == Some(k.as_slice()) {
return Err(Error::DuplicateKey);
}
rows = rows.checked_add(1).ok_or(Error::TooLarge)?;
if range_min.is_none() { range_min = Some(k.clone()); }
range_max = Some(k.clone());
prev_key = Some(k.clone());
let rec = if is_marker {
let m: [u8; 12] = v.as_slice().try_into().map_err(|_| {
Error::Io(invalid_scratch("overflow marker scratch value is not exactly 12 bytes"))
})?;
crate::btree::enc_leaf_marker(&k, &m)
} else {
enc_leaf(&k, &v, pool.compact_cells())
};
let need = rec.len() + 4;
if need > capacity {
return Err(Error::TooLarge);
}
let fits = matches!(&cur, Some((_, _, _, used)) if used + need <= usable);
if !fits {
let no = sink.reserve()?;
flush_leaf(&mut cur, &mut level, no, sink)?;
if first_leaf.is_none() { first_leaf = Some(no); }
last_leaf = Some(no);
cur = Some((no, Vec::new(), k.clone(), 0));
}
if let Some((_, recs, _, used)) = cur.as_mut() { recs.push(rec); *used += need; }
}
flush_leaf(&mut cur, &mut level, last_next, sink)?;
level.seal()?;
if level.count == 0 {
let no = sink.reserve()?;
let scratch = build_scratch(PageKind::Leaf, tree_id, no, &[])?;
sink.push(no, &scratch)?;
sink.finish()?;
return Ok(PackedRange {
root: no,
first_leaf: no,
last_leaf: no,
rows: 0,
min: None,
max: None,
});
}
while level.count > 1 {
let mut up = Separators::create(scratch_dir, pack, &mut seq)?;
{
let mut r = level.reader()?;
let mut seen = 0u64;
let mut pending = std::collections::VecDeque::new();
if let Some(first) = read_sep(&mut r, level.max_key_len)? {
seen += 1;
pending.push_back(first);
}
while let Some((first, child0)) = pending.pop_front() {
let no = sink.reserve()?;
let mut recs: Vec<(Vec<u8>, u32, Vec<u8>)> = Vec::new();
let mut used = 0usize;
loop {
let next = if let Some(queued) = pending.pop_front() {
Some(queued)
} else {
let read = read_sep(&mut r, level.max_key_len)?;
if read.is_some() { seen += 1; }
read
};
let Some((k, child)) = next else { break };
let rec = enc_interior(&k, child);
let need = rec.len() + 4;
if need > capacity { return Err(Error::TooLarge); }
if used + need > usable && !recs.is_empty() {
if seen == level.count && used + need <= capacity {
used += need;
recs.push((k, child, rec));
continue;
}
if seen == level.count {
if recs.len() == 1 { return Err(Error::TooLarge); }
let (moved_key, moved_child, _) = recs.pop().unwrap();
pending.push_back((moved_key, moved_child));
pending.push_back((k, child));
break;
}
pending.push_back((k, child));
break;
}
used += need;
recs.push((k, child, rec));
}
let encoded: Vec<Vec<u8>> = recs.into_iter().map(|(_, _, rec)| rec).collect();
let mut scratch = build_scratch(PageKind::Interior, tree_id, no, &encoded)?;
{
let mut p = PageMut::reopen(&mut scratch);
p.set_child0(child0);
p.finalise(0);
}
sink.push(no, &scratch)?;
up.push(&first, no)?;
}
if seen != level.count { return Err(short_level(seen, level.count)); }
}
up.seal()?;
level = up;
}
let mut r = level.reader()?;
let root = match read_sep(&mut r, level.max_key_len)? {
Some((_, no)) => no,
None => return Err(short_level(0, level.count)),
};
if read_sep(&mut r, level.max_key_len)?.is_some() {
return Err(Error::Io(std::io::Error::new(
std::io::ErrorKind::InvalidData,
"the final separator level held more than the one entry it counted",
)));
}
sink.finish()?;
Ok(PackedRange {
root,
first_leaf: first_leaf.expect("a non-empty level has a first leaf"),
last_leaf: last_leaf.expect("a non-empty level has a last leaf"),
rows,
min: range_min,
max: range_max,
})
}
pub fn pack_tree<I>(pool: &BufferPool, tree_id: u16, sorted: I, fill: f32, scratch_dir: &Path) -> Result<u32>
where I: Iterator<Item = Result<(Vec<u8>, Vec<u8>, bool)>> {
Ok(pack_range(pool, tree_id, sorted, fill, scratch_dir, 0)?.root)
}
pub fn pack_tree_pooled<I>(pool: &BufferPool, tree_id: u16, sorted: I, fill: f32, scratch_dir: &Path)
-> Result<(u32, u64)>
where I: Iterator<Item = Result<(Vec<u8>, Vec<u8>, bool)>> {
let mut peek = sorted.peekable();
if peek.peek().is_none() { return Ok((0, 0)); }
let packed = pack_range_pooled(pool, tree_id, peek, fill, scratch_dir, 0)?;
Ok((packed.root, packed.rows))
}