use std::collections::{HashMap, VecDeque};
use std::fs;
use std::io::{self, Write};
use std::path::{Path, PathBuf};
use std::sync::atomic::{AtomicU64, Ordering};
use std::sync::{Arc, Condvar, Mutex};
use std::time::Instant;
use sha2::{Digest, Sha256};
use crate::cache::KvCache;
use crate::kv_block::BlockHash;
use crate::kv_signature::{
CacheSignature, KvBlock, KvDtype, UnverifiedBlock, BLOCK_FORMAT_VERSION,
READABLE_FORMAT_VERSIONS,
};
const MAGIC: &[u8; 8] = b"FRXKVBLK";
const PREFIX_LEN: usize = 8 + 4 + 4 + 8 + 32;
pub const BLOCK_FILE_EXT: &str = "kvb";
const TMP_DIR: &str = ".tmp";
const DTYPE_F32: u32 = 0;
fn dtype_code(dtype: KvDtype) -> u32 {
match dtype {
KvDtype::F32 => DTYPE_F32,
}
}
fn dtype_from_code(code: u32) -> Option<KvDtype> {
match code {
DTYPE_F32 => Some(KvDtype::F32),
_ => None,
}
}
fn dtype_width(dtype: KvDtype) -> usize {
match dtype {
KvDtype::F32 => 4,
}
}
#[derive(Clone, Debug, PartialEq, Eq)]
pub enum BlockFormatError {
TooShort { len: usize },
BadMagic,
UnsupportedFormat {
found: u32,
readable: &'static [u32],
},
Truncated { expected: u64, actual: u64 },
ChecksumMismatch,
Malformed(&'static str),
UnknownDtype(u32),
}
impl std::fmt::Display for BlockFormatError {
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
match self {
BlockFormatError::TooShort { len } => write!(
f,
"KV block file is {len} bytes, shorter than the {PREFIX_LEN}-byte header prefix"
),
BlockFormatError::BadMagic => write!(f, "KV block file has the wrong magic"),
BlockFormatError::UnsupportedFormat { found, readable } => write!(
f,
"KV block file format version {found} is not readable by this build (readable: {readable:?})"
),
BlockFormatError::Truncated { expected, actual } => write!(
f,
"KV block file declares {expected} bytes but is {actual}; refusing a torn file"
),
BlockFormatError::ChecksumMismatch => {
write!(f, "KV block file failed its SHA-256 checksum")
}
BlockFormatError::Malformed(what) => {
write!(f, "KV block file is malformed: {what}")
}
BlockFormatError::UnknownDtype(code) => {
write!(f, "KV block file has unknown dtype code {code}")
}
}
}
}
impl std::error::Error for BlockFormatError {}
pub fn encode_block(hash: &BlockHash, block: &KvBlock) -> Vec<u8> {
let sig = block.signature();
let mut header = Vec::with_capacity(64 + sig.model.len());
header.extend_from_slice(hash.as_bytes());
header.extend_from_slice(&(sig.n_layers as u32).to_le_bytes());
header.extend_from_slice(&(sig.n_kv_heads as u32).to_le_bytes());
header.extend_from_slice(&(sig.head_dim as u32).to_le_bytes());
header.extend_from_slice(&(sig.tokens as u32).to_le_bytes());
header.extend_from_slice(&dtype_code(sig.dtype).to_le_bytes());
header.extend_from_slice(&(sig.model.len() as u32).to_le_bytes());
header.extend_from_slice(sig.model.as_bytes());
let mut body = Vec::with_capacity(body_len(sig) as usize);
for layer in block.layers() {
for value in &layer.k {
body.extend_from_slice(&value.to_le_bytes());
}
for value in &layer.v {
body.extend_from_slice(&value.to_le_bytes());
}
}
let mut digest = Sha256::new();
digest.update(&header);
digest.update(&body);
let digest: [u8; 32] = digest.finalize().into();
let mut out = Vec::with_capacity(PREFIX_LEN + header.len() + body.len());
out.extend_from_slice(MAGIC);
out.extend_from_slice(&BLOCK_FORMAT_VERSION.to_le_bytes());
out.extend_from_slice(&(header.len() as u32).to_le_bytes());
out.extend_from_slice(&(body.len() as u64).to_le_bytes());
out.extend_from_slice(&digest);
out.extend_from_slice(&header);
out.extend_from_slice(&body);
out
}
fn body_len(sig: &CacheSignature) -> u64 {
let per_layer = sig.tokens as u64
* sig.n_kv_heads as u64
* sig.head_dim as u64
* dtype_width(sig.dtype) as u64;
per_layer * 2 * sig.n_layers as u64
}
pub fn encoded_len(sig: &CacheSignature) -> u64 {
let header = 32 + 4 * 6 + sig.model.len() as u64;
PREFIX_LEN as u64 + header + body_len(sig)
}
#[derive(Debug)]
pub struct DecodedBlock {
pub hash: BlockHash,
pub block: UnverifiedBlock,
}
pub fn decode_block(bytes: &[u8]) -> Result<DecodedBlock, BlockFormatError> {
if bytes.len() < PREFIX_LEN {
return Err(BlockFormatError::TooShort { len: bytes.len() });
}
if &bytes[..8] != MAGIC {
return Err(BlockFormatError::BadMagic);
}
let version = u32::from_le_bytes(bytes[8..12].try_into().unwrap());
if !READABLE_FORMAT_VERSIONS.contains(&version) {
return Err(BlockFormatError::UnsupportedFormat {
found: version,
readable: READABLE_FORMAT_VERSIONS,
});
}
let header_len = u32::from_le_bytes(bytes[12..16].try_into().unwrap()) as u64;
let body_len = u64::from_le_bytes(bytes[16..24].try_into().unwrap());
let declared = PREFIX_LEN as u64 + header_len + body_len;
if declared != bytes.len() as u64 {
return Err(BlockFormatError::Truncated {
expected: declared,
actual: bytes.len() as u64,
});
}
let digest_recorded = &bytes[24..PREFIX_LEN];
let mut digest = Sha256::new();
digest.update(&bytes[PREFIX_LEN..]);
let digest: [u8; 32] = digest.finalize().into();
if digest != digest_recorded {
return Err(BlockFormatError::ChecksumMismatch);
}
let header = &bytes[PREFIX_LEN..PREFIX_LEN + header_len as usize];
let body = &bytes[PREFIX_LEN + header_len as usize..];
if header.len() < 32 + 4 * 6 {
return Err(BlockFormatError::Malformed(
"header shorter than its fields",
));
}
let mut hash = [0u8; 32];
hash.copy_from_slice(&header[..32]);
let hash = BlockHash::from_bytes(hash);
let field = |i: usize| u32::from_le_bytes(header[32 + i * 4..36 + i * 4].try_into().unwrap());
let n_layers = field(0) as usize;
let n_kv_heads = field(1) as usize;
let head_dim = field(2) as usize;
let tokens = field(3) as usize;
let dtype_code = field(4);
let model_len = field(5) as usize;
let dtype = dtype_from_code(dtype_code).ok_or(BlockFormatError::UnknownDtype(dtype_code))?;
if header.len() != 32 + 4 * 6 + model_len {
return Err(BlockFormatError::Malformed(
"model name length disagrees with header",
));
}
let model = std::str::from_utf8(&header[32 + 4 * 6..])
.map_err(|_| BlockFormatError::Malformed("model name is not UTF-8"))?
.to_string();
if n_layers == 0 || n_kv_heads == 0 || head_dim == 0 {
return Err(BlockFormatError::Malformed(
"zero layers, heads, or head dim",
));
}
let per_layer_elems = tokens
.checked_mul(n_kv_heads)
.and_then(|n| n.checked_mul(head_dim))
.ok_or(BlockFormatError::Malformed("layer size overflows"))?;
let expected_body = (per_layer_elems as u64)
.checked_mul(2 * n_layers as u64)
.and_then(|n| n.checked_mul(dtype_width(dtype) as u64))
.ok_or(BlockFormatError::Malformed("body size overflows"))?;
if expected_body != body.len() as u64 {
return Err(BlockFormatError::Malformed(
"body does not match declared dims",
));
}
let mut layers = Vec::with_capacity(n_layers);
let mut offset = 0usize;
for _ in 0..n_layers {
let k = read_f32(&body[offset..offset + per_layer_elems * 4]);
offset += per_layer_elems * 4;
let v = read_f32(&body[offset..offset + per_layer_elems * 4]);
offset += per_layer_elems * 4;
let mut cache = KvCache::new(n_kv_heads, head_dim);
cache.k = k;
cache.v = v;
cache.seq_len = tokens;
layers.push(cache);
}
let signature = CacheSignature {
format_version: version,
model,
n_layers,
n_kv_heads,
head_dim,
dtype,
tokens,
};
Ok(DecodedBlock {
hash,
block: UnverifiedBlock::new(Some(signature), layers),
})
}
fn read_f32(bytes: &[u8]) -> Vec<f32> {
bytes
.chunks_exact(4)
.map(|c| f32::from_le_bytes(c.try_into().unwrap()))
.collect()
}
#[derive(Clone, Debug, PartialEq, Eq)]
pub enum StoreError {
Io {
op: &'static str,
path: PathBuf,
message: String,
},
MissingPayload { hash: BlockHash },
}
impl std::fmt::Display for StoreError {
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
match self {
StoreError::Io { op, path, message } => {
write!(
f,
"KV block store failed to {op} {}: {message}",
path.display()
)
}
StoreError::MissingPayload { hash } => write!(
f,
"KV block store index names {hash:?} but it has neither a file nor a buffered \
payload; the write-ordering invariant was violated"
),
}
}
}
impl std::error::Error for StoreError {}
fn io_err(op: &'static str, path: &Path, err: io::Error) -> StoreError {
StoreError::Io {
op,
path: path.to_path_buf(),
message: err.to_string(),
}
}
pub type FreeSpaceProbe = Arc<dyn Fn(&Path) -> Option<u64> + Send + Sync>;
#[cfg(unix)]
#[allow(clippy::unnecessary_cast)] fn platform_free_bytes(path: &Path) -> Option<u64> {
use std::os::unix::ffi::OsStrExt;
let c_path = std::ffi::CString::new(path.as_os_str().as_bytes()).ok()?;
let stat = unsafe {
let mut stat: libc::statvfs = std::mem::zeroed();
if libc::statvfs(c_path.as_ptr(), &mut stat) != 0 {
return None;
}
stat
};
let block = if stat.f_frsize > 0 {
stat.f_frsize as u64
} else {
stat.f_bsize as u64
};
Some((stat.f_bavail as u64).saturating_mul(block))
}
#[cfg(not(unix))]
fn platform_free_bytes(_path: &Path) -> Option<u64> {
None
}
struct FreeSpace {
checked_at: Option<Instant>,
bytes: Option<u64>,
}
#[derive(Clone)]
pub struct DiskConfig {
pub root: PathBuf,
pub max_bytes: u64,
pub shard_chars: usize,
pub queue_capacity: usize,
pub writer_threads: usize,
pub reader_threads: usize,
pub prefetch_capacity: usize,
pub reserve_bytes: u64,
pub free_space_ttl: std::time::Duration,
pub free_space_probe: FreeSpaceProbe,
}
impl std::fmt::Debug for DiskConfig {
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
f.debug_struct("DiskConfig")
.field("root", &self.root)
.field("max_bytes", &self.max_bytes)
.field("shard_chars", &self.shard_chars)
.field("queue_capacity", &self.queue_capacity)
.field("writer_threads", &self.writer_threads)
.field("reader_threads", &self.reader_threads)
.field("prefetch_capacity", &self.prefetch_capacity)
.field("reserve_bytes", &self.reserve_bytes)
.field("free_space_ttl", &self.free_space_ttl)
.finish_non_exhaustive()
}
}
impl DiskConfig {
pub fn new(root: impl Into<PathBuf>) -> Self {
DiskConfig {
root: root.into(),
max_bytes: 1 << 30,
shard_chars: 2,
queue_capacity: 64,
writer_threads: 1,
reader_threads: 2,
prefetch_capacity: 64,
reserve_bytes: 1 << 30,
free_space_ttl: std::time::Duration::from_secs(2),
free_space_probe: Arc::new(platform_free_bytes),
}
}
pub fn with_max_bytes(mut self, max_bytes: u64) -> Self {
self.max_bytes = max_bytes;
self
}
pub fn with_shard_chars(mut self, shard_chars: usize) -> Self {
self.shard_chars = shard_chars.clamp(1, 8);
self
}
pub fn with_queue_capacity(mut self, queue_capacity: usize) -> Self {
self.queue_capacity = queue_capacity;
self
}
pub fn with_writer_threads(mut self, writer_threads: usize) -> Self {
self.writer_threads = writer_threads;
self
}
pub fn with_reader_threads(mut self, reader_threads: usize) -> Self {
self.reader_threads = reader_threads;
self
}
pub fn with_prefetch_capacity(mut self, prefetch_capacity: usize) -> Self {
self.prefetch_capacity = prefetch_capacity;
self
}
pub fn with_reserve_bytes(mut self, reserve_bytes: u64) -> Self {
self.reserve_bytes = reserve_bytes;
self
}
pub fn with_free_space_ttl(mut self, free_space_ttl: std::time::Duration) -> Self {
self.free_space_ttl = free_space_ttl;
self
}
pub fn with_free_space_probe(mut self, probe: FreeSpaceProbe) -> Self {
self.free_space_probe = probe;
self
}
}
#[derive(Default)]
struct Stats {
writes: AtomicU64,
queued_writes: AtomicU64,
inline_writes: AtomicU64,
write_failures: AtomicU64,
write_skipped: AtomicU64,
write_nanos: AtomicU64,
write_raced_eviction: AtomicU64,
hits: AtomicU64,
buffer_hits: AtomicU64,
misses: AtomicU64,
corrupt: AtomicU64,
incompatible: AtomicU64,
read_nanos: AtomicU64,
evictions: AtomicU64,
evicted_bytes: AtomicU64,
prefetch_issued: AtomicU64,
prefetch_dropped: AtomicU64,
prefetch_hits: AtomicU64,
prefetch_waits: AtomicU64,
async_reads: AtomicU64,
enospc: AtomicU64,
space_clamped: AtomicU64,
}
#[derive(Clone, Debug, Default, PartialEq, Eq)]
pub struct DiskStats {
pub blocks: usize,
pub bytes: u64,
pub disk_bytes: u64,
pub queue_depth: usize,
pub writes: u64,
pub queued_writes: u64,
pub inline_writes: u64,
pub write_failures: u64,
pub write_skipped: u64,
pub write_raced_eviction: u64,
pub write_nanos: u64,
pub hits: u64,
pub buffer_hits: u64,
pub misses: u64,
pub corrupt: u64,
pub incompatible: u64,
pub read_nanos: u64,
pub evictions: u64,
pub evicted_bytes: u64,
pub prefetch_issued: u64,
pub prefetch_dropped: u64,
pub prefetch_hits: u64,
pub prefetch_waits: u64,
pub async_reads: u64,
pub staged_blocks: usize,
pub enospc: u64,
pub space_clamped: u64,
pub effective_capacity: u64,
}
struct Entry {
bytes: u64,
last_used: u64,
published: bool,
generation: u64,
}
struct Index {
entries: HashMap<BlockHash, Entry>,
bytes: u64,
disk_bytes: u64,
clock: u64,
}
impl Index {
fn touch(&mut self) -> u64 {
self.clock += 1;
self.clock
}
fn insert_entry(&mut self, hash: BlockHash, entry: Entry) {
let bytes = entry.bytes;
if let Some(previous) = self.entries.insert(hash, entry) {
self.uncharge(&previous);
}
self.bytes += bytes;
}
fn remove_entry(&mut self, hash: &BlockHash) -> Option<Entry> {
let entry = self.entries.remove(hash)?;
self.uncharge(&entry);
Some(entry)
}
fn uncharge(&mut self, entry: &Entry) {
self.bytes -= entry.bytes;
if entry.published {
self.disk_bytes -= entry.bytes;
}
}
}
enum Source {
Disk(PathBuf),
Buffer(Arc<KvBlock>),
}
struct Buffered {
generation: u64,
block: Arc<KvBlock>,
}
#[derive(Clone, Copy)]
struct WriteJob {
hash: BlockHash,
generation: u64,
}
struct QueueState {
jobs: VecDeque<WriteJob>,
running: usize,
shutdown: bool,
}
struct WriteQueue {
state: Mutex<QueueState>,
ready: Condvar,
idle: Condvar,
capacity: usize,
}
impl WriteQueue {
fn new(capacity: usize) -> Self {
WriteQueue {
state: Mutex::new(QueueState {
jobs: VecDeque::new(),
running: 0,
shutdown: false,
}),
ready: Condvar::new(),
idle: Condvar::new(),
capacity: capacity.max(1),
}
}
fn lock(&self) -> std::sync::MutexGuard<'_, QueueState> {
self.state.lock().expect("kv disk write queue poisoned")
}
fn try_push(&self, job: WriteJob) -> bool {
let mut state = self.lock();
if state.shutdown || state.jobs.len() >= self.capacity {
return false;
}
state.jobs.push_back(job);
self.ready.notify_one();
true
}
fn pop_blocking(&self) -> Option<WriteJob> {
let mut state = self.lock();
loop {
if let Some(job) = state.jobs.pop_front() {
state.running += 1;
return Some(job);
}
if state.shutdown {
return None;
}
state = self
.ready
.wait(state)
.expect("kv disk write queue poisoned");
}
}
fn pop_now(&self) -> Option<WriteJob> {
let mut state = self.lock();
let job = state.jobs.pop_front()?;
state.running += 1;
Some(job)
}
fn finish(&self) {
let mut state = self.lock();
state.running -= 1;
self.idle.notify_all();
}
fn shutdown(&self) {
let mut state = self.lock();
state.shutdown = true;
self.ready.notify_all();
}
fn depth(&self) -> usize {
self.lock().jobs.len()
}
}
pub type ReadOutcome = Result<Option<Arc<KvBlock>>, StoreError>;
struct ReadSlot {
expected: CacheSignature,
outcome: Mutex<Option<ReadOutcome>>,
done: Condvar,
}
impl ReadSlot {
fn pending(expected: CacheSignature) -> Arc<Self> {
Arc::new(ReadSlot {
expected,
outcome: Mutex::new(None),
done: Condvar::new(),
})
}
fn ready(expected: CacheSignature, outcome: ReadOutcome) -> Arc<Self> {
Arc::new(ReadSlot {
expected,
outcome: Mutex::new(Some(outcome)),
done: Condvar::new(),
})
}
fn is_ready(&self) -> bool {
self.outcome
.lock()
.expect("kv disk read slot poisoned")
.is_some()
}
fn fulfil(&self, outcome: ReadOutcome) {
let mut slot = self.outcome.lock().expect("kv disk read slot poisoned");
*slot = Some(outcome);
self.done.notify_all();
}
fn wait(&self) -> ReadOutcome {
let mut slot = self.outcome.lock().expect("kv disk read slot poisoned");
loop {
if let Some(outcome) = slot.as_ref() {
return outcome.clone();
}
slot = self.done.wait(slot).expect("kv disk read slot poisoned");
}
}
}
struct ReadJob {
hash: BlockHash,
path: PathBuf,
slot: Arc<ReadSlot>,
}
struct ReadQueueState {
jobs: VecDeque<ReadJob>,
shutdown: bool,
}
struct ReadQueue {
state: Mutex<ReadQueueState>,
ready: Condvar,
capacity: usize,
}
impl ReadQueue {
fn new(capacity: usize) -> Self {
ReadQueue {
state: Mutex::new(ReadQueueState {
jobs: VecDeque::new(),
shutdown: false,
}),
ready: Condvar::new(),
capacity: capacity.max(1),
}
}
fn lock(&self) -> std::sync::MutexGuard<'_, ReadQueueState> {
self.state.lock().expect("kv disk read queue poisoned")
}
fn try_push(&self, job: ReadJob, demand: bool) -> bool {
let mut state = self.lock();
if state.shutdown || state.jobs.len() >= self.capacity {
return false;
}
if demand {
state.jobs.push_front(job);
} else {
state.jobs.push_back(job);
}
self.ready.notify_one();
true
}
fn pop_blocking(&self) -> Option<ReadJob> {
let mut state = self.lock();
loop {
if let Some(job) = state.jobs.pop_front() {
return Some(job);
}
if state.shutdown {
return None;
}
state = self.ready.wait(state).expect("kv disk read queue poisoned");
}
}
fn shutdown(&self) {
let mut state = self.lock();
state.shutdown = true;
self.ready.notify_all();
}
}
pub struct ReadHandle {
shared: Arc<Shared>,
hash: BlockHash,
slot: Arc<ReadSlot>,
staged: bool,
}
impl ReadHandle {
pub fn is_ready(&self) -> bool {
self.slot.is_ready()
}
pub fn try_claim(&self) -> Option<ReadOutcome> {
if !self.slot.is_ready() {
return None;
}
Some(self.claim())
}
pub fn wait(self) -> ReadOutcome {
self.claim()
}
fn claim(&self) -> ReadOutcome {
let outcome = self.slot.wait();
if self.staged {
self.shared.unstage(&self.hash, &self.slot);
}
outcome
}
}
impl std::fmt::Debug for ReadHandle {
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
f.debug_struct("ReadHandle")
.field("hash", &self.hash)
.field("ready", &self.is_ready())
.finish()
}
}
#[cfg(test)]
type Hook = Arc<dyn Fn(&BlockHash) + Send + Sync>;
#[cfg(test)]
#[derive(Clone, Copy, Debug, Default, PartialEq, Eq)]
enum WriteOrder {
#[default]
BufferThenIndex,
IndexBeforeBuffer,
DropBufferBeforeMarking,
}
#[cfg(test)]
#[derive(Default)]
struct Hooks {
order: Mutex<WriteOrder>,
after_rename: Mutex<Option<Hook>>,
in_put_window: Mutex<Option<Hook>>,
in_publish_window: Mutex<Option<Hook>>,
fail_with_enospc: std::sync::atomic::AtomicBool,
}
#[cfg(test)]
impl Hooks {
fn fire(slot: &Mutex<Option<Hook>>, hash: &BlockHash) {
let hook = slot.lock().expect("kv disk hook poisoned").clone();
if let Some(hook) = hook {
hook(hash);
}
}
}
struct Shared {
root: PathBuf,
shard_chars: usize,
max_bytes: u64,
index: Mutex<Index>,
buffer: Mutex<HashMap<BlockHash, Buffered>>,
queue: WriteQueue,
reads: ReadQueue,
staging: Mutex<HashMap<BlockHash, Arc<ReadSlot>>>,
prefetch_capacity: usize,
has_readers: bool,
reserve_bytes: u64,
free_space_ttl: std::time::Duration,
free_space_probe: FreeSpaceProbe,
free_space: Mutex<FreeSpace>,
stats: Stats,
seq: AtomicU64,
generation: AtomicU64,
#[cfg(test)]
hooks: Hooks,
}
pub struct DiskKvStore {
shared: Arc<Shared>,
writers: Vec<std::thread::JoinHandle<()>>,
readers: Vec<std::thread::JoinHandle<()>>,
}
impl Drop for DiskKvStore {
fn drop(&mut self) {
self.shared.queue.shutdown();
self.shared.reads.shutdown();
for writer in self.writers.drain(..) {
let _ = writer.join();
}
for reader in self.readers.drain(..) {
let _ = reader.join();
}
}
}
impl DiskKvStore {
pub fn open(config: DiskConfig) -> Result<Self, StoreError> {
let root = config.root.clone();
fs::create_dir_all(&root).map_err(|e| io_err("create", &root, e))?;
let tmp = root.join(TMP_DIR);
fs::create_dir_all(&tmp).map_err(|e| io_err("create", &tmp, e))?;
let shared = Arc::new(Shared {
root,
shard_chars: config.shard_chars.clamp(1, 8),
max_bytes: config.max_bytes,
index: Mutex::new(Index {
entries: HashMap::new(),
bytes: 0,
disk_bytes: 0,
clock: 0,
}),
buffer: Mutex::new(HashMap::new()),
queue: WriteQueue::new(config.queue_capacity),
reads: ReadQueue::new(config.queue_capacity),
staging: Mutex::new(HashMap::new()),
prefetch_capacity: config.prefetch_capacity,
has_readers: config.reader_threads > 0,
reserve_bytes: config.reserve_bytes,
free_space_ttl: config.free_space_ttl,
free_space_probe: Arc::clone(&config.free_space_probe),
free_space: Mutex::new(FreeSpace {
checked_at: None,
bytes: None,
}),
stats: Stats::default(),
seq: AtomicU64::new(0),
generation: AtomicU64::new(0),
#[cfg(test)]
hooks: Hooks::default(),
});
let mut writers = Vec::with_capacity(config.writer_threads);
for n in 0..config.writer_threads {
let shared = Arc::clone(&shared);
let handle = std::thread::Builder::new()
.name(format!("ferrox-kv-write-{n}"))
.spawn(move || {
while let Some(job) = shared.queue.pop_blocking() {
shared.run_job(job);
shared.queue.finish();
}
})
.map_err(|e| io_err("spawn writer for", &config.root, e))?;
writers.push(handle);
}
let mut readers = Vec::with_capacity(config.reader_threads);
for n in 0..config.reader_threads {
let shared = Arc::clone(&shared);
let handle = std::thread::Builder::new()
.name(format!("ferrox-kv-read-{n}"))
.spawn(move || {
while let Some(job) = shared.reads.pop_blocking() {
shared.stats.async_reads.fetch_add(1, Ordering::Relaxed);
let outcome = shared.read_timed(&job.path, &job.hash, &job.slot.expected);
job.slot.fulfil(outcome);
}
})
.map_err(|e| io_err("spawn reader for", &config.root, e))?;
readers.push(handle);
}
Ok(DiskKvStore {
shared,
writers,
readers,
})
}
pub fn root(&self) -> &Path {
&self.shared.root
}
pub fn put(&self, hash: BlockHash, block: KvBlock) -> Result<(), StoreError> {
self.shared.put(hash, block, false)
}
pub fn put_blocking(&self, hash: BlockHash, block: KvBlock) -> Result<(), StoreError> {
self.shared.put(hash, block, true)
}
pub fn flush(&self) {
loop {
if let Some(job) = self.shared.queue.pop_now() {
self.shared.run_job(job);
self.shared.queue.finish();
continue;
}
let state = self.shared.queue.lock();
if state.jobs.is_empty() && state.running == 0 {
return;
}
let _ = self
.shared
.queue
.idle
.wait_timeout(state, std::time::Duration::from_millis(1));
}
}
pub fn get(&self, hash: &BlockHash, expected: &CacheSignature) -> ReadOutcome {
self.shared.read_async(hash, expected, true).wait()
}
pub fn read_async(&self, hash: &BlockHash, expected: &CacheSignature) -> ReadHandle {
self.shared.read_async(hash, expected, true)
}
pub fn prefetch(&self, hashes: &[BlockHash], expected: &CacheSignature) {
self.shared.prefetch(hashes, expected);
}
pub fn clear_prefetch(&self) {
self.shared
.staging
.lock()
.expect("kv disk staging poisoned")
.clear();
}
pub fn reindex(&self) -> Result<usize, StoreError> {
self.shared.reindex()
}
pub fn remove(&self, hash: &BlockHash) {
self.shared.quarantine(hash);
}
pub fn contains(&self, hash: &BlockHash) -> bool {
let index = self.shared.index.lock().expect("kv disk index poisoned");
index.entries.contains_key(hash)
}
pub fn capacity(&self) -> u64 {
self.shared.max_bytes
}
pub fn effective_capacity(&self) -> u64 {
let used = {
let index = self.shared.index.lock().expect("kv disk index poisoned");
index.bytes
};
self.shared.effective_capacity(used)
}
pub fn block_path(&self, hash: &BlockHash) -> PathBuf {
self.shared.block_path(hash)
}
pub fn stats(&self) -> DiskStats {
self.shared.stats()
}
}
impl Shared {
fn next_generation(&self) -> u64 {
self.generation.fetch_add(1, Ordering::SeqCst) + 1
}
fn put(&self, hash: BlockHash, block: KvBlock, inline: bool) -> Result<(), StoreError> {
let bytes = encoded_len(block.signature());
let block = Arc::new(block);
let generation = self.next_generation();
#[cfg(test)]
let index_first = *self.hooks.order.lock().expect("kv disk hook poisoned")
== WriteOrder::IndexBeforeBuffer;
#[cfg(not(test))]
let index_first = false;
if index_first {
self.reserve(hash, bytes, generation);
#[cfg(test)]
Hooks::fire(&self.hooks.in_put_window, &hash);
self.buffer_block(hash, generation, Arc::clone(&block));
} else {
self.buffer_block(hash, generation, Arc::clone(&block));
#[cfg(test)]
Hooks::fire(&self.hooks.in_put_window, &hash);
self.reserve(hash, bytes, generation);
}
let job = WriteJob { hash, generation };
if !inline && self.queue.try_push(job) {
self.stats.queued_writes.fetch_add(1, Ordering::Relaxed);
return Ok(());
}
if !inline {
self.stats.inline_writes.fetch_add(1, Ordering::Relaxed);
}
self.run_write(job, block)
}
fn buffer_block(&self, hash: BlockHash, generation: u64, block: Arc<KvBlock>) {
self.buffer
.lock()
.expect("kv disk buffer poisoned")
.insert(hash, Buffered { generation, block });
}
fn release_buffer(&self, hash: &BlockHash, generation: u64) {
let mut buffer = self.buffer.lock().expect("kv disk buffer poisoned");
if buffer.get(hash).is_some_and(|b| b.generation == generation) {
buffer.remove(hash);
}
}
fn buffered(&self, hash: &BlockHash, generation: u64) -> Option<Arc<KvBlock>> {
let buffer = self.buffer.lock().expect("kv disk buffer poisoned");
buffer
.get(hash)
.filter(|b| b.generation == generation)
.map(|b| Arc::clone(&b.block))
}
fn run_job(&self, job: WriteJob) {
match self.buffered(&job.hash, job.generation) {
Some(block) => {
let _ = self.run_write(job, block);
}
None => {
self.stats.write_skipped.fetch_add(1, Ordering::Relaxed);
}
}
}
fn run_write(&self, job: WriteJob, block: Arc<KvBlock>) -> Result<(), StoreError> {
let started = Instant::now();
let result = self.write_and_publish(&job.hash, &block, job.generation);
self.stats
.write_nanos
.fetch_add(started.elapsed().as_nanos() as u64, Ordering::Relaxed);
match result {
Ok(()) => {
self.stats.writes.fetch_add(1, Ordering::Relaxed);
Ok(())
}
Err(err) => {
self.stats.write_failures.fetch_add(1, Ordering::Relaxed);
self.abandon(&job.hash, job.generation);
Err(err)
}
}
}
fn reserve(&self, hash: BlockHash, bytes: u64, generation: u64) {
let mut index = self.index.lock().expect("kv disk index poisoned");
let last_used = index.touch();
index.insert_entry(
hash,
Entry {
bytes,
last_used,
published: false,
generation,
},
);
let victims = self.collect_victims(&mut index, Some(&hash));
drop(index);
self.discard(victims);
}
fn abandon(&self, hash: &BlockHash, generation: u64) {
{
let mut index = self.index.lock().expect("kv disk index poisoned");
if index
.entries
.get(hash)
.is_some_and(|e| e.generation == generation)
{
index.remove_entry(hash);
}
}
self.release_buffer(hash, generation);
}
fn write_and_publish(
&self,
hash: &BlockHash,
block: &KvBlock,
generation: u64,
) -> Result<(), StoreError> {
let bytes = encode_block(hash, block);
let final_path = self.block_path(hash);
let shard = final_path.parent().expect("block path has a parent");
fs::create_dir_all(shard).map_err(|e| io_err("create", shard, e))?;
let tmp_path = self.tmp_path(hash);
{
let mut file =
fs::File::create(&tmp_path).map_err(|e| io_err("create", &tmp_path, e))?;
#[cfg(test)]
let written = if self.hooks.fail_with_enospc.load(Ordering::Relaxed) {
Err(io::Error::from(io::ErrorKind::StorageFull))
} else {
file.write_all(&bytes)
};
#[cfg(not(test))]
let written = file.write_all(&bytes);
if let Err(e) = written {
let _ = fs::remove_file(&tmp_path);
self.note_if_enospc(&e);
return Err(io_err("write", &tmp_path, e));
}
if let Err(e) = file.sync_all() {
let _ = fs::remove_file(&tmp_path);
self.note_if_enospc(&e);
return Err(io_err("sync", &tmp_path, e));
}
}
fs::rename(&tmp_path, &final_path).map_err(|e| {
let _ = fs::remove_file(&tmp_path);
io_err("publish", &final_path, e)
})?;
#[cfg(test)]
Hooks::fire(&self.hooks.after_rename, hash);
#[cfg(test)]
let drop_buffer_first = *self.hooks.order.lock().expect("kv disk hook poisoned")
== WriteOrder::DropBufferBeforeMarking;
#[cfg(not(test))]
let drop_buffer_first = false;
let survived = if drop_buffer_first {
self.release_buffer(hash, generation);
#[cfg(test)]
Hooks::fire(&self.hooks.in_publish_window, hash);
self.mark_published(hash, generation)
} else {
let survived = self.mark_published(hash, generation);
#[cfg(test)]
Hooks::fire(&self.hooks.in_publish_window, hash);
self.release_buffer(hash, generation);
survived
};
if !survived {
let _ = fs::remove_file(&final_path);
self.stats
.write_raced_eviction
.fetch_add(1, Ordering::Relaxed);
}
Ok(())
}
fn mark_published(&self, hash: &BlockHash, generation: u64) -> bool {
let mut index = self.index.lock().expect("kv disk index poisoned");
match index.entries.get_mut(hash) {
Some(entry) if entry.generation == generation => {
if !entry.published {
entry.published = true;
let bytes = entry.bytes;
index.disk_bytes += bytes;
}
true
}
_ => false,
}
}
fn source(&self, hash: &BlockHash) -> Result<Option<Source>, StoreError> {
let mut index = self.index.lock().expect("kv disk index poisoned");
let clock = index.clock + 1;
let Some(entry) = index.entries.get_mut(hash) else {
return Ok(None);
};
entry.last_used = clock;
let published = entry.published;
index.clock = clock;
if published {
return Ok(Some(Source::Disk(self.block_path(hash))));
}
let buffered = self
.buffer
.lock()
.expect("kv disk buffer poisoned")
.get(hash)
.map(|b| Arc::clone(&b.block));
match buffered {
Some(block) => Ok(Some(Source::Buffer(block))),
None => Err(StoreError::MissingPayload { hash: *hash }),
}
}
fn read_async(
self: &Arc<Self>,
hash: &BlockHash,
expected: &CacheSignature,
demand: bool,
) -> ReadHandle {
let staged = {
let staging = self.staging.lock().expect("kv disk staging poisoned");
staging
.get(hash)
.filter(|slot| &slot.expected == expected)
.map(Arc::clone)
};
if let Some(slot) = staged {
if demand {
let counter = if slot.is_ready() {
&self.stats.prefetch_hits
} else {
&self.stats.prefetch_waits
};
counter.fetch_add(1, Ordering::Relaxed);
}
return self.handle(*hash, slot, true);
}
let ready = |outcome: ReadOutcome| ReadHandle {
shared: Arc::clone(self),
hash: *hash,
slot: ReadSlot::ready(expected.clone(), outcome),
staged: false,
};
let path = match self.source(hash) {
Err(err) => return ready(Err(err)),
Ok(None) => {
self.stats.misses.fetch_add(1, Ordering::Relaxed);
return ready(Ok(None));
}
Ok(Some(Source::Buffer(block))) => {
if block.signature() != expected {
self.stats.incompatible.fetch_add(1, Ordering::Relaxed);
return ready(Ok(None));
}
self.stats.buffer_hits.fetch_add(1, Ordering::Relaxed);
return ready(Ok(Some(block)));
}
Ok(Some(Source::Disk(path))) => path,
};
let slot = ReadSlot::pending(expected.clone());
let dispatched = self.has_readers && {
self.staging
.lock()
.expect("kv disk staging poisoned")
.insert(*hash, Arc::clone(&slot));
let job = ReadJob {
hash: *hash,
path: path.clone(),
slot: Arc::clone(&slot),
};
let pushed = self.reads.try_push(job, demand);
if !pushed {
self.unstage(hash, &slot);
}
pushed
};
if dispatched {
return self.handle(*hash, slot, true);
}
slot.fulfil(self.read_timed(&path, hash, expected));
self.handle(*hash, slot, false)
}
fn handle(self: &Arc<Self>, hash: BlockHash, slot: Arc<ReadSlot>, staged: bool) -> ReadHandle {
ReadHandle {
shared: Arc::clone(self),
hash,
slot,
staged,
}
}
fn unstage(&self, hash: &BlockHash, slot: &Arc<ReadSlot>) {
let mut staging = self.staging.lock().expect("kv disk staging poisoned");
if staging.get(hash).is_some_and(|s| Arc::ptr_eq(s, slot)) {
staging.remove(hash);
}
}
fn prefetch(self: &Arc<Self>, hashes: &[BlockHash], expected: &CacheSignature) {
if !self.has_readers {
self.stats
.prefetch_dropped
.fetch_add(hashes.len() as u64, Ordering::Relaxed);
return;
}
for hash in hashes {
let room = {
let staging = self.staging.lock().expect("kv disk staging poisoned");
!staging.contains_key(hash) && staging.len() < self.prefetch_capacity
};
if !room {
self.stats.prefetch_dropped.fetch_add(1, Ordering::Relaxed);
continue;
}
let handle = self.read_async(hash, expected, false);
if handle.staged {
self.stats.prefetch_issued.fetch_add(1, Ordering::Relaxed);
} else {
self.stats.prefetch_dropped.fetch_add(1, Ordering::Relaxed);
}
drop(handle);
}
}
fn effective_capacity(&self, used: u64) -> u64 {
let Some(free) = self.free_bytes() else {
return self.max_bytes;
};
let headroom = free as i128 - self.reserve_bytes as i128;
let allowed = (used as i128 + headroom).max(0) as u64;
self.max_bytes.min(allowed)
}
fn free_bytes(&self) -> Option<u64> {
let mut cache = self.free_space.lock().expect("kv disk free space poisoned");
if let Some(checked_at) = cache.checked_at {
if checked_at.elapsed() < self.free_space_ttl {
return cache.bytes;
}
}
let bytes = (self.free_space_probe)(&self.root);
cache.checked_at = Some(Instant::now());
cache.bytes = bytes;
bytes
}
fn note_enospc(&self) {
self.stats.enospc.fetch_add(1, Ordering::Relaxed);
{
let mut cache = self.free_space.lock().expect("kv disk free space poisoned");
cache.checked_at = None;
cache.bytes = None;
}
let victims = {
let mut index = self.index.lock().expect("kv disk index poisoned");
self.collect_victims(&mut index, None)
};
self.discard(victims);
}
fn note_if_enospc(&self, err: &io::Error) {
if err.kind() == io::ErrorKind::StorageFull {
self.note_enospc();
}
}
fn read_timed(&self, path: &Path, hash: &BlockHash, expected: &CacheSignature) -> ReadOutcome {
let started = Instant::now();
let outcome = self.read_verified(path, hash, expected);
self.stats
.read_nanos
.fetch_add(started.elapsed().as_nanos() as u64, Ordering::Relaxed);
outcome
}
fn read_verified(
&self,
path: &Path,
hash: &BlockHash,
expected: &CacheSignature,
) -> Result<Option<Arc<KvBlock>>, StoreError> {
let bytes = match fs::read(path) {
Ok(bytes) => bytes,
Err(e) if e.kind() == io::ErrorKind::NotFound => {
self.stats.misses.fetch_add(1, Ordering::Relaxed);
self.drop_entry(hash);
return Ok(None);
}
Err(e) => return Err(io_err("read", path, e)),
};
let decoded = match decode_block(&bytes) {
Ok(decoded) => decoded,
Err(_) => {
self.stats.corrupt.fetch_add(1, Ordering::Relaxed);
self.quarantine(hash);
return Ok(None);
}
};
if &decoded.hash != hash {
self.stats.corrupt.fetch_add(1, Ordering::Relaxed);
self.quarantine(hash);
return Ok(None);
}
match decoded.block.verify(expected) {
Ok(block) => {
self.stats.hits.fetch_add(1, Ordering::Relaxed);
Ok(Some(Arc::new(block)))
}
Err(_) => {
self.stats.incompatible.fetch_add(1, Ordering::Relaxed);
Ok(None)
}
}
}
fn quarantine(&self, hash: &BlockHash) {
self.drop_entry(hash);
self.buffer
.lock()
.expect("kv disk buffer poisoned")
.remove(hash);
let _ = fs::remove_file(self.block_path(hash));
}
fn drop_entry(&self, hash: &BlockHash) {
let mut index = self.index.lock().expect("kv disk index poisoned");
index.remove_entry(hash);
}
fn reindex(&self) -> Result<usize, StoreError> {
let tmp = self.root.join(TMP_DIR);
if let Ok(entries) = fs::read_dir(&tmp) {
for entry in entries.flatten() {
let _ = fs::remove_file(entry.path());
}
}
let mut found = Vec::new();
let shards = fs::read_dir(&self.root).map_err(|e| io_err("read", &self.root, e))?;
for shard in shards.flatten() {
if !shard.file_type().map(|t| t.is_dir()).unwrap_or(false) {
continue;
}
if shard.file_name() == TMP_DIR {
continue;
}
let Ok(files) = fs::read_dir(shard.path()) else {
continue;
};
for file in files.flatten() {
let path = file.path();
if path.extension().and_then(|e| e.to_str()) != Some(BLOCK_FILE_EXT) {
continue;
}
let Some(hash) = path
.file_stem()
.and_then(|s| s.to_str())
.and_then(parse_hex_hash)
else {
continue;
};
let Ok(meta) = file.metadata() else { continue };
found.push((hash, meta.len()));
}
}
let mut adopted = 0;
let mut index = self.index.lock().expect("kv disk index poisoned");
for (hash, bytes) in found {
if index.entries.contains_key(&hash) {
continue;
}
let last_used = index.touch();
index.insert_entry(
hash,
Entry {
bytes,
last_used,
published: true,
generation: self.next_generation(),
},
);
index.disk_bytes += bytes;
adopted += 1;
}
let victims = self.collect_victims(&mut index, None);
drop(index);
self.discard(victims);
Ok(adopted)
}
fn collect_victims(&self, index: &mut Index, protect: Option<&BlockHash>) -> Vec<Victim> {
let budget = self.effective_capacity(index.disk_bytes);
if budget < self.max_bytes {
self.stats.space_clamped.fetch_add(1, Ordering::Relaxed);
}
if index.bytes <= budget {
return Vec::new();
}
let mut candidates: Vec<(u64, BlockHash)> = index
.entries
.iter()
.filter(|(hash, _)| Some(*hash) != protect)
.map(|(hash, entry)| (entry.last_used, *hash))
.collect();
candidates.sort_unstable();
let mut victims = Vec::new();
for (_, hash) in candidates {
if index.bytes <= budget {
break;
}
if let Some(entry) = index.remove_entry(&hash) {
self.stats.evictions.fetch_add(1, Ordering::Relaxed);
self.stats
.evicted_bytes
.fetch_add(entry.bytes, Ordering::Relaxed);
victims.push(Victim {
hash,
published: entry.published,
});
}
}
victims
}
fn discard(&self, victims: Vec<Victim>) {
if victims.is_empty() {
return;
}
{
let mut buffer = self.buffer.lock().expect("kv disk buffer poisoned");
for victim in &victims {
buffer.remove(&victim.hash);
}
}
for victim in victims {
if victim.published {
let _ = fs::remove_file(self.block_path(&victim.hash));
}
}
}
fn stats(&self) -> DiskStats {
let index = self.index.lock().expect("kv disk index poisoned");
let stats = &self.stats;
DiskStats {
blocks: index.entries.len(),
bytes: index.bytes,
queue_depth: self.queue.depth(),
writes: stats.writes.load(Ordering::Relaxed),
queued_writes: stats.queued_writes.load(Ordering::Relaxed),
inline_writes: stats.inline_writes.load(Ordering::Relaxed),
write_failures: stats.write_failures.load(Ordering::Relaxed),
write_skipped: stats.write_skipped.load(Ordering::Relaxed),
write_raced_eviction: stats.write_raced_eviction.load(Ordering::Relaxed),
write_nanos: stats.write_nanos.load(Ordering::Relaxed),
hits: stats.hits.load(Ordering::Relaxed),
buffer_hits: stats.buffer_hits.load(Ordering::Relaxed),
misses: stats.misses.load(Ordering::Relaxed),
corrupt: stats.corrupt.load(Ordering::Relaxed),
incompatible: stats.incompatible.load(Ordering::Relaxed),
read_nanos: stats.read_nanos.load(Ordering::Relaxed),
evictions: stats.evictions.load(Ordering::Relaxed),
evicted_bytes: stats.evicted_bytes.load(Ordering::Relaxed),
prefetch_issued: stats.prefetch_issued.load(Ordering::Relaxed),
prefetch_dropped: stats.prefetch_dropped.load(Ordering::Relaxed),
prefetch_hits: stats.prefetch_hits.load(Ordering::Relaxed),
prefetch_waits: stats.prefetch_waits.load(Ordering::Relaxed),
async_reads: stats.async_reads.load(Ordering::Relaxed),
staged_blocks: self.staging.lock().expect("kv disk staging poisoned").len(),
enospc: stats.enospc.load(Ordering::Relaxed),
space_clamped: stats.space_clamped.load(Ordering::Relaxed),
effective_capacity: self.effective_capacity(index.disk_bytes),
disk_bytes: index.disk_bytes,
}
}
fn block_path(&self, hash: &BlockHash) -> PathBuf {
self.root
.join(hash.shard_prefix(self.shard_chars))
.join(format!("{}.{BLOCK_FILE_EXT}", hash.to_hex()))
}
fn tmp_path(&self, hash: &BlockHash) -> PathBuf {
let n = self.seq.fetch_add(1, Ordering::Relaxed);
self.root.join(TMP_DIR).join(format!(
"{}.{}.{n}.tmp",
hash.shard_prefix(16),
std::process::id()
))
}
}
struct Victim {
hash: BlockHash,
published: bool,
}
fn parse_hex_hash(text: &str) -> Option<BlockHash> {
if text.len() != 64 {
return None;
}
let mut out = [0u8; 32];
for (i, byte) in out.iter_mut().enumerate() {
let hi = text.as_bytes()[i * 2] as char;
let lo = text.as_bytes()[i * 2 + 1] as char;
*byte = ((hi.to_digit(16)? << 4) | lo.to_digit(16)?) as u8;
}
Some(BlockHash::from_bytes(out))
}
#[cfg(test)]
mod tests {
use super::*;
use crate::kv_block::BlockHasher;
use std::sync::atomic::AtomicUsize;
struct TempDir(PathBuf);
impl TempDir {
fn new(tag: &str) -> Self {
static N: AtomicU64 = AtomicU64::new(0);
let path = std::env::temp_dir().join(format!(
"ferrox-kvdisk-{tag}-{}-{}",
std::process::id(),
N.fetch_add(1, Ordering::Relaxed)
));
let _ = fs::remove_dir_all(&path);
fs::create_dir_all(&path).expect("temp dir");
TempDir(path)
}
fn path(&self) -> &Path {
&self.0
}
}
impl Drop for TempDir {
fn drop(&mut self) {
let _ = fs::remove_dir_all(&self.0);
}
}
fn layer(n_kv_heads: usize, head_dim: usize, tokens: usize, fill: f32) -> KvCache {
let mut cache = KvCache::new(n_kv_heads, head_dim);
for t in 0..tokens {
let k = vec![fill + t as f32; n_kv_heads * head_dim];
let v = vec![fill - t as f32; n_kv_heads * head_dim];
cache.push(&k, &v).expect("unpooled push cannot fail");
}
cache
}
fn block(model: &str, n_layers: usize, tokens: usize, fill: f32) -> KvBlock {
let layers = (0..n_layers)
.map(|l| layer(2, 4, tokens, fill + l as f32 * 100.0))
.collect();
KvBlock::stamp(model, layers).expect("stamp")
}
fn expected(model: &str, n_layers: usize, tokens: usize) -> CacheSignature {
CacheSignature::expected(model, n_layers, 2, 4, tokens)
}
fn hash(n: usize) -> BlockHash {
BlockHasher::new("model-a", &[] as &[&str]).chain(&[n, n + 1], 2)[0]
}
fn plenty() -> FreeSpaceProbe {
Arc::new(|_: &Path| Some(1 << 40))
}
fn store(dir: &TempDir, max_bytes: u64) -> DiskKvStore {
DiskKvStore::open(
DiskConfig::new(dir.path())
.with_max_bytes(max_bytes)
.with_writer_threads(0)
.with_free_space_probe(plenty()),
)
.expect("open")
}
fn put_now(store: &DiskKvStore, hash: BlockHash, block: KvBlock) {
store.put_blocking(hash, block).expect("put");
}
#[test]
fn a_block_round_trips_through_a_file() {
let dir = TempDir::new("roundtrip");
let store = store(&dir, 1 << 20);
let h = hash(1);
let written = block("model-a", 3, 4, 1.0);
let copy = block("model-a", 3, 4, 1.0);
put_now(&store, h, written);
let read = store
.get(&h, &expected("model-a", 3, 4))
.expect("get")
.expect("the block just written must be found");
assert_eq!(read.layers().len(), 3);
for (a, b) in read.layers().iter().zip(copy.layers()) {
assert_eq!(a.k, b.k);
assert_eq!(a.v, b.v);
assert_eq!(a.seq_len, b.seq_len);
}
let stats = store.stats();
assert_eq!(stats.hits, 1);
assert_eq!(stats.writes, 1);
assert_eq!(stats.blocks, 1);
assert!(stats.read_nanos > 0, "a read must be timed");
assert!(stats.write_nanos > 0, "a write must be timed");
}
#[test]
fn the_accounted_size_is_the_real_file_size() {
let dir = TempDir::new("size");
let store = store(&dir, 1 << 20);
let h = hash(2);
let written = block("model-a", 2, 8, 0.25);
let predicted = encoded_len(written.signature());
put_now(&store, h, written);
let on_disk = fs::metadata(store.block_path(&h)).expect("stat").len();
assert_eq!(
predicted, on_disk,
"the budget charges what the file really costs"
);
assert_eq!(store.stats().bytes, on_disk);
}
#[test]
fn blocks_are_sharded_by_hash_prefix() {
let dir = TempDir::new("shard");
let store = DiskKvStore::open(
DiskConfig::new(dir.path())
.with_shard_chars(2)
.with_writer_threads(0)
.with_free_space_probe(plenty()),
)
.expect("open");
let h = hash(3);
put_now(&store, h, block("model-a", 1, 2, 1.0));
let path = store.block_path(&h);
assert_eq!(
path.parent()
.unwrap()
.file_name()
.unwrap()
.to_str()
.unwrap(),
&h.to_hex()[..2]
);
assert!(path.exists());
}
#[test]
fn a_truncated_file_is_refused_at_every_cut_point() {
let h = hash(4);
let bytes = encode_block(&h, &block("model-a", 2, 4, 3.0));
assert!(bytes.len() > PREFIX_LEN + 16);
let err = decode_block(&bytes[..PREFIX_LEN - 1]).expect_err("short file");
assert_eq!(
err,
BlockFormatError::TooShort {
len: PREFIX_LEN - 1
}
);
for cut in [PREFIX_LEN, PREFIX_LEN + 8, bytes.len() - 4, bytes.len() - 1] {
let err = decode_block(&bytes[..cut]).expect_err("truncated file");
assert_eq!(
err,
BlockFormatError::Truncated {
expected: bytes.len() as u64,
actual: cut as u64,
},
"a file cut at {cut} must be refused"
);
}
let mut flipped = bytes.clone();
let last = flipped.len() - 1;
flipped[last] ^= 0xff;
assert_eq!(
decode_block(&flipped).expect_err("altered file"),
BlockFormatError::ChecksumMismatch
);
let mut alien = bytes;
alien[0] = b'X';
assert_eq!(
decode_block(&alien).expect_err("foreign file"),
BlockFormatError::BadMagic
);
}
#[test]
fn a_torn_file_on_disk_is_a_miss_and_is_quarantined() {
let dir = TempDir::new("torn");
let store = store(&dir, 1 << 20);
let h = hash(5);
put_now(&store, h, block("model-a", 2, 4, 1.0));
let path = store.block_path(&h);
let full = fs::read(&path).expect("read back");
fs::write(&path, &full[..full.len() / 2]).expect("truncate");
let got = store.get(&h, &expected("model-a", 2, 4)).expect("get");
assert!(got.is_none(), "a torn block must not be returned");
assert_eq!(store.stats().corrupt, 1);
assert!(!path.exists(), "a torn block must not be left to trip over");
assert!(!store.contains(&h));
}
#[test]
fn an_unreadable_format_version_is_refused() {
let h = hash(6);
let mut bytes = encode_block(&h, &block("model-a", 1, 2, 1.0));
bytes[8..12].copy_from_slice(&99u32.to_le_bytes());
let mut digest = Sha256::new();
digest.update(&bytes[PREFIX_LEN..]);
let digest: [u8; 32] = digest.finalize().into();
bytes[24..PREFIX_LEN].copy_from_slice(&digest);
assert_eq!(
decode_block(&bytes).expect_err("unknown version"),
BlockFormatError::UnsupportedFormat {
found: 99,
readable: READABLE_FORMAT_VERSIONS,
}
);
}
#[test]
fn a_block_from_a_different_config_is_a_miss_not_a_hit() {
let dir = TempDir::new("config");
let store = store(&dir, 1 << 20);
let h = hash(7);
put_now(&store, h, block("model-a", 2, 4, 1.0));
assert!(store
.get(&h, &expected("model-b", 2, 4))
.expect("get")
.is_none());
assert!(store
.get(&h, &CacheSignature::expected("model-a", 2, 8, 4, 4))
.expect("get")
.is_none());
assert_eq!(store.stats().incompatible, 2);
assert_eq!(store.stats().hits, 0);
assert!(store
.get(&h, &expected("model-a", 2, 4))
.expect("get")
.is_some());
}
#[test]
fn a_file_stored_under_the_wrong_name_is_rejected() {
let dir = TempDir::new("misfiled");
let store = store(&dir, 1 << 20);
let (a, b) = (hash(8), hash(9));
put_now(&store, a, block("model-a", 1, 2, 1.0));
put_now(&store, b, block("model-a", 1, 2, 2.0));
let bytes = fs::read(store.block_path(&b)).expect("read b");
fs::write(store.block_path(&a), bytes).expect("misfile");
assert!(store
.get(&a, &expected("model-a", 1, 2))
.expect("get")
.is_none());
assert_eq!(store.stats().corrupt, 1);
}
#[test]
fn eviction_keeps_the_store_inside_its_budget() {
let dir = TempDir::new("evict");
let one = encoded_len(block("model-a", 1, 4, 1.0).signature());
let store = store(&dir, one * 2 + 8);
let hashes: Vec<BlockHash> = (0..4).map(|i| hash(20 + i)).collect();
for (i, h) in hashes.iter().enumerate() {
put_now(&store, *h, block("model-a", 1, 4, i as f32));
assert!(
store.stats().bytes <= store.capacity(),
"the store must never sit over budget"
);
}
let stats = store.stats();
assert_eq!(stats.blocks, 2);
assert_eq!(stats.evictions, 2);
assert!(stats.evicted_bytes >= one * 2);
for h in &hashes[..2] {
assert!(!store.contains(h));
assert!(
!store.block_path(h).exists(),
"an evicted file must be deleted"
);
}
for h in &hashes[2..] {
assert!(store.contains(h));
}
}
#[test]
fn a_read_makes_a_block_the_least_likely_eviction_victim() {
let dir = TempDir::new("lru");
let one = encoded_len(block("model-a", 1, 4, 1.0).signature());
let store = store(&dir, one * 2 + 8);
let (a, b, c) = (hash(30), hash(31), hash(32));
put_now(&store, a, block("model-a", 1, 4, 1.0));
put_now(&store, b, block("model-a", 1, 4, 2.0));
assert!(store.get(&a, &expected("model-a", 1, 4)).unwrap().is_some());
put_now(&store, c, block("model-a", 1, 4, 3.0));
assert!(store.contains(&a), "a recently read block must survive");
assert!(!store.contains(&b));
assert!(store.contains(&c));
}
#[test]
fn a_block_evicted_mid_write_does_not_leave_its_file_behind() {
let dir = TempDir::new("raced");
let store = store(&dir, 1 << 20);
let h = hash(40);
{
let evicting = Arc::clone(&store.shared);
let mut hook = store
.shared
.hooks
.after_rename
.lock()
.expect("hook lock poisoned");
*hook = Some(Arc::new(move |hash: &BlockHash| {
evicting.drop_entry(hash);
}));
}
put_now(&store, h, block("model-a", 1, 4, 1.0));
assert!(
!store.block_path(&h).exists(),
"a file published for an entry that no longer exists must be withdrawn"
);
assert!(!store.contains(&h));
let stats = store.stats();
assert_eq!(stats.write_raced_eviction, 1);
assert_eq!(stats.bytes, 0, "no bytes may be left unaccounted");
}
#[test]
fn no_temp_files_survive_a_successful_write() {
let dir = TempDir::new("tmp");
let store = store(&dir, 1 << 20);
for i in 0..4 {
put_now(&store, hash(50 + i), block("model-a", 1, 2, i as f32));
}
let leftovers: Vec<_> = fs::read_dir(dir.path().join(TMP_DIR))
.expect("tmp dir")
.flatten()
.collect();
assert!(
leftovers.is_empty(),
"temp files must not accumulate: {leftovers:?}"
);
}
#[test]
fn a_new_store_reattaches_to_what_the_previous_one_published() {
let dir = TempDir::new("restart");
let h = hash(60);
{
let store = store(&dir, 1 << 20);
put_now(&store, h, block("model-a", 2, 4, 7.0));
}
let orphan = dir.path().join(TMP_DIR).join("dead.tmp");
fs::write(&orphan, b"half a block").expect("orphan");
let reopened = store(&dir, 1 << 20);
assert!(
!reopened.contains(&h),
"reattaching must be an explicit step, not a side effect of open()"
);
assert_eq!(reopened.reindex().expect("reindex"), 1);
assert!(reopened.contains(&h));
assert!(!orphan.exists(), "an unpublished temp file must be swept");
let read = reopened
.get(&h, &expected("model-a", 2, 4))
.expect("get")
.expect("a block written before the restart must still be readable");
assert_eq!(read.tokens(), 4);
}
#[test]
fn reindex_evicts_down_to_the_budget() {
let dir = TempDir::new("reindex-evict");
let one = encoded_len(block("model-a", 1, 4, 1.0).signature());
{
let store = store(&dir, 1 << 20);
for i in 0..4 {
put_now(&store, hash(70 + i), block("model-a", 1, 4, i as f32));
}
}
let small = store(&dir, one * 2 + 8);
small.reindex().expect("reindex");
let stats = small.stats();
assert_eq!(stats.blocks, 2, "a shrunken budget must bind on restart");
assert!(stats.bytes <= small.capacity());
}
#[test]
fn an_absent_block_is_a_plain_miss() {
let dir = TempDir::new("miss");
let store = store(&dir, 1 << 20);
assert!(store
.get(&hash(80), &expected("model-a", 1, 2))
.expect("get")
.is_none());
assert_eq!(store.stats().misses, 1);
assert_eq!(store.stats().corrupt, 0);
}
#[test]
fn a_file_deleted_behind_the_stores_back_is_a_miss() {
let dir = TempDir::new("vanished");
let store = store(&dir, 1 << 20);
let h = hash(90);
put_now(&store, h, block("model-a", 1, 2, 1.0));
fs::remove_file(store.block_path(&h)).expect("remove");
assert!(store
.get(&h, &expected("model-a", 1, 2))
.expect("get")
.is_none());
assert!(!store.contains(&h));
assert_eq!(store.stats().bytes, 0);
}
#[test]
fn rewriting_a_block_does_not_double_charge_it() {
let dir = TempDir::new("rewrite");
let store = store(&dir, 1 << 20);
let h = hash(100);
put_now(&store, h, block("model-a", 1, 4, 1.0));
let once = store.stats().bytes;
put_now(&store, h, block("model-a", 1, 4, 1.0));
assert_eq!(store.stats().bytes, once);
assert_eq!(store.stats().blocks, 1);
}
#[test]
fn hex_names_round_trip() {
let h = hash(110);
assert_eq!(parse_hex_hash(&h.to_hex()), Some(h));
assert_eq!(parse_hex_hash("nothex"), None);
assert_eq!(parse_hex_hash(&"z".repeat(64)), None);
}
fn probe_window(order: WriteOrder, publish_window: bool) -> (usize, usize) {
let dir = TempDir::new("ordering");
let store = store(&dir, 1 << 20);
*store.shared.hooks.order.lock().unwrap() = order;
let violations = Arc::new(AtomicUsize::new(0));
let served = Arc::new(AtomicUsize::new(0));
let reader = Arc::clone(&store.shared);
let v = Arc::clone(&violations);
let s = Arc::clone(&served);
let hook: Hook = Arc::new(move |hash: &BlockHash| {
match reader
.read_async(hash, &expected("model-a", 1, 4), true)
.wait()
{
Ok(Some(_)) => {
s.fetch_add(1, Ordering::Relaxed);
}
Ok(None) => {}
Err(StoreError::MissingPayload { .. }) => {
v.fetch_add(1, Ordering::Relaxed);
}
Err(other) => panic!("unexpected store error: {other}"),
}
});
let slot = if publish_window {
&store.shared.hooks.in_publish_window
} else {
&store.shared.hooks.in_put_window
};
*slot.lock().unwrap() = Some(hook);
put_now(&store, hash(200), block("model-a", 1, 4, 1.0));
(
violations.load(Ordering::Relaxed),
served.load(Ordering::Relaxed),
)
}
#[test]
fn a_reader_never_sees_an_index_hit_with_no_payload() {
let (violations, _) = probe_window(WriteOrder::BufferThenIndex, false);
assert_eq!(violations, 0, "admission window must be safe");
let (violations, served) = probe_window(WriteOrder::BufferThenIndex, true);
assert_eq!(violations, 0, "publication window must be safe");
assert_eq!(
served, 1,
"the reader must actually have reached the block, or this test proves nothing"
);
}
#[test]
fn indexing_before_buffering_is_caught() {
let (violations, _) = probe_window(WriteOrder::IndexBeforeBuffer, false);
assert_eq!(
violations, 1,
"index-then-buffer must be detected as an invariant violation"
);
}
#[test]
fn releasing_the_buffer_before_publishing_is_caught() {
let (violations, _) = probe_window(WriteOrder::DropBufferBeforeMarking, true);
assert_eq!(
violations, 1,
"release-then-mark must be detected as an invariant violation"
);
}
#[test]
fn concurrent_readers_never_see_an_index_hit_with_no_payload() {
let dir = TempDir::new("concurrent");
let one = encoded_len(block("model-a", 1, 4, 1.0).signature());
let store = Arc::new(
DiskKvStore::open(
DiskConfig::new(dir.path())
.with_max_bytes(one * 8)
.with_queue_capacity(4)
.with_writer_threads(2)
.with_free_space_probe(plenty()),
)
.expect("open"),
);
let hashes: Vec<BlockHash> = (0..16).map(|i| hash(300 + i)).collect();
let violations = Arc::new(AtomicUsize::new(0));
let stop = Arc::new(std::sync::atomic::AtomicBool::new(false));
let readers: Vec<_> = (0..4)
.map(|_| {
let store = Arc::clone(&store);
let hashes = hashes.clone();
let violations = Arc::clone(&violations);
let stop = Arc::clone(&stop);
std::thread::spawn(move || {
let want = expected("model-a", 1, 4);
while !stop.load(Ordering::Relaxed) {
for h in &hashes {
match store.get(h, &want) {
Ok(_) => {}
Err(StoreError::MissingPayload { .. }) => {
violations.fetch_add(1, Ordering::Relaxed);
}
Err(other) => panic!("unexpected store error: {other}"),
}
}
}
})
})
.collect();
for round in 0..4 {
for (i, h) in hashes.iter().enumerate() {
store
.put(*h, block("model-a", 1, 4, (round * 16 + i) as f32))
.expect("put");
}
}
store.flush();
stop.store(true, Ordering::Relaxed);
for reader in readers {
reader.join().expect("reader thread");
}
assert_eq!(
violations.load(Ordering::Relaxed),
0,
"no reader may ever see an index hit with no payload"
);
let stats = store.stats();
assert!(
stats.buffer_hits > 0,
"readers must have caught blocks still in the write buffer, \
or this test never entered the window"
);
assert!(stats.evictions > 0, "the budget must have bound");
assert!(stats.bytes <= store.capacity());
}
#[test]
fn a_queued_block_is_readable_before_it_reaches_disk() {
let dir = TempDir::new("buffered");
let store = store(&dir, 1 << 20);
let h = hash(400);
store.put(h, block("model-a", 1, 4, 1.0)).expect("put");
assert!(
!store.block_path(&h).exists(),
"nothing has been written yet"
);
let got = store
.get(&h, &expected("model-a", 1, 4))
.expect("get")
.expect("a queued block must be readable immediately");
assert_eq!(got.tokens(), 4);
assert_eq!(store.stats().buffer_hits, 1);
store.flush();
assert!(store.block_path(&h).exists(), "flush must publish it");
assert!(store
.get(&h, &expected("model-a", 1, 4))
.expect("get")
.is_some());
assert_eq!(store.stats().hits, 1, "and now it comes off the disk");
}
#[test]
fn a_full_queue_writes_inline_rather_than_dropping_the_block() {
let dir = TempDir::new("backpressure");
let store = DiskKvStore::open(
DiskConfig::new(dir.path())
.with_queue_capacity(2)
.with_writer_threads(0)
.with_free_space_probe(plenty()),
)
.expect("open");
let hashes: Vec<BlockHash> = (0..5).map(|i| hash(500 + i)).collect();
for (i, h) in hashes.iter().enumerate() {
store
.put(*h, block("model-a", 1, 4, i as f32))
.expect("put");
}
let stats = store.stats();
assert_eq!(stats.queued_writes, 2, "the queue holds exactly its cap");
assert_eq!(stats.inline_writes, 3, "the rest fall back to this thread");
assert_eq!(stats.writes, 3, "and the fallbacks really wrote");
let want = expected("model-a", 1, 4);
for h in &hashes {
assert!(
store.get(h, &want).expect("get").is_some(),
"no block may be lost to a full queue"
);
}
store.flush();
for h in &hashes {
assert!(store.block_path(h).exists(), "flush publishes the rest");
}
}
#[test]
fn a_queued_write_evicted_before_it_runs_is_skipped() {
let dir = TempDir::new("skipped");
let store = store(&dir, 1 << 20);
let h = hash(600);
store.put(h, block("model-a", 1, 4, 1.0)).expect("put");
store.remove(&h);
store.flush();
let stats = store.stats();
assert_eq!(stats.write_skipped, 1);
assert_eq!(stats.writes, 0);
assert!(!store.block_path(&h).exists());
assert_eq!(stats.bytes, 0);
}
#[test]
fn a_superseded_queued_write_does_not_overwrite_the_newer_block() {
let dir = TempDir::new("superseded");
let store = store(&dir, 1 << 20);
let h = hash(700);
store.put(h, block("model-a", 1, 4, 1.0)).expect("put");
store.put(h, block("model-a", 1, 4, 9.0)).expect("put");
store.flush();
let got = store
.get(&h, &expected("model-a", 1, 4))
.expect("get")
.expect("hit");
assert_eq!(
got.layers()[0].k[0],
9.0,
"the newer block must win, not whichever write ran last"
);
assert_eq!(store.stats().write_skipped, 1);
assert_eq!(store.stats().blocks, 1);
}
fn reading_store(dir: &TempDir, readers: usize) -> DiskKvStore {
DiskKvStore::open(
DiskConfig::new(dir.path())
.with_writer_threads(0)
.with_reader_threads(readers)
.with_free_space_probe(plenty()),
)
.expect("open")
}
fn wait_staged(store: &DiskKvStore, hash: &BlockHash) {
for _ in 0..2000 {
let ready = store
.shared
.staging
.lock()
.unwrap()
.get(hash)
.map(|slot| slot.is_ready())
.unwrap_or(false);
if ready {
return;
}
std::thread::sleep(std::time::Duration::from_millis(1));
}
panic!("prefetch never completed");
}
#[test]
fn a_prefetched_block_is_already_read_when_the_request_arrives() {
let dir = TempDir::new("prefetch");
let store = reading_store(&dir, 1);
let h = hash(800);
put_now(&store, h, block("model-a", 2, 4, 5.0));
let want = expected("model-a", 2, 4);
store.prefetch(&[h], &want);
wait_staged(&store, &h);
fs::remove_file(store.block_path(&h)).expect("remove");
let got = store
.get(&h, &want)
.expect("get")
.expect("the prefetch already had it");
assert_eq!(got.tokens(), 4);
assert_eq!(got.layers()[0].k[0], 5.0);
let stats = store.stats();
assert_eq!(
stats.prefetch_hits, 1,
"the request must have found it ready"
);
assert_eq!(
stats.hits, 1,
"and the file must have been read exactly once"
);
assert_eq!(stats.async_reads, 1, "on a reader thread, not the caller's");
assert_eq!(stats.staged_blocks, 0, "a claimed read leaves staging");
}
#[test]
fn a_whole_chain_can_be_read_ahead_in_one_call() {
let dir = TempDir::new("chain");
let store = reading_store(&dir, 2);
let hashes: Vec<BlockHash> = (0..4).map(|i| hash(810 + i)).collect();
for (i, h) in hashes.iter().enumerate() {
put_now(&store, *h, block("model-a", 1, 4, i as f32));
}
let want = expected("model-a", 1, 4);
store.prefetch(&hashes, &want);
for h in &hashes {
wait_staged(&store, h);
fs::remove_file(store.block_path(h)).expect("remove");
}
for (i, h) in hashes.iter().enumerate() {
let got = store.get(h, &want).expect("get").expect("read ahead");
assert_eq!(got.layers()[0].k[0], i as f32);
}
let stats = store.stats();
assert_eq!(stats.prefetch_issued, 4);
assert_eq!(stats.prefetch_hits, 4);
assert_eq!(stats.hits, 4, "four blocks, four reads, none repeated");
}
#[test]
fn a_request_joins_a_read_already_running_rather_than_repeating_it() {
let dir = TempDir::new("join");
let store = reading_store(&dir, 1);
let h = hash(820);
put_now(&store, h, block("model-a", 1, 4, 1.0));
let want = expected("model-a", 1, 4);
store.prefetch(&[h], &want);
let got = store.get(&h, &want).expect("get").expect("hit");
assert_eq!(got.tokens(), 4);
let stats = store.stats();
assert_eq!(
stats.prefetch_hits + stats.prefetch_waits,
1,
"the request either found the read done or waited for it"
);
assert_eq!(stats.hits, 1, "one physical read, whichever way it went");
}
#[test]
fn a_staged_read_is_not_reused_by_a_reader_that_wants_another_shape() {
let dir = TempDir::new("staged-shape");
let store = reading_store(&dir, 1);
let h = hash(830);
put_now(&store, h, block("model-a", 1, 4, 1.0));
store.prefetch(&[h], &expected("model-a", 1, 4));
wait_staged(&store, &h);
let got = store.get(&h, &expected("model-b", 1, 4)).expect("get");
assert!(got.is_none(), "a different model must not be served");
assert_eq!(store.stats().incompatible, 1);
assert_eq!(
store.stats().prefetch_hits,
0,
"the staged answer was for another expectation and must not be claimed"
);
}
#[test]
fn prefetching_is_bounded() {
let dir = TempDir::new("prefetch-bound");
let store = DiskKvStore::open(
DiskConfig::new(dir.path())
.with_writer_threads(0)
.with_reader_threads(1)
.with_prefetch_capacity(2)
.with_free_space_probe(plenty()),
)
.expect("open");
let hashes: Vec<BlockHash> = (0..6).map(|i| hash(840 + i)).collect();
for (i, h) in hashes.iter().enumerate() {
put_now(&store, *h, block("model-a", 1, 4, i as f32));
}
store.prefetch(&hashes, &expected("model-a", 1, 4));
let stats = store.stats();
assert!(
stats.staged_blocks <= 2,
"staging must respect its cap, got {}",
stats.staged_blocks
);
assert!(
stats.prefetch_dropped >= 4,
"the refusals must be visible, got {}",
stats.prefetch_dropped
);
let want = expected("model-a", 1, 4);
for (i, h) in hashes.iter().enumerate() {
let got = store.get(h, &want).expect("get").expect("hit");
assert_eq!(got.layers()[0].k[0], i as f32);
}
}
#[test]
fn without_reader_threads_reads_run_on_the_caller() {
let dir = TempDir::new("no-readers");
let store = reading_store(&dir, 0);
let h = hash(850);
put_now(&store, h, block("model-a", 1, 4, 1.0));
let want = expected("model-a", 1, 4);
store.prefetch(&[h], &want);
assert_eq!(store.stats().prefetch_dropped, 1);
assert_eq!(store.stats().staged_blocks, 0);
assert!(store.get(&h, &want).expect("get").is_some());
let stats = store.stats();
assert_eq!(stats.hits, 1);
assert_eq!(stats.async_reads, 0);
}
#[test]
fn a_read_handle_can_be_polled_to_completion() {
let dir = TempDir::new("handle");
let store = reading_store(&dir, 1);
let h = hash(860);
put_now(&store, h, block("model-a", 1, 4, 2.0));
let want = expected("model-a", 1, 4);
let handle = store.read_async(&h, &want);
for _ in 0..2000 {
if let Some(outcome) = handle.try_claim() {
let got = outcome.expect("read").expect("hit");
assert_eq!(got.layers()[0].k[0], 2.0);
assert_eq!(store.stats().staged_blocks, 0);
return;
}
std::thread::sleep(std::time::Duration::from_millis(1));
}
panic!("read never completed");
}
#[test]
fn a_miss_is_answered_without_dispatching_a_read() {
let dir = TempDir::new("ready-miss");
let store = reading_store(&dir, 1);
let handle = store.read_async(&hash(870), &expected("model-a", 1, 4));
assert!(handle.is_ready(), "a miss must not cost a thread hop");
assert!(handle.wait().expect("read").is_none());
assert_eq!(store.stats().async_reads, 0);
}
#[test]
fn clearing_the_prefetch_releases_staged_blocks() {
let dir = TempDir::new("clear");
let store = reading_store(&dir, 1);
let h = hash(880);
put_now(&store, h, block("model-a", 1, 4, 1.0));
store.prefetch(&[h], &expected("model-a", 1, 4));
wait_staged(&store, &h);
assert_eq!(store.stats().staged_blocks, 1);
store.clear_prefetch();
assert_eq!(store.stats().staged_blocks, 0);
}
fn budgeted_store(
dir: &TempDir,
max_bytes: u64,
reserve: u64,
ttl: std::time::Duration,
probe: FreeSpaceProbe,
) -> DiskKvStore {
DiskKvStore::open(
DiskConfig::new(dir.path())
.with_max_bytes(max_bytes)
.with_reserve_bytes(reserve)
.with_free_space_ttl(ttl)
.with_free_space_probe(probe)
.with_writer_threads(0)
.with_reader_threads(0),
)
.expect("open")
}
fn dir_bytes(root: &Path) -> u64 {
let mut total = 0;
let Ok(entries) = fs::read_dir(root) else {
return 0;
};
for entry in entries.flatten() {
let path = entry.path();
if path.is_dir() {
total += dir_bytes(&path);
} else if let Ok(meta) = entry.metadata() {
total += meta.len();
}
}
total
}
#[test]
fn the_ceiling_falls_when_the_filesystem_fills_up() {
let dir = TempDir::new("budget");
let one = encoded_len(block("model-a", 1, 4, 1.0).signature());
let device = Arc::new(AtomicU64::new(1 << 40));
let probe: FreeSpaceProbe = {
let device = Arc::clone(&device);
let root = dir.path().to_path_buf();
Arc::new(move |_: &Path| {
Some(
device
.load(Ordering::Relaxed)
.saturating_sub(dir_bytes(&root)),
)
})
};
let reserve = one * 2;
let store = budgeted_store(&dir, one * 100, reserve, std::time::Duration::ZERO, probe);
for i in 0..4 {
put_now(&store, hash(900 + i), block("model-a", 1, 4, i as f32));
}
assert_eq!(store.stats().blocks, 4);
assert_eq!(store.stats().evictions, 0, "nothing binds yet");
assert_eq!(
store.effective_capacity(),
store.capacity(),
"with a terabyte free the configured budget is the ceiling"
);
device.store(one * 6, Ordering::Relaxed);
assert_eq!(
store.effective_capacity(),
one * 4,
"the ceiling must follow the device down to total - reserve"
);
for i in 0..6 {
put_now(&store, hash(910 + i), block("model-a", 1, 4, i as f32));
let on_disk = dir_bytes(dir.path());
assert!(
on_disk + reserve <= one * 6,
"the store must hand the device its reserve back before the \
filesystem has to: {on_disk} bytes used of {}, {reserve} reserved",
one * 6
);
}
let stats = store.stats();
assert_eq!(stats.blocks, 4, "settled at total - reserve");
assert!(stats.evictions >= 6, "got {}", stats.evictions);
assert!(stats.space_clamped > 0, "the clamp must be visible");
assert!(
stats.bytes < one * 100,
"far under the configured budget it never reached"
);
assert_eq!(
stats.disk_bytes,
dir_bytes(dir.path()),
"the store's idea of its disk footprint must be the real one"
);
}
#[test]
fn the_free_space_reading_is_cached_for_its_ttl() {
let dir = TempDir::new("ttl");
let calls = Arc::new(AtomicUsize::new(0));
let probe: FreeSpaceProbe = {
let calls = Arc::clone(&calls);
Arc::new(move |_: &Path| {
calls.fetch_add(1, Ordering::Relaxed);
Some(1 << 40)
})
};
let store = budgeted_store(&dir, 1 << 20, 0, std::time::Duration::from_secs(60), probe);
for _ in 0..5 {
store.effective_capacity();
}
for i in 0..3 {
put_now(&store, hash(920 + i), block("model-a", 1, 4, i as f32));
}
assert_eq!(
calls.load(Ordering::Relaxed),
1,
"a TTL'd reading must not be re-taken per operation"
);
}
#[test]
fn enospc_throws_away_the_cached_free_space() {
let dir = TempDir::new("enospc");
let calls = Arc::new(AtomicUsize::new(0));
let probe: FreeSpaceProbe = {
let calls = Arc::clone(&calls);
Arc::new(move |_: &Path| {
calls.fetch_add(1, Ordering::Relaxed);
Some(1 << 40)
})
};
let store = budgeted_store(
&dir,
1 << 20,
0,
std::time::Duration::from_secs(3600),
probe,
);
let h = hash(930);
put_now(&store, h, block("model-a", 1, 4, 1.0));
store.effective_capacity();
assert_eq!(calls.load(Ordering::Relaxed), 1);
store
.shared
.hooks
.fail_with_enospc
.store(true, Ordering::Relaxed);
let full = hash(931);
let err = store
.put_blocking(full, block("model-a", 1, 4, 2.0))
.expect_err("a full filesystem must be reported, not swallowed");
assert!(matches!(err, StoreError::Io { .. }), "{err}");
let stats = store.stats();
assert_eq!(stats.enospc, 1);
assert!(
calls.load(Ordering::Relaxed) > 1,
"ENOSPC must invalidate the cached reading immediately"
);
assert!(
!store.contains(&full),
"a block that could not be written must not be indexed"
);
assert_eq!(stats.write_failures, 1);
let leftovers: Vec<_> = fs::read_dir(dir.path().join(TMP_DIR))
.expect("tmp dir")
.flatten()
.collect();
assert!(
leftovers.is_empty(),
"a failed write must clean up after itself: {leftovers:?}"
);
assert!(store
.get(&h, &expected("model-a", 1, 4))
.expect("get")
.is_some());
store
.shared
.hooks
.fail_with_enospc
.store(false, Ordering::Relaxed);
put_now(&store, full, block("model-a", 1, 4, 2.0));
assert!(store.contains(&full));
}
#[test]
fn an_unmeasurable_filesystem_falls_back_to_the_configured_budget() {
let dir = TempDir::new("unknowable");
let store = budgeted_store(
&dir,
1 << 20,
1 << 30,
std::time::Duration::ZERO,
Arc::new(|_: &Path| None),
);
assert_eq!(store.effective_capacity(), 1 << 20);
for i in 0..3 {
put_now(&store, hash(940 + i), block("model-a", 1, 4, i as f32));
}
assert_eq!(store.stats().blocks, 3);
assert_eq!(store.stats().evictions, 0);
}
#[test]
#[cfg(unix)]
fn the_platform_probe_measures_a_real_filesystem() {
let dir = TempDir::new("statvfs");
let free = platform_free_bytes(dir.path()).expect("statvfs on a directory that exists");
assert!(free > 0, "a writable temp dir with zero bytes free?");
assert!(
platform_free_bytes(&dir.path().join("no-such-dir")).is_none(),
"a path that does not exist cannot report free space"
);
}
}