Skip to main content

ferrox_core/
kv_disk.rs

1//! The disk tier for the KV prefix cache: where a block goes so that a
2//! prefix survives eviction from RAM, and a process restart.
3//!
4//! One file per block, named by its
5//! [`BlockHash`](crate::kv_block::BlockHash) and sharded into
6//! subdirectories by a hex prefix of that hash, so a machine that has
7//! cached a million prefixes never puts a million entries in one
8//! directory. The payload is the block's per-layer K and V tensors,
9//! flattened, in the cache's own dtype -- **no re-encoding, no
10//! compression**: a KV block is already dense float data, and a
11//! compressor would only spend CPU on the request path to lose.
12//!
13//! # What a reader is protected from
14//!
15//! A cache file outlives the process that wrote it, so every failure
16//! mode here is "someone else's bytes":
17//!
18//! - **A torn write.** The publish is temp-file + `fsync` + `rename`,
19//!   which is atomic within a directory on every filesystem ferrox
20//!   targets -- a reader sees the whole file or no file. But a crash
21//!   mid-`write` to the *temp* file, a truncated copy, or a partial
22//!   restore from a backup can still leave a short file lying around,
23//!   so the format records its own total length and a SHA-256 of its
24//!   body. A file that does not match is refused
25//!   ([`BlockFormatError`]), never partially deserialized.
26//! - **A different build.** The format is versioned with an explicit
27//!   readable-set; an unknown version is refused rather than guessed
28//!   at.
29//! - **A different model or config.** That is
30//!   [`kv_signature`](crate::kv_signature)'s job, and this module does
31//!   not duplicate it: a decoded file becomes an
32//!   [`UnverifiedBlock`], and only
33//!   [`UnverifiedBlock::verify`] against the reader's own expectation
34//!   produces a usable [`KvBlock`].
35//!
36//! # The write-ordering invariant
37//!
38//! A write is accepted on one thread and finished on another, so for a
39//! while a block is "in the store" without being on disk. The rule that
40//! makes that safe, and the one every step below is ordered around:
41//!
42//! > **buffer -> index -> queue.** A concurrent reader must never see
43//! > an index hit for a block that has neither a file nor a buffered
44//! > payload.
45//!
46//! So the payload is reachable *before* anything claims the block
47//! exists, and on the way out the file is published *before* the
48//! buffered copy is released. A reader holds the index lock while it
49//! consults the buffer, because "this block is not on disk yet" and
50//! "here is its payload" have to be one decision -- as two, the writer
51//! can publish and release in between and the reader finds nothing.
52//! When that invariant does break, the reader gets
53//! [`StoreError::MissingPayload`] rather than a quiet miss: a
54//! correctness bug that degrades into a cache miss is a bug nobody ever
55//! finds.
56//!
57//! The queue is bounded and **never drops**: a full queue makes the
58//! caller write the block itself ([`DiskStats::inline_writes`] counts
59//! it). Dropping writes silently would be indistinguishable from a cold
60//! cache later.
61//!
62//! # Layout
63//!
64//! ```text
65//! <root>/.tmp/<hash>.<pid>.<n>.tmp     in-progress writes
66//! <root>/<hh>/<full-hex>.kvb           published blocks (hh = shard prefix)
67//! ```
68//!
69//! # File format (version 2)
70//!
71//! ```text
72//! magic           8   b"FRXKVBLK"
73//! format_version  4   u32 LE, checked against READABLE_FORMAT_VERSIONS
74//! header_len      4   u32 LE
75//! body_len        8   u64 LE
76//! digest         32   SHA-256 over header || body
77//! header  header_len  block hash, dims, dtype, block layout, model identity
78//! body      body_len  per layer: all K elements, then all V elements
79//! ```
80//!
81//! Version 2 added the two block-layout fields -- block size and
82//! sliding window -- and version 1 was dropped from the readable set
83//! rather than being read with the window assumed absent. A v1 file
84//! cannot say what window it was cut under, and "it did not say" is not
85//! "there was none": see [`kv_swa`](crate::kv_swa) for what a
86//! mis-aligned block does to an answer. A restart onto this build
87//! therefore starts from a cold cache once, and the old files are
88//! evicted as unreadable rather than reinterpreted.
89//!
90//! The digest covers header and body but not the fixed prefix, so the
91//! lengths are checked against the real file size *before* anything is
92//! hashed or parsed -- a 4 GB `body_len` on a 200-byte file is rejected
93//! by arithmetic, not by allocating.
94
95use std::collections::{HashMap, VecDeque};
96use std::fs;
97use std::io::{self, Write};
98use std::path::{Path, PathBuf};
99use std::sync::atomic::{AtomicU64, Ordering};
100use std::sync::{Arc, Condvar, Mutex};
101use std::time::Instant;
102
103use sha2::{Digest, Sha256};
104
105use crate::cache::KvCache;
106use crate::kv_block::BlockHash;
107use crate::kv_signature::{
108    CacheSignature, KvBlock, KvDtype, UnverifiedBlock, BLOCK_FORMAT_VERSION,
109    READABLE_FORMAT_VERSIONS,
110};
111use crate::kv_swa::{BlockLayout, BlockLayoutError};
112
113const MAGIC: &[u8; 8] = b"FRXKVBLK";
114/// Fixed u32 fields in a block header, after the 32-byte hash:
115/// n_layers, n_kv_heads, head_dim, tokens, dtype, block_size,
116/// sliding_window, model_len.
117const HEADER_FIELDS: usize = 8;
118/// magic + version + header_len + body_len + digest.
119const PREFIX_LEN: usize = 8 + 4 + 4 + 8 + 32;
120/// Extension of a published block file.
121pub const BLOCK_FILE_EXT: &str = "kvb";
122/// Subdirectory holding in-progress writes. Not a valid shard name --
123/// shard directories are lowercase hex, and `.` is not a hex digit --
124/// so it can never collide with one.
125const TMP_DIR: &str = ".tmp";
126
127const DTYPE_F32: u32 = 0;
128
129fn dtype_code(dtype: KvDtype) -> u32 {
130    match dtype {
131        KvDtype::F32 => DTYPE_F32,
132    }
133}
134
135fn dtype_from_code(code: u32) -> Option<KvDtype> {
136    match code {
137        DTYPE_F32 => Some(KvDtype::F32),
138        _ => None,
139    }
140}
141
142/// Encodes a sliding window as a u32, with 0 meaning "no window".
143/// `BlockLayout` refuses a zero window, so the two cases cannot
144/// collide.
145fn window_code(window: Option<usize>) -> u32 {
146    window.unwrap_or(0) as u32
147}
148
149fn window_from_code(code: u32) -> Option<usize> {
150    if code == 0 {
151        None
152    } else {
153        Some(code as usize)
154    }
155}
156
157fn dtype_width(dtype: KvDtype) -> usize {
158    match dtype {
159        KvDtype::F32 => 4,
160    }
161}
162
163/// Why a block file was refused. Every variant means "these bytes are
164/// not a block this build can read", and none of them is recoverable by
165/// reading harder.
166#[derive(Clone, Debug, PartialEq, Eq)]
167pub enum BlockFormatError {
168    /// Shorter than the fixed prefix: there is not even a header to
169    /// check.
170    TooShort { len: usize },
171    /// Not a ferrox block file at all.
172    BadMagic,
173    /// Written by a build whose layout this one does not know.
174    UnsupportedFormat {
175        found: u32,
176        readable: &'static [u32],
177    },
178    /// The file's own declared length does not match the bytes present
179    /// -- a half-written or truncated file.
180    Truncated { expected: u64, actual: u64 },
181    /// Right length, wrong bytes: bit rot, an interrupted overwrite, or
182    /// a file that was edited.
183    ChecksumMismatch,
184    /// Structurally impossible content: a dimension of zero, a body
185    /// that cannot hold the tensors the header describes.
186    Malformed(&'static str),
187    /// A dtype code this build has no reader for.
188    UnknownDtype(u32),
189    /// The recorded block layout is not one any correct writer could
190    /// have produced -- a block size that does not divide the sliding
191    /// window it was cut against.
192    BadLayout(BlockLayoutError),
193}
194
195impl std::fmt::Display for BlockFormatError {
196    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
197        match self {
198            BlockFormatError::TooShort { len } => write!(
199                f,
200                "KV block file is {len} bytes, shorter than the {PREFIX_LEN}-byte header prefix"
201            ),
202            BlockFormatError::BadMagic => write!(f, "KV block file has the wrong magic"),
203            BlockFormatError::UnsupportedFormat { found, readable } => write!(
204                f,
205                "KV block file format version {found} is not readable by this build (readable: {readable:?})"
206            ),
207            BlockFormatError::Truncated { expected, actual } => write!(
208                f,
209                "KV block file declares {expected} bytes but is {actual}; refusing a torn file"
210            ),
211            BlockFormatError::ChecksumMismatch => {
212                write!(f, "KV block file failed its SHA-256 checksum")
213            }
214            BlockFormatError::Malformed(what) => {
215                write!(f, "KV block file is malformed: {what}")
216            }
217            BlockFormatError::UnknownDtype(code) => {
218                write!(f, "KV block file has unknown dtype code {code}")
219            }
220            BlockFormatError::BadLayout(err) => {
221                write!(f, "KV block file records an impossible block layout: {err}")
222            }
223        }
224    }
225}
226
227impl std::error::Error for BlockFormatError {}
228
229/// Serializes a block. The `hash` is stored inside the file as well as
230/// in its name, so a block found under a wrong or renamed path can
231/// still be checked against the identity it claims.
232pub fn encode_block(hash: &BlockHash, block: &KvBlock) -> Vec<u8> {
233    let sig = block.signature();
234    let mut header = Vec::with_capacity(64 + sig.model.len());
235    header.extend_from_slice(hash.as_bytes());
236    header.extend_from_slice(&(sig.n_layers as u32).to_le_bytes());
237    header.extend_from_slice(&(sig.n_kv_heads as u32).to_le_bytes());
238    header.extend_from_slice(&(sig.head_dim as u32).to_le_bytes());
239    header.extend_from_slice(&(sig.tokens as u32).to_le_bytes());
240    header.extend_from_slice(&dtype_code(sig.dtype).to_le_bytes());
241    header.extend_from_slice(&(sig.layout.block_size() as u32).to_le_bytes());
242    // 0 encodes "no sliding window". A real window is never 0 --
243    // `BlockLayout` refuses `Some(0)` precisely so this encoding is
244    // unambiguous.
245    header.extend_from_slice(&(window_code(sig.layout.sliding_window())).to_le_bytes());
246    header.extend_from_slice(&(sig.model.len() as u32).to_le_bytes());
247    header.extend_from_slice(sig.model.as_bytes());
248
249    let mut body = Vec::with_capacity(body_len(sig) as usize);
250    for layer in block.layers() {
251        for value in &layer.k {
252            body.extend_from_slice(&value.to_le_bytes());
253        }
254        for value in &layer.v {
255            body.extend_from_slice(&value.to_le_bytes());
256        }
257    }
258
259    let mut digest = Sha256::new();
260    digest.update(&header);
261    digest.update(&body);
262    let digest: [u8; 32] = digest.finalize().into();
263
264    let mut out = Vec::with_capacity(PREFIX_LEN + header.len() + body.len());
265    out.extend_from_slice(MAGIC);
266    out.extend_from_slice(&BLOCK_FORMAT_VERSION.to_le_bytes());
267    out.extend_from_slice(&(header.len() as u32).to_le_bytes());
268    out.extend_from_slice(&(body.len() as u64).to_le_bytes());
269    out.extend_from_slice(&digest);
270    out.extend_from_slice(&header);
271    out.extend_from_slice(&body);
272    out
273}
274
275/// Bytes the body of a block with this signature occupies. Lets the
276/// store charge a block against its budget without serializing it
277/// first.
278fn body_len(sig: &CacheSignature) -> u64 {
279    let per_layer = sig.tokens as u64
280        * sig.n_kv_heads as u64
281        * sig.head_dim as u64
282        * dtype_width(sig.dtype) as u64;
283    // K and V.
284    per_layer * 2 * sig.n_layers as u64
285}
286
287/// Total on-disk size of a block with this signature, header included.
288pub fn encoded_len(sig: &CacheSignature) -> u64 {
289    let header = 32 + 4 * HEADER_FIELDS as u64 + sig.model.len() as u64;
290    PREFIX_LEN as u64 + header + body_len(sig)
291}
292
293/// The identity and payload recovered from a block file. The signature
294/// is deliberately *unverified*: use
295/// [`UnverifiedBlock::verify`](crate::kv_signature::UnverifiedBlock::verify)
296/// to turn it into a block this process may use.
297#[derive(Debug)]
298pub struct DecodedBlock {
299    /// The hash the file claims to be stored under.
300    pub hash: BlockHash,
301    pub block: UnverifiedBlock,
302}
303
304/// Parses a block file. Checks, in order: length, magic, format
305/// version, declared-vs-actual size, checksum, then structure. Nothing
306/// is allocated from a length field until that length has been checked
307/// against the bytes actually present.
308pub fn decode_block(bytes: &[u8]) -> Result<DecodedBlock, BlockFormatError> {
309    if bytes.len() < PREFIX_LEN {
310        return Err(BlockFormatError::TooShort { len: bytes.len() });
311    }
312    if &bytes[..8] != MAGIC {
313        return Err(BlockFormatError::BadMagic);
314    }
315    let version = u32::from_le_bytes(bytes[8..12].try_into().unwrap());
316    if !READABLE_FORMAT_VERSIONS.contains(&version) {
317        return Err(BlockFormatError::UnsupportedFormat {
318            found: version,
319            readable: READABLE_FORMAT_VERSIONS,
320        });
321    }
322    let header_len = u32::from_le_bytes(bytes[12..16].try_into().unwrap()) as u64;
323    let body_len = u64::from_le_bytes(bytes[16..24].try_into().unwrap());
324    let declared = PREFIX_LEN as u64 + header_len + body_len;
325    if declared != bytes.len() as u64 {
326        return Err(BlockFormatError::Truncated {
327            expected: declared,
328            actual: bytes.len() as u64,
329        });
330    }
331    let digest_recorded = &bytes[24..PREFIX_LEN];
332    let mut digest = Sha256::new();
333    digest.update(&bytes[PREFIX_LEN..]);
334    let digest: [u8; 32] = digest.finalize().into();
335    if digest != digest_recorded {
336        return Err(BlockFormatError::ChecksumMismatch);
337    }
338
339    let header = &bytes[PREFIX_LEN..PREFIX_LEN + header_len as usize];
340    let body = &bytes[PREFIX_LEN + header_len as usize..];
341    if header.len() < 32 + 4 * HEADER_FIELDS {
342        return Err(BlockFormatError::Malformed(
343            "header shorter than its fields",
344        ));
345    }
346    let mut hash = [0u8; 32];
347    hash.copy_from_slice(&header[..32]);
348    let hash = BlockHash::from_bytes(hash);
349    let field = |i: usize| u32::from_le_bytes(header[32 + i * 4..36 + i * 4].try_into().unwrap());
350    let n_layers = field(0) as usize;
351    let n_kv_heads = field(1) as usize;
352    let head_dim = field(2) as usize;
353    let tokens = field(3) as usize;
354    let dtype_code = field(4);
355    let block_size = field(5) as usize;
356    let window_code = field(6);
357    let model_len = field(7) as usize;
358    let dtype = dtype_from_code(dtype_code).ok_or(BlockFormatError::UnknownDtype(dtype_code))?;
359    if header.len() != 32 + 4 * HEADER_FIELDS + model_len {
360        return Err(BlockFormatError::Malformed(
361            "model name length disagrees with header",
362        ));
363    }
364    let model = std::str::from_utf8(&header[32 + 4 * HEADER_FIELDS..])
365        .map_err(|_| BlockFormatError::Malformed("model name is not UTF-8"))?
366        .to_string();
367    if n_layers == 0 || n_kv_heads == 0 || head_dim == 0 {
368        return Err(BlockFormatError::Malformed(
369            "zero layers, heads, or head dim",
370        ));
371    }
372    // A file whose recorded layout is not a layout at all -- a block
373    // size that does not divide its window -- is refused here rather
374    // than reconstructed into a `BlockLayout` that could not have been
375    // built by any correct writer.
376    let layout = BlockLayout::new(block_size, window_from_code(window_code))
377        .map_err(BlockFormatError::BadLayout)?;
378
379    let per_layer_elems = tokens
380        .checked_mul(n_kv_heads)
381        .and_then(|n| n.checked_mul(head_dim))
382        .ok_or(BlockFormatError::Malformed("layer size overflows"))?;
383    let expected_body = (per_layer_elems as u64)
384        .checked_mul(2 * n_layers as u64)
385        .and_then(|n| n.checked_mul(dtype_width(dtype) as u64))
386        .ok_or(BlockFormatError::Malformed("body size overflows"))?;
387    if expected_body != body.len() as u64 {
388        return Err(BlockFormatError::Malformed(
389            "body does not match declared dims",
390        ));
391    }
392
393    let mut layers = Vec::with_capacity(n_layers);
394    let mut offset = 0usize;
395    for _ in 0..n_layers {
396        let k = read_f32(&body[offset..offset + per_layer_elems * 4]);
397        offset += per_layer_elems * 4;
398        let v = read_f32(&body[offset..offset + per_layer_elems * 4]);
399        offset += per_layer_elems * 4;
400        let mut cache = KvCache::new(n_kv_heads, head_dim);
401        cache.k = k;
402        cache.v = v;
403        // The rows were just written by hand, so this is the one place
404        // a position count is set rather than counted by `push`.
405        cache.set_positions(tokens);
406        layers.push(cache);
407    }
408
409    let signature = CacheSignature {
410        format_version: version,
411        model,
412        n_layers,
413        n_kv_heads,
414        head_dim,
415        dtype,
416        tokens,
417        layout,
418    };
419    Ok(DecodedBlock {
420        hash,
421        block: UnverifiedBlock::new(Some(signature), layers),
422    })
423}
424
425fn read_f32(bytes: &[u8]) -> Vec<f32> {
426    bytes
427        .as_chunks::<4>()
428        .0
429        .iter()
430        .map(|c| f32::from_le_bytes(*c))
431        .collect()
432}
433
434/// Something went wrong reaching the disk tier. Corruption and
435/// incompatibility are **not** here: those are misses, reported through
436/// [`DiskStats`], because a caller's only sane response to either is to
437/// recompute the prefix.
438#[derive(Clone, Debug, PartialEq, Eq)]
439pub enum StoreError {
440    Io {
441        op: &'static str,
442        path: PathBuf,
443        message: String,
444    },
445    /// The index named a block whose payload is nowhere: not on disk,
446    /// not in the write buffer. This is not a cache miss -- it is the
447    /// write-ordering invariant (buffer -> index -> queue) having been
448    /// violated, i.e. a bug in this module, and it is surfaced rather
449    /// than smoothed into a miss so a test can fail on it.
450    MissingPayload { hash: BlockHash },
451}
452
453impl std::fmt::Display for StoreError {
454    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
455        match self {
456            StoreError::Io { op, path, message } => {
457                write!(
458                    f,
459                    "KV block store failed to {op} {}: {message}",
460                    path.display()
461                )
462            }
463            StoreError::MissingPayload { hash } => write!(
464                f,
465                "KV block store index names {hash:?} but it has neither a file nor a buffered \
466                 payload; the write-ordering invariant was violated"
467            ),
468        }
469    }
470}
471
472impl std::error::Error for StoreError {}
473
474fn io_err(op: &'static str, path: &Path, err: io::Error) -> StoreError {
475    StoreError::Io {
476        op,
477        path: path.to_path_buf(),
478        message: err.to_string(),
479    }
480}
481
482/// Asks how many bytes are still free on the filesystem holding a
483/// path. `None` means "cannot tell", and the store then trusts only its
484/// configured byte budget.
485///
486/// Injectable so the budget's behaviour under a nearly-full disk is
487/// testable without actually filling one.
488pub type FreeSpaceProbe = Arc<dyn Fn(&Path) -> Option<u64> + Send + Sync>;
489
490/// Free space via `statvfs`. `f_bavail` (blocks available to an
491/// unprivileged process), not `f_bfree`, because the reserved blocks a
492/// filesystem keeps back are not ours to spend.
493#[cfg(unix)]
494#[allow(clippy::unnecessary_cast)] // field widths differ across unixes
495fn platform_free_bytes(path: &Path) -> Option<u64> {
496    use std::os::unix::ffi::OsStrExt;
497    let c_path = std::ffi::CString::new(path.as_os_str().as_bytes()).ok()?;
498    // SAFETY: `c_path` is a NUL-terminated path that outlives the call,
499    // and `stat` is a correctly sized, writable `statvfs`.
500    let stat = unsafe {
501        let mut stat: libc::statvfs = std::mem::zeroed();
502        if libc::statvfs(c_path.as_ptr(), &mut stat) != 0 {
503            return None;
504        }
505        stat
506    };
507    let block = if stat.f_frsize > 0 {
508        stat.f_frsize as u64
509    } else {
510        stat.f_bsize as u64
511    };
512    Some((stat.f_bavail as u64).saturating_mul(block))
513}
514
515#[cfg(not(unix))]
516fn platform_free_bytes(_path: &Path) -> Option<u64> {
517    None
518}
519
520/// A free-space reading and when it was taken. `statvfs` is a syscall
521/// per call and the answer changes slowly, so it is cached -- but only
522/// for a TTL, and any `ENOSPC` throws it away immediately, because at
523/// that moment the cached number is known to be a lie.
524struct FreeSpace {
525    checked_at: Option<Instant>,
526    bytes: Option<u64>,
527}
528
529/// How the store is sized, laid out, and how much writing it will do
530/// off the calling thread.
531#[derive(Clone)]
532pub struct DiskConfig {
533    /// Directory the store owns. Created if absent.
534    pub root: PathBuf,
535    /// Byte budget for blocks the store is accounting for. Eviction
536    /// keeps the store at or under this.
537    pub max_bytes: u64,
538    /// Hex characters of the hash used as the shard subdirectory name.
539    /// 2 gives 256 shards, which keeps directory sizes sane well past a
540    /// million blocks.
541    pub shard_chars: usize,
542    /// Writes that may be waiting for a writer thread at once. When it
543    /// is full, [`DiskKvStore::put`] writes on the calling thread
544    /// instead of dropping the block -- backpressure, not loss.
545    pub queue_capacity: usize,
546    /// Background writer threads. `0` is legitimate and means every
547    /// write happens on the thread that asked for it.
548    pub writer_threads: usize,
549    /// Background reader threads, serving prefetches and any demand
550    /// read that does not want to block its own thread. `0` means every
551    /// read happens on the thread that asked for it, and a prefetch is
552    /// a no-op.
553    pub reader_threads: usize,
554    /// Blocks that may sit in the prefetch staging area at once. A
555    /// prefetch is a hint, so this is a hard refusal rather than
556    /// backpressure: reading ahead must never be the thing that runs
557    /// the process out of memory.
558    pub prefetch_capacity: usize,
559    /// Bytes to leave free on the filesystem. The store evicts to stay
560    /// this far away from a full disk, rather than letting `ENOSPC` be
561    /// the mechanism that tells it to stop.
562    pub reserve_bytes: u64,
563    /// How long a free-space reading is trusted before it is taken
564    /// again.
565    pub free_space_ttl: std::time::Duration,
566    /// How free space is measured. Defaults to `statvfs` on unix, and
567    /// to "cannot tell" elsewhere.
568    pub free_space_probe: FreeSpaceProbe,
569}
570
571impl std::fmt::Debug for DiskConfig {
572    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
573        f.debug_struct("DiskConfig")
574            .field("root", &self.root)
575            .field("max_bytes", &self.max_bytes)
576            .field("shard_chars", &self.shard_chars)
577            .field("queue_capacity", &self.queue_capacity)
578            .field("writer_threads", &self.writer_threads)
579            .field("reader_threads", &self.reader_threads)
580            .field("prefetch_capacity", &self.prefetch_capacity)
581            .field("reserve_bytes", &self.reserve_bytes)
582            .field("free_space_ttl", &self.free_space_ttl)
583            .finish_non_exhaustive()
584    }
585}
586
587impl DiskConfig {
588    pub fn new(root: impl Into<PathBuf>) -> Self {
589        DiskConfig {
590            root: root.into(),
591            max_bytes: 1 << 30,
592            shard_chars: 2,
593            queue_capacity: 64,
594            writer_threads: 1,
595            reader_threads: 2,
596            prefetch_capacity: 64,
597            reserve_bytes: 1 << 30,
598            free_space_ttl: std::time::Duration::from_secs(2),
599            free_space_probe: Arc::new(platform_free_bytes),
600        }
601    }
602
603    pub fn with_max_bytes(mut self, max_bytes: u64) -> Self {
604        self.max_bytes = max_bytes;
605        self
606    }
607
608    pub fn with_shard_chars(mut self, shard_chars: usize) -> Self {
609        self.shard_chars = shard_chars.clamp(1, 8);
610        self
611    }
612
613    pub fn with_queue_capacity(mut self, queue_capacity: usize) -> Self {
614        self.queue_capacity = queue_capacity;
615        self
616    }
617
618    pub fn with_writer_threads(mut self, writer_threads: usize) -> Self {
619        self.writer_threads = writer_threads;
620        self
621    }
622
623    pub fn with_reader_threads(mut self, reader_threads: usize) -> Self {
624        self.reader_threads = reader_threads;
625        self
626    }
627
628    pub fn with_prefetch_capacity(mut self, prefetch_capacity: usize) -> Self {
629        self.prefetch_capacity = prefetch_capacity;
630        self
631    }
632
633    pub fn with_reserve_bytes(mut self, reserve_bytes: u64) -> Self {
634        self.reserve_bytes = reserve_bytes;
635        self
636    }
637
638    pub fn with_free_space_ttl(mut self, free_space_ttl: std::time::Duration) -> Self {
639        self.free_space_ttl = free_space_ttl;
640        self
641    }
642
643    pub fn with_free_space_probe(mut self, probe: FreeSpaceProbe) -> Self {
644        self.free_space_probe = probe;
645        self
646    }
647}
648
649#[derive(Default)]
650struct Stats {
651    writes: AtomicU64,
652    queued_writes: AtomicU64,
653    /// Writes that ran on the calling thread because the queue was
654    /// full. The plan's "count the fallbacks": a store that is
655    /// permanently inline is a store whose queue is too small or whose
656    /// disk is too slow, and that is invisible without this.
657    inline_writes: AtomicU64,
658    write_failures: AtomicU64,
659    /// Queued writes whose block was evicted (or superseded) before a
660    /// writer thread reached it.
661    write_skipped: AtomicU64,
662    write_nanos: AtomicU64,
663    /// Writes whose block was evicted while it was being written, so
664    /// the published file was withdrawn again.
665    write_raced_eviction: AtomicU64,
666    hits: AtomicU64,
667    /// Reads served from the write buffer, before the block reached
668    /// disk. These are what make the write path asynchronous *and*
669    /// immediately visible.
670    buffer_hits: AtomicU64,
671    misses: AtomicU64,
672    /// Files that failed [`decode_block`] and were quarantined.
673    corrupt: AtomicU64,
674    /// Blocks that decoded cleanly but do not match this reader's
675    /// signature expectation.
676    incompatible: AtomicU64,
677    read_nanos: AtomicU64,
678    evictions: AtomicU64,
679    evicted_bytes: AtomicU64,
680    /// Reads handed to a reader thread by [`DiskKvStore::prefetch`].
681    prefetch_issued: AtomicU64,
682    /// Prefetches refused: staging full, no reader threads, or the
683    /// block was already staged or in flight.
684    prefetch_dropped: AtomicU64,
685    /// Demand reads that found a prefetch already *finished*. This is
686    /// the one that says the read-ahead paid for itself: the request
687    /// did no I/O at all.
688    prefetch_hits: AtomicU64,
689    /// Demand reads that found a prefetch still running and waited for
690    /// it instead of issuing a second read of the same file.
691    prefetch_waits: AtomicU64,
692    /// Reads that ran on a reader thread rather than the caller's.
693    async_reads: AtomicU64,
694    /// Writes that failed because the filesystem was full. Non-zero
695    /// means the budget lost the race it exists to win.
696    enospc: AtomicU64,
697    /// Eviction passes whose ceiling came from free disk rather than
698    /// from `max_bytes`.
699    space_clamped: AtomicU64,
700}
701
702/// A snapshot of the tier's behaviour. Note the two time-valued fields:
703/// hit *rate* alone cannot tell an operator whether a disk hit was
704/// cheaper than recomputing the prefix, which is the only question that
705/// decides whether the tier is worth having.
706#[derive(Clone, Debug, Default, PartialEq, Eq)]
707pub struct DiskStats {
708    pub blocks: usize,
709    /// Everything the store has accepted, published or still buffered.
710    pub bytes: u64,
711    /// The subset that has actually reached the filesystem.
712    pub disk_bytes: u64,
713    pub queue_depth: usize,
714    pub writes: u64,
715    pub queued_writes: u64,
716    pub inline_writes: u64,
717    pub write_failures: u64,
718    pub write_skipped: u64,
719    pub write_raced_eviction: u64,
720    pub write_nanos: u64,
721    pub hits: u64,
722    pub buffer_hits: u64,
723    pub misses: u64,
724    pub corrupt: u64,
725    pub incompatible: u64,
726    pub read_nanos: u64,
727    pub evictions: u64,
728    pub evicted_bytes: u64,
729    pub prefetch_issued: u64,
730    pub prefetch_dropped: u64,
731    pub prefetch_hits: u64,
732    pub prefetch_waits: u64,
733    pub async_reads: u64,
734    pub staged_blocks: usize,
735    pub enospc: u64,
736    pub space_clamped: u64,
737    /// The ceiling eviction is currently working to: `max_bytes`, or
738    /// less when free disk says so.
739    pub effective_capacity: u64,
740}
741
742struct Entry {
743    bytes: u64,
744    last_used: u64,
745    /// False between admitting the block and its file landing under its
746    /// final name. While it is false the payload lives in the write
747    /// buffer, and a reader is served from there.
748    published: bool,
749    /// Bumped every time this hash is admitted, so a write that
750    /// finishes after its entry was evicted and re-created cannot mark
751    /// the *new* entry published, and a queued job whose block has been
752    /// superseded can tell.
753    generation: u64,
754}
755
756struct Index {
757    entries: HashMap<BlockHash, Entry>,
758    /// Everything the store has accepted, published or still buffered.
759    bytes: u64,
760    /// The subset that is actually occupying filesystem space. The
761    /// free-space budget reasons about this one: a block that has been
762    /// accepted but not written has not taken any disk yet, and
763    /// charging it twice (once here, once as space the device has
764    /// already lost) makes the ceiling oscillate.
765    disk_bytes: u64,
766    clock: u64,
767}
768
769impl Index {
770    fn touch(&mut self) -> u64 {
771        self.clock += 1;
772        self.clock
773    }
774
775    fn insert_entry(&mut self, hash: BlockHash, entry: Entry) {
776        let bytes = entry.bytes;
777        if let Some(previous) = self.entries.insert(hash, entry) {
778            self.uncharge(&previous);
779        }
780        self.bytes += bytes;
781    }
782
783    fn remove_entry(&mut self, hash: &BlockHash) -> Option<Entry> {
784        let entry = self.entries.remove(hash)?;
785        self.uncharge(&entry);
786        Some(entry)
787    }
788
789    fn uncharge(&mut self, entry: &Entry) {
790        self.bytes -= entry.bytes;
791        if entry.published {
792            self.disk_bytes -= entry.bytes;
793        }
794    }
795}
796
797/// Where a block's payload is, once the index has been consulted.
798enum Source {
799    Disk(PathBuf),
800    Buffer(Arc<KvBlock>),
801}
802
803/// A block that has been accepted but whose file is not on disk yet.
804struct Buffered {
805    generation: u64,
806    block: Arc<KvBlock>,
807}
808
809#[derive(Clone, Copy)]
810struct WriteJob {
811    hash: BlockHash,
812    generation: u64,
813}
814
815struct QueueState {
816    jobs: VecDeque<WriteJob>,
817    running: usize,
818    shutdown: bool,
819}
820
821/// A bounded queue of pending block writes.
822///
823/// Bounded, and **never lossy**: `try_push` refusing is the caller's
824/// signal to write the block itself, not to drop it. A dropped write is
825/// indistinguishable from a cache miss later, which is exactly the kind
826/// of silent degradation that makes a cache tier impossible to trust.
827struct WriteQueue {
828    state: Mutex<QueueState>,
829    ready: Condvar,
830    idle: Condvar,
831    capacity: usize,
832}
833
834impl WriteQueue {
835    fn new(capacity: usize) -> Self {
836        WriteQueue {
837            state: Mutex::new(QueueState {
838                jobs: VecDeque::new(),
839                running: 0,
840                shutdown: false,
841            }),
842            ready: Condvar::new(),
843            idle: Condvar::new(),
844            capacity: capacity.max(1),
845        }
846    }
847
848    fn lock(&self) -> std::sync::MutexGuard<'_, QueueState> {
849        self.state.lock().expect("kv disk write queue poisoned")
850    }
851
852    /// `false` means "full, or shutting down" -- write it yourself.
853    fn try_push(&self, job: WriteJob) -> bool {
854        let mut state = self.lock();
855        if state.shutdown || state.jobs.len() >= self.capacity {
856            return false;
857        }
858        state.jobs.push_back(job);
859        self.ready.notify_one();
860        true
861    }
862
863    fn pop_blocking(&self) -> Option<WriteJob> {
864        let mut state = self.lock();
865        loop {
866            if let Some(job) = state.jobs.pop_front() {
867                state.running += 1;
868                return Some(job);
869            }
870            if state.shutdown {
871                return None;
872            }
873            state = self
874                .ready
875                .wait(state)
876                .expect("kv disk write queue poisoned");
877        }
878    }
879
880    fn pop_now(&self) -> Option<WriteJob> {
881        let mut state = self.lock();
882        let job = state.jobs.pop_front()?;
883        state.running += 1;
884        Some(job)
885    }
886
887    fn finish(&self) {
888        let mut state = self.lock();
889        state.running -= 1;
890        self.idle.notify_all();
891    }
892
893    fn shutdown(&self) {
894        let mut state = self.lock();
895        state.shutdown = true;
896        self.ready.notify_all();
897    }
898
899    fn depth(&self) -> usize {
900        self.lock().jobs.len()
901    }
902}
903
904/// What a read produced. `Ok(None)` is a miss (absent, corrupt, or
905/// incompatible); `Err` is I/O, or the write-ordering invariant
906/// breaking.
907pub type ReadOutcome = Result<Option<Arc<KvBlock>>, StoreError>;
908
909/// A read in progress, or one that has finished and is waiting to be
910/// claimed. Shared between the thread that asked for it, the thread
911/// that runs it, and any later request for the same block.
912struct ReadSlot {
913    /// The signature this read was issued for. A staged result is only
914    /// reusable by a reader that wants the *same* shape -- otherwise a
915    /// prefetch issued under one config would hand its answer to a
916    /// request under another.
917    expected: CacheSignature,
918    outcome: Mutex<Option<ReadOutcome>>,
919    done: Condvar,
920}
921
922impl ReadSlot {
923    fn pending(expected: CacheSignature) -> Arc<Self> {
924        Arc::new(ReadSlot {
925            expected,
926            outcome: Mutex::new(None),
927            done: Condvar::new(),
928        })
929    }
930
931    fn ready(expected: CacheSignature, outcome: ReadOutcome) -> Arc<Self> {
932        Arc::new(ReadSlot {
933            expected,
934            outcome: Mutex::new(Some(outcome)),
935            done: Condvar::new(),
936        })
937    }
938
939    fn is_ready(&self) -> bool {
940        self.outcome
941            .lock()
942            .expect("kv disk read slot poisoned")
943            .is_some()
944    }
945
946    fn fulfil(&self, outcome: ReadOutcome) {
947        let mut slot = self.outcome.lock().expect("kv disk read slot poisoned");
948        *slot = Some(outcome);
949        self.done.notify_all();
950    }
951
952    fn wait(&self) -> ReadOutcome {
953        let mut slot = self.outcome.lock().expect("kv disk read slot poisoned");
954        loop {
955            if let Some(outcome) = slot.as_ref() {
956                return outcome.clone();
957            }
958            slot = self.done.wait(slot).expect("kv disk read slot poisoned");
959        }
960    }
961}
962
963struct ReadJob {
964    hash: BlockHash,
965    path: PathBuf,
966    slot: Arc<ReadSlot>,
967}
968
969struct ReadQueueState {
970    jobs: VecDeque<ReadJob>,
971    shutdown: bool,
972}
973
974/// Pending disk reads. Unlike the write queue this one *is* allowed to
975/// refuse work -- a prefetch is a hint, and a refused hint costs a
976/// later request one read. A refused *demand* read is not dropped
977/// either: it runs on the calling thread.
978struct ReadQueue {
979    state: Mutex<ReadQueueState>,
980    ready: Condvar,
981    capacity: usize,
982}
983
984impl ReadQueue {
985    fn new(capacity: usize) -> Self {
986        ReadQueue {
987            state: Mutex::new(ReadQueueState {
988                jobs: VecDeque::new(),
989                shutdown: false,
990            }),
991            ready: Condvar::new(),
992            capacity: capacity.max(1),
993        }
994    }
995
996    fn lock(&self) -> std::sync::MutexGuard<'_, ReadQueueState> {
997        self.state.lock().expect("kv disk read queue poisoned")
998    }
999
1000    /// `demand` jumps the queue: a request that is waiting must not sit
1001    /// behind speculative read-ahead.
1002    fn try_push(&self, job: ReadJob, demand: bool) -> bool {
1003        let mut state = self.lock();
1004        if state.shutdown || state.jobs.len() >= self.capacity {
1005            return false;
1006        }
1007        if demand {
1008            state.jobs.push_front(job);
1009        } else {
1010            state.jobs.push_back(job);
1011        }
1012        self.ready.notify_one();
1013        true
1014    }
1015
1016    fn pop_blocking(&self) -> Option<ReadJob> {
1017        let mut state = self.lock();
1018        loop {
1019            if let Some(job) = state.jobs.pop_front() {
1020                return Some(job);
1021            }
1022            if state.shutdown {
1023                return None;
1024            }
1025            state = self.ready.wait(state).expect("kv disk read queue poisoned");
1026        }
1027    }
1028
1029    fn shutdown(&self) {
1030        let mut state = self.lock();
1031        state.shutdown = true;
1032        self.ready.notify_all();
1033    }
1034}
1035
1036/// A read that may not have finished yet.
1037///
1038/// The disk tier is asynchronous **by construction**, not as a later
1039/// retrofit: [`DiskKvStore::get`] is this handle plus a `wait`, so
1040/// there is no synchronous read path that a prefetch has to work
1041/// around.
1042pub struct ReadHandle {
1043    shared: Arc<Shared>,
1044    hash: BlockHash,
1045    slot: Arc<ReadSlot>,
1046    /// Whether this handle's slot is registered in the staging map and
1047    /// should be removed once claimed.
1048    staged: bool,
1049}
1050
1051impl ReadHandle {
1052    /// True if the block is already in hand -- a memory hit, a miss, or
1053    /// a prefetch that has landed.
1054    pub fn is_ready(&self) -> bool {
1055        self.slot.is_ready()
1056    }
1057
1058    /// The result, without blocking. Returns `None` if the read is
1059    /// still running.
1060    pub fn try_claim(&self) -> Option<ReadOutcome> {
1061        if !self.slot.is_ready() {
1062            return None;
1063        }
1064        Some(self.claim())
1065    }
1066
1067    /// Blocks until the read finishes.
1068    pub fn wait(self) -> ReadOutcome {
1069        self.claim()
1070    }
1071
1072    fn claim(&self) -> ReadOutcome {
1073        let outcome = self.slot.wait();
1074        if self.staged {
1075            self.shared.unstage(&self.hash, &self.slot);
1076        }
1077        outcome
1078    }
1079}
1080
1081impl std::fmt::Debug for ReadHandle {
1082    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
1083        f.debug_struct("ReadHandle")
1084            .field("hash", &self.hash)
1085            .field("ready", &self.is_ready())
1086            .finish()
1087    }
1088}
1089
1090#[cfg(test)]
1091type Hook = Arc<dyn Fn(&BlockHash) + Send + Sync>;
1092
1093/// Which order the write path uses. Production is
1094/// `BufferThenIndex`; the other two exist so a test can prove the
1095/// invariant test is not vacuous -- a concurrency test that passes on
1096/// broken code proves nothing.
1097#[cfg(test)]
1098#[derive(Clone, Copy, Debug, Default, PartialEq, Eq)]
1099enum WriteOrder {
1100    #[default]
1101    BufferThenIndex,
1102    /// Index the block before buffering it: a reader can then hit the
1103    /// index for a block with no file and no payload.
1104    IndexBeforeBuffer,
1105    /// Release the buffered payload before marking the file published:
1106    /// same hole, at the other end of the write.
1107    DropBufferBeforeMarking,
1108}
1109
1110#[cfg(test)]
1111#[derive(Default)]
1112struct Hooks {
1113    order: Mutex<WriteOrder>,
1114    /// Called after the rename and before the post-publish eviction
1115    /// re-check, so a test can evict the block in exactly the window
1116    /// the re-check exists to cover.
1117    after_rename: Mutex<Option<Hook>>,
1118    /// Called inside the window between the two steps of admission.
1119    in_put_window: Mutex<Option<Hook>>,
1120    /// Called inside the window between the two steps of publication.
1121    in_publish_window: Mutex<Option<Hook>>,
1122    /// Makes the next write fail as though the filesystem were full,
1123    /// which is the one failure that cannot be arranged for real
1124    /// without filling a real disk.
1125    fail_with_enospc: std::sync::atomic::AtomicBool,
1126}
1127
1128#[cfg(test)]
1129impl Hooks {
1130    fn fire(slot: &Mutex<Option<Hook>>, hash: &BlockHash) {
1131        // Cloned out before calling, so a hook that re-enters the store
1132        // cannot deadlock on the hook slot itself.
1133        let hook = slot.lock().expect("kv disk hook poisoned").clone();
1134        if let Some(hook) = hook {
1135            hook(hash);
1136        }
1137    }
1138}
1139
1140/// Everything the store's threads share. Deliberately holds no join
1141/// handles: the writer threads hold an `Arc<Shared>`, so a `Drop` here
1142/// that joined them could run *on* a writer thread and deadlock. The
1143/// handles live in [`DiskKvStore`], which is not `Clone`.
1144struct Shared {
1145    root: PathBuf,
1146    shard_chars: usize,
1147    max_bytes: u64,
1148    index: Mutex<Index>,
1149    /// Blocks accepted but not yet on disk.
1150    ///
1151    /// **Lock order: `index`, then `buffer`, never the reverse.** A
1152    /// reader decides "the index says this block exists" and "here is
1153    /// its payload" as one atomic act; otherwise the writer could mark
1154    /// a block published and drop its buffered copy in between, and the
1155    /// reader would find nothing.
1156    buffer: Mutex<HashMap<BlockHash, Buffered>>,
1157    queue: WriteQueue,
1158    reads: ReadQueue,
1159    /// Reads in flight or finished and unclaimed, keyed by block. One
1160    /// slot per block, so a demand read that arrives while a prefetch
1161    /// is running joins it instead of reading the same file twice.
1162    staging: Mutex<HashMap<BlockHash, Arc<ReadSlot>>>,
1163    prefetch_capacity: usize,
1164    has_readers: bool,
1165    reserve_bytes: u64,
1166    free_space_ttl: std::time::Duration,
1167    free_space_probe: FreeSpaceProbe,
1168    /// **Lock order: `index`, then `free_space`.** Eviction reads the
1169    /// budget while holding the index.
1170    free_space: Mutex<FreeSpace>,
1171    stats: Stats,
1172    seq: AtomicU64,
1173    generation: AtomicU64,
1174    #[cfg(test)]
1175    hooks: Hooks,
1176}
1177
1178/// A content-addressed block store on disk, with its own writer
1179/// threads.
1180///
1181/// Not `Clone` on purpose -- it owns the writer threads and joins them
1182/// when it drops. Share it as an `Arc<DiskKvStore>`; every method takes
1183/// `&self`.
1184pub struct DiskKvStore {
1185    shared: Arc<Shared>,
1186    writers: Vec<std::thread::JoinHandle<()>>,
1187    readers: Vec<std::thread::JoinHandle<()>>,
1188}
1189
1190impl Drop for DiskKvStore {
1191    /// Stops accepting queued work and joins the worker threads. Blocks
1192    /// already queued but not started are **not** written: they were
1193    /// never durable, and a shutdown that waits for an arbitrarily deep
1194    /// queue is worse than a cold cache. Call [`Self::flush`] first if
1195    /// they matter.
1196    fn drop(&mut self) {
1197        self.shared.queue.shutdown();
1198        self.shared.reads.shutdown();
1199        for writer in self.writers.drain(..) {
1200            let _ = writer.join();
1201        }
1202        for reader in self.readers.drain(..) {
1203            let _ = reader.join();
1204        }
1205    }
1206}
1207
1208impl DiskKvStore {
1209    /// Creates the store's directories and starts its writer threads.
1210    /// Does **not** scan `root` for pre-existing blocks: reattaching to
1211    /// a store left by a previous process is [`Self::reindex`], an
1212    /// explicit step, because it costs a directory walk and a caller
1213    /// may prefer to start cold.
1214    pub fn open(config: DiskConfig) -> Result<Self, StoreError> {
1215        let root = config.root.clone();
1216        fs::create_dir_all(&root).map_err(|e| io_err("create", &root, e))?;
1217        let tmp = root.join(TMP_DIR);
1218        fs::create_dir_all(&tmp).map_err(|e| io_err("create", &tmp, e))?;
1219        let shared = Arc::new(Shared {
1220            root,
1221            shard_chars: config.shard_chars.clamp(1, 8),
1222            max_bytes: config.max_bytes,
1223            index: Mutex::new(Index {
1224                entries: HashMap::new(),
1225                bytes: 0,
1226                disk_bytes: 0,
1227                clock: 0,
1228            }),
1229            buffer: Mutex::new(HashMap::new()),
1230            queue: WriteQueue::new(config.queue_capacity),
1231            reads: ReadQueue::new(config.queue_capacity),
1232            staging: Mutex::new(HashMap::new()),
1233            prefetch_capacity: config.prefetch_capacity,
1234            has_readers: config.reader_threads > 0,
1235            reserve_bytes: config.reserve_bytes,
1236            free_space_ttl: config.free_space_ttl,
1237            free_space_probe: Arc::clone(&config.free_space_probe),
1238            free_space: Mutex::new(FreeSpace {
1239                checked_at: None,
1240                bytes: None,
1241            }),
1242            stats: Stats::default(),
1243            seq: AtomicU64::new(0),
1244            generation: AtomicU64::new(0),
1245            #[cfg(test)]
1246            hooks: Hooks::default(),
1247        });
1248        let mut writers = Vec::with_capacity(config.writer_threads);
1249        for n in 0..config.writer_threads {
1250            let shared = Arc::clone(&shared);
1251            let handle = std::thread::Builder::new()
1252                .name(format!("ferrox-kv-write-{n}"))
1253                .spawn(move || {
1254                    while let Some(job) = shared.queue.pop_blocking() {
1255                        shared.run_job(job);
1256                        shared.queue.finish();
1257                    }
1258                })
1259                .map_err(|e| io_err("spawn writer for", &config.root, e))?;
1260            writers.push(handle);
1261        }
1262        let mut readers = Vec::with_capacity(config.reader_threads);
1263        for n in 0..config.reader_threads {
1264            let shared = Arc::clone(&shared);
1265            let handle = std::thread::Builder::new()
1266                .name(format!("ferrox-kv-read-{n}"))
1267                .spawn(move || {
1268                    while let Some(job) = shared.reads.pop_blocking() {
1269                        shared.stats.async_reads.fetch_add(1, Ordering::Relaxed);
1270                        let outcome = shared.read_timed(&job.path, &job.hash, &job.slot.expected);
1271                        job.slot.fulfil(outcome);
1272                    }
1273                })
1274                .map_err(|e| io_err("spawn reader for", &config.root, e))?;
1275            readers.push(handle);
1276        }
1277        Ok(DiskKvStore {
1278            shared,
1279            writers,
1280            readers,
1281        })
1282    }
1283
1284    pub fn root(&self) -> &Path {
1285        &self.shared.root
1286    }
1287
1288    /// Accepts a block: buffered, indexed, then queued. Returns as soon
1289    /// as the block is *visible* -- a reader can have it immediately,
1290    /// whether or not it has reached disk.
1291    ///
1292    /// If the write queue is full the block is written on this thread
1293    /// rather than dropped, and the fallback is counted.
1294    pub fn put(&self, hash: BlockHash, block: KvBlock) -> Result<(), StoreError> {
1295        self.shared.put(hash, block, false)
1296    }
1297
1298    /// Like [`Self::put`], but always writes on the calling thread.
1299    pub fn put_blocking(&self, hash: BlockHash, block: KvBlock) -> Result<(), StoreError> {
1300        self.shared.put(hash, block, true)
1301    }
1302
1303    /// Runs every pending write to completion, on this thread if no
1304    /// writer thread gets there first. Returns once the queue is empty
1305    /// and nothing is in flight.
1306    pub fn flush(&self) {
1307        loop {
1308            if let Some(job) = self.shared.queue.pop_now() {
1309                self.shared.run_job(job);
1310                self.shared.queue.finish();
1311                continue;
1312            }
1313            let state = self.shared.queue.lock();
1314            if state.jobs.is_empty() && state.running == 0 {
1315                return;
1316            }
1317            // Timed, so a store with no writer threads and a job pushed
1318            // by another thread cannot park here forever.
1319            let _ = self
1320                .shared
1321                .queue
1322                .idle
1323                .wait_timeout(state, std::time::Duration::from_millis(1));
1324        }
1325    }
1326
1327    /// Looks a block up, verifying it against `expected` before
1328    /// returning it. Blocks until the answer is in hand.
1329    ///
1330    /// This is exactly [`Self::read_async`] plus a wait -- the disk
1331    /// tier has no separate synchronous read path for a prefetch to
1332    /// have to work around later.
1333    pub fn get(&self, hash: &BlockHash, expected: &CacheSignature) -> ReadOutcome {
1334        self.shared.read_async(hash, expected, true).wait()
1335    }
1336
1337    /// Starts a read and returns immediately. The handle can be polled
1338    /// with [`ReadHandle::try_claim`] or waited on.
1339    pub fn read_async(&self, hash: &BlockHash, expected: &CacheSignature) -> ReadHandle {
1340        self.shared.read_async(hash, expected, true)
1341    }
1342
1343    /// Reads `hashes` ahead of anyone asking for them, on the reader
1344    /// threads, and returns without waiting.
1345    ///
1346    /// Intended for a prefix chain the moment its hashes are known:
1347    /// every block the request is about to want, read while the tokens
1348    /// before it are still being processed. Blocks already staged, in
1349    /// flight, or in memory are skipped; so is everything past
1350    /// `prefetch_capacity`, because reading ahead must never be the
1351    /// thing that exhausts memory.
1352    pub fn prefetch(&self, hashes: &[BlockHash], expected: &CacheSignature) {
1353        self.shared.prefetch(hashes, expected);
1354    }
1355
1356    /// Drops any staged read-ahead results. A caller that abandons a
1357    /// request it prefetched for should say so rather than leave the
1358    /// blocks occupying staging until something else needs the room.
1359    pub fn clear_prefetch(&self) {
1360        self.shared
1361            .staging
1362            .lock()
1363            .expect("kv disk staging poisoned")
1364            .clear();
1365    }
1366
1367    /// Adopts the blocks already under `root`, as a restart would.
1368    pub fn reindex(&self) -> Result<usize, StoreError> {
1369        self.shared.reindex()
1370    }
1371
1372    /// Removes a block from the index, the write buffer, and the disk.
1373    pub fn remove(&self, hash: &BlockHash) {
1374        self.shared.quarantine(hash);
1375    }
1376
1377    /// True if the index holds an entry for `hash` -- on disk or still
1378    /// buffered. Says nothing about whether its contents will verify.
1379    pub fn contains(&self, hash: &BlockHash) -> bool {
1380        let index = self.shared.index.lock().expect("kv disk index poisoned");
1381        index.entries.contains_key(hash)
1382    }
1383
1384    /// Byte ceiling an operator configured.
1385    pub fn capacity(&self) -> u64 {
1386        self.shared.max_bytes
1387    }
1388
1389    /// The ceiling eviction is actually working to right now: the
1390    /// configured one, or less when free disk says so.
1391    pub fn effective_capacity(&self) -> u64 {
1392        let used = {
1393            let index = self.shared.index.lock().expect("kv disk index poisoned");
1394            index.bytes
1395        };
1396        self.shared.effective_capacity(used)
1397    }
1398
1399    /// Where a block's file lives (or would). Useful to an operator
1400    /// tracing one block; the file may not exist yet, or at all.
1401    pub fn block_path(&self, hash: &BlockHash) -> PathBuf {
1402        self.shared.block_path(hash)
1403    }
1404
1405    pub fn stats(&self) -> DiskStats {
1406        self.shared.stats()
1407    }
1408}
1409
1410impl Shared {
1411    fn next_generation(&self) -> u64 {
1412        self.generation.fetch_add(1, Ordering::SeqCst) + 1
1413    }
1414
1415    /// The write path, in the order the plan requires:
1416    ///
1417    /// 1. **buffer** -- the payload is reachable before anything claims
1418    ///    it exists;
1419    /// 2. **index** -- now it is claimed to exist, and is subject to
1420    ///    eviction and budget accounting;
1421    /// 3. **queue** -- only now does the write get scheduled.
1422    ///
1423    /// Reversing 1 and 2 lets a reader hit the index for a block with
1424    /// no file and no payload. Reversing 2 and 3 would be harmless but
1425    /// pointless: a queued write whose block is not indexed cannot be
1426    /// evicted or accounted for.
1427    fn put(&self, hash: BlockHash, block: KvBlock, inline: bool) -> Result<(), StoreError> {
1428        let bytes = encoded_len(block.signature());
1429        let block = Arc::new(block);
1430        let generation = self.next_generation();
1431
1432        #[cfg(test)]
1433        let index_first = *self.hooks.order.lock().expect("kv disk hook poisoned")
1434            == WriteOrder::IndexBeforeBuffer;
1435        #[cfg(not(test))]
1436        let index_first = false;
1437
1438        if index_first {
1439            self.reserve(hash, bytes, generation);
1440            #[cfg(test)]
1441            Hooks::fire(&self.hooks.in_put_window, &hash);
1442            self.buffer_block(hash, generation, Arc::clone(&block));
1443        } else {
1444            self.buffer_block(hash, generation, Arc::clone(&block));
1445            #[cfg(test)]
1446            Hooks::fire(&self.hooks.in_put_window, &hash);
1447            self.reserve(hash, bytes, generation);
1448        }
1449
1450        let job = WriteJob { hash, generation };
1451        if !inline && self.queue.try_push(job) {
1452            self.stats.queued_writes.fetch_add(1, Ordering::Relaxed);
1453            return Ok(());
1454        }
1455        if !inline {
1456            self.stats.inline_writes.fetch_add(1, Ordering::Relaxed);
1457        }
1458        self.run_write(job, block)
1459    }
1460
1461    fn buffer_block(&self, hash: BlockHash, generation: u64, block: Arc<KvBlock>) {
1462        self.buffer
1463            .lock()
1464            .expect("kv disk buffer poisoned")
1465            .insert(hash, Buffered { generation, block });
1466    }
1467
1468    /// Drops a buffered payload, but only if it is still the one this
1469    /// generation put there -- a newer `put` for the same hash owns the
1470    /// slot now.
1471    fn release_buffer(&self, hash: &BlockHash, generation: u64) {
1472        let mut buffer = self.buffer.lock().expect("kv disk buffer poisoned");
1473        if buffer.get(hash).is_some_and(|b| b.generation == generation) {
1474            buffer.remove(hash);
1475        }
1476    }
1477
1478    fn buffered(&self, hash: &BlockHash, generation: u64) -> Option<Arc<KvBlock>> {
1479        let buffer = self.buffer.lock().expect("kv disk buffer poisoned");
1480        buffer
1481            .get(hash)
1482            .filter(|b| b.generation == generation)
1483            .map(|b| Arc::clone(&b.block))
1484    }
1485
1486    /// Runs a queued write. A block whose buffered payload is gone was
1487    /// evicted or superseded while it waited, and is skipped rather
1488    /// than resurrected.
1489    fn run_job(&self, job: WriteJob) {
1490        match self.buffered(&job.hash, job.generation) {
1491            Some(block) => {
1492                let _ = self.run_write(job, block);
1493            }
1494            None => {
1495                self.stats.write_skipped.fetch_add(1, Ordering::Relaxed);
1496            }
1497        }
1498    }
1499
1500    fn run_write(&self, job: WriteJob, block: Arc<KvBlock>) -> Result<(), StoreError> {
1501        let started = Instant::now();
1502        let result = self.write_and_publish(&job.hash, &block, job.generation);
1503        self.stats
1504            .write_nanos
1505            .fetch_add(started.elapsed().as_nanos() as u64, Ordering::Relaxed);
1506        match result {
1507            Ok(()) => {
1508                self.stats.writes.fetch_add(1, Ordering::Relaxed);
1509                Ok(())
1510            }
1511            Err(err) => {
1512                self.stats.write_failures.fetch_add(1, Ordering::Relaxed);
1513                self.abandon(&job.hash, job.generation);
1514                Err(err)
1515            }
1516        }
1517    }
1518
1519    /// Reserves (or re-reserves) an index entry, charges its bytes, and
1520    /// evicts whatever that pushes over budget.
1521    fn reserve(&self, hash: BlockHash, bytes: u64, generation: u64) {
1522        let mut index = self.index.lock().expect("kv disk index poisoned");
1523        let last_used = index.touch();
1524        index.insert_entry(
1525            hash,
1526            Entry {
1527                bytes,
1528                last_used,
1529                published: false,
1530                generation,
1531            },
1532        );
1533        let victims = self.collect_victims(&mut index, Some(&hash));
1534        drop(index);
1535        self.discard(victims);
1536    }
1537
1538    /// Drops an entry that will never be published (a failed write).
1539    /// Index first, then the buffer: an entry that is gone from the
1540    /// index is unreachable, so no reader can be looking for the
1541    /// payload we are about to free.
1542    fn abandon(&self, hash: &BlockHash, generation: u64) {
1543        {
1544            let mut index = self.index.lock().expect("kv disk index poisoned");
1545            if index
1546                .entries
1547                .get(hash)
1548                .is_some_and(|e| e.generation == generation)
1549            {
1550                index.remove_entry(hash);
1551            }
1552        }
1553        self.release_buffer(hash, generation);
1554    }
1555
1556    fn write_and_publish(
1557        &self,
1558        hash: &BlockHash,
1559        block: &KvBlock,
1560        generation: u64,
1561    ) -> Result<(), StoreError> {
1562        let bytes = encode_block(hash, block);
1563        let final_path = self.block_path(hash);
1564        let shard = final_path.parent().expect("block path has a parent");
1565        fs::create_dir_all(shard).map_err(|e| io_err("create", shard, e))?;
1566        let tmp_path = self.tmp_path(hash);
1567        {
1568            let mut file =
1569                fs::File::create(&tmp_path).map_err(|e| io_err("create", &tmp_path, e))?;
1570            #[cfg(test)]
1571            let written = if self.hooks.fail_with_enospc.load(Ordering::Relaxed) {
1572                Err(io::Error::from(io::ErrorKind::StorageFull))
1573            } else {
1574                file.write_all(&bytes)
1575            };
1576            #[cfg(not(test))]
1577            let written = file.write_all(&bytes);
1578            if let Err(e) = written {
1579                let _ = fs::remove_file(&tmp_path);
1580                self.note_if_enospc(&e);
1581                return Err(io_err("write", &tmp_path, e));
1582            }
1583            // Without this the rename can be durable while the contents
1584            // are not, which is exactly how a zero-length "published"
1585            // block file appears after a power loss.
1586            if let Err(e) = file.sync_all() {
1587                let _ = fs::remove_file(&tmp_path);
1588                self.note_if_enospc(&e);
1589                return Err(io_err("sync", &tmp_path, e));
1590            }
1591        }
1592        fs::rename(&tmp_path, &final_path).map_err(|e| {
1593            let _ = fs::remove_file(&tmp_path);
1594            io_err("publish", &final_path, e)
1595        })?;
1596
1597        #[cfg(test)]
1598        Hooks::fire(&self.hooks.after_rename, hash);
1599
1600        #[cfg(test)]
1601        let drop_buffer_first = *self.hooks.order.lock().expect("kv disk hook poisoned")
1602            == WriteOrder::DropBufferBeforeMarking;
1603        #[cfg(not(test))]
1604        let drop_buffer_first = false;
1605
1606        // Mark published, *then* release the buffered payload. In
1607        // between, a reader either sees "published" and reads the file
1608        // (which exists) or sees "buffered" and reads the buffer (which
1609        // still holds it). Releasing first opens a window where neither
1610        // is true.
1611        let survived = if drop_buffer_first {
1612            self.release_buffer(hash, generation);
1613            #[cfg(test)]
1614            Hooks::fire(&self.hooks.in_publish_window, hash);
1615            self.mark_published(hash, generation)
1616        } else {
1617            let survived = self.mark_published(hash, generation);
1618            #[cfg(test)]
1619            Hooks::fire(&self.hooks.in_publish_window, hash);
1620            self.release_buffer(hash, generation);
1621            survived
1622        };
1623
1624        if !survived {
1625            // Evicted (or superseded) mid-write. The eviction already
1626            // released this entry's bytes and could not delete a file
1627            // that did not exist yet, so the file is ours to withdraw.
1628            let _ = fs::remove_file(&final_path);
1629            self.stats
1630                .write_raced_eviction
1631                .fetch_add(1, Ordering::Relaxed);
1632        }
1633        Ok(())
1634    }
1635
1636    fn mark_published(&self, hash: &BlockHash, generation: u64) -> bool {
1637        let mut index = self.index.lock().expect("kv disk index poisoned");
1638        match index.entries.get_mut(hash) {
1639            Some(entry) if entry.generation == generation => {
1640                if !entry.published {
1641                    entry.published = true;
1642                    let bytes = entry.bytes;
1643                    index.disk_bytes += bytes;
1644                }
1645                true
1646            }
1647            _ => false,
1648        }
1649    }
1650
1651    /// Where a block's payload currently is, decided under the index
1652    /// lock. `Ok(None)` is an outright miss.
1653    fn source(&self, hash: &BlockHash) -> Result<Option<Source>, StoreError> {
1654        let mut index = self.index.lock().expect("kv disk index poisoned");
1655        let clock = index.clock + 1;
1656        let Some(entry) = index.entries.get_mut(hash) else {
1657            return Ok(None);
1658        };
1659        entry.last_used = clock;
1660        let published = entry.published;
1661        index.clock = clock;
1662        if published {
1663            return Ok(Some(Source::Disk(self.block_path(hash))));
1664        }
1665        // The index lock is deliberately still held: "not published"
1666        // and "here is the buffered payload" must be one decision. Two
1667        // decisions leave a gap for the writer to publish and release
1668        // in between.
1669        let buffered = self
1670            .buffer
1671            .lock()
1672            .expect("kv disk buffer poisoned")
1673            .get(hash)
1674            .map(|b| Arc::clone(&b.block));
1675        match buffered {
1676            Some(block) => Ok(Some(Source::Buffer(block))),
1677            None => Err(StoreError::MissingPayload { hash: *hash }),
1678        }
1679    }
1680
1681    /// Begins a read.
1682    ///
1683    /// `Ok(None)` in the eventual outcome covers every "you will have
1684    /// to recompute this" case -- absent, corrupt on disk, or built for
1685    /// a different config -- because they are the same answer to the
1686    /// caller. The counters in [`DiskStats`] tell them apart. `Err` is
1687    /// I/O that failed, or the write-ordering invariant breaking.
1688    ///
1689    /// Anything answerable from memory (a miss, or a block still in the
1690    /// write buffer) is answered here and comes back already ready; no
1691    /// thread is involved. Only a real file read is dispatched, and a
1692    /// `demand` read jumps ahead of queued prefetches -- or, if the
1693    /// queue is full, runs on this thread rather than queueing behind
1694    /// speculative work.
1695    fn read_async(
1696        self: &Arc<Self>,
1697        hash: &BlockHash,
1698        expected: &CacheSignature,
1699        demand: bool,
1700    ) -> ReadHandle {
1701        // An in-flight or finished read for the same block and the same
1702        // expectation is the answer -- never a second read of the same
1703        // file.
1704        let staged = {
1705            let staging = self.staging.lock().expect("kv disk staging poisoned");
1706            staging
1707                .get(hash)
1708                .filter(|slot| &slot.expected == expected)
1709                .map(Arc::clone)
1710        };
1711        if let Some(slot) = staged {
1712            if demand {
1713                let counter = if slot.is_ready() {
1714                    &self.stats.prefetch_hits
1715                } else {
1716                    &self.stats.prefetch_waits
1717                };
1718                counter.fetch_add(1, Ordering::Relaxed);
1719            }
1720            return self.handle(*hash, slot, true);
1721        }
1722
1723        let ready = |outcome: ReadOutcome| ReadHandle {
1724            shared: Arc::clone(self),
1725            hash: *hash,
1726            slot: ReadSlot::ready(expected.clone(), outcome),
1727            staged: false,
1728        };
1729
1730        let path = match self.source(hash) {
1731            Err(err) => return ready(Err(err)),
1732            Ok(None) => {
1733                self.stats.misses.fetch_add(1, Ordering::Relaxed);
1734                return ready(Ok(None));
1735            }
1736            Ok(Some(Source::Buffer(block))) => {
1737                if block.signature() != expected {
1738                    self.stats.incompatible.fetch_add(1, Ordering::Relaxed);
1739                    return ready(Ok(None));
1740                }
1741                self.stats.buffer_hits.fetch_add(1, Ordering::Relaxed);
1742                return ready(Ok(Some(block)));
1743            }
1744            Ok(Some(Source::Disk(path))) => path,
1745        };
1746
1747        let slot = ReadSlot::pending(expected.clone());
1748        let dispatched = self.has_readers && {
1749            self.staging
1750                .lock()
1751                .expect("kv disk staging poisoned")
1752                .insert(*hash, Arc::clone(&slot));
1753            let job = ReadJob {
1754                hash: *hash,
1755                path: path.clone(),
1756                slot: Arc::clone(&slot),
1757            };
1758            let pushed = self.reads.try_push(job, demand);
1759            if !pushed {
1760                self.unstage(hash, &slot);
1761            }
1762            pushed
1763        };
1764        if dispatched {
1765            return self.handle(*hash, slot, true);
1766        }
1767        slot.fulfil(self.read_timed(&path, hash, expected));
1768        self.handle(*hash, slot, false)
1769    }
1770
1771    fn handle(self: &Arc<Self>, hash: BlockHash, slot: Arc<ReadSlot>, staged: bool) -> ReadHandle {
1772        ReadHandle {
1773            shared: Arc::clone(self),
1774            hash,
1775            slot,
1776            staged,
1777        }
1778    }
1779
1780    /// Removes a staging entry, but only if it is still the slot the
1781    /// claimant was holding -- a newer read for the same block owns the
1782    /// entry now.
1783    fn unstage(&self, hash: &BlockHash, slot: &Arc<ReadSlot>) {
1784        let mut staging = self.staging.lock().expect("kv disk staging poisoned");
1785        if staging.get(hash).is_some_and(|s| Arc::ptr_eq(s, slot)) {
1786            staging.remove(hash);
1787        }
1788    }
1789
1790    fn prefetch(self: &Arc<Self>, hashes: &[BlockHash], expected: &CacheSignature) {
1791        if !self.has_readers {
1792            self.stats
1793                .prefetch_dropped
1794                .fetch_add(hashes.len() as u64, Ordering::Relaxed);
1795            return;
1796        }
1797        for hash in hashes {
1798            let room = {
1799                let staging = self.staging.lock().expect("kv disk staging poisoned");
1800                !staging.contains_key(hash) && staging.len() < self.prefetch_capacity
1801            };
1802            if !room {
1803                self.stats.prefetch_dropped.fetch_add(1, Ordering::Relaxed);
1804                continue;
1805            }
1806            let handle = self.read_async(hash, expected, false);
1807            if handle.staged {
1808                self.stats.prefetch_issued.fetch_add(1, Ordering::Relaxed);
1809            } else {
1810                // Answered from memory, or the read queue was full and
1811                // it ran here. Either way there is nothing staged to
1812                // claim later, so drop the answer: a prefetch is a
1813                // hint, and re-reading is cheaper than holding blocks
1814                // nobody asked for.
1815                self.stats.prefetch_dropped.fetch_add(1, Ordering::Relaxed);
1816            }
1817            // Dropped unclaimed on purpose: a staged slot stays in
1818            // staging for whoever asks for the block next, which is the
1819            // entire point of reading it early.
1820            drop(handle);
1821        }
1822    }
1823
1824    /// The ceiling eviction works to, given how many bytes the store
1825    /// currently holds.
1826    ///
1827    /// `max_bytes` is what an operator asked for; free disk is what the
1828    /// machine can actually give. The store may occupy what it already
1829    /// occupies plus whatever the filesystem still has, less a reserve:
1830    ///
1831    /// ```text
1832    /// effective = min(max_bytes, used + free - reserve)
1833    /// ```
1834    ///
1835    /// As free space falls below the reserve the ceiling drops below
1836    /// `used`, so the next write evicts instead of pushing the
1837    /// filesystem to `ENOSPC`. That is the whole point: the cache
1838    /// should give the disk back before the disk takes it back.
1839    fn effective_capacity(&self, used: u64) -> u64 {
1840        let Some(free) = self.free_bytes() else {
1841            return self.max_bytes;
1842        };
1843        // `free - reserve` is deliberately signed. Clamping it at zero
1844        // would make the ceiling equal `used` exactly when the disk is
1845        // fullest -- the store would sit still and let the filesystem
1846        // do the refusing, which is the failure this whole budget
1847        // exists to avoid.
1848        let headroom = free as i128 - self.reserve_bytes as i128;
1849        let allowed = (used as i128 + headroom).max(0) as u64;
1850        self.max_bytes.min(allowed)
1851    }
1852
1853    /// Free space, re-measured at most once per TTL. A `statvfs` per
1854    /// block write would be a syscall on the write path for a number
1855    /// that moves slowly.
1856    fn free_bytes(&self) -> Option<u64> {
1857        let mut cache = self.free_space.lock().expect("kv disk free space poisoned");
1858        if let Some(checked_at) = cache.checked_at {
1859            if checked_at.elapsed() < self.free_space_ttl {
1860                return cache.bytes;
1861            }
1862        }
1863        let bytes = (self.free_space_probe)(&self.root);
1864        cache.checked_at = Some(Instant::now());
1865        cache.bytes = bytes;
1866        bytes
1867    }
1868
1869    /// The filesystem said "full". Whatever the cached free-space
1870    /// reading says, it is wrong *now*, so it is thrown away rather
1871    /// than left to expire -- and an eviction pass runs immediately, so
1872    /// the next write has somewhere to go.
1873    fn note_enospc(&self) {
1874        self.stats.enospc.fetch_add(1, Ordering::Relaxed);
1875        {
1876            let mut cache = self.free_space.lock().expect("kv disk free space poisoned");
1877            cache.checked_at = None;
1878            cache.bytes = None;
1879        }
1880        let victims = {
1881            let mut index = self.index.lock().expect("kv disk index poisoned");
1882            self.collect_victims(&mut index, None)
1883        };
1884        self.discard(victims);
1885    }
1886
1887    /// Recognises the one I/O error the budget exists to prevent.
1888    fn note_if_enospc(&self, err: &io::Error) {
1889        if err.kind() == io::ErrorKind::StorageFull {
1890            self.note_enospc();
1891        }
1892    }
1893
1894    fn read_timed(&self, path: &Path, hash: &BlockHash, expected: &CacheSignature) -> ReadOutcome {
1895        let started = Instant::now();
1896        let outcome = self.read_verified(path, hash, expected);
1897        self.stats
1898            .read_nanos
1899            .fetch_add(started.elapsed().as_nanos() as u64, Ordering::Relaxed);
1900        outcome
1901    }
1902
1903    fn read_verified(
1904        &self,
1905        path: &Path,
1906        hash: &BlockHash,
1907        expected: &CacheSignature,
1908    ) -> Result<Option<Arc<KvBlock>>, StoreError> {
1909        let bytes = match fs::read(path) {
1910            Ok(bytes) => bytes,
1911            Err(e) if e.kind() == io::ErrorKind::NotFound => {
1912                // The file vanished under us -- an eviction between the
1913                // index lookup and the open, or an external cleaner.
1914                self.stats.misses.fetch_add(1, Ordering::Relaxed);
1915                self.drop_entry(hash);
1916                return Ok(None);
1917            }
1918            Err(e) => return Err(io_err("read", path, e)),
1919        };
1920        let decoded = match decode_block(&bytes) {
1921            Ok(decoded) => decoded,
1922            Err(_) => {
1923                self.stats.corrupt.fetch_add(1, Ordering::Relaxed);
1924                self.quarantine(hash);
1925                return Ok(None);
1926            }
1927        };
1928        if &decoded.hash != hash {
1929            // The file under this name is some other block. Treat it
1930            // exactly like corruption: the name is the identity.
1931            self.stats.corrupt.fetch_add(1, Ordering::Relaxed);
1932            self.quarantine(hash);
1933            return Ok(None);
1934        }
1935        match decoded.block.verify(expected) {
1936            Ok(block) => {
1937                self.stats.hits.fetch_add(1, Ordering::Relaxed);
1938                Ok(Some(Arc::new(block)))
1939            }
1940            Err(_) => {
1941                self.stats.incompatible.fetch_add(1, Ordering::Relaxed);
1942                Ok(None)
1943            }
1944        }
1945    }
1946
1947    fn quarantine(&self, hash: &BlockHash) {
1948        self.drop_entry(hash);
1949        self.buffer
1950            .lock()
1951            .expect("kv disk buffer poisoned")
1952            .remove(hash);
1953        let _ = fs::remove_file(self.block_path(hash));
1954    }
1955
1956    fn drop_entry(&self, hash: &BlockHash) {
1957        let mut index = self.index.lock().expect("kv disk index poisoned");
1958        index.remove_entry(hash);
1959    }
1960
1961    fn reindex(&self) -> Result<usize, StoreError> {
1962        let tmp = self.root.join(TMP_DIR);
1963        if let Ok(entries) = fs::read_dir(&tmp) {
1964            for entry in entries.flatten() {
1965                let _ = fs::remove_file(entry.path());
1966            }
1967        }
1968        let mut found = Vec::new();
1969        let shards = fs::read_dir(&self.root).map_err(|e| io_err("read", &self.root, e))?;
1970        for shard in shards.flatten() {
1971            if !shard.file_type().map(|t| t.is_dir()).unwrap_or(false) {
1972                continue;
1973            }
1974            if shard.file_name() == TMP_DIR {
1975                continue;
1976            }
1977            let Ok(files) = fs::read_dir(shard.path()) else {
1978                continue;
1979            };
1980            for file in files.flatten() {
1981                let path = file.path();
1982                if path.extension().and_then(|e| e.to_str()) != Some(BLOCK_FILE_EXT) {
1983                    continue;
1984                }
1985                let Some(hash) = path
1986                    .file_stem()
1987                    .and_then(|s| s.to_str())
1988                    .and_then(parse_hex_hash)
1989                else {
1990                    continue;
1991                };
1992                let Ok(meta) = file.metadata() else { continue };
1993                found.push((hash, meta.len()));
1994            }
1995        }
1996        let mut adopted = 0;
1997        let mut index = self.index.lock().expect("kv disk index poisoned");
1998        for (hash, bytes) in found {
1999            if index.entries.contains_key(&hash) {
2000                continue;
2001            }
2002            let last_used = index.touch();
2003            index.insert_entry(
2004                hash,
2005                Entry {
2006                    bytes,
2007                    last_used,
2008                    published: true,
2009                    generation: self.next_generation(),
2010                },
2011            );
2012            index.disk_bytes += bytes;
2013            adopted += 1;
2014        }
2015        let victims = self.collect_victims(&mut index, None);
2016        drop(index);
2017        self.discard(victims);
2018        Ok(adopted)
2019    }
2020
2021    /// Picks least-recently-used entries until the store fits its
2022    /// budget, removing them from the index and returning them for
2023    /// disposal. Index first, payload second -- an entry that is gone
2024    /// from the index is unreachable, whereas a file deleted while the
2025    /// index still points at it would be a hit that fails to open.
2026    fn collect_victims(&self, index: &mut Index, protect: Option<&BlockHash>) -> Vec<Victim> {
2027        // Computed once per pass, against the pre-eviction size. It has
2028        // to be: the ceiling is a function of how much the store
2029        // already occupies, so re-evaluating it as blocks leave would
2030        // make eviction chase its own tail down to empty.
2031        let budget = self.effective_capacity(index.disk_bytes);
2032        if budget < self.max_bytes {
2033            self.stats.space_clamped.fetch_add(1, Ordering::Relaxed);
2034        }
2035        if index.bytes <= budget {
2036            return Vec::new();
2037        }
2038        let mut candidates: Vec<(u64, BlockHash)> = index
2039            .entries
2040            .iter()
2041            .filter(|(hash, _)| Some(*hash) != protect)
2042            .map(|(hash, entry)| (entry.last_used, *hash))
2043            .collect();
2044        candidates.sort_unstable();
2045        let mut victims = Vec::new();
2046        for (_, hash) in candidates {
2047            if index.bytes <= budget {
2048                break;
2049            }
2050            if let Some(entry) = index.remove_entry(&hash) {
2051                self.stats.evictions.fetch_add(1, Ordering::Relaxed);
2052                self.stats
2053                    .evicted_bytes
2054                    .fetch_add(entry.bytes, Ordering::Relaxed);
2055                victims.push(Victim {
2056                    hash,
2057                    published: entry.published,
2058                });
2059            }
2060        }
2061        victims
2062    }
2063
2064    /// Frees what eviction removed from the index: the buffered payload
2065    /// and, if it got that far, the file.
2066    fn discard(&self, victims: Vec<Victim>) {
2067        if victims.is_empty() {
2068            return;
2069        }
2070        {
2071            let mut buffer = self.buffer.lock().expect("kv disk buffer poisoned");
2072            for victim in &victims {
2073                buffer.remove(&victim.hash);
2074            }
2075        }
2076        for victim in victims {
2077            if victim.published {
2078                let _ = fs::remove_file(self.block_path(&victim.hash));
2079            }
2080        }
2081    }
2082
2083    fn stats(&self) -> DiskStats {
2084        let index = self.index.lock().expect("kv disk index poisoned");
2085        let stats = &self.stats;
2086        DiskStats {
2087            blocks: index.entries.len(),
2088            bytes: index.bytes,
2089            queue_depth: self.queue.depth(),
2090            writes: stats.writes.load(Ordering::Relaxed),
2091            queued_writes: stats.queued_writes.load(Ordering::Relaxed),
2092            inline_writes: stats.inline_writes.load(Ordering::Relaxed),
2093            write_failures: stats.write_failures.load(Ordering::Relaxed),
2094            write_skipped: stats.write_skipped.load(Ordering::Relaxed),
2095            write_raced_eviction: stats.write_raced_eviction.load(Ordering::Relaxed),
2096            write_nanos: stats.write_nanos.load(Ordering::Relaxed),
2097            hits: stats.hits.load(Ordering::Relaxed),
2098            buffer_hits: stats.buffer_hits.load(Ordering::Relaxed),
2099            misses: stats.misses.load(Ordering::Relaxed),
2100            corrupt: stats.corrupt.load(Ordering::Relaxed),
2101            incompatible: stats.incompatible.load(Ordering::Relaxed),
2102            read_nanos: stats.read_nanos.load(Ordering::Relaxed),
2103            evictions: stats.evictions.load(Ordering::Relaxed),
2104            evicted_bytes: stats.evicted_bytes.load(Ordering::Relaxed),
2105            prefetch_issued: stats.prefetch_issued.load(Ordering::Relaxed),
2106            prefetch_dropped: stats.prefetch_dropped.load(Ordering::Relaxed),
2107            prefetch_hits: stats.prefetch_hits.load(Ordering::Relaxed),
2108            prefetch_waits: stats.prefetch_waits.load(Ordering::Relaxed),
2109            async_reads: stats.async_reads.load(Ordering::Relaxed),
2110            staged_blocks: self.staging.lock().expect("kv disk staging poisoned").len(),
2111            enospc: stats.enospc.load(Ordering::Relaxed),
2112            space_clamped: stats.space_clamped.load(Ordering::Relaxed),
2113            effective_capacity: self.effective_capacity(index.disk_bytes),
2114            disk_bytes: index.disk_bytes,
2115        }
2116    }
2117
2118    fn block_path(&self, hash: &BlockHash) -> PathBuf {
2119        self.root
2120            .join(hash.shard_prefix(self.shard_chars))
2121            .join(format!("{}.{BLOCK_FILE_EXT}", hash.to_hex()))
2122    }
2123
2124    fn tmp_path(&self, hash: &BlockHash) -> PathBuf {
2125        let n = self.seq.fetch_add(1, Ordering::Relaxed);
2126        self.root.join(TMP_DIR).join(format!(
2127            "{}.{}.{n}.tmp",
2128            hash.shard_prefix(16),
2129            std::process::id()
2130        ))
2131    }
2132}
2133
2134/// An entry eviction has already removed from the index, awaiting
2135/// disposal of its payload.
2136struct Victim {
2137    hash: BlockHash,
2138    published: bool,
2139}
2140
2141fn parse_hex_hash(text: &str) -> Option<BlockHash> {
2142    if text.len() != 64 {
2143        return None;
2144    }
2145    let mut out = [0u8; 32];
2146    for (i, byte) in out.iter_mut().enumerate() {
2147        let hi = text.as_bytes()[i * 2] as char;
2148        let lo = text.as_bytes()[i * 2 + 1] as char;
2149        *byte = ((hi.to_digit(16)? << 4) | lo.to_digit(16)?) as u8;
2150    }
2151    Some(BlockHash::from_bytes(out))
2152}
2153
2154#[cfg(test)]
2155mod tests {
2156    use super::*;
2157    use crate::kv_block::BlockHasher;
2158    use std::sync::atomic::AtomicUsize;
2159
2160    /// A throwaway directory that removes itself, so a failing test
2161    /// does not leave block files behind. `std::env::temp_dir` plus the
2162    /// pid and a counter: the workspace has no `tempfile` dependency
2163    /// and this needs eight lines.
2164    struct TempDir(PathBuf);
2165
2166    impl TempDir {
2167        fn new(tag: &str) -> Self {
2168            static N: AtomicU64 = AtomicU64::new(0);
2169            let path = std::env::temp_dir().join(format!(
2170                "ferrox-kvdisk-{tag}-{}-{}",
2171                std::process::id(),
2172                N.fetch_add(1, Ordering::Relaxed)
2173            ));
2174            let _ = fs::remove_dir_all(&path);
2175            fs::create_dir_all(&path).expect("temp dir");
2176            TempDir(path)
2177        }
2178
2179        fn path(&self) -> &Path {
2180            &self.0
2181        }
2182    }
2183
2184    impl Drop for TempDir {
2185        fn drop(&mut self) {
2186            let _ = fs::remove_dir_all(&self.0);
2187        }
2188    }
2189
2190    fn layer(n_kv_heads: usize, head_dim: usize, tokens: usize, fill: f32) -> KvCache {
2191        let mut cache = KvCache::new(n_kv_heads, head_dim);
2192        for t in 0..tokens {
2193            let k = vec![fill + t as f32; n_kv_heads * head_dim];
2194            let v = vec![fill - t as f32; n_kv_heads * head_dim];
2195            cache.push(&k, &v).expect("unpooled push cannot fail");
2196        }
2197        cache
2198    }
2199
2200    /// A full-causal layout whose block size is the block's own token
2201    /// depth -- the default for every test that is not about SWA.
2202    fn flat(tokens: usize) -> BlockLayout {
2203        BlockLayout::full_attention(tokens).expect("positive block size")
2204    }
2205
2206    fn block(model: &str, n_layers: usize, tokens: usize, fill: f32) -> KvBlock {
2207        block_with_layout(model, n_layers, tokens, fill, flat(tokens))
2208    }
2209
2210    fn block_with_layout(
2211        model: &str,
2212        n_layers: usize,
2213        tokens: usize,
2214        fill: f32,
2215        layout: BlockLayout,
2216    ) -> KvBlock {
2217        let layers = (0..n_layers)
2218            .map(|l| layer(2, 4, tokens, fill + l as f32 * 100.0))
2219            .collect();
2220        KvBlock::stamp(model, layout, layers).expect("stamp")
2221    }
2222
2223    fn expected(model: &str, n_layers: usize, tokens: usize) -> CacheSignature {
2224        CacheSignature::expected(model, flat(tokens), n_layers, 2, 4, tokens)
2225    }
2226
2227    fn hash(n: usize) -> BlockHash {
2228        BlockHasher::new("model-a", &[] as &[&str]).chain(&[n, n + 1], 2)[0]
2229    }
2230
2231    /// A free-space probe that reports a terabyte. Tests must not
2232    /// depend on how full the machine's real disk happens to be -- that
2233    /// is exactly the variable the budget tests below control on
2234    /// purpose.
2235    fn plenty() -> FreeSpaceProbe {
2236        Arc::new(|_: &Path| Some(1 << 40))
2237    }
2238
2239    /// Default test store: everything on the calling thread, so a test
2240    /// that does not care about the writer pool never races it.
2241    fn store(dir: &TempDir, max_bytes: u64) -> DiskKvStore {
2242        DiskKvStore::open(
2243            DiskConfig::new(dir.path())
2244                .with_max_bytes(max_bytes)
2245                .with_writer_threads(0)
2246                .with_free_space_probe(plenty()),
2247        )
2248        .expect("open")
2249    }
2250
2251    fn put_now(store: &DiskKvStore, hash: BlockHash, block: KvBlock) {
2252        store.put_blocking(hash, block).expect("put");
2253    }
2254
2255    #[test]
2256    fn a_block_round_trips_through_a_file() {
2257        let dir = TempDir::new("roundtrip");
2258        let store = store(&dir, 1 << 20);
2259        let h = hash(1);
2260        let written = block("model-a", 3, 4, 1.0);
2261        let copy = block("model-a", 3, 4, 1.0);
2262        put_now(&store, h, written);
2263
2264        let read = store
2265            .get(&h, &expected("model-a", 3, 4))
2266            .expect("get")
2267            .expect("the block just written must be found");
2268        assert_eq!(read.layers().len(), 3);
2269        for (a, b) in read.layers().iter().zip(copy.layers()) {
2270            assert_eq!(a.k, b.k);
2271            assert_eq!(a.v, b.v);
2272            assert_eq!(a.positions(), b.positions());
2273        }
2274        let stats = store.stats();
2275        assert_eq!(stats.hits, 1);
2276        assert_eq!(stats.writes, 1);
2277        assert_eq!(stats.blocks, 1);
2278        assert!(stats.read_nanos > 0, "a read must be timed");
2279        assert!(stats.write_nanos > 0, "a write must be timed");
2280    }
2281
2282    #[test]
2283    fn the_accounted_size_is_the_real_file_size() {
2284        let dir = TempDir::new("size");
2285        let store = store(&dir, 1 << 20);
2286        let h = hash(2);
2287        let written = block("model-a", 2, 8, 0.25);
2288        let predicted = encoded_len(written.signature());
2289        put_now(&store, h, written);
2290        let on_disk = fs::metadata(store.block_path(&h)).expect("stat").len();
2291        assert_eq!(
2292            predicted, on_disk,
2293            "the budget charges what the file really costs"
2294        );
2295        assert_eq!(store.stats().bytes, on_disk);
2296    }
2297
2298    #[test]
2299    fn blocks_are_sharded_by_hash_prefix() {
2300        let dir = TempDir::new("shard");
2301        let store = DiskKvStore::open(
2302            DiskConfig::new(dir.path())
2303                .with_shard_chars(2)
2304                .with_writer_threads(0)
2305                .with_free_space_probe(plenty()),
2306        )
2307        .expect("open");
2308        let h = hash(3);
2309        put_now(&store, h, block("model-a", 1, 2, 1.0));
2310        let path = store.block_path(&h);
2311        assert_eq!(
2312            path.parent()
2313                .unwrap()
2314                .file_name()
2315                .unwrap()
2316                .to_str()
2317                .unwrap(),
2318            &h.to_hex()[..2]
2319        );
2320        assert!(path.exists());
2321    }
2322
2323    /// The crash-safety case, stated as the failure it prevents: a file
2324    /// cut short must be *refused*, not parsed into whatever the
2325    /// remaining bytes happen to say.
2326    #[test]
2327    fn a_truncated_file_is_refused_at_every_cut_point() {
2328        let h = hash(4);
2329        let bytes = encode_block(&h, &block("model-a", 2, 4, 3.0));
2330        assert!(bytes.len() > PREFIX_LEN + 16);
2331
2332        // Cut inside the fixed prefix: not even a header to read.
2333        let err = decode_block(&bytes[..PREFIX_LEN - 1]).expect_err("short file");
2334        assert_eq!(
2335            err,
2336            BlockFormatError::TooShort {
2337                len: PREFIX_LEN - 1
2338            }
2339        );
2340
2341        // Cut inside the body: the declared length no longer matches.
2342        for cut in [PREFIX_LEN, PREFIX_LEN + 8, bytes.len() - 4, bytes.len() - 1] {
2343            let err = decode_block(&bytes[..cut]).expect_err("truncated file");
2344            assert_eq!(
2345                err,
2346                BlockFormatError::Truncated {
2347                    expected: bytes.len() as u64,
2348                    actual: cut as u64,
2349                },
2350                "a file cut at {cut} must be refused"
2351            );
2352        }
2353
2354        // A file that is the right length but has been altered.
2355        let mut flipped = bytes.clone();
2356        let last = flipped.len() - 1;
2357        flipped[last] ^= 0xff;
2358        assert_eq!(
2359            decode_block(&flipped).expect_err("altered file"),
2360            BlockFormatError::ChecksumMismatch
2361        );
2362
2363        // And a file that is not one of ours at all.
2364        let mut alien = bytes;
2365        alien[0] = b'X';
2366        assert_eq!(
2367            decode_block(&alien).expect_err("foreign file"),
2368            BlockFormatError::BadMagic
2369        );
2370    }
2371
2372    /// Truncation through the whole store, not just the decoder: a torn
2373    /// file is a miss and is removed, so the next request recomputes
2374    /// instead of tripping over it forever.
2375    #[test]
2376    fn a_torn_file_on_disk_is_a_miss_and_is_quarantined() {
2377        let dir = TempDir::new("torn");
2378        let store = store(&dir, 1 << 20);
2379        let h = hash(5);
2380        put_now(&store, h, block("model-a", 2, 4, 1.0));
2381        let path = store.block_path(&h);
2382
2383        // Simulate a write that died half way.
2384        let full = fs::read(&path).expect("read back");
2385        fs::write(&path, &full[..full.len() / 2]).expect("truncate");
2386
2387        let got = store.get(&h, &expected("model-a", 2, 4)).expect("get");
2388        assert!(got.is_none(), "a torn block must not be returned");
2389        assert_eq!(store.stats().corrupt, 1);
2390        assert!(!path.exists(), "a torn block must not be left to trip over");
2391        assert!(!store.contains(&h));
2392    }
2393
2394    #[test]
2395    fn an_unreadable_format_version_is_refused() {
2396        let h = hash(6);
2397        let mut bytes = encode_block(&h, &block("model-a", 1, 2, 1.0));
2398        bytes[8..12].copy_from_slice(&99u32.to_le_bytes());
2399        // Re-checksum so the version is the only thing wrong.
2400        let mut digest = Sha256::new();
2401        digest.update(&bytes[PREFIX_LEN..]);
2402        let digest: [u8; 32] = digest.finalize().into();
2403        bytes[24..PREFIX_LEN].copy_from_slice(&digest);
2404        assert_eq!(
2405            decode_block(&bytes).expect_err("unknown version"),
2406            BlockFormatError::UnsupportedFormat {
2407                found: 99,
2408                readable: READABLE_FORMAT_VERSIONS,
2409            }
2410        );
2411    }
2412
2413    /// The signature discipline is not re-implemented here, but it is
2414    /// enforced here: a block written under one config must not be
2415    /// handed to a reader expecting another.
2416    #[test]
2417    fn a_block_from_a_different_config_is_a_miss_not_a_hit() {
2418        let dir = TempDir::new("config");
2419        let store = store(&dir, 1 << 20);
2420        let h = hash(7);
2421        put_now(&store, h, block("model-a", 2, 4, 1.0));
2422
2423        assert!(store
2424            .get(&h, &expected("model-b", 2, 4))
2425            .expect("get")
2426            .is_none());
2427        assert!(store
2428            .get(
2429                &h,
2430                &CacheSignature::expected("model-a", flat(4), 2, 8, 4, 4)
2431            )
2432            .expect("get")
2433            .is_none());
2434        assert_eq!(store.stats().incompatible, 2);
2435        assert_eq!(store.stats().hits, 0);
2436        // Still readable by a reader that does match: an incompatible
2437        // read is not destructive.
2438        assert!(store
2439            .get(&h, &expected("model-a", 2, 4))
2440            .expect("get")
2441            .is_some());
2442    }
2443
2444    /// A file whose name says one block and whose contents say another
2445    /// is treated as corruption. The name is the identity; a store that
2446    /// trusted the contents instead would serve a prefix under the
2447    /// wrong hash, which is the silent-wrong-answer case the hashing
2448    /// exists to prevent.
2449    #[test]
2450    fn a_file_stored_under_the_wrong_name_is_rejected() {
2451        let dir = TempDir::new("misfiled");
2452        let store = store(&dir, 1 << 20);
2453        let (a, b) = (hash(8), hash(9));
2454        put_now(&store, a, block("model-a", 1, 2, 1.0));
2455        put_now(&store, b, block("model-a", 1, 2, 2.0));
2456        // Put b's bytes under a's name.
2457        let bytes = fs::read(store.block_path(&b)).expect("read b");
2458        fs::write(store.block_path(&a), bytes).expect("misfile");
2459
2460        assert!(store
2461            .get(&a, &expected("model-a", 1, 2))
2462            .expect("get")
2463            .is_none());
2464        assert_eq!(store.stats().corrupt, 1);
2465    }
2466
2467    #[test]
2468    fn eviction_keeps_the_store_inside_its_budget() {
2469        let dir = TempDir::new("evict");
2470        let one = encoded_len(block("model-a", 1, 4, 1.0).signature());
2471        // Room for two blocks and change, never three.
2472        let store = store(&dir, one * 2 + 8);
2473        let hashes: Vec<BlockHash> = (0..4).map(|i| hash(20 + i)).collect();
2474        for (i, h) in hashes.iter().enumerate() {
2475            put_now(&store, *h, block("model-a", 1, 4, i as f32));
2476            assert!(
2477                store.stats().bytes <= store.capacity(),
2478                "the store must never sit over budget"
2479            );
2480        }
2481        let stats = store.stats();
2482        assert_eq!(stats.blocks, 2);
2483        assert_eq!(stats.evictions, 2);
2484        assert!(stats.evicted_bytes >= one * 2);
2485        // The two oldest are gone, from the index and from the disk.
2486        for h in &hashes[..2] {
2487            assert!(!store.contains(h));
2488            assert!(
2489                !store.block_path(h).exists(),
2490                "an evicted file must be deleted"
2491            );
2492        }
2493        for h in &hashes[2..] {
2494            assert!(store.contains(h));
2495        }
2496    }
2497
2498    #[test]
2499    fn a_read_makes_a_block_the_least_likely_eviction_victim() {
2500        let dir = TempDir::new("lru");
2501        let one = encoded_len(block("model-a", 1, 4, 1.0).signature());
2502        let store = store(&dir, one * 2 + 8);
2503        let (a, b, c) = (hash(30), hash(31), hash(32));
2504        put_now(&store, a, block("model-a", 1, 4, 1.0));
2505        put_now(&store, b, block("model-a", 1, 4, 2.0));
2506        // Touch `a`, so `b` is now the oldest.
2507        assert!(store.get(&a, &expected("model-a", 1, 4)).unwrap().is_some());
2508        put_now(&store, c, block("model-a", 1, 4, 3.0));
2509
2510        assert!(store.contains(&a), "a recently read block must survive");
2511        assert!(!store.contains(&b));
2512        assert!(store.contains(&c));
2513    }
2514
2515    /// The post-rename re-check. The block is evicted in the window
2516    /// between the rename and the index update -- exactly the window
2517    /// the re-check exists for. Without it the file stays on disk
2518    /// forever with nothing accounting for its bytes, and the store
2519    /// drifts over budget one raced write at a time.
2520    #[test]
2521    fn a_block_evicted_mid_write_does_not_leave_its_file_behind() {
2522        let dir = TempDir::new("raced");
2523        let store = store(&dir, 1 << 20);
2524        let h = hash(40);
2525        {
2526            let evicting = Arc::clone(&store.shared);
2527            let mut hook = store
2528                .shared
2529                .hooks
2530                .after_rename
2531                .lock()
2532                .expect("hook lock poisoned");
2533            *hook = Some(Arc::new(move |hash: &BlockHash| {
2534                // Whoever evicts cannot delete a file that does not
2535                // exist yet; the writer must notice and withdraw it.
2536                evicting.drop_entry(hash);
2537            }));
2538        }
2539        put_now(&store, h, block("model-a", 1, 4, 1.0));
2540
2541        assert!(
2542            !store.block_path(&h).exists(),
2543            "a file published for an entry that no longer exists must be withdrawn"
2544        );
2545        assert!(!store.contains(&h));
2546        let stats = store.stats();
2547        assert_eq!(stats.write_raced_eviction, 1);
2548        assert_eq!(stats.bytes, 0, "no bytes may be left unaccounted");
2549    }
2550
2551    #[test]
2552    fn no_temp_files_survive_a_successful_write() {
2553        let dir = TempDir::new("tmp");
2554        let store = store(&dir, 1 << 20);
2555        for i in 0..4 {
2556            put_now(&store, hash(50 + i), block("model-a", 1, 2, i as f32));
2557        }
2558        let leftovers: Vec<_> = fs::read_dir(dir.path().join(TMP_DIR))
2559            .expect("tmp dir")
2560            .flatten()
2561            .collect();
2562        assert!(
2563            leftovers.is_empty(),
2564            "temp files must not accumulate: {leftovers:?}"
2565        );
2566    }
2567
2568    /// Survival across a restart is the whole point of the tier: a
2569    /// second store opened on the same directory finds what the first
2570    /// one wrote, and sweeps away temp files that never published.
2571    #[test]
2572    fn a_new_store_reattaches_to_what_the_previous_one_published() {
2573        let dir = TempDir::new("restart");
2574        let h = hash(60);
2575        {
2576            let store = store(&dir, 1 << 20);
2577            put_now(&store, h, block("model-a", 2, 4, 7.0));
2578        }
2579        // A write that died before publishing.
2580        let orphan = dir.path().join(TMP_DIR).join("dead.tmp");
2581        fs::write(&orphan, b"half a block").expect("orphan");
2582
2583        let reopened = store(&dir, 1 << 20);
2584        assert!(
2585            !reopened.contains(&h),
2586            "reattaching must be an explicit step, not a side effect of open()"
2587        );
2588        assert_eq!(reopened.reindex().expect("reindex"), 1);
2589        assert!(reopened.contains(&h));
2590        assert!(!orphan.exists(), "an unpublished temp file must be swept");
2591
2592        let read = reopened
2593            .get(&h, &expected("model-a", 2, 4))
2594            .expect("get")
2595            .expect("a block written before the restart must still be readable");
2596        assert_eq!(read.tokens(), 4);
2597    }
2598
2599    /// `kv-swa-block-alignment`, at the layer that makes it dangerous.
2600    ///
2601    /// The disk tier is the thing that carries a block past the death
2602    /// of the process that computed it, so a window change between two
2603    /// runs is not a hypothetical: run 1 fills the cache at window 128,
2604    /// somebody edits the config, run 2 reindexes the same directory
2605    /// and asks for the same prefix. The tokens match, the hash matches,
2606    /// the tensors are the right shape -- everything except the mask the
2607    /// state was computed under. The block must be refused, and counted
2608    /// as incompatible so an operator can see why their hit rate went
2609    /// to zero.
2610    #[test]
2611    fn a_block_written_under_one_window_is_not_served_to_a_reader_expecting_another() {
2612        let dir = TempDir::new("swa-window-restart");
2613        let h = hash(90);
2614        let window_128 = BlockLayout::new(4, Some(128)).expect("4 divides 128");
2615        let window_256 = BlockLayout::new(4, Some(256)).expect("4 divides 256");
2616        {
2617            let store = store(&dir, 1 << 20);
2618            put_now(
2619                &store,
2620                h,
2621                block_with_layout("model-a", 2, 4, 7.0, window_128),
2622            );
2623        }
2624
2625        let reopened = store(&dir, 1 << 20);
2626        assert_eq!(reopened.reindex().expect("reindex"), 1);
2627        assert!(reopened.contains(&h), "the file is there");
2628
2629        let after_window_change = reopened
2630            .get(
2631                &h,
2632                &CacheSignature::expected("model-a", window_256, 2, 2, 4, 4),
2633            )
2634            .expect("a config change is a miss, not an I/O error");
2635        assert!(
2636            after_window_change.is_none(),
2637            "a block cut against a 128 window must not be handed to a 256-window reader"
2638        );
2639        assert_eq!(reopened.stats().incompatible, 1);
2640        assert_eq!(reopened.stats().hits, 0);
2641
2642        // Unchanged config still hits: the guard invalidates on change,
2643        // not on principle.
2644        let same = reopened
2645            .get(
2646                &h,
2647                &CacheSignature::expected("model-a", window_128, 2, 2, 4, 4),
2648            )
2649            .expect("get")
2650            .expect("the same window must still hit");
2651        assert_eq!(same.tokens(), 4);
2652        assert_eq!(same.layout(), window_128);
2653    }
2654
2655    /// The window survives the round trip through the file at all --
2656    /// without this, the test above would pass for the wrong reason
2657    /// (every decoded block reporting "no window" and every reader
2658    /// expecting one missing).
2659    #[test]
2660    fn the_block_layout_round_trips_through_the_file_format() {
2661        let sliding = BlockLayout::new(4, Some(512)).expect("4 divides 512");
2662        let h = hash(91);
2663        let bytes = encode_block(&h, &block_with_layout("model-a", 2, 4, 1.0, sliding));
2664        let decoded = decode_block(&bytes).expect("decode");
2665        let sig = decoded.block.signature.as_ref().expect("signature");
2666        assert_eq!(sig.layout, sliding);
2667        assert_eq!(sig.layout.sliding_window(), Some(512));
2668        assert_eq!(sig.layout.block_size(), 4);
2669
2670        // And a full-causal block round-trips as full-causal, not as
2671        // "window 0".
2672        let bytes = encode_block(&h, &block("model-a", 2, 4, 1.0));
2673        let decoded = decode_block(&bytes).expect("decode");
2674        let sig = decoded.block.signature.as_ref().expect("signature");
2675        assert_eq!(sig.layout.sliding_window(), None);
2676    }
2677
2678    /// A file whose header records a block size that does not divide
2679    /// its window could not have been written by any correct build, so
2680    /// it is refused at parse time rather than reconstructed into a
2681    /// layout and checked later.
2682    #[test]
2683    fn a_file_recording_a_mis_aligned_layout_is_refused_at_parse_time() {
2684        let h = hash(92);
2685        let mut bytes = encode_block(&h, &block("model-a", 2, 4, 1.0));
2686        // Header sits after the fixed prefix; the window is the 7th u32
2687        // after the 32-byte hash. Block size is 4, so 6 is not a
2688        // multiple of it.
2689        let window_at = PREFIX_LEN + 32 + 4 * 6;
2690        bytes[window_at..window_at + 4].copy_from_slice(&6u32.to_le_bytes());
2691        // Re-checksum, so the file fails on its layout rather than on
2692        // its digest -- the point is that a *valid* file with an
2693        // impossible layout is still refused.
2694        let mut digest = Sha256::new();
2695        digest.update(&bytes[PREFIX_LEN..]);
2696        let digest: [u8; 32] = digest.finalize().into();
2697        bytes[24..PREFIX_LEN].copy_from_slice(&digest);
2698
2699        let err = decode_block(&bytes).expect_err("6 is not a multiple of 4");
2700        assert!(
2701            matches!(err, BlockFormatError::BadLayout(_)),
2702            "expected a layout refusal, got {err}"
2703        );
2704    }
2705
2706    #[test]
2707    fn reindex_evicts_down_to_the_budget() {
2708        let dir = TempDir::new("reindex-evict");
2709        let one = encoded_len(block("model-a", 1, 4, 1.0).signature());
2710        {
2711            let store = store(&dir, 1 << 20);
2712            for i in 0..4 {
2713                put_now(&store, hash(70 + i), block("model-a", 1, 4, i as f32));
2714            }
2715        }
2716        let small = store(&dir, one * 2 + 8);
2717        small.reindex().expect("reindex");
2718        let stats = small.stats();
2719        assert_eq!(stats.blocks, 2, "a shrunken budget must bind on restart");
2720        assert!(stats.bytes <= small.capacity());
2721    }
2722
2723    #[test]
2724    fn an_absent_block_is_a_plain_miss() {
2725        let dir = TempDir::new("miss");
2726        let store = store(&dir, 1 << 20);
2727        assert!(store
2728            .get(&hash(80), &expected("model-a", 1, 2))
2729            .expect("get")
2730            .is_none());
2731        assert_eq!(store.stats().misses, 1);
2732        assert_eq!(store.stats().corrupt, 0);
2733    }
2734
2735    /// A hit whose file has been deleted behind the store's back (an
2736    /// external cleaner, a `rm -rf` on the shard) is a miss, and the
2737    /// stale entry is dropped rather than left to fail forever.
2738    #[test]
2739    fn a_file_deleted_behind_the_stores_back_is_a_miss() {
2740        let dir = TempDir::new("vanished");
2741        let store = store(&dir, 1 << 20);
2742        let h = hash(90);
2743        put_now(&store, h, block("model-a", 1, 2, 1.0));
2744        fs::remove_file(store.block_path(&h)).expect("remove");
2745        assert!(store
2746            .get(&h, &expected("model-a", 1, 2))
2747            .expect("get")
2748            .is_none());
2749        assert!(!store.contains(&h));
2750        assert_eq!(store.stats().bytes, 0);
2751    }
2752
2753    #[test]
2754    fn rewriting_a_block_does_not_double_charge_it() {
2755        let dir = TempDir::new("rewrite");
2756        let store = store(&dir, 1 << 20);
2757        let h = hash(100);
2758        put_now(&store, h, block("model-a", 1, 4, 1.0));
2759        let once = store.stats().bytes;
2760        put_now(&store, h, block("model-a", 1, 4, 1.0));
2761        assert_eq!(store.stats().bytes, once);
2762        assert_eq!(store.stats().blocks, 1);
2763    }
2764
2765    #[test]
2766    fn hex_names_round_trip() {
2767        let h = hash(110);
2768        assert_eq!(parse_hex_hash(&h.to_hex()), Some(h));
2769        assert_eq!(parse_hex_hash("nothex"), None);
2770        assert_eq!(parse_hex_hash(&"z".repeat(64)), None);
2771    }
2772
2773    // ---------------------------------------------------------------
2774    // Write ordering: buffer -> index -> queue
2775    // ---------------------------------------------------------------
2776
2777    /// Runs one `put` with a reader firing inside the window between
2778    /// the two ordered steps of the write path, and reports what the
2779    /// reader saw. Deterministic on purpose: a sleep-and-hope
2780    /// concurrency test that passes tells you nothing about a window it
2781    /// may simply have missed.
2782    ///
2783    /// Returns `(missing_payload_errors, served)`.
2784    fn probe_window(order: WriteOrder, publish_window: bool) -> (usize, usize) {
2785        let dir = TempDir::new("ordering");
2786        let store = store(&dir, 1 << 20);
2787        *store.shared.hooks.order.lock().unwrap() = order;
2788
2789        let violations = Arc::new(AtomicUsize::new(0));
2790        let served = Arc::new(AtomicUsize::new(0));
2791        let reader = Arc::clone(&store.shared);
2792        let v = Arc::clone(&violations);
2793        let s = Arc::clone(&served);
2794        let hook: Hook = Arc::new(move |hash: &BlockHash| {
2795            match reader
2796                .read_async(hash, &expected("model-a", 1, 4), true)
2797                .wait()
2798            {
2799                Ok(Some(_)) => {
2800                    s.fetch_add(1, Ordering::Relaxed);
2801                }
2802                // Not indexed yet: an honest miss, the reader simply
2803                // recomputes.
2804                Ok(None) => {}
2805                Err(StoreError::MissingPayload { .. }) => {
2806                    v.fetch_add(1, Ordering::Relaxed);
2807                }
2808                Err(other) => panic!("unexpected store error: {other}"),
2809            }
2810        });
2811        let slot = if publish_window {
2812            &store.shared.hooks.in_publish_window
2813        } else {
2814            &store.shared.hooks.in_put_window
2815        };
2816        *slot.lock().unwrap() = Some(hook);
2817
2818        put_now(&store, hash(200), block("model-a", 1, 4, 1.0));
2819        (
2820            violations.load(Ordering::Relaxed),
2821            served.load(Ordering::Relaxed),
2822        )
2823    }
2824
2825    /// The invariant. A reader that looks inside either window -- after
2826    /// the block is buffered but before it is indexed, and after it is
2827    /// published but before the buffer is released -- either misses
2828    /// cleanly or gets the block. It never gets an index hit with
2829    /// nothing behind it.
2830    #[test]
2831    fn a_reader_never_sees_an_index_hit_with_no_payload() {
2832        let (violations, _) = probe_window(WriteOrder::BufferThenIndex, false);
2833        assert_eq!(violations, 0, "admission window must be safe");
2834
2835        let (violations, served) = probe_window(WriteOrder::BufferThenIndex, true);
2836        assert_eq!(violations, 0, "publication window must be safe");
2837        assert_eq!(
2838            served, 1,
2839            "the reader must actually have reached the block, or this test proves nothing"
2840        );
2841    }
2842
2843    /// The proof that the test above is not vacuous: with the two steps
2844    /// of admission reversed -- index first, buffer second, which is
2845    /// the natural way to write it -- the very same reader hits an
2846    /// index entry for a block with no file and no payload.
2847    #[test]
2848    fn indexing_before_buffering_is_caught() {
2849        let (violations, _) = probe_window(WriteOrder::IndexBeforeBuffer, false);
2850        assert_eq!(
2851            violations, 1,
2852            "index-then-buffer must be detected as an invariant violation"
2853        );
2854    }
2855
2856    /// The other end of the write, and the subtler half: releasing the
2857    /// buffered payload before marking the file published leaves the
2858    /// same gap.
2859    #[test]
2860    fn releasing_the_buffer_before_publishing_is_caught() {
2861        let (violations, _) = probe_window(WriteOrder::DropBufferBeforeMarking, true);
2862        assert_eq!(
2863            violations, 1,
2864            "release-then-mark must be detected as an invariant violation"
2865        );
2866    }
2867
2868    /// The same invariant under real concurrency rather than a hook:
2869    /// writers queueing blocks while readers hammer them. This one can
2870    /// only ever *sample* the windows, which is why the deterministic
2871    /// probes above exist -- but it also exercises the writer threads,
2872    /// the queue, and eviction all at once.
2873    #[test]
2874    fn concurrent_readers_never_see_an_index_hit_with_no_payload() {
2875        let dir = TempDir::new("concurrent");
2876        let one = encoded_len(block("model-a", 1, 4, 1.0).signature());
2877        let store = Arc::new(
2878            DiskKvStore::open(
2879                DiskConfig::new(dir.path())
2880                    // Tight enough that eviction runs constantly.
2881                    .with_max_bytes(one * 8)
2882                    .with_queue_capacity(4)
2883                    .with_writer_threads(2)
2884                    .with_free_space_probe(plenty()),
2885            )
2886            .expect("open"),
2887        );
2888        let hashes: Vec<BlockHash> = (0..16).map(|i| hash(300 + i)).collect();
2889        let violations = Arc::new(AtomicUsize::new(0));
2890        let stop = Arc::new(std::sync::atomic::AtomicBool::new(false));
2891
2892        let readers: Vec<_> = (0..4)
2893            .map(|_| {
2894                let store = Arc::clone(&store);
2895                let hashes = hashes.clone();
2896                let violations = Arc::clone(&violations);
2897                let stop = Arc::clone(&stop);
2898                std::thread::spawn(move || {
2899                    let want = expected("model-a", 1, 4);
2900                    while !stop.load(Ordering::Relaxed) {
2901                        for h in &hashes {
2902                            match store.get(h, &want) {
2903                                Ok(_) => {}
2904                                Err(StoreError::MissingPayload { .. }) => {
2905                                    violations.fetch_add(1, Ordering::Relaxed);
2906                                }
2907                                Err(other) => panic!("unexpected store error: {other}"),
2908                            }
2909                        }
2910                    }
2911                })
2912            })
2913            .collect();
2914
2915        for round in 0..4 {
2916            for (i, h) in hashes.iter().enumerate() {
2917                store
2918                    .put(*h, block("model-a", 1, 4, (round * 16 + i) as f32))
2919                    .expect("put");
2920            }
2921        }
2922        store.flush();
2923        stop.store(true, Ordering::Relaxed);
2924        for reader in readers {
2925            reader.join().expect("reader thread");
2926        }
2927
2928        assert_eq!(
2929            violations.load(Ordering::Relaxed),
2930            0,
2931            "no reader may ever see an index hit with no payload"
2932        );
2933        let stats = store.stats();
2934        assert!(
2935            stats.buffer_hits > 0,
2936            "readers must have caught blocks still in the write buffer, \
2937             or this test never entered the window"
2938        );
2939        assert!(stats.evictions > 0, "the budget must have bound");
2940        assert!(stats.bytes <= store.capacity());
2941    }
2942
2943    /// A block is readable the instant `put` returns, before any writer
2944    /// thread has touched it. This is what makes the queue safe to use
2945    /// on the request path.
2946    #[test]
2947    fn a_queued_block_is_readable_before_it_reaches_disk() {
2948        let dir = TempDir::new("buffered");
2949        // No writer threads: nothing can publish until we flush.
2950        let store = store(&dir, 1 << 20);
2951        let h = hash(400);
2952        store.put(h, block("model-a", 1, 4, 1.0)).expect("put");
2953
2954        assert!(
2955            !store.block_path(&h).exists(),
2956            "nothing has been written yet"
2957        );
2958        let got = store
2959            .get(&h, &expected("model-a", 1, 4))
2960            .expect("get")
2961            .expect("a queued block must be readable immediately");
2962        assert_eq!(got.tokens(), 4);
2963        assert_eq!(store.stats().buffer_hits, 1);
2964
2965        store.flush();
2966        assert!(store.block_path(&h).exists(), "flush must publish it");
2967        assert!(store
2968            .get(&h, &expected("model-a", 1, 4))
2969            .expect("get")
2970            .is_some());
2971        assert_eq!(store.stats().hits, 1, "and now it comes off the disk");
2972    }
2973
2974    /// Backpressure, not loss: when the queue is full the block is
2975    /// written on the calling thread. Nothing is dropped, and the
2976    /// fallback is counted so an operator can see a queue that is too
2977    /// small.
2978    #[test]
2979    fn a_full_queue_writes_inline_rather_than_dropping_the_block() {
2980        let dir = TempDir::new("backpressure");
2981        let store = DiskKvStore::open(
2982            DiskConfig::new(dir.path())
2983                .with_queue_capacity(2)
2984                // Nothing drains the queue, so it stays full.
2985                .with_writer_threads(0)
2986                .with_free_space_probe(plenty()),
2987        )
2988        .expect("open");
2989
2990        let hashes: Vec<BlockHash> = (0..5).map(|i| hash(500 + i)).collect();
2991        for (i, h) in hashes.iter().enumerate() {
2992            store
2993                .put(*h, block("model-a", 1, 4, i as f32))
2994                .expect("put");
2995        }
2996        let stats = store.stats();
2997        assert_eq!(stats.queued_writes, 2, "the queue holds exactly its cap");
2998        assert_eq!(stats.inline_writes, 3, "the rest fall back to this thread");
2999        assert_eq!(stats.writes, 3, "and the fallbacks really wrote");
3000
3001        // Every block is readable regardless of which path it took --
3002        // the point of "fall back" instead of "drop".
3003        let want = expected("model-a", 1, 4);
3004        for h in &hashes {
3005            assert!(
3006                store.get(h, &want).expect("get").is_some(),
3007                "no block may be lost to a full queue"
3008            );
3009        }
3010        store.flush();
3011        for h in &hashes {
3012            assert!(store.block_path(h).exists(), "flush publishes the rest");
3013        }
3014    }
3015
3016    /// A queued write whose block was evicted before a writer reached
3017    /// it is skipped, not resurrected: eviction has already released
3018    /// its bytes, so writing it would put the store over budget with a
3019    /// file nothing accounts for.
3020    #[test]
3021    fn a_queued_write_evicted_before_it_runs_is_skipped() {
3022        let dir = TempDir::new("skipped");
3023        let store = store(&dir, 1 << 20);
3024        let h = hash(600);
3025        store.put(h, block("model-a", 1, 4, 1.0)).expect("put");
3026        store.remove(&h);
3027        store.flush();
3028
3029        let stats = store.stats();
3030        assert_eq!(stats.write_skipped, 1);
3031        assert_eq!(stats.writes, 0);
3032        assert!(!store.block_path(&h).exists());
3033        assert_eq!(stats.bytes, 0);
3034    }
3035
3036    /// A second `put` for the same hash supersedes the first: the
3037    /// queued job for the older generation finds a payload that is no
3038    /// longer its own and skips, rather than overwriting the newer
3039    /// block with the older one.
3040    #[test]
3041    fn a_superseded_queued_write_does_not_overwrite_the_newer_block() {
3042        let dir = TempDir::new("superseded");
3043        let store = store(&dir, 1 << 20);
3044        let h = hash(700);
3045        store.put(h, block("model-a", 1, 4, 1.0)).expect("put");
3046        store.put(h, block("model-a", 1, 4, 9.0)).expect("put");
3047        store.flush();
3048
3049        let got = store
3050            .get(&h, &expected("model-a", 1, 4))
3051            .expect("get")
3052            .expect("hit");
3053        assert_eq!(
3054            got.layers()[0].k[0],
3055            9.0,
3056            "the newer block must win, not whichever write ran last"
3057        );
3058        assert_eq!(store.stats().write_skipped, 1);
3059        assert_eq!(store.stats().blocks, 1);
3060    }
3061
3062    // ---------------------------------------------------------------
3063    // Asynchronous and prefetched reads
3064    // ---------------------------------------------------------------
3065
3066    /// A store that reads on its own threads, writes on the caller's
3067    /// (so a test's writes are done when `put_blocking` returns).
3068    fn reading_store(dir: &TempDir, readers: usize) -> DiskKvStore {
3069        DiskKvStore::open(
3070            DiskConfig::new(dir.path())
3071                .with_writer_threads(0)
3072                .with_reader_threads(readers)
3073                .with_free_space_probe(plenty()),
3074        )
3075        .expect("open")
3076    }
3077
3078    /// Blocks until a prefetch for `hash` has landed in staging.
3079    fn wait_staged(store: &DiskKvStore, hash: &BlockHash) {
3080        for _ in 0..2000 {
3081            let ready = store
3082                .shared
3083                .staging
3084                .lock()
3085                .unwrap()
3086                .get(hash)
3087                .map(|slot| slot.is_ready())
3088                .unwrap_or(false);
3089            if ready {
3090                return;
3091            }
3092            std::thread::sleep(std::time::Duration::from_millis(1));
3093        }
3094        panic!("prefetch never completed");
3095    }
3096
3097    /// The point of prefetching, proved the only way it can be: the
3098    /// block file is **deleted** after the prefetch and before the
3099    /// request. The request still gets its block, so the read
3100    /// demonstrably happened before it was asked for -- not on its
3101    /// thread, not on its clock.
3102    #[test]
3103    fn a_prefetched_block_is_already_read_when_the_request_arrives() {
3104        let dir = TempDir::new("prefetch");
3105        let store = reading_store(&dir, 1);
3106        let h = hash(800);
3107        put_now(&store, h, block("model-a", 2, 4, 5.0));
3108        let want = expected("model-a", 2, 4);
3109
3110        store.prefetch(&[h], &want);
3111        wait_staged(&store, &h);
3112        fs::remove_file(store.block_path(&h)).expect("remove");
3113
3114        let got = store
3115            .get(&h, &want)
3116            .expect("get")
3117            .expect("the prefetch already had it");
3118        assert_eq!(got.tokens(), 4);
3119        assert_eq!(got.layers()[0].k[0], 5.0);
3120
3121        let stats = store.stats();
3122        assert_eq!(
3123            stats.prefetch_hits, 1,
3124            "the request must have found it ready"
3125        );
3126        assert_eq!(
3127            stats.hits, 1,
3128            "and the file must have been read exactly once"
3129        );
3130        assert_eq!(stats.async_reads, 1, "on a reader thread, not the caller's");
3131        assert_eq!(stats.staged_blocks, 0, "a claimed read leaves staging");
3132    }
3133
3134    /// A whole prefix chain read ahead in one call, which is how a
3135    /// prefix cache would use this: the request's blocks are known
3136    /// before the request needs them.
3137    #[test]
3138    fn a_whole_chain_can_be_read_ahead_in_one_call() {
3139        let dir = TempDir::new("chain");
3140        let store = reading_store(&dir, 2);
3141        let hashes: Vec<BlockHash> = (0..4).map(|i| hash(810 + i)).collect();
3142        for (i, h) in hashes.iter().enumerate() {
3143            put_now(&store, *h, block("model-a", 1, 4, i as f32));
3144        }
3145        let want = expected("model-a", 1, 4);
3146
3147        store.prefetch(&hashes, &want);
3148        for h in &hashes {
3149            wait_staged(&store, h);
3150            fs::remove_file(store.block_path(h)).expect("remove");
3151        }
3152        for (i, h) in hashes.iter().enumerate() {
3153            let got = store.get(h, &want).expect("get").expect("read ahead");
3154            assert_eq!(got.layers()[0].k[0], i as f32);
3155        }
3156        let stats = store.stats();
3157        assert_eq!(stats.prefetch_issued, 4);
3158        assert_eq!(stats.prefetch_hits, 4);
3159        assert_eq!(stats.hits, 4, "four blocks, four reads, none repeated");
3160    }
3161
3162    /// Whether or not the prefetch has landed, the request never reads
3163    /// the same file twice: it joins the read in flight.
3164    #[test]
3165    fn a_request_joins_a_read_already_running_rather_than_repeating_it() {
3166        let dir = TempDir::new("join");
3167        let store = reading_store(&dir, 1);
3168        let h = hash(820);
3169        put_now(&store, h, block("model-a", 1, 4, 1.0));
3170        let want = expected("model-a", 1, 4);
3171
3172        store.prefetch(&[h], &want);
3173        // Deliberately no wait: this races the reader thread, and the
3174        // assertion holds either way.
3175        let got = store.get(&h, &want).expect("get").expect("hit");
3176        assert_eq!(got.tokens(), 4);
3177        let stats = store.stats();
3178        assert_eq!(
3179            stats.prefetch_hits + stats.prefetch_waits,
3180            1,
3181            "the request either found the read done or waited for it"
3182        );
3183        assert_eq!(stats.hits, 1, "one physical read, whichever way it went");
3184    }
3185
3186    /// A staged read belongs to the shape it was issued for. Handing a
3187    /// prefetch's answer to a request that wants a different layout
3188    /// would defeat the whole signature discipline, so the staged slot
3189    /// is only reused on an exact match.
3190    #[test]
3191    fn a_staged_read_is_not_reused_by_a_reader_that_wants_another_shape() {
3192        let dir = TempDir::new("staged-shape");
3193        let store = reading_store(&dir, 1);
3194        let h = hash(830);
3195        put_now(&store, h, block("model-a", 1, 4, 1.0));
3196
3197        store.prefetch(&[h], &expected("model-a", 1, 4));
3198        wait_staged(&store, &h);
3199
3200        let got = store.get(&h, &expected("model-b", 1, 4)).expect("get");
3201        assert!(got.is_none(), "a different model must not be served");
3202        assert_eq!(store.stats().incompatible, 1);
3203        assert_eq!(
3204            store.stats().prefetch_hits,
3205            0,
3206            "the staged answer was for another expectation and must not be claimed"
3207        );
3208    }
3209
3210    /// Reading ahead is a hint and must never be the thing that
3211    /// exhausts memory: past the staging cap, prefetches are refused
3212    /// and counted rather than queued.
3213    #[test]
3214    fn prefetching_is_bounded() {
3215        let dir = TempDir::new("prefetch-bound");
3216        let store = DiskKvStore::open(
3217            DiskConfig::new(dir.path())
3218                .with_writer_threads(0)
3219                .with_reader_threads(1)
3220                .with_prefetch_capacity(2)
3221                .with_free_space_probe(plenty()),
3222        )
3223        .expect("open");
3224        let hashes: Vec<BlockHash> = (0..6).map(|i| hash(840 + i)).collect();
3225        for (i, h) in hashes.iter().enumerate() {
3226            put_now(&store, *h, block("model-a", 1, 4, i as f32));
3227        }
3228
3229        store.prefetch(&hashes, &expected("model-a", 1, 4));
3230        let stats = store.stats();
3231        assert!(
3232            stats.staged_blocks <= 2,
3233            "staging must respect its cap, got {}",
3234            stats.staged_blocks
3235        );
3236        assert!(
3237            stats.prefetch_dropped >= 4,
3238            "the refusals must be visible, got {}",
3239            stats.prefetch_dropped
3240        );
3241
3242        // Refusing a hint costs a read, never an answer.
3243        let want = expected("model-a", 1, 4);
3244        for (i, h) in hashes.iter().enumerate() {
3245            let got = store.get(h, &want).expect("get").expect("hit");
3246            assert_eq!(got.layers()[0].k[0], i as f32);
3247        }
3248    }
3249
3250    /// With no reader threads every read happens on the calling thread
3251    /// and a prefetch is an honest no-op -- counted, not pretended.
3252    #[test]
3253    fn without_reader_threads_reads_run_on_the_caller() {
3254        let dir = TempDir::new("no-readers");
3255        let store = reading_store(&dir, 0);
3256        let h = hash(850);
3257        put_now(&store, h, block("model-a", 1, 4, 1.0));
3258        let want = expected("model-a", 1, 4);
3259
3260        store.prefetch(&[h], &want);
3261        assert_eq!(store.stats().prefetch_dropped, 1);
3262        assert_eq!(store.stats().staged_blocks, 0);
3263
3264        assert!(store.get(&h, &want).expect("get").is_some());
3265        let stats = store.stats();
3266        assert_eq!(stats.hits, 1);
3267        assert_eq!(stats.async_reads, 0);
3268    }
3269
3270    /// The handle is pollable: a caller that has other work to do can
3271    /// start the read, do the work, and collect it.
3272    #[test]
3273    fn a_read_handle_can_be_polled_to_completion() {
3274        let dir = TempDir::new("handle");
3275        let store = reading_store(&dir, 1);
3276        let h = hash(860);
3277        put_now(&store, h, block("model-a", 1, 4, 2.0));
3278        let want = expected("model-a", 1, 4);
3279
3280        let handle = store.read_async(&h, &want);
3281        for _ in 0..2000 {
3282            if let Some(outcome) = handle.try_claim() {
3283                let got = outcome.expect("read").expect("hit");
3284                assert_eq!(got.layers()[0].k[0], 2.0);
3285                assert_eq!(store.stats().staged_blocks, 0);
3286                return;
3287            }
3288            std::thread::sleep(std::time::Duration::from_millis(1));
3289        }
3290        panic!("read never completed");
3291    }
3292
3293    /// A miss needs no thread at all: it is decided under the index
3294    /// lock and comes back already answered.
3295    #[test]
3296    fn a_miss_is_answered_without_dispatching_a_read() {
3297        let dir = TempDir::new("ready-miss");
3298        let store = reading_store(&dir, 1);
3299        let handle = store.read_async(&hash(870), &expected("model-a", 1, 4));
3300        assert!(handle.is_ready(), "a miss must not cost a thread hop");
3301        assert!(handle.wait().expect("read").is_none());
3302        assert_eq!(store.stats().async_reads, 0);
3303    }
3304
3305    /// Abandoning a request drops what was read for it, rather than
3306    /// leaving the blocks parked in staging.
3307    #[test]
3308    fn clearing_the_prefetch_releases_staged_blocks() {
3309        let dir = TempDir::new("clear");
3310        let store = reading_store(&dir, 1);
3311        let h = hash(880);
3312        put_now(&store, h, block("model-a", 1, 4, 1.0));
3313        store.prefetch(&[h], &expected("model-a", 1, 4));
3314        wait_staged(&store, &h);
3315        assert_eq!(store.stats().staged_blocks, 1);
3316        store.clear_prefetch();
3317        assert_eq!(store.stats().staged_blocks, 0);
3318    }
3319
3320    // ---------------------------------------------------------------
3321    // The disk budget
3322    // ---------------------------------------------------------------
3323
3324    /// A store whose free-space reading the test controls.
3325    fn budgeted_store(
3326        dir: &TempDir,
3327        max_bytes: u64,
3328        reserve: u64,
3329        ttl: std::time::Duration,
3330        probe: FreeSpaceProbe,
3331    ) -> DiskKvStore {
3332        DiskKvStore::open(
3333            DiskConfig::new(dir.path())
3334                .with_max_bytes(max_bytes)
3335                .with_reserve_bytes(reserve)
3336                .with_free_space_ttl(ttl)
3337                .with_free_space_probe(probe)
3338                .with_writer_threads(0)
3339                .with_reader_threads(0),
3340        )
3341        .expect("open")
3342    }
3343
3344    /// Bytes really occupying the store's directory. The device model
3345    /// below is driven by this rather than by the store's own
3346    /// accounting, so the test cannot pass by agreeing with itself.
3347    fn dir_bytes(root: &Path) -> u64 {
3348        let mut total = 0;
3349        let Ok(entries) = fs::read_dir(root) else {
3350            return 0;
3351        };
3352        for entry in entries.flatten() {
3353            let path = entry.path();
3354            if path.is_dir() {
3355                total += dir_bytes(&path);
3356            } else if let Ok(meta) = entry.metadata() {
3357                total += meta.len();
3358            }
3359        }
3360        total
3361    }
3362
3363    /// The budget the plan asks for: the configured ceiling is an upper
3364    /// bound, and real free disk can lower it. When the filesystem
3365    /// fills up under a store that is nowhere near its byte budget,
3366    /// eviction is what fires -- not `ENOSPC`.
3367    ///
3368    /// The probe models a real device: free space is what the modelled
3369    /// device has left after the store's actual files, so evicting
3370    /// really does give space back and the ceiling has a fixed point
3371    /// instead of chasing itself down to empty.
3372    #[test]
3373    fn the_ceiling_falls_when_the_filesystem_fills_up() {
3374        let dir = TempDir::new("budget");
3375        let one = encoded_len(block("model-a", 1, 4, 1.0).signature());
3376        let device = Arc::new(AtomicU64::new(1 << 40));
3377        let probe: FreeSpaceProbe = {
3378            let device = Arc::clone(&device);
3379            let root = dir.path().to_path_buf();
3380            Arc::new(move |_: &Path| {
3381                Some(
3382                    device
3383                        .load(Ordering::Relaxed)
3384                        .saturating_sub(dir_bytes(&root)),
3385                )
3386            })
3387        };
3388        // Room for a hundred blocks by byte budget; two blocks' worth
3389        // of headroom demanded on the device.
3390        let reserve = one * 2;
3391        let store = budgeted_store(&dir, one * 100, reserve, std::time::Duration::ZERO, probe);
3392
3393        for i in 0..4 {
3394            put_now(&store, hash(900 + i), block("model-a", 1, 4, i as f32));
3395        }
3396        assert_eq!(store.stats().blocks, 4);
3397        assert_eq!(store.stats().evictions, 0, "nothing binds yet");
3398        assert_eq!(
3399            store.effective_capacity(),
3400            store.capacity(),
3401            "with a terabyte free the configured budget is the ceiling"
3402        );
3403
3404        // The device turns out to be small -- or something else on it
3405        // grew. Six blocks total, of which two must stay free.
3406        device.store(one * 6, Ordering::Relaxed);
3407        assert_eq!(
3408            store.effective_capacity(),
3409            one * 4,
3410            "the ceiling must follow the device down to total - reserve"
3411        );
3412
3413        for i in 0..6 {
3414            put_now(&store, hash(910 + i), block("model-a", 1, 4, i as f32));
3415            let on_disk = dir_bytes(dir.path());
3416            assert!(
3417                on_disk + reserve <= one * 6,
3418                "the store must hand the device its reserve back before the \
3419                 filesystem has to: {on_disk} bytes used of {}, {reserve} reserved",
3420                one * 6
3421            );
3422        }
3423
3424        let stats = store.stats();
3425        assert_eq!(stats.blocks, 4, "settled at total - reserve");
3426        assert!(stats.evictions >= 6, "got {}", stats.evictions);
3427        assert!(stats.space_clamped > 0, "the clamp must be visible");
3428        assert!(
3429            stats.bytes < one * 100,
3430            "far under the configured budget it never reached"
3431        );
3432        assert_eq!(
3433            stats.disk_bytes,
3434            dir_bytes(dir.path()),
3435            "the store's idea of its disk footprint must be the real one"
3436        );
3437    }
3438
3439    /// `statvfs` is a syscall, and free space moves slowly. It is read
3440    /// at most once per TTL rather than once per block written.
3441    #[test]
3442    fn the_free_space_reading_is_cached_for_its_ttl() {
3443        let dir = TempDir::new("ttl");
3444        let calls = Arc::new(AtomicUsize::new(0));
3445        let probe: FreeSpaceProbe = {
3446            let calls = Arc::clone(&calls);
3447            Arc::new(move |_: &Path| {
3448                calls.fetch_add(1, Ordering::Relaxed);
3449                Some(1 << 40)
3450            })
3451        };
3452        let store = budgeted_store(&dir, 1 << 20, 0, std::time::Duration::from_secs(60), probe);
3453
3454        for _ in 0..5 {
3455            store.effective_capacity();
3456        }
3457        for i in 0..3 {
3458            put_now(&store, hash(920 + i), block("model-a", 1, 4, i as f32));
3459        }
3460        assert_eq!(
3461            calls.load(Ordering::Relaxed),
3462            1,
3463            "a TTL'd reading must not be re-taken per operation"
3464        );
3465    }
3466
3467    /// The invalidation that matters: at the moment the filesystem says
3468    /// "full", the cached reading is known to be wrong, so it is thrown
3469    /// away rather than left to expire. And the store does not keep an
3470    /// index entry for a block it could not write.
3471    #[test]
3472    fn enospc_throws_away_the_cached_free_space() {
3473        let dir = TempDir::new("enospc");
3474        let calls = Arc::new(AtomicUsize::new(0));
3475        let probe: FreeSpaceProbe = {
3476            let calls = Arc::clone(&calls);
3477            Arc::new(move |_: &Path| {
3478                calls.fetch_add(1, Ordering::Relaxed);
3479                Some(1 << 40)
3480            })
3481        };
3482        let store = budgeted_store(
3483            &dir,
3484            1 << 20,
3485            0,
3486            // Long enough that nothing but the ENOSPC can re-take it.
3487            std::time::Duration::from_secs(3600),
3488            probe,
3489        );
3490        let h = hash(930);
3491        put_now(&store, h, block("model-a", 1, 4, 1.0));
3492        store.effective_capacity();
3493        assert_eq!(calls.load(Ordering::Relaxed), 1);
3494
3495        store
3496            .shared
3497            .hooks
3498            .fail_with_enospc
3499            .store(true, Ordering::Relaxed);
3500        let full = hash(931);
3501        let err = store
3502            .put_blocking(full, block("model-a", 1, 4, 2.0))
3503            .expect_err("a full filesystem must be reported, not swallowed");
3504        assert!(matches!(err, StoreError::Io { .. }), "{err}");
3505
3506        let stats = store.stats();
3507        assert_eq!(stats.enospc, 1);
3508        assert!(
3509            calls.load(Ordering::Relaxed) > 1,
3510            "ENOSPC must invalidate the cached reading immediately"
3511        );
3512        assert!(
3513            !store.contains(&full),
3514            "a block that could not be written must not be indexed"
3515        );
3516        assert_eq!(stats.write_failures, 1);
3517        let leftovers: Vec<_> = fs::read_dir(dir.path().join(TMP_DIR))
3518            .expect("tmp dir")
3519            .flatten()
3520            .collect();
3521        assert!(
3522            leftovers.is_empty(),
3523            "a failed write must clean up after itself: {leftovers:?}"
3524        );
3525
3526        // The block written before the failure is untouched.
3527        assert!(store
3528            .get(&h, &expected("model-a", 1, 4))
3529            .expect("get")
3530            .is_some());
3531
3532        // And once there is room again, writing resumes.
3533        store
3534            .shared
3535            .hooks
3536            .fail_with_enospc
3537            .store(false, Ordering::Relaxed);
3538        put_now(&store, full, block("model-a", 1, 4, 2.0));
3539        assert!(store.contains(&full));
3540    }
3541
3542    /// A filesystem that cannot be measured is not treated as full, and
3543    /// not treated as infinite either: the configured budget is simply
3544    /// the only ceiling.
3545    #[test]
3546    fn an_unmeasurable_filesystem_falls_back_to_the_configured_budget() {
3547        let dir = TempDir::new("unknowable");
3548        let store = budgeted_store(
3549            &dir,
3550            1 << 20,
3551            1 << 30,
3552            std::time::Duration::ZERO,
3553            Arc::new(|_: &Path| None),
3554        );
3555        assert_eq!(store.effective_capacity(), 1 << 20);
3556        for i in 0..3 {
3557            put_now(&store, hash(940 + i), block("model-a", 1, 4, i as f32));
3558        }
3559        assert_eq!(store.stats().blocks, 3);
3560        assert_eq!(store.stats().evictions, 0);
3561    }
3562
3563    /// The real probe, exercised: it is `unsafe` FFI, and a binding
3564    /// that silently returned nonsense would disable the whole budget
3565    /// without failing anything else.
3566    #[test]
3567    #[cfg(unix)]
3568    fn the_platform_probe_measures_a_real_filesystem() {
3569        let dir = TempDir::new("statvfs");
3570        let free = platform_free_bytes(dir.path()).expect("statvfs on a directory that exists");
3571        assert!(free > 0, "a writable temp dir with zero bytes free?");
3572        assert!(
3573            platform_free_bytes(&dir.path().join("no-such-dir")).is_none(),
3574            "a path that does not exist cannot report free space"
3575        );
3576    }
3577}