1use 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";
114const HEADER_FIELDS: usize = 8;
118const PREFIX_LEN: usize = 8 + 4 + 4 + 8 + 32;
120pub const BLOCK_FILE_EXT: &str = "kvb";
122const 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
142fn 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#[derive(Clone, Debug, PartialEq, Eq)]
167pub enum BlockFormatError {
168 TooShort { len: usize },
171 BadMagic,
173 UnsupportedFormat {
175 found: u32,
176 readable: &'static [u32],
177 },
178 Truncated { expected: u64, actual: u64 },
181 ChecksumMismatch,
184 Malformed(&'static str),
187 UnknownDtype(u32),
189 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
229pub 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 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
275fn 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 per_layer * 2 * sig.n_layers as u64
285}
286
287pub 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#[derive(Debug)]
298pub struct DecodedBlock {
299 pub hash: BlockHash,
301 pub block: UnverifiedBlock,
302}
303
304pub 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 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 cache.seq_len = tokens;
404 layers.push(cache);
405 }
406
407 let signature = CacheSignature {
408 format_version: version,
409 model,
410 n_layers,
411 n_kv_heads,
412 head_dim,
413 dtype,
414 tokens,
415 layout,
416 };
417 Ok(DecodedBlock {
418 hash,
419 block: UnverifiedBlock::new(Some(signature), layers),
420 })
421}
422
423fn read_f32(bytes: &[u8]) -> Vec<f32> {
424 bytes
425 .as_chunks::<4>()
426 .0
427 .iter()
428 .map(|c| f32::from_le_bytes(*c))
429 .collect()
430}
431
432#[derive(Clone, Debug, PartialEq, Eq)]
437pub enum StoreError {
438 Io {
439 op: &'static str,
440 path: PathBuf,
441 message: String,
442 },
443 MissingPayload { hash: BlockHash },
449}
450
451impl std::fmt::Display for StoreError {
452 fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
453 match self {
454 StoreError::Io { op, path, message } => {
455 write!(
456 f,
457 "KV block store failed to {op} {}: {message}",
458 path.display()
459 )
460 }
461 StoreError::MissingPayload { hash } => write!(
462 f,
463 "KV block store index names {hash:?} but it has neither a file nor a buffered \
464 payload; the write-ordering invariant was violated"
465 ),
466 }
467 }
468}
469
470impl std::error::Error for StoreError {}
471
472fn io_err(op: &'static str, path: &Path, err: io::Error) -> StoreError {
473 StoreError::Io {
474 op,
475 path: path.to_path_buf(),
476 message: err.to_string(),
477 }
478}
479
480pub type FreeSpaceProbe = Arc<dyn Fn(&Path) -> Option<u64> + Send + Sync>;
487
488#[cfg(unix)]
492#[allow(clippy::unnecessary_cast)] fn platform_free_bytes(path: &Path) -> Option<u64> {
494 use std::os::unix::ffi::OsStrExt;
495 let c_path = std::ffi::CString::new(path.as_os_str().as_bytes()).ok()?;
496 let stat = unsafe {
499 let mut stat: libc::statvfs = std::mem::zeroed();
500 if libc::statvfs(c_path.as_ptr(), &mut stat) != 0 {
501 return None;
502 }
503 stat
504 };
505 let block = if stat.f_frsize > 0 {
506 stat.f_frsize as u64
507 } else {
508 stat.f_bsize as u64
509 };
510 Some((stat.f_bavail as u64).saturating_mul(block))
511}
512
513#[cfg(not(unix))]
514fn platform_free_bytes(_path: &Path) -> Option<u64> {
515 None
516}
517
518struct FreeSpace {
523 checked_at: Option<Instant>,
524 bytes: Option<u64>,
525}
526
527#[derive(Clone)]
530pub struct DiskConfig {
531 pub root: PathBuf,
533 pub max_bytes: u64,
536 pub shard_chars: usize,
540 pub queue_capacity: usize,
544 pub writer_threads: usize,
547 pub reader_threads: usize,
552 pub prefetch_capacity: usize,
557 pub reserve_bytes: u64,
561 pub free_space_ttl: std::time::Duration,
564 pub free_space_probe: FreeSpaceProbe,
567}
568
569impl std::fmt::Debug for DiskConfig {
570 fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
571 f.debug_struct("DiskConfig")
572 .field("root", &self.root)
573 .field("max_bytes", &self.max_bytes)
574 .field("shard_chars", &self.shard_chars)
575 .field("queue_capacity", &self.queue_capacity)
576 .field("writer_threads", &self.writer_threads)
577 .field("reader_threads", &self.reader_threads)
578 .field("prefetch_capacity", &self.prefetch_capacity)
579 .field("reserve_bytes", &self.reserve_bytes)
580 .field("free_space_ttl", &self.free_space_ttl)
581 .finish_non_exhaustive()
582 }
583}
584
585impl DiskConfig {
586 pub fn new(root: impl Into<PathBuf>) -> Self {
587 DiskConfig {
588 root: root.into(),
589 max_bytes: 1 << 30,
590 shard_chars: 2,
591 queue_capacity: 64,
592 writer_threads: 1,
593 reader_threads: 2,
594 prefetch_capacity: 64,
595 reserve_bytes: 1 << 30,
596 free_space_ttl: std::time::Duration::from_secs(2),
597 free_space_probe: Arc::new(platform_free_bytes),
598 }
599 }
600
601 pub fn with_max_bytes(mut self, max_bytes: u64) -> Self {
602 self.max_bytes = max_bytes;
603 self
604 }
605
606 pub fn with_shard_chars(mut self, shard_chars: usize) -> Self {
607 self.shard_chars = shard_chars.clamp(1, 8);
608 self
609 }
610
611 pub fn with_queue_capacity(mut self, queue_capacity: usize) -> Self {
612 self.queue_capacity = queue_capacity;
613 self
614 }
615
616 pub fn with_writer_threads(mut self, writer_threads: usize) -> Self {
617 self.writer_threads = writer_threads;
618 self
619 }
620
621 pub fn with_reader_threads(mut self, reader_threads: usize) -> Self {
622 self.reader_threads = reader_threads;
623 self
624 }
625
626 pub fn with_prefetch_capacity(mut self, prefetch_capacity: usize) -> Self {
627 self.prefetch_capacity = prefetch_capacity;
628 self
629 }
630
631 pub fn with_reserve_bytes(mut self, reserve_bytes: u64) -> Self {
632 self.reserve_bytes = reserve_bytes;
633 self
634 }
635
636 pub fn with_free_space_ttl(mut self, free_space_ttl: std::time::Duration) -> Self {
637 self.free_space_ttl = free_space_ttl;
638 self
639 }
640
641 pub fn with_free_space_probe(mut self, probe: FreeSpaceProbe) -> Self {
642 self.free_space_probe = probe;
643 self
644 }
645}
646
647#[derive(Default)]
648struct Stats {
649 writes: AtomicU64,
650 queued_writes: AtomicU64,
651 inline_writes: AtomicU64,
656 write_failures: AtomicU64,
657 write_skipped: AtomicU64,
660 write_nanos: AtomicU64,
661 write_raced_eviction: AtomicU64,
664 hits: AtomicU64,
665 buffer_hits: AtomicU64,
669 misses: AtomicU64,
670 corrupt: AtomicU64,
672 incompatible: AtomicU64,
675 read_nanos: AtomicU64,
676 evictions: AtomicU64,
677 evicted_bytes: AtomicU64,
678 prefetch_issued: AtomicU64,
680 prefetch_dropped: AtomicU64,
683 prefetch_hits: AtomicU64,
687 prefetch_waits: AtomicU64,
690 async_reads: AtomicU64,
692 enospc: AtomicU64,
695 space_clamped: AtomicU64,
698}
699
700#[derive(Clone, Debug, Default, PartialEq, Eq)]
705pub struct DiskStats {
706 pub blocks: usize,
707 pub bytes: u64,
709 pub disk_bytes: u64,
711 pub queue_depth: usize,
712 pub writes: u64,
713 pub queued_writes: u64,
714 pub inline_writes: u64,
715 pub write_failures: u64,
716 pub write_skipped: u64,
717 pub write_raced_eviction: u64,
718 pub write_nanos: u64,
719 pub hits: u64,
720 pub buffer_hits: u64,
721 pub misses: u64,
722 pub corrupt: u64,
723 pub incompatible: u64,
724 pub read_nanos: u64,
725 pub evictions: u64,
726 pub evicted_bytes: u64,
727 pub prefetch_issued: u64,
728 pub prefetch_dropped: u64,
729 pub prefetch_hits: u64,
730 pub prefetch_waits: u64,
731 pub async_reads: u64,
732 pub staged_blocks: usize,
733 pub enospc: u64,
734 pub space_clamped: u64,
735 pub effective_capacity: u64,
738}
739
740struct Entry {
741 bytes: u64,
742 last_used: u64,
743 published: bool,
747 generation: u64,
752}
753
754struct Index {
755 entries: HashMap<BlockHash, Entry>,
756 bytes: u64,
758 disk_bytes: u64,
764 clock: u64,
765}
766
767impl Index {
768 fn touch(&mut self) -> u64 {
769 self.clock += 1;
770 self.clock
771 }
772
773 fn insert_entry(&mut self, hash: BlockHash, entry: Entry) {
774 let bytes = entry.bytes;
775 if let Some(previous) = self.entries.insert(hash, entry) {
776 self.uncharge(&previous);
777 }
778 self.bytes += bytes;
779 }
780
781 fn remove_entry(&mut self, hash: &BlockHash) -> Option<Entry> {
782 let entry = self.entries.remove(hash)?;
783 self.uncharge(&entry);
784 Some(entry)
785 }
786
787 fn uncharge(&mut self, entry: &Entry) {
788 self.bytes -= entry.bytes;
789 if entry.published {
790 self.disk_bytes -= entry.bytes;
791 }
792 }
793}
794
795enum Source {
797 Disk(PathBuf),
798 Buffer(Arc<KvBlock>),
799}
800
801struct Buffered {
803 generation: u64,
804 block: Arc<KvBlock>,
805}
806
807#[derive(Clone, Copy)]
808struct WriteJob {
809 hash: BlockHash,
810 generation: u64,
811}
812
813struct QueueState {
814 jobs: VecDeque<WriteJob>,
815 running: usize,
816 shutdown: bool,
817}
818
819struct WriteQueue {
826 state: Mutex<QueueState>,
827 ready: Condvar,
828 idle: Condvar,
829 capacity: usize,
830}
831
832impl WriteQueue {
833 fn new(capacity: usize) -> Self {
834 WriteQueue {
835 state: Mutex::new(QueueState {
836 jobs: VecDeque::new(),
837 running: 0,
838 shutdown: false,
839 }),
840 ready: Condvar::new(),
841 idle: Condvar::new(),
842 capacity: capacity.max(1),
843 }
844 }
845
846 fn lock(&self) -> std::sync::MutexGuard<'_, QueueState> {
847 self.state.lock().expect("kv disk write queue poisoned")
848 }
849
850 fn try_push(&self, job: WriteJob) -> bool {
852 let mut state = self.lock();
853 if state.shutdown || state.jobs.len() >= self.capacity {
854 return false;
855 }
856 state.jobs.push_back(job);
857 self.ready.notify_one();
858 true
859 }
860
861 fn pop_blocking(&self) -> Option<WriteJob> {
862 let mut state = self.lock();
863 loop {
864 if let Some(job) = state.jobs.pop_front() {
865 state.running += 1;
866 return Some(job);
867 }
868 if state.shutdown {
869 return None;
870 }
871 state = self
872 .ready
873 .wait(state)
874 .expect("kv disk write queue poisoned");
875 }
876 }
877
878 fn pop_now(&self) -> Option<WriteJob> {
879 let mut state = self.lock();
880 let job = state.jobs.pop_front()?;
881 state.running += 1;
882 Some(job)
883 }
884
885 fn finish(&self) {
886 let mut state = self.lock();
887 state.running -= 1;
888 self.idle.notify_all();
889 }
890
891 fn shutdown(&self) {
892 let mut state = self.lock();
893 state.shutdown = true;
894 self.ready.notify_all();
895 }
896
897 fn depth(&self) -> usize {
898 self.lock().jobs.len()
899 }
900}
901
902pub type ReadOutcome = Result<Option<Arc<KvBlock>>, StoreError>;
906
907struct ReadSlot {
911 expected: CacheSignature,
916 outcome: Mutex<Option<ReadOutcome>>,
917 done: Condvar,
918}
919
920impl ReadSlot {
921 fn pending(expected: CacheSignature) -> Arc<Self> {
922 Arc::new(ReadSlot {
923 expected,
924 outcome: Mutex::new(None),
925 done: Condvar::new(),
926 })
927 }
928
929 fn ready(expected: CacheSignature, outcome: ReadOutcome) -> Arc<Self> {
930 Arc::new(ReadSlot {
931 expected,
932 outcome: Mutex::new(Some(outcome)),
933 done: Condvar::new(),
934 })
935 }
936
937 fn is_ready(&self) -> bool {
938 self.outcome
939 .lock()
940 .expect("kv disk read slot poisoned")
941 .is_some()
942 }
943
944 fn fulfil(&self, outcome: ReadOutcome) {
945 let mut slot = self.outcome.lock().expect("kv disk read slot poisoned");
946 *slot = Some(outcome);
947 self.done.notify_all();
948 }
949
950 fn wait(&self) -> ReadOutcome {
951 let mut slot = self.outcome.lock().expect("kv disk read slot poisoned");
952 loop {
953 if let Some(outcome) = slot.as_ref() {
954 return outcome.clone();
955 }
956 slot = self.done.wait(slot).expect("kv disk read slot poisoned");
957 }
958 }
959}
960
961struct ReadJob {
962 hash: BlockHash,
963 path: PathBuf,
964 slot: Arc<ReadSlot>,
965}
966
967struct ReadQueueState {
968 jobs: VecDeque<ReadJob>,
969 shutdown: bool,
970}
971
972struct ReadQueue {
977 state: Mutex<ReadQueueState>,
978 ready: Condvar,
979 capacity: usize,
980}
981
982impl ReadQueue {
983 fn new(capacity: usize) -> Self {
984 ReadQueue {
985 state: Mutex::new(ReadQueueState {
986 jobs: VecDeque::new(),
987 shutdown: false,
988 }),
989 ready: Condvar::new(),
990 capacity: capacity.max(1),
991 }
992 }
993
994 fn lock(&self) -> std::sync::MutexGuard<'_, ReadQueueState> {
995 self.state.lock().expect("kv disk read queue poisoned")
996 }
997
998 fn try_push(&self, job: ReadJob, demand: bool) -> bool {
1001 let mut state = self.lock();
1002 if state.shutdown || state.jobs.len() >= self.capacity {
1003 return false;
1004 }
1005 if demand {
1006 state.jobs.push_front(job);
1007 } else {
1008 state.jobs.push_back(job);
1009 }
1010 self.ready.notify_one();
1011 true
1012 }
1013
1014 fn pop_blocking(&self) -> Option<ReadJob> {
1015 let mut state = self.lock();
1016 loop {
1017 if let Some(job) = state.jobs.pop_front() {
1018 return Some(job);
1019 }
1020 if state.shutdown {
1021 return None;
1022 }
1023 state = self.ready.wait(state).expect("kv disk read queue poisoned");
1024 }
1025 }
1026
1027 fn shutdown(&self) {
1028 let mut state = self.lock();
1029 state.shutdown = true;
1030 self.ready.notify_all();
1031 }
1032}
1033
1034pub struct ReadHandle {
1041 shared: Arc<Shared>,
1042 hash: BlockHash,
1043 slot: Arc<ReadSlot>,
1044 staged: bool,
1047}
1048
1049impl ReadHandle {
1050 pub fn is_ready(&self) -> bool {
1053 self.slot.is_ready()
1054 }
1055
1056 pub fn try_claim(&self) -> Option<ReadOutcome> {
1059 if !self.slot.is_ready() {
1060 return None;
1061 }
1062 Some(self.claim())
1063 }
1064
1065 pub fn wait(self) -> ReadOutcome {
1067 self.claim()
1068 }
1069
1070 fn claim(&self) -> ReadOutcome {
1071 let outcome = self.slot.wait();
1072 if self.staged {
1073 self.shared.unstage(&self.hash, &self.slot);
1074 }
1075 outcome
1076 }
1077}
1078
1079impl std::fmt::Debug for ReadHandle {
1080 fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
1081 f.debug_struct("ReadHandle")
1082 .field("hash", &self.hash)
1083 .field("ready", &self.is_ready())
1084 .finish()
1085 }
1086}
1087
1088#[cfg(test)]
1089type Hook = Arc<dyn Fn(&BlockHash) + Send + Sync>;
1090
1091#[cfg(test)]
1096#[derive(Clone, Copy, Debug, Default, PartialEq, Eq)]
1097enum WriteOrder {
1098 #[default]
1099 BufferThenIndex,
1100 IndexBeforeBuffer,
1103 DropBufferBeforeMarking,
1106}
1107
1108#[cfg(test)]
1109#[derive(Default)]
1110struct Hooks {
1111 order: Mutex<WriteOrder>,
1112 after_rename: Mutex<Option<Hook>>,
1116 in_put_window: Mutex<Option<Hook>>,
1118 in_publish_window: Mutex<Option<Hook>>,
1120 fail_with_enospc: std::sync::atomic::AtomicBool,
1124}
1125
1126#[cfg(test)]
1127impl Hooks {
1128 fn fire(slot: &Mutex<Option<Hook>>, hash: &BlockHash) {
1129 let hook = slot.lock().expect("kv disk hook poisoned").clone();
1132 if let Some(hook) = hook {
1133 hook(hash);
1134 }
1135 }
1136}
1137
1138struct Shared {
1143 root: PathBuf,
1144 shard_chars: usize,
1145 max_bytes: u64,
1146 index: Mutex<Index>,
1147 buffer: Mutex<HashMap<BlockHash, Buffered>>,
1155 queue: WriteQueue,
1156 reads: ReadQueue,
1157 staging: Mutex<HashMap<BlockHash, Arc<ReadSlot>>>,
1161 prefetch_capacity: usize,
1162 has_readers: bool,
1163 reserve_bytes: u64,
1164 free_space_ttl: std::time::Duration,
1165 free_space_probe: FreeSpaceProbe,
1166 free_space: Mutex<FreeSpace>,
1169 stats: Stats,
1170 seq: AtomicU64,
1171 generation: AtomicU64,
1172 #[cfg(test)]
1173 hooks: Hooks,
1174}
1175
1176pub struct DiskKvStore {
1183 shared: Arc<Shared>,
1184 writers: Vec<std::thread::JoinHandle<()>>,
1185 readers: Vec<std::thread::JoinHandle<()>>,
1186}
1187
1188impl Drop for DiskKvStore {
1189 fn drop(&mut self) {
1195 self.shared.queue.shutdown();
1196 self.shared.reads.shutdown();
1197 for writer in self.writers.drain(..) {
1198 let _ = writer.join();
1199 }
1200 for reader in self.readers.drain(..) {
1201 let _ = reader.join();
1202 }
1203 }
1204}
1205
1206impl DiskKvStore {
1207 pub fn open(config: DiskConfig) -> Result<Self, StoreError> {
1213 let root = config.root.clone();
1214 fs::create_dir_all(&root).map_err(|e| io_err("create", &root, e))?;
1215 let tmp = root.join(TMP_DIR);
1216 fs::create_dir_all(&tmp).map_err(|e| io_err("create", &tmp, e))?;
1217 let shared = Arc::new(Shared {
1218 root,
1219 shard_chars: config.shard_chars.clamp(1, 8),
1220 max_bytes: config.max_bytes,
1221 index: Mutex::new(Index {
1222 entries: HashMap::new(),
1223 bytes: 0,
1224 disk_bytes: 0,
1225 clock: 0,
1226 }),
1227 buffer: Mutex::new(HashMap::new()),
1228 queue: WriteQueue::new(config.queue_capacity),
1229 reads: ReadQueue::new(config.queue_capacity),
1230 staging: Mutex::new(HashMap::new()),
1231 prefetch_capacity: config.prefetch_capacity,
1232 has_readers: config.reader_threads > 0,
1233 reserve_bytes: config.reserve_bytes,
1234 free_space_ttl: config.free_space_ttl,
1235 free_space_probe: Arc::clone(&config.free_space_probe),
1236 free_space: Mutex::new(FreeSpace {
1237 checked_at: None,
1238 bytes: None,
1239 }),
1240 stats: Stats::default(),
1241 seq: AtomicU64::new(0),
1242 generation: AtomicU64::new(0),
1243 #[cfg(test)]
1244 hooks: Hooks::default(),
1245 });
1246 let mut writers = Vec::with_capacity(config.writer_threads);
1247 for n in 0..config.writer_threads {
1248 let shared = Arc::clone(&shared);
1249 let handle = std::thread::Builder::new()
1250 .name(format!("ferrox-kv-write-{n}"))
1251 .spawn(move || {
1252 while let Some(job) = shared.queue.pop_blocking() {
1253 shared.run_job(job);
1254 shared.queue.finish();
1255 }
1256 })
1257 .map_err(|e| io_err("spawn writer for", &config.root, e))?;
1258 writers.push(handle);
1259 }
1260 let mut readers = Vec::with_capacity(config.reader_threads);
1261 for n in 0..config.reader_threads {
1262 let shared = Arc::clone(&shared);
1263 let handle = std::thread::Builder::new()
1264 .name(format!("ferrox-kv-read-{n}"))
1265 .spawn(move || {
1266 while let Some(job) = shared.reads.pop_blocking() {
1267 shared.stats.async_reads.fetch_add(1, Ordering::Relaxed);
1268 let outcome = shared.read_timed(&job.path, &job.hash, &job.slot.expected);
1269 job.slot.fulfil(outcome);
1270 }
1271 })
1272 .map_err(|e| io_err("spawn reader for", &config.root, e))?;
1273 readers.push(handle);
1274 }
1275 Ok(DiskKvStore {
1276 shared,
1277 writers,
1278 readers,
1279 })
1280 }
1281
1282 pub fn root(&self) -> &Path {
1283 &self.shared.root
1284 }
1285
1286 pub fn put(&self, hash: BlockHash, block: KvBlock) -> Result<(), StoreError> {
1293 self.shared.put(hash, block, false)
1294 }
1295
1296 pub fn put_blocking(&self, hash: BlockHash, block: KvBlock) -> Result<(), StoreError> {
1298 self.shared.put(hash, block, true)
1299 }
1300
1301 pub fn flush(&self) {
1305 loop {
1306 if let Some(job) = self.shared.queue.pop_now() {
1307 self.shared.run_job(job);
1308 self.shared.queue.finish();
1309 continue;
1310 }
1311 let state = self.shared.queue.lock();
1312 if state.jobs.is_empty() && state.running == 0 {
1313 return;
1314 }
1315 let _ = self
1318 .shared
1319 .queue
1320 .idle
1321 .wait_timeout(state, std::time::Duration::from_millis(1));
1322 }
1323 }
1324
1325 pub fn get(&self, hash: &BlockHash, expected: &CacheSignature) -> ReadOutcome {
1332 self.shared.read_async(hash, expected, true).wait()
1333 }
1334
1335 pub fn read_async(&self, hash: &BlockHash, expected: &CacheSignature) -> ReadHandle {
1338 self.shared.read_async(hash, expected, true)
1339 }
1340
1341 pub fn prefetch(&self, hashes: &[BlockHash], expected: &CacheSignature) {
1351 self.shared.prefetch(hashes, expected);
1352 }
1353
1354 pub fn clear_prefetch(&self) {
1358 self.shared
1359 .staging
1360 .lock()
1361 .expect("kv disk staging poisoned")
1362 .clear();
1363 }
1364
1365 pub fn reindex(&self) -> Result<usize, StoreError> {
1367 self.shared.reindex()
1368 }
1369
1370 pub fn remove(&self, hash: &BlockHash) {
1372 self.shared.quarantine(hash);
1373 }
1374
1375 pub fn contains(&self, hash: &BlockHash) -> bool {
1378 let index = self.shared.index.lock().expect("kv disk index poisoned");
1379 index.entries.contains_key(hash)
1380 }
1381
1382 pub fn capacity(&self) -> u64 {
1384 self.shared.max_bytes
1385 }
1386
1387 pub fn effective_capacity(&self) -> u64 {
1390 let used = {
1391 let index = self.shared.index.lock().expect("kv disk index poisoned");
1392 index.bytes
1393 };
1394 self.shared.effective_capacity(used)
1395 }
1396
1397 pub fn block_path(&self, hash: &BlockHash) -> PathBuf {
1400 self.shared.block_path(hash)
1401 }
1402
1403 pub fn stats(&self) -> DiskStats {
1404 self.shared.stats()
1405 }
1406}
1407
1408impl Shared {
1409 fn next_generation(&self) -> u64 {
1410 self.generation.fetch_add(1, Ordering::SeqCst) + 1
1411 }
1412
1413 fn put(&self, hash: BlockHash, block: KvBlock, inline: bool) -> Result<(), StoreError> {
1426 let bytes = encoded_len(block.signature());
1427 let block = Arc::new(block);
1428 let generation = self.next_generation();
1429
1430 #[cfg(test)]
1431 let index_first = *self.hooks.order.lock().expect("kv disk hook poisoned")
1432 == WriteOrder::IndexBeforeBuffer;
1433 #[cfg(not(test))]
1434 let index_first = false;
1435
1436 if index_first {
1437 self.reserve(hash, bytes, generation);
1438 #[cfg(test)]
1439 Hooks::fire(&self.hooks.in_put_window, &hash);
1440 self.buffer_block(hash, generation, Arc::clone(&block));
1441 } else {
1442 self.buffer_block(hash, generation, Arc::clone(&block));
1443 #[cfg(test)]
1444 Hooks::fire(&self.hooks.in_put_window, &hash);
1445 self.reserve(hash, bytes, generation);
1446 }
1447
1448 let job = WriteJob { hash, generation };
1449 if !inline && self.queue.try_push(job) {
1450 self.stats.queued_writes.fetch_add(1, Ordering::Relaxed);
1451 return Ok(());
1452 }
1453 if !inline {
1454 self.stats.inline_writes.fetch_add(1, Ordering::Relaxed);
1455 }
1456 self.run_write(job, block)
1457 }
1458
1459 fn buffer_block(&self, hash: BlockHash, generation: u64, block: Arc<KvBlock>) {
1460 self.buffer
1461 .lock()
1462 .expect("kv disk buffer poisoned")
1463 .insert(hash, Buffered { generation, block });
1464 }
1465
1466 fn release_buffer(&self, hash: &BlockHash, generation: u64) {
1470 let mut buffer = self.buffer.lock().expect("kv disk buffer poisoned");
1471 if buffer.get(hash).is_some_and(|b| b.generation == generation) {
1472 buffer.remove(hash);
1473 }
1474 }
1475
1476 fn buffered(&self, hash: &BlockHash, generation: u64) -> Option<Arc<KvBlock>> {
1477 let buffer = self.buffer.lock().expect("kv disk buffer poisoned");
1478 buffer
1479 .get(hash)
1480 .filter(|b| b.generation == generation)
1481 .map(|b| Arc::clone(&b.block))
1482 }
1483
1484 fn run_job(&self, job: WriteJob) {
1488 match self.buffered(&job.hash, job.generation) {
1489 Some(block) => {
1490 let _ = self.run_write(job, block);
1491 }
1492 None => {
1493 self.stats.write_skipped.fetch_add(1, Ordering::Relaxed);
1494 }
1495 }
1496 }
1497
1498 fn run_write(&self, job: WriteJob, block: Arc<KvBlock>) -> Result<(), StoreError> {
1499 let started = Instant::now();
1500 let result = self.write_and_publish(&job.hash, &block, job.generation);
1501 self.stats
1502 .write_nanos
1503 .fetch_add(started.elapsed().as_nanos() as u64, Ordering::Relaxed);
1504 match result {
1505 Ok(()) => {
1506 self.stats.writes.fetch_add(1, Ordering::Relaxed);
1507 Ok(())
1508 }
1509 Err(err) => {
1510 self.stats.write_failures.fetch_add(1, Ordering::Relaxed);
1511 self.abandon(&job.hash, job.generation);
1512 Err(err)
1513 }
1514 }
1515 }
1516
1517 fn reserve(&self, hash: BlockHash, bytes: u64, generation: u64) {
1520 let mut index = self.index.lock().expect("kv disk index poisoned");
1521 let last_used = index.touch();
1522 index.insert_entry(
1523 hash,
1524 Entry {
1525 bytes,
1526 last_used,
1527 published: false,
1528 generation,
1529 },
1530 );
1531 let victims = self.collect_victims(&mut index, Some(&hash));
1532 drop(index);
1533 self.discard(victims);
1534 }
1535
1536 fn abandon(&self, hash: &BlockHash, generation: u64) {
1541 {
1542 let mut index = self.index.lock().expect("kv disk index poisoned");
1543 if index
1544 .entries
1545 .get(hash)
1546 .is_some_and(|e| e.generation == generation)
1547 {
1548 index.remove_entry(hash);
1549 }
1550 }
1551 self.release_buffer(hash, generation);
1552 }
1553
1554 fn write_and_publish(
1555 &self,
1556 hash: &BlockHash,
1557 block: &KvBlock,
1558 generation: u64,
1559 ) -> Result<(), StoreError> {
1560 let bytes = encode_block(hash, block);
1561 let final_path = self.block_path(hash);
1562 let shard = final_path.parent().expect("block path has a parent");
1563 fs::create_dir_all(shard).map_err(|e| io_err("create", shard, e))?;
1564 let tmp_path = self.tmp_path(hash);
1565 {
1566 let mut file =
1567 fs::File::create(&tmp_path).map_err(|e| io_err("create", &tmp_path, e))?;
1568 #[cfg(test)]
1569 let written = if self.hooks.fail_with_enospc.load(Ordering::Relaxed) {
1570 Err(io::Error::from(io::ErrorKind::StorageFull))
1571 } else {
1572 file.write_all(&bytes)
1573 };
1574 #[cfg(not(test))]
1575 let written = file.write_all(&bytes);
1576 if let Err(e) = written {
1577 let _ = fs::remove_file(&tmp_path);
1578 self.note_if_enospc(&e);
1579 return Err(io_err("write", &tmp_path, e));
1580 }
1581 if let Err(e) = file.sync_all() {
1585 let _ = fs::remove_file(&tmp_path);
1586 self.note_if_enospc(&e);
1587 return Err(io_err("sync", &tmp_path, e));
1588 }
1589 }
1590 fs::rename(&tmp_path, &final_path).map_err(|e| {
1591 let _ = fs::remove_file(&tmp_path);
1592 io_err("publish", &final_path, e)
1593 })?;
1594
1595 #[cfg(test)]
1596 Hooks::fire(&self.hooks.after_rename, hash);
1597
1598 #[cfg(test)]
1599 let drop_buffer_first = *self.hooks.order.lock().expect("kv disk hook poisoned")
1600 == WriteOrder::DropBufferBeforeMarking;
1601 #[cfg(not(test))]
1602 let drop_buffer_first = false;
1603
1604 let survived = if drop_buffer_first {
1610 self.release_buffer(hash, generation);
1611 #[cfg(test)]
1612 Hooks::fire(&self.hooks.in_publish_window, hash);
1613 self.mark_published(hash, generation)
1614 } else {
1615 let survived = self.mark_published(hash, generation);
1616 #[cfg(test)]
1617 Hooks::fire(&self.hooks.in_publish_window, hash);
1618 self.release_buffer(hash, generation);
1619 survived
1620 };
1621
1622 if !survived {
1623 let _ = fs::remove_file(&final_path);
1627 self.stats
1628 .write_raced_eviction
1629 .fetch_add(1, Ordering::Relaxed);
1630 }
1631 Ok(())
1632 }
1633
1634 fn mark_published(&self, hash: &BlockHash, generation: u64) -> bool {
1635 let mut index = self.index.lock().expect("kv disk index poisoned");
1636 match index.entries.get_mut(hash) {
1637 Some(entry) if entry.generation == generation => {
1638 if !entry.published {
1639 entry.published = true;
1640 let bytes = entry.bytes;
1641 index.disk_bytes += bytes;
1642 }
1643 true
1644 }
1645 _ => false,
1646 }
1647 }
1648
1649 fn source(&self, hash: &BlockHash) -> Result<Option<Source>, StoreError> {
1652 let mut index = self.index.lock().expect("kv disk index poisoned");
1653 let clock = index.clock + 1;
1654 let Some(entry) = index.entries.get_mut(hash) else {
1655 return Ok(None);
1656 };
1657 entry.last_used = clock;
1658 let published = entry.published;
1659 index.clock = clock;
1660 if published {
1661 return Ok(Some(Source::Disk(self.block_path(hash))));
1662 }
1663 let buffered = self
1668 .buffer
1669 .lock()
1670 .expect("kv disk buffer poisoned")
1671 .get(hash)
1672 .map(|b| Arc::clone(&b.block));
1673 match buffered {
1674 Some(block) => Ok(Some(Source::Buffer(block))),
1675 None => Err(StoreError::MissingPayload { hash: *hash }),
1676 }
1677 }
1678
1679 fn read_async(
1694 self: &Arc<Self>,
1695 hash: &BlockHash,
1696 expected: &CacheSignature,
1697 demand: bool,
1698 ) -> ReadHandle {
1699 let staged = {
1703 let staging = self.staging.lock().expect("kv disk staging poisoned");
1704 staging
1705 .get(hash)
1706 .filter(|slot| &slot.expected == expected)
1707 .map(Arc::clone)
1708 };
1709 if let Some(slot) = staged {
1710 if demand {
1711 let counter = if slot.is_ready() {
1712 &self.stats.prefetch_hits
1713 } else {
1714 &self.stats.prefetch_waits
1715 };
1716 counter.fetch_add(1, Ordering::Relaxed);
1717 }
1718 return self.handle(*hash, slot, true);
1719 }
1720
1721 let ready = |outcome: ReadOutcome| ReadHandle {
1722 shared: Arc::clone(self),
1723 hash: *hash,
1724 slot: ReadSlot::ready(expected.clone(), outcome),
1725 staged: false,
1726 };
1727
1728 let path = match self.source(hash) {
1729 Err(err) => return ready(Err(err)),
1730 Ok(None) => {
1731 self.stats.misses.fetch_add(1, Ordering::Relaxed);
1732 return ready(Ok(None));
1733 }
1734 Ok(Some(Source::Buffer(block))) => {
1735 if block.signature() != expected {
1736 self.stats.incompatible.fetch_add(1, Ordering::Relaxed);
1737 return ready(Ok(None));
1738 }
1739 self.stats.buffer_hits.fetch_add(1, Ordering::Relaxed);
1740 return ready(Ok(Some(block)));
1741 }
1742 Ok(Some(Source::Disk(path))) => path,
1743 };
1744
1745 let slot = ReadSlot::pending(expected.clone());
1746 let dispatched = self.has_readers && {
1747 self.staging
1748 .lock()
1749 .expect("kv disk staging poisoned")
1750 .insert(*hash, Arc::clone(&slot));
1751 let job = ReadJob {
1752 hash: *hash,
1753 path: path.clone(),
1754 slot: Arc::clone(&slot),
1755 };
1756 let pushed = self.reads.try_push(job, demand);
1757 if !pushed {
1758 self.unstage(hash, &slot);
1759 }
1760 pushed
1761 };
1762 if dispatched {
1763 return self.handle(*hash, slot, true);
1764 }
1765 slot.fulfil(self.read_timed(&path, hash, expected));
1766 self.handle(*hash, slot, false)
1767 }
1768
1769 fn handle(self: &Arc<Self>, hash: BlockHash, slot: Arc<ReadSlot>, staged: bool) -> ReadHandle {
1770 ReadHandle {
1771 shared: Arc::clone(self),
1772 hash,
1773 slot,
1774 staged,
1775 }
1776 }
1777
1778 fn unstage(&self, hash: &BlockHash, slot: &Arc<ReadSlot>) {
1782 let mut staging = self.staging.lock().expect("kv disk staging poisoned");
1783 if staging.get(hash).is_some_and(|s| Arc::ptr_eq(s, slot)) {
1784 staging.remove(hash);
1785 }
1786 }
1787
1788 fn prefetch(self: &Arc<Self>, hashes: &[BlockHash], expected: &CacheSignature) {
1789 if !self.has_readers {
1790 self.stats
1791 .prefetch_dropped
1792 .fetch_add(hashes.len() as u64, Ordering::Relaxed);
1793 return;
1794 }
1795 for hash in hashes {
1796 let room = {
1797 let staging = self.staging.lock().expect("kv disk staging poisoned");
1798 !staging.contains_key(hash) && staging.len() < self.prefetch_capacity
1799 };
1800 if !room {
1801 self.stats.prefetch_dropped.fetch_add(1, Ordering::Relaxed);
1802 continue;
1803 }
1804 let handle = self.read_async(hash, expected, false);
1805 if handle.staged {
1806 self.stats.prefetch_issued.fetch_add(1, Ordering::Relaxed);
1807 } else {
1808 self.stats.prefetch_dropped.fetch_add(1, Ordering::Relaxed);
1814 }
1815 drop(handle);
1819 }
1820 }
1821
1822 fn effective_capacity(&self, used: u64) -> u64 {
1838 let Some(free) = self.free_bytes() else {
1839 return self.max_bytes;
1840 };
1841 let headroom = free as i128 - self.reserve_bytes as i128;
1847 let allowed = (used as i128 + headroom).max(0) as u64;
1848 self.max_bytes.min(allowed)
1849 }
1850
1851 fn free_bytes(&self) -> Option<u64> {
1855 let mut cache = self.free_space.lock().expect("kv disk free space poisoned");
1856 if let Some(checked_at) = cache.checked_at {
1857 if checked_at.elapsed() < self.free_space_ttl {
1858 return cache.bytes;
1859 }
1860 }
1861 let bytes = (self.free_space_probe)(&self.root);
1862 cache.checked_at = Some(Instant::now());
1863 cache.bytes = bytes;
1864 bytes
1865 }
1866
1867 fn note_enospc(&self) {
1872 self.stats.enospc.fetch_add(1, Ordering::Relaxed);
1873 {
1874 let mut cache = self.free_space.lock().expect("kv disk free space poisoned");
1875 cache.checked_at = None;
1876 cache.bytes = None;
1877 }
1878 let victims = {
1879 let mut index = self.index.lock().expect("kv disk index poisoned");
1880 self.collect_victims(&mut index, None)
1881 };
1882 self.discard(victims);
1883 }
1884
1885 fn note_if_enospc(&self, err: &io::Error) {
1887 if err.kind() == io::ErrorKind::StorageFull {
1888 self.note_enospc();
1889 }
1890 }
1891
1892 fn read_timed(&self, path: &Path, hash: &BlockHash, expected: &CacheSignature) -> ReadOutcome {
1893 let started = Instant::now();
1894 let outcome = self.read_verified(path, hash, expected);
1895 self.stats
1896 .read_nanos
1897 .fetch_add(started.elapsed().as_nanos() as u64, Ordering::Relaxed);
1898 outcome
1899 }
1900
1901 fn read_verified(
1902 &self,
1903 path: &Path,
1904 hash: &BlockHash,
1905 expected: &CacheSignature,
1906 ) -> Result<Option<Arc<KvBlock>>, StoreError> {
1907 let bytes = match fs::read(path) {
1908 Ok(bytes) => bytes,
1909 Err(e) if e.kind() == io::ErrorKind::NotFound => {
1910 self.stats.misses.fetch_add(1, Ordering::Relaxed);
1913 self.drop_entry(hash);
1914 return Ok(None);
1915 }
1916 Err(e) => return Err(io_err("read", path, e)),
1917 };
1918 let decoded = match decode_block(&bytes) {
1919 Ok(decoded) => decoded,
1920 Err(_) => {
1921 self.stats.corrupt.fetch_add(1, Ordering::Relaxed);
1922 self.quarantine(hash);
1923 return Ok(None);
1924 }
1925 };
1926 if &decoded.hash != hash {
1927 self.stats.corrupt.fetch_add(1, Ordering::Relaxed);
1930 self.quarantine(hash);
1931 return Ok(None);
1932 }
1933 match decoded.block.verify(expected) {
1934 Ok(block) => {
1935 self.stats.hits.fetch_add(1, Ordering::Relaxed);
1936 Ok(Some(Arc::new(block)))
1937 }
1938 Err(_) => {
1939 self.stats.incompatible.fetch_add(1, Ordering::Relaxed);
1940 Ok(None)
1941 }
1942 }
1943 }
1944
1945 fn quarantine(&self, hash: &BlockHash) {
1946 self.drop_entry(hash);
1947 self.buffer
1948 .lock()
1949 .expect("kv disk buffer poisoned")
1950 .remove(hash);
1951 let _ = fs::remove_file(self.block_path(hash));
1952 }
1953
1954 fn drop_entry(&self, hash: &BlockHash) {
1955 let mut index = self.index.lock().expect("kv disk index poisoned");
1956 index.remove_entry(hash);
1957 }
1958
1959 fn reindex(&self) -> Result<usize, StoreError> {
1960 let tmp = self.root.join(TMP_DIR);
1961 if let Ok(entries) = fs::read_dir(&tmp) {
1962 for entry in entries.flatten() {
1963 let _ = fs::remove_file(entry.path());
1964 }
1965 }
1966 let mut found = Vec::new();
1967 let shards = fs::read_dir(&self.root).map_err(|e| io_err("read", &self.root, e))?;
1968 for shard in shards.flatten() {
1969 if !shard.file_type().map(|t| t.is_dir()).unwrap_or(false) {
1970 continue;
1971 }
1972 if shard.file_name() == TMP_DIR {
1973 continue;
1974 }
1975 let Ok(files) = fs::read_dir(shard.path()) else {
1976 continue;
1977 };
1978 for file in files.flatten() {
1979 let path = file.path();
1980 if path.extension().and_then(|e| e.to_str()) != Some(BLOCK_FILE_EXT) {
1981 continue;
1982 }
1983 let Some(hash) = path
1984 .file_stem()
1985 .and_then(|s| s.to_str())
1986 .and_then(parse_hex_hash)
1987 else {
1988 continue;
1989 };
1990 let Ok(meta) = file.metadata() else { continue };
1991 found.push((hash, meta.len()));
1992 }
1993 }
1994 let mut adopted = 0;
1995 let mut index = self.index.lock().expect("kv disk index poisoned");
1996 for (hash, bytes) in found {
1997 if index.entries.contains_key(&hash) {
1998 continue;
1999 }
2000 let last_used = index.touch();
2001 index.insert_entry(
2002 hash,
2003 Entry {
2004 bytes,
2005 last_used,
2006 published: true,
2007 generation: self.next_generation(),
2008 },
2009 );
2010 index.disk_bytes += bytes;
2011 adopted += 1;
2012 }
2013 let victims = self.collect_victims(&mut index, None);
2014 drop(index);
2015 self.discard(victims);
2016 Ok(adopted)
2017 }
2018
2019 fn collect_victims(&self, index: &mut Index, protect: Option<&BlockHash>) -> Vec<Victim> {
2025 let budget = self.effective_capacity(index.disk_bytes);
2030 if budget < self.max_bytes {
2031 self.stats.space_clamped.fetch_add(1, Ordering::Relaxed);
2032 }
2033 if index.bytes <= budget {
2034 return Vec::new();
2035 }
2036 let mut candidates: Vec<(u64, BlockHash)> = index
2037 .entries
2038 .iter()
2039 .filter(|(hash, _)| Some(*hash) != protect)
2040 .map(|(hash, entry)| (entry.last_used, *hash))
2041 .collect();
2042 candidates.sort_unstable();
2043 let mut victims = Vec::new();
2044 for (_, hash) in candidates {
2045 if index.bytes <= budget {
2046 break;
2047 }
2048 if let Some(entry) = index.remove_entry(&hash) {
2049 self.stats.evictions.fetch_add(1, Ordering::Relaxed);
2050 self.stats
2051 .evicted_bytes
2052 .fetch_add(entry.bytes, Ordering::Relaxed);
2053 victims.push(Victim {
2054 hash,
2055 published: entry.published,
2056 });
2057 }
2058 }
2059 victims
2060 }
2061
2062 fn discard(&self, victims: Vec<Victim>) {
2065 if victims.is_empty() {
2066 return;
2067 }
2068 {
2069 let mut buffer = self.buffer.lock().expect("kv disk buffer poisoned");
2070 for victim in &victims {
2071 buffer.remove(&victim.hash);
2072 }
2073 }
2074 for victim in victims {
2075 if victim.published {
2076 let _ = fs::remove_file(self.block_path(&victim.hash));
2077 }
2078 }
2079 }
2080
2081 fn stats(&self) -> DiskStats {
2082 let index = self.index.lock().expect("kv disk index poisoned");
2083 let stats = &self.stats;
2084 DiskStats {
2085 blocks: index.entries.len(),
2086 bytes: index.bytes,
2087 queue_depth: self.queue.depth(),
2088 writes: stats.writes.load(Ordering::Relaxed),
2089 queued_writes: stats.queued_writes.load(Ordering::Relaxed),
2090 inline_writes: stats.inline_writes.load(Ordering::Relaxed),
2091 write_failures: stats.write_failures.load(Ordering::Relaxed),
2092 write_skipped: stats.write_skipped.load(Ordering::Relaxed),
2093 write_raced_eviction: stats.write_raced_eviction.load(Ordering::Relaxed),
2094 write_nanos: stats.write_nanos.load(Ordering::Relaxed),
2095 hits: stats.hits.load(Ordering::Relaxed),
2096 buffer_hits: stats.buffer_hits.load(Ordering::Relaxed),
2097 misses: stats.misses.load(Ordering::Relaxed),
2098 corrupt: stats.corrupt.load(Ordering::Relaxed),
2099 incompatible: stats.incompatible.load(Ordering::Relaxed),
2100 read_nanos: stats.read_nanos.load(Ordering::Relaxed),
2101 evictions: stats.evictions.load(Ordering::Relaxed),
2102 evicted_bytes: stats.evicted_bytes.load(Ordering::Relaxed),
2103 prefetch_issued: stats.prefetch_issued.load(Ordering::Relaxed),
2104 prefetch_dropped: stats.prefetch_dropped.load(Ordering::Relaxed),
2105 prefetch_hits: stats.prefetch_hits.load(Ordering::Relaxed),
2106 prefetch_waits: stats.prefetch_waits.load(Ordering::Relaxed),
2107 async_reads: stats.async_reads.load(Ordering::Relaxed),
2108 staged_blocks: self.staging.lock().expect("kv disk staging poisoned").len(),
2109 enospc: stats.enospc.load(Ordering::Relaxed),
2110 space_clamped: stats.space_clamped.load(Ordering::Relaxed),
2111 effective_capacity: self.effective_capacity(index.disk_bytes),
2112 disk_bytes: index.disk_bytes,
2113 }
2114 }
2115
2116 fn block_path(&self, hash: &BlockHash) -> PathBuf {
2117 self.root
2118 .join(hash.shard_prefix(self.shard_chars))
2119 .join(format!("{}.{BLOCK_FILE_EXT}", hash.to_hex()))
2120 }
2121
2122 fn tmp_path(&self, hash: &BlockHash) -> PathBuf {
2123 let n = self.seq.fetch_add(1, Ordering::Relaxed);
2124 self.root.join(TMP_DIR).join(format!(
2125 "{}.{}.{n}.tmp",
2126 hash.shard_prefix(16),
2127 std::process::id()
2128 ))
2129 }
2130}
2131
2132struct Victim {
2135 hash: BlockHash,
2136 published: bool,
2137}
2138
2139fn parse_hex_hash(text: &str) -> Option<BlockHash> {
2140 if text.len() != 64 {
2141 return None;
2142 }
2143 let mut out = [0u8; 32];
2144 for (i, byte) in out.iter_mut().enumerate() {
2145 let hi = text.as_bytes()[i * 2] as char;
2146 let lo = text.as_bytes()[i * 2 + 1] as char;
2147 *byte = ((hi.to_digit(16)? << 4) | lo.to_digit(16)?) as u8;
2148 }
2149 Some(BlockHash::from_bytes(out))
2150}
2151
2152#[cfg(test)]
2153mod tests {
2154 use super::*;
2155 use crate::kv_block::BlockHasher;
2156 use std::sync::atomic::AtomicUsize;
2157
2158 struct TempDir(PathBuf);
2163
2164 impl TempDir {
2165 fn new(tag: &str) -> Self {
2166 static N: AtomicU64 = AtomicU64::new(0);
2167 let path = std::env::temp_dir().join(format!(
2168 "ferrox-kvdisk-{tag}-{}-{}",
2169 std::process::id(),
2170 N.fetch_add(1, Ordering::Relaxed)
2171 ));
2172 let _ = fs::remove_dir_all(&path);
2173 fs::create_dir_all(&path).expect("temp dir");
2174 TempDir(path)
2175 }
2176
2177 fn path(&self) -> &Path {
2178 &self.0
2179 }
2180 }
2181
2182 impl Drop for TempDir {
2183 fn drop(&mut self) {
2184 let _ = fs::remove_dir_all(&self.0);
2185 }
2186 }
2187
2188 fn layer(n_kv_heads: usize, head_dim: usize, tokens: usize, fill: f32) -> KvCache {
2189 let mut cache = KvCache::new(n_kv_heads, head_dim);
2190 for t in 0..tokens {
2191 let k = vec![fill + t as f32; n_kv_heads * head_dim];
2192 let v = vec![fill - t as f32; n_kv_heads * head_dim];
2193 cache.push(&k, &v).expect("unpooled push cannot fail");
2194 }
2195 cache
2196 }
2197
2198 fn flat(tokens: usize) -> BlockLayout {
2201 BlockLayout::full_attention(tokens).expect("positive block size")
2202 }
2203
2204 fn block(model: &str, n_layers: usize, tokens: usize, fill: f32) -> KvBlock {
2205 block_with_layout(model, n_layers, tokens, fill, flat(tokens))
2206 }
2207
2208 fn block_with_layout(
2209 model: &str,
2210 n_layers: usize,
2211 tokens: usize,
2212 fill: f32,
2213 layout: BlockLayout,
2214 ) -> KvBlock {
2215 let layers = (0..n_layers)
2216 .map(|l| layer(2, 4, tokens, fill + l as f32 * 100.0))
2217 .collect();
2218 KvBlock::stamp(model, layout, layers).expect("stamp")
2219 }
2220
2221 fn expected(model: &str, n_layers: usize, tokens: usize) -> CacheSignature {
2222 CacheSignature::expected(model, flat(tokens), n_layers, 2, 4, tokens)
2223 }
2224
2225 fn hash(n: usize) -> BlockHash {
2226 BlockHasher::new("model-a", &[] as &[&str]).chain(&[n, n + 1], 2)[0]
2227 }
2228
2229 fn plenty() -> FreeSpaceProbe {
2234 Arc::new(|_: &Path| Some(1 << 40))
2235 }
2236
2237 fn store(dir: &TempDir, max_bytes: u64) -> DiskKvStore {
2240 DiskKvStore::open(
2241 DiskConfig::new(dir.path())
2242 .with_max_bytes(max_bytes)
2243 .with_writer_threads(0)
2244 .with_free_space_probe(plenty()),
2245 )
2246 .expect("open")
2247 }
2248
2249 fn put_now(store: &DiskKvStore, hash: BlockHash, block: KvBlock) {
2250 store.put_blocking(hash, block).expect("put");
2251 }
2252
2253 #[test]
2254 fn a_block_round_trips_through_a_file() {
2255 let dir = TempDir::new("roundtrip");
2256 let store = store(&dir, 1 << 20);
2257 let h = hash(1);
2258 let written = block("model-a", 3, 4, 1.0);
2259 let copy = block("model-a", 3, 4, 1.0);
2260 put_now(&store, h, written);
2261
2262 let read = store
2263 .get(&h, &expected("model-a", 3, 4))
2264 .expect("get")
2265 .expect("the block just written must be found");
2266 assert_eq!(read.layers().len(), 3);
2267 for (a, b) in read.layers().iter().zip(copy.layers()) {
2268 assert_eq!(a.k, b.k);
2269 assert_eq!(a.v, b.v);
2270 assert_eq!(a.seq_len, b.seq_len);
2271 }
2272 let stats = store.stats();
2273 assert_eq!(stats.hits, 1);
2274 assert_eq!(stats.writes, 1);
2275 assert_eq!(stats.blocks, 1);
2276 assert!(stats.read_nanos > 0, "a read must be timed");
2277 assert!(stats.write_nanos > 0, "a write must be timed");
2278 }
2279
2280 #[test]
2281 fn the_accounted_size_is_the_real_file_size() {
2282 let dir = TempDir::new("size");
2283 let store = store(&dir, 1 << 20);
2284 let h = hash(2);
2285 let written = block("model-a", 2, 8, 0.25);
2286 let predicted = encoded_len(written.signature());
2287 put_now(&store, h, written);
2288 let on_disk = fs::metadata(store.block_path(&h)).expect("stat").len();
2289 assert_eq!(
2290 predicted, on_disk,
2291 "the budget charges what the file really costs"
2292 );
2293 assert_eq!(store.stats().bytes, on_disk);
2294 }
2295
2296 #[test]
2297 fn blocks_are_sharded_by_hash_prefix() {
2298 let dir = TempDir::new("shard");
2299 let store = DiskKvStore::open(
2300 DiskConfig::new(dir.path())
2301 .with_shard_chars(2)
2302 .with_writer_threads(0)
2303 .with_free_space_probe(plenty()),
2304 )
2305 .expect("open");
2306 let h = hash(3);
2307 put_now(&store, h, block("model-a", 1, 2, 1.0));
2308 let path = store.block_path(&h);
2309 assert_eq!(
2310 path.parent()
2311 .unwrap()
2312 .file_name()
2313 .unwrap()
2314 .to_str()
2315 .unwrap(),
2316 &h.to_hex()[..2]
2317 );
2318 assert!(path.exists());
2319 }
2320
2321 #[test]
2325 fn a_truncated_file_is_refused_at_every_cut_point() {
2326 let h = hash(4);
2327 let bytes = encode_block(&h, &block("model-a", 2, 4, 3.0));
2328 assert!(bytes.len() > PREFIX_LEN + 16);
2329
2330 let err = decode_block(&bytes[..PREFIX_LEN - 1]).expect_err("short file");
2332 assert_eq!(
2333 err,
2334 BlockFormatError::TooShort {
2335 len: PREFIX_LEN - 1
2336 }
2337 );
2338
2339 for cut in [PREFIX_LEN, PREFIX_LEN + 8, bytes.len() - 4, bytes.len() - 1] {
2341 let err = decode_block(&bytes[..cut]).expect_err("truncated file");
2342 assert_eq!(
2343 err,
2344 BlockFormatError::Truncated {
2345 expected: bytes.len() as u64,
2346 actual: cut as u64,
2347 },
2348 "a file cut at {cut} must be refused"
2349 );
2350 }
2351
2352 let mut flipped = bytes.clone();
2354 let last = flipped.len() - 1;
2355 flipped[last] ^= 0xff;
2356 assert_eq!(
2357 decode_block(&flipped).expect_err("altered file"),
2358 BlockFormatError::ChecksumMismatch
2359 );
2360
2361 let mut alien = bytes;
2363 alien[0] = b'X';
2364 assert_eq!(
2365 decode_block(&alien).expect_err("foreign file"),
2366 BlockFormatError::BadMagic
2367 );
2368 }
2369
2370 #[test]
2374 fn a_torn_file_on_disk_is_a_miss_and_is_quarantined() {
2375 let dir = TempDir::new("torn");
2376 let store = store(&dir, 1 << 20);
2377 let h = hash(5);
2378 put_now(&store, h, block("model-a", 2, 4, 1.0));
2379 let path = store.block_path(&h);
2380
2381 let full = fs::read(&path).expect("read back");
2383 fs::write(&path, &full[..full.len() / 2]).expect("truncate");
2384
2385 let got = store.get(&h, &expected("model-a", 2, 4)).expect("get");
2386 assert!(got.is_none(), "a torn block must not be returned");
2387 assert_eq!(store.stats().corrupt, 1);
2388 assert!(!path.exists(), "a torn block must not be left to trip over");
2389 assert!(!store.contains(&h));
2390 }
2391
2392 #[test]
2393 fn an_unreadable_format_version_is_refused() {
2394 let h = hash(6);
2395 let mut bytes = encode_block(&h, &block("model-a", 1, 2, 1.0));
2396 bytes[8..12].copy_from_slice(&99u32.to_le_bytes());
2397 let mut digest = Sha256::new();
2399 digest.update(&bytes[PREFIX_LEN..]);
2400 let digest: [u8; 32] = digest.finalize().into();
2401 bytes[24..PREFIX_LEN].copy_from_slice(&digest);
2402 assert_eq!(
2403 decode_block(&bytes).expect_err("unknown version"),
2404 BlockFormatError::UnsupportedFormat {
2405 found: 99,
2406 readable: READABLE_FORMAT_VERSIONS,
2407 }
2408 );
2409 }
2410
2411 #[test]
2415 fn a_block_from_a_different_config_is_a_miss_not_a_hit() {
2416 let dir = TempDir::new("config");
2417 let store = store(&dir, 1 << 20);
2418 let h = hash(7);
2419 put_now(&store, h, block("model-a", 2, 4, 1.0));
2420
2421 assert!(store
2422 .get(&h, &expected("model-b", 2, 4))
2423 .expect("get")
2424 .is_none());
2425 assert!(store
2426 .get(
2427 &h,
2428 &CacheSignature::expected("model-a", flat(4), 2, 8, 4, 4)
2429 )
2430 .expect("get")
2431 .is_none());
2432 assert_eq!(store.stats().incompatible, 2);
2433 assert_eq!(store.stats().hits, 0);
2434 assert!(store
2437 .get(&h, &expected("model-a", 2, 4))
2438 .expect("get")
2439 .is_some());
2440 }
2441
2442 #[test]
2448 fn a_file_stored_under_the_wrong_name_is_rejected() {
2449 let dir = TempDir::new("misfiled");
2450 let store = store(&dir, 1 << 20);
2451 let (a, b) = (hash(8), hash(9));
2452 put_now(&store, a, block("model-a", 1, 2, 1.0));
2453 put_now(&store, b, block("model-a", 1, 2, 2.0));
2454 let bytes = fs::read(store.block_path(&b)).expect("read b");
2456 fs::write(store.block_path(&a), bytes).expect("misfile");
2457
2458 assert!(store
2459 .get(&a, &expected("model-a", 1, 2))
2460 .expect("get")
2461 .is_none());
2462 assert_eq!(store.stats().corrupt, 1);
2463 }
2464
2465 #[test]
2466 fn eviction_keeps_the_store_inside_its_budget() {
2467 let dir = TempDir::new("evict");
2468 let one = encoded_len(block("model-a", 1, 4, 1.0).signature());
2469 let store = store(&dir, one * 2 + 8);
2471 let hashes: Vec<BlockHash> = (0..4).map(|i| hash(20 + i)).collect();
2472 for (i, h) in hashes.iter().enumerate() {
2473 put_now(&store, *h, block("model-a", 1, 4, i as f32));
2474 assert!(
2475 store.stats().bytes <= store.capacity(),
2476 "the store must never sit over budget"
2477 );
2478 }
2479 let stats = store.stats();
2480 assert_eq!(stats.blocks, 2);
2481 assert_eq!(stats.evictions, 2);
2482 assert!(stats.evicted_bytes >= one * 2);
2483 for h in &hashes[..2] {
2485 assert!(!store.contains(h));
2486 assert!(
2487 !store.block_path(h).exists(),
2488 "an evicted file must be deleted"
2489 );
2490 }
2491 for h in &hashes[2..] {
2492 assert!(store.contains(h));
2493 }
2494 }
2495
2496 #[test]
2497 fn a_read_makes_a_block_the_least_likely_eviction_victim() {
2498 let dir = TempDir::new("lru");
2499 let one = encoded_len(block("model-a", 1, 4, 1.0).signature());
2500 let store = store(&dir, one * 2 + 8);
2501 let (a, b, c) = (hash(30), hash(31), hash(32));
2502 put_now(&store, a, block("model-a", 1, 4, 1.0));
2503 put_now(&store, b, block("model-a", 1, 4, 2.0));
2504 assert!(store.get(&a, &expected("model-a", 1, 4)).unwrap().is_some());
2506 put_now(&store, c, block("model-a", 1, 4, 3.0));
2507
2508 assert!(store.contains(&a), "a recently read block must survive");
2509 assert!(!store.contains(&b));
2510 assert!(store.contains(&c));
2511 }
2512
2513 #[test]
2519 fn a_block_evicted_mid_write_does_not_leave_its_file_behind() {
2520 let dir = TempDir::new("raced");
2521 let store = store(&dir, 1 << 20);
2522 let h = hash(40);
2523 {
2524 let evicting = Arc::clone(&store.shared);
2525 let mut hook = store
2526 .shared
2527 .hooks
2528 .after_rename
2529 .lock()
2530 .expect("hook lock poisoned");
2531 *hook = Some(Arc::new(move |hash: &BlockHash| {
2532 evicting.drop_entry(hash);
2535 }));
2536 }
2537 put_now(&store, h, block("model-a", 1, 4, 1.0));
2538
2539 assert!(
2540 !store.block_path(&h).exists(),
2541 "a file published for an entry that no longer exists must be withdrawn"
2542 );
2543 assert!(!store.contains(&h));
2544 let stats = store.stats();
2545 assert_eq!(stats.write_raced_eviction, 1);
2546 assert_eq!(stats.bytes, 0, "no bytes may be left unaccounted");
2547 }
2548
2549 #[test]
2550 fn no_temp_files_survive_a_successful_write() {
2551 let dir = TempDir::new("tmp");
2552 let store = store(&dir, 1 << 20);
2553 for i in 0..4 {
2554 put_now(&store, hash(50 + i), block("model-a", 1, 2, i as f32));
2555 }
2556 let leftovers: Vec<_> = fs::read_dir(dir.path().join(TMP_DIR))
2557 .expect("tmp dir")
2558 .flatten()
2559 .collect();
2560 assert!(
2561 leftovers.is_empty(),
2562 "temp files must not accumulate: {leftovers:?}"
2563 );
2564 }
2565
2566 #[test]
2570 fn a_new_store_reattaches_to_what_the_previous_one_published() {
2571 let dir = TempDir::new("restart");
2572 let h = hash(60);
2573 {
2574 let store = store(&dir, 1 << 20);
2575 put_now(&store, h, block("model-a", 2, 4, 7.0));
2576 }
2577 let orphan = dir.path().join(TMP_DIR).join("dead.tmp");
2579 fs::write(&orphan, b"half a block").expect("orphan");
2580
2581 let reopened = store(&dir, 1 << 20);
2582 assert!(
2583 !reopened.contains(&h),
2584 "reattaching must be an explicit step, not a side effect of open()"
2585 );
2586 assert_eq!(reopened.reindex().expect("reindex"), 1);
2587 assert!(reopened.contains(&h));
2588 assert!(!orphan.exists(), "an unpublished temp file must be swept");
2589
2590 let read = reopened
2591 .get(&h, &expected("model-a", 2, 4))
2592 .expect("get")
2593 .expect("a block written before the restart must still be readable");
2594 assert_eq!(read.tokens(), 4);
2595 }
2596
2597 #[test]
2609 fn a_block_written_under_one_window_is_not_served_to_a_reader_expecting_another() {
2610 let dir = TempDir::new("swa-window-restart");
2611 let h = hash(90);
2612 let window_128 = BlockLayout::new(4, Some(128)).expect("4 divides 128");
2613 let window_256 = BlockLayout::new(4, Some(256)).expect("4 divides 256");
2614 {
2615 let store = store(&dir, 1 << 20);
2616 put_now(
2617 &store,
2618 h,
2619 block_with_layout("model-a", 2, 4, 7.0, window_128),
2620 );
2621 }
2622
2623 let reopened = store(&dir, 1 << 20);
2624 assert_eq!(reopened.reindex().expect("reindex"), 1);
2625 assert!(reopened.contains(&h), "the file is there");
2626
2627 let after_window_change = reopened
2628 .get(
2629 &h,
2630 &CacheSignature::expected("model-a", window_256, 2, 2, 4, 4),
2631 )
2632 .expect("a config change is a miss, not an I/O error");
2633 assert!(
2634 after_window_change.is_none(),
2635 "a block cut against a 128 window must not be handed to a 256-window reader"
2636 );
2637 assert_eq!(reopened.stats().incompatible, 1);
2638 assert_eq!(reopened.stats().hits, 0);
2639
2640 let same = reopened
2643 .get(
2644 &h,
2645 &CacheSignature::expected("model-a", window_128, 2, 2, 4, 4),
2646 )
2647 .expect("get")
2648 .expect("the same window must still hit");
2649 assert_eq!(same.tokens(), 4);
2650 assert_eq!(same.layout(), window_128);
2651 }
2652
2653 #[test]
2658 fn the_block_layout_round_trips_through_the_file_format() {
2659 let sliding = BlockLayout::new(4, Some(512)).expect("4 divides 512");
2660 let h = hash(91);
2661 let bytes = encode_block(&h, &block_with_layout("model-a", 2, 4, 1.0, sliding));
2662 let decoded = decode_block(&bytes).expect("decode");
2663 let sig = decoded.block.signature.as_ref().expect("signature");
2664 assert_eq!(sig.layout, sliding);
2665 assert_eq!(sig.layout.sliding_window(), Some(512));
2666 assert_eq!(sig.layout.block_size(), 4);
2667
2668 let bytes = encode_block(&h, &block("model-a", 2, 4, 1.0));
2671 let decoded = decode_block(&bytes).expect("decode");
2672 let sig = decoded.block.signature.as_ref().expect("signature");
2673 assert_eq!(sig.layout.sliding_window(), None);
2674 }
2675
2676 #[test]
2681 fn a_file_recording_a_mis_aligned_layout_is_refused_at_parse_time() {
2682 let h = hash(92);
2683 let mut bytes = encode_block(&h, &block("model-a", 2, 4, 1.0));
2684 let window_at = PREFIX_LEN + 32 + 4 * 6;
2688 bytes[window_at..window_at + 4].copy_from_slice(&6u32.to_le_bytes());
2689 let mut digest = Sha256::new();
2693 digest.update(&bytes[PREFIX_LEN..]);
2694 let digest: [u8; 32] = digest.finalize().into();
2695 bytes[24..PREFIX_LEN].copy_from_slice(&digest);
2696
2697 let err = decode_block(&bytes).expect_err("6 is not a multiple of 4");
2698 assert!(
2699 matches!(err, BlockFormatError::BadLayout(_)),
2700 "expected a layout refusal, got {err}"
2701 );
2702 }
2703
2704 #[test]
2705 fn reindex_evicts_down_to_the_budget() {
2706 let dir = TempDir::new("reindex-evict");
2707 let one = encoded_len(block("model-a", 1, 4, 1.0).signature());
2708 {
2709 let store = store(&dir, 1 << 20);
2710 for i in 0..4 {
2711 put_now(&store, hash(70 + i), block("model-a", 1, 4, i as f32));
2712 }
2713 }
2714 let small = store(&dir, one * 2 + 8);
2715 small.reindex().expect("reindex");
2716 let stats = small.stats();
2717 assert_eq!(stats.blocks, 2, "a shrunken budget must bind on restart");
2718 assert!(stats.bytes <= small.capacity());
2719 }
2720
2721 #[test]
2722 fn an_absent_block_is_a_plain_miss() {
2723 let dir = TempDir::new("miss");
2724 let store = store(&dir, 1 << 20);
2725 assert!(store
2726 .get(&hash(80), &expected("model-a", 1, 2))
2727 .expect("get")
2728 .is_none());
2729 assert_eq!(store.stats().misses, 1);
2730 assert_eq!(store.stats().corrupt, 0);
2731 }
2732
2733 #[test]
2737 fn a_file_deleted_behind_the_stores_back_is_a_miss() {
2738 let dir = TempDir::new("vanished");
2739 let store = store(&dir, 1 << 20);
2740 let h = hash(90);
2741 put_now(&store, h, block("model-a", 1, 2, 1.0));
2742 fs::remove_file(store.block_path(&h)).expect("remove");
2743 assert!(store
2744 .get(&h, &expected("model-a", 1, 2))
2745 .expect("get")
2746 .is_none());
2747 assert!(!store.contains(&h));
2748 assert_eq!(store.stats().bytes, 0);
2749 }
2750
2751 #[test]
2752 fn rewriting_a_block_does_not_double_charge_it() {
2753 let dir = TempDir::new("rewrite");
2754 let store = store(&dir, 1 << 20);
2755 let h = hash(100);
2756 put_now(&store, h, block("model-a", 1, 4, 1.0));
2757 let once = store.stats().bytes;
2758 put_now(&store, h, block("model-a", 1, 4, 1.0));
2759 assert_eq!(store.stats().bytes, once);
2760 assert_eq!(store.stats().blocks, 1);
2761 }
2762
2763 #[test]
2764 fn hex_names_round_trip() {
2765 let h = hash(110);
2766 assert_eq!(parse_hex_hash(&h.to_hex()), Some(h));
2767 assert_eq!(parse_hex_hash("nothex"), None);
2768 assert_eq!(parse_hex_hash(&"z".repeat(64)), None);
2769 }
2770
2771 fn probe_window(order: WriteOrder, publish_window: bool) -> (usize, usize) {
2783 let dir = TempDir::new("ordering");
2784 let store = store(&dir, 1 << 20);
2785 *store.shared.hooks.order.lock().unwrap() = order;
2786
2787 let violations = Arc::new(AtomicUsize::new(0));
2788 let served = Arc::new(AtomicUsize::new(0));
2789 let reader = Arc::clone(&store.shared);
2790 let v = Arc::clone(&violations);
2791 let s = Arc::clone(&served);
2792 let hook: Hook = Arc::new(move |hash: &BlockHash| {
2793 match reader
2794 .read_async(hash, &expected("model-a", 1, 4), true)
2795 .wait()
2796 {
2797 Ok(Some(_)) => {
2798 s.fetch_add(1, Ordering::Relaxed);
2799 }
2800 Ok(None) => {}
2803 Err(StoreError::MissingPayload { .. }) => {
2804 v.fetch_add(1, Ordering::Relaxed);
2805 }
2806 Err(other) => panic!("unexpected store error: {other}"),
2807 }
2808 });
2809 let slot = if publish_window {
2810 &store.shared.hooks.in_publish_window
2811 } else {
2812 &store.shared.hooks.in_put_window
2813 };
2814 *slot.lock().unwrap() = Some(hook);
2815
2816 put_now(&store, hash(200), block("model-a", 1, 4, 1.0));
2817 (
2818 violations.load(Ordering::Relaxed),
2819 served.load(Ordering::Relaxed),
2820 )
2821 }
2822
2823 #[test]
2829 fn a_reader_never_sees_an_index_hit_with_no_payload() {
2830 let (violations, _) = probe_window(WriteOrder::BufferThenIndex, false);
2831 assert_eq!(violations, 0, "admission window must be safe");
2832
2833 let (violations, served) = probe_window(WriteOrder::BufferThenIndex, true);
2834 assert_eq!(violations, 0, "publication window must be safe");
2835 assert_eq!(
2836 served, 1,
2837 "the reader must actually have reached the block, or this test proves nothing"
2838 );
2839 }
2840
2841 #[test]
2846 fn indexing_before_buffering_is_caught() {
2847 let (violations, _) = probe_window(WriteOrder::IndexBeforeBuffer, false);
2848 assert_eq!(
2849 violations, 1,
2850 "index-then-buffer must be detected as an invariant violation"
2851 );
2852 }
2853
2854 #[test]
2858 fn releasing_the_buffer_before_publishing_is_caught() {
2859 let (violations, _) = probe_window(WriteOrder::DropBufferBeforeMarking, true);
2860 assert_eq!(
2861 violations, 1,
2862 "release-then-mark must be detected as an invariant violation"
2863 );
2864 }
2865
2866 #[test]
2872 fn concurrent_readers_never_see_an_index_hit_with_no_payload() {
2873 let dir = TempDir::new("concurrent");
2874 let one = encoded_len(block("model-a", 1, 4, 1.0).signature());
2875 let store = Arc::new(
2876 DiskKvStore::open(
2877 DiskConfig::new(dir.path())
2878 .with_max_bytes(one * 8)
2880 .with_queue_capacity(4)
2881 .with_writer_threads(2)
2882 .with_free_space_probe(plenty()),
2883 )
2884 .expect("open"),
2885 );
2886 let hashes: Vec<BlockHash> = (0..16).map(|i| hash(300 + i)).collect();
2887 let violations = Arc::new(AtomicUsize::new(0));
2888 let stop = Arc::new(std::sync::atomic::AtomicBool::new(false));
2889
2890 let readers: Vec<_> = (0..4)
2891 .map(|_| {
2892 let store = Arc::clone(&store);
2893 let hashes = hashes.clone();
2894 let violations = Arc::clone(&violations);
2895 let stop = Arc::clone(&stop);
2896 std::thread::spawn(move || {
2897 let want = expected("model-a", 1, 4);
2898 while !stop.load(Ordering::Relaxed) {
2899 for h in &hashes {
2900 match store.get(h, &want) {
2901 Ok(_) => {}
2902 Err(StoreError::MissingPayload { .. }) => {
2903 violations.fetch_add(1, Ordering::Relaxed);
2904 }
2905 Err(other) => panic!("unexpected store error: {other}"),
2906 }
2907 }
2908 }
2909 })
2910 })
2911 .collect();
2912
2913 for round in 0..4 {
2914 for (i, h) in hashes.iter().enumerate() {
2915 store
2916 .put(*h, block("model-a", 1, 4, (round * 16 + i) as f32))
2917 .expect("put");
2918 }
2919 }
2920 store.flush();
2921 stop.store(true, Ordering::Relaxed);
2922 for reader in readers {
2923 reader.join().expect("reader thread");
2924 }
2925
2926 assert_eq!(
2927 violations.load(Ordering::Relaxed),
2928 0,
2929 "no reader may ever see an index hit with no payload"
2930 );
2931 let stats = store.stats();
2932 assert!(
2933 stats.buffer_hits > 0,
2934 "readers must have caught blocks still in the write buffer, \
2935 or this test never entered the window"
2936 );
2937 assert!(stats.evictions > 0, "the budget must have bound");
2938 assert!(stats.bytes <= store.capacity());
2939 }
2940
2941 #[test]
2945 fn a_queued_block_is_readable_before_it_reaches_disk() {
2946 let dir = TempDir::new("buffered");
2947 let store = store(&dir, 1 << 20);
2949 let h = hash(400);
2950 store.put(h, block("model-a", 1, 4, 1.0)).expect("put");
2951
2952 assert!(
2953 !store.block_path(&h).exists(),
2954 "nothing has been written yet"
2955 );
2956 let got = store
2957 .get(&h, &expected("model-a", 1, 4))
2958 .expect("get")
2959 .expect("a queued block must be readable immediately");
2960 assert_eq!(got.tokens(), 4);
2961 assert_eq!(store.stats().buffer_hits, 1);
2962
2963 store.flush();
2964 assert!(store.block_path(&h).exists(), "flush must publish it");
2965 assert!(store
2966 .get(&h, &expected("model-a", 1, 4))
2967 .expect("get")
2968 .is_some());
2969 assert_eq!(store.stats().hits, 1, "and now it comes off the disk");
2970 }
2971
2972 #[test]
2977 fn a_full_queue_writes_inline_rather_than_dropping_the_block() {
2978 let dir = TempDir::new("backpressure");
2979 let store = DiskKvStore::open(
2980 DiskConfig::new(dir.path())
2981 .with_queue_capacity(2)
2982 .with_writer_threads(0)
2984 .with_free_space_probe(plenty()),
2985 )
2986 .expect("open");
2987
2988 let hashes: Vec<BlockHash> = (0..5).map(|i| hash(500 + i)).collect();
2989 for (i, h) in hashes.iter().enumerate() {
2990 store
2991 .put(*h, block("model-a", 1, 4, i as f32))
2992 .expect("put");
2993 }
2994 let stats = store.stats();
2995 assert_eq!(stats.queued_writes, 2, "the queue holds exactly its cap");
2996 assert_eq!(stats.inline_writes, 3, "the rest fall back to this thread");
2997 assert_eq!(stats.writes, 3, "and the fallbacks really wrote");
2998
2999 let want = expected("model-a", 1, 4);
3002 for h in &hashes {
3003 assert!(
3004 store.get(h, &want).expect("get").is_some(),
3005 "no block may be lost to a full queue"
3006 );
3007 }
3008 store.flush();
3009 for h in &hashes {
3010 assert!(store.block_path(h).exists(), "flush publishes the rest");
3011 }
3012 }
3013
3014 #[test]
3019 fn a_queued_write_evicted_before_it_runs_is_skipped() {
3020 let dir = TempDir::new("skipped");
3021 let store = store(&dir, 1 << 20);
3022 let h = hash(600);
3023 store.put(h, block("model-a", 1, 4, 1.0)).expect("put");
3024 store.remove(&h);
3025 store.flush();
3026
3027 let stats = store.stats();
3028 assert_eq!(stats.write_skipped, 1);
3029 assert_eq!(stats.writes, 0);
3030 assert!(!store.block_path(&h).exists());
3031 assert_eq!(stats.bytes, 0);
3032 }
3033
3034 #[test]
3039 fn a_superseded_queued_write_does_not_overwrite_the_newer_block() {
3040 let dir = TempDir::new("superseded");
3041 let store = store(&dir, 1 << 20);
3042 let h = hash(700);
3043 store.put(h, block("model-a", 1, 4, 1.0)).expect("put");
3044 store.put(h, block("model-a", 1, 4, 9.0)).expect("put");
3045 store.flush();
3046
3047 let got = store
3048 .get(&h, &expected("model-a", 1, 4))
3049 .expect("get")
3050 .expect("hit");
3051 assert_eq!(
3052 got.layers()[0].k[0],
3053 9.0,
3054 "the newer block must win, not whichever write ran last"
3055 );
3056 assert_eq!(store.stats().write_skipped, 1);
3057 assert_eq!(store.stats().blocks, 1);
3058 }
3059
3060 fn reading_store(dir: &TempDir, readers: usize) -> DiskKvStore {
3067 DiskKvStore::open(
3068 DiskConfig::new(dir.path())
3069 .with_writer_threads(0)
3070 .with_reader_threads(readers)
3071 .with_free_space_probe(plenty()),
3072 )
3073 .expect("open")
3074 }
3075
3076 fn wait_staged(store: &DiskKvStore, hash: &BlockHash) {
3078 for _ in 0..2000 {
3079 let ready = store
3080 .shared
3081 .staging
3082 .lock()
3083 .unwrap()
3084 .get(hash)
3085 .map(|slot| slot.is_ready())
3086 .unwrap_or(false);
3087 if ready {
3088 return;
3089 }
3090 std::thread::sleep(std::time::Duration::from_millis(1));
3091 }
3092 panic!("prefetch never completed");
3093 }
3094
3095 #[test]
3101 fn a_prefetched_block_is_already_read_when_the_request_arrives() {
3102 let dir = TempDir::new("prefetch");
3103 let store = reading_store(&dir, 1);
3104 let h = hash(800);
3105 put_now(&store, h, block("model-a", 2, 4, 5.0));
3106 let want = expected("model-a", 2, 4);
3107
3108 store.prefetch(&[h], &want);
3109 wait_staged(&store, &h);
3110 fs::remove_file(store.block_path(&h)).expect("remove");
3111
3112 let got = store
3113 .get(&h, &want)
3114 .expect("get")
3115 .expect("the prefetch already had it");
3116 assert_eq!(got.tokens(), 4);
3117 assert_eq!(got.layers()[0].k[0], 5.0);
3118
3119 let stats = store.stats();
3120 assert_eq!(
3121 stats.prefetch_hits, 1,
3122 "the request must have found it ready"
3123 );
3124 assert_eq!(
3125 stats.hits, 1,
3126 "and the file must have been read exactly once"
3127 );
3128 assert_eq!(stats.async_reads, 1, "on a reader thread, not the caller's");
3129 assert_eq!(stats.staged_blocks, 0, "a claimed read leaves staging");
3130 }
3131
3132 #[test]
3136 fn a_whole_chain_can_be_read_ahead_in_one_call() {
3137 let dir = TempDir::new("chain");
3138 let store = reading_store(&dir, 2);
3139 let hashes: Vec<BlockHash> = (0..4).map(|i| hash(810 + i)).collect();
3140 for (i, h) in hashes.iter().enumerate() {
3141 put_now(&store, *h, block("model-a", 1, 4, i as f32));
3142 }
3143 let want = expected("model-a", 1, 4);
3144
3145 store.prefetch(&hashes, &want);
3146 for h in &hashes {
3147 wait_staged(&store, h);
3148 fs::remove_file(store.block_path(h)).expect("remove");
3149 }
3150 for (i, h) in hashes.iter().enumerate() {
3151 let got = store.get(h, &want).expect("get").expect("read ahead");
3152 assert_eq!(got.layers()[0].k[0], i as f32);
3153 }
3154 let stats = store.stats();
3155 assert_eq!(stats.prefetch_issued, 4);
3156 assert_eq!(stats.prefetch_hits, 4);
3157 assert_eq!(stats.hits, 4, "four blocks, four reads, none repeated");
3158 }
3159
3160 #[test]
3163 fn a_request_joins_a_read_already_running_rather_than_repeating_it() {
3164 let dir = TempDir::new("join");
3165 let store = reading_store(&dir, 1);
3166 let h = hash(820);
3167 put_now(&store, h, block("model-a", 1, 4, 1.0));
3168 let want = expected("model-a", 1, 4);
3169
3170 store.prefetch(&[h], &want);
3171 let got = store.get(&h, &want).expect("get").expect("hit");
3174 assert_eq!(got.tokens(), 4);
3175 let stats = store.stats();
3176 assert_eq!(
3177 stats.prefetch_hits + stats.prefetch_waits,
3178 1,
3179 "the request either found the read done or waited for it"
3180 );
3181 assert_eq!(stats.hits, 1, "one physical read, whichever way it went");
3182 }
3183
3184 #[test]
3189 fn a_staged_read_is_not_reused_by_a_reader_that_wants_another_shape() {
3190 let dir = TempDir::new("staged-shape");
3191 let store = reading_store(&dir, 1);
3192 let h = hash(830);
3193 put_now(&store, h, block("model-a", 1, 4, 1.0));
3194
3195 store.prefetch(&[h], &expected("model-a", 1, 4));
3196 wait_staged(&store, &h);
3197
3198 let got = store.get(&h, &expected("model-b", 1, 4)).expect("get");
3199 assert!(got.is_none(), "a different model must not be served");
3200 assert_eq!(store.stats().incompatible, 1);
3201 assert_eq!(
3202 store.stats().prefetch_hits,
3203 0,
3204 "the staged answer was for another expectation and must not be claimed"
3205 );
3206 }
3207
3208 #[test]
3212 fn prefetching_is_bounded() {
3213 let dir = TempDir::new("prefetch-bound");
3214 let store = DiskKvStore::open(
3215 DiskConfig::new(dir.path())
3216 .with_writer_threads(0)
3217 .with_reader_threads(1)
3218 .with_prefetch_capacity(2)
3219 .with_free_space_probe(plenty()),
3220 )
3221 .expect("open");
3222 let hashes: Vec<BlockHash> = (0..6).map(|i| hash(840 + i)).collect();
3223 for (i, h) in hashes.iter().enumerate() {
3224 put_now(&store, *h, block("model-a", 1, 4, i as f32));
3225 }
3226
3227 store.prefetch(&hashes, &expected("model-a", 1, 4));
3228 let stats = store.stats();
3229 assert!(
3230 stats.staged_blocks <= 2,
3231 "staging must respect its cap, got {}",
3232 stats.staged_blocks
3233 );
3234 assert!(
3235 stats.prefetch_dropped >= 4,
3236 "the refusals must be visible, got {}",
3237 stats.prefetch_dropped
3238 );
3239
3240 let want = expected("model-a", 1, 4);
3242 for (i, h) in hashes.iter().enumerate() {
3243 let got = store.get(h, &want).expect("get").expect("hit");
3244 assert_eq!(got.layers()[0].k[0], i as f32);
3245 }
3246 }
3247
3248 #[test]
3251 fn without_reader_threads_reads_run_on_the_caller() {
3252 let dir = TempDir::new("no-readers");
3253 let store = reading_store(&dir, 0);
3254 let h = hash(850);
3255 put_now(&store, h, block("model-a", 1, 4, 1.0));
3256 let want = expected("model-a", 1, 4);
3257
3258 store.prefetch(&[h], &want);
3259 assert_eq!(store.stats().prefetch_dropped, 1);
3260 assert_eq!(store.stats().staged_blocks, 0);
3261
3262 assert!(store.get(&h, &want).expect("get").is_some());
3263 let stats = store.stats();
3264 assert_eq!(stats.hits, 1);
3265 assert_eq!(stats.async_reads, 0);
3266 }
3267
3268 #[test]
3271 fn a_read_handle_can_be_polled_to_completion() {
3272 let dir = TempDir::new("handle");
3273 let store = reading_store(&dir, 1);
3274 let h = hash(860);
3275 put_now(&store, h, block("model-a", 1, 4, 2.0));
3276 let want = expected("model-a", 1, 4);
3277
3278 let handle = store.read_async(&h, &want);
3279 for _ in 0..2000 {
3280 if let Some(outcome) = handle.try_claim() {
3281 let got = outcome.expect("read").expect("hit");
3282 assert_eq!(got.layers()[0].k[0], 2.0);
3283 assert_eq!(store.stats().staged_blocks, 0);
3284 return;
3285 }
3286 std::thread::sleep(std::time::Duration::from_millis(1));
3287 }
3288 panic!("read never completed");
3289 }
3290
3291 #[test]
3294 fn a_miss_is_answered_without_dispatching_a_read() {
3295 let dir = TempDir::new("ready-miss");
3296 let store = reading_store(&dir, 1);
3297 let handle = store.read_async(&hash(870), &expected("model-a", 1, 4));
3298 assert!(handle.is_ready(), "a miss must not cost a thread hop");
3299 assert!(handle.wait().expect("read").is_none());
3300 assert_eq!(store.stats().async_reads, 0);
3301 }
3302
3303 #[test]
3306 fn clearing_the_prefetch_releases_staged_blocks() {
3307 let dir = TempDir::new("clear");
3308 let store = reading_store(&dir, 1);
3309 let h = hash(880);
3310 put_now(&store, h, block("model-a", 1, 4, 1.0));
3311 store.prefetch(&[h], &expected("model-a", 1, 4));
3312 wait_staged(&store, &h);
3313 assert_eq!(store.stats().staged_blocks, 1);
3314 store.clear_prefetch();
3315 assert_eq!(store.stats().staged_blocks, 0);
3316 }
3317
3318 fn budgeted_store(
3324 dir: &TempDir,
3325 max_bytes: u64,
3326 reserve: u64,
3327 ttl: std::time::Duration,
3328 probe: FreeSpaceProbe,
3329 ) -> DiskKvStore {
3330 DiskKvStore::open(
3331 DiskConfig::new(dir.path())
3332 .with_max_bytes(max_bytes)
3333 .with_reserve_bytes(reserve)
3334 .with_free_space_ttl(ttl)
3335 .with_free_space_probe(probe)
3336 .with_writer_threads(0)
3337 .with_reader_threads(0),
3338 )
3339 .expect("open")
3340 }
3341
3342 fn dir_bytes(root: &Path) -> u64 {
3346 let mut total = 0;
3347 let Ok(entries) = fs::read_dir(root) else {
3348 return 0;
3349 };
3350 for entry in entries.flatten() {
3351 let path = entry.path();
3352 if path.is_dir() {
3353 total += dir_bytes(&path);
3354 } else if let Ok(meta) = entry.metadata() {
3355 total += meta.len();
3356 }
3357 }
3358 total
3359 }
3360
3361 #[test]
3371 fn the_ceiling_falls_when_the_filesystem_fills_up() {
3372 let dir = TempDir::new("budget");
3373 let one = encoded_len(block("model-a", 1, 4, 1.0).signature());
3374 let device = Arc::new(AtomicU64::new(1 << 40));
3375 let probe: FreeSpaceProbe = {
3376 let device = Arc::clone(&device);
3377 let root = dir.path().to_path_buf();
3378 Arc::new(move |_: &Path| {
3379 Some(
3380 device
3381 .load(Ordering::Relaxed)
3382 .saturating_sub(dir_bytes(&root)),
3383 )
3384 })
3385 };
3386 let reserve = one * 2;
3389 let store = budgeted_store(&dir, one * 100, reserve, std::time::Duration::ZERO, probe);
3390
3391 for i in 0..4 {
3392 put_now(&store, hash(900 + i), block("model-a", 1, 4, i as f32));
3393 }
3394 assert_eq!(store.stats().blocks, 4);
3395 assert_eq!(store.stats().evictions, 0, "nothing binds yet");
3396 assert_eq!(
3397 store.effective_capacity(),
3398 store.capacity(),
3399 "with a terabyte free the configured budget is the ceiling"
3400 );
3401
3402 device.store(one * 6, Ordering::Relaxed);
3405 assert_eq!(
3406 store.effective_capacity(),
3407 one * 4,
3408 "the ceiling must follow the device down to total - reserve"
3409 );
3410
3411 for i in 0..6 {
3412 put_now(&store, hash(910 + i), block("model-a", 1, 4, i as f32));
3413 let on_disk = dir_bytes(dir.path());
3414 assert!(
3415 on_disk + reserve <= one * 6,
3416 "the store must hand the device its reserve back before the \
3417 filesystem has to: {on_disk} bytes used of {}, {reserve} reserved",
3418 one * 6
3419 );
3420 }
3421
3422 let stats = store.stats();
3423 assert_eq!(stats.blocks, 4, "settled at total - reserve");
3424 assert!(stats.evictions >= 6, "got {}", stats.evictions);
3425 assert!(stats.space_clamped > 0, "the clamp must be visible");
3426 assert!(
3427 stats.bytes < one * 100,
3428 "far under the configured budget it never reached"
3429 );
3430 assert_eq!(
3431 stats.disk_bytes,
3432 dir_bytes(dir.path()),
3433 "the store's idea of its disk footprint must be the real one"
3434 );
3435 }
3436
3437 #[test]
3440 fn the_free_space_reading_is_cached_for_its_ttl() {
3441 let dir = TempDir::new("ttl");
3442 let calls = Arc::new(AtomicUsize::new(0));
3443 let probe: FreeSpaceProbe = {
3444 let calls = Arc::clone(&calls);
3445 Arc::new(move |_: &Path| {
3446 calls.fetch_add(1, Ordering::Relaxed);
3447 Some(1 << 40)
3448 })
3449 };
3450 let store = budgeted_store(&dir, 1 << 20, 0, std::time::Duration::from_secs(60), probe);
3451
3452 for _ in 0..5 {
3453 store.effective_capacity();
3454 }
3455 for i in 0..3 {
3456 put_now(&store, hash(920 + i), block("model-a", 1, 4, i as f32));
3457 }
3458 assert_eq!(
3459 calls.load(Ordering::Relaxed),
3460 1,
3461 "a TTL'd reading must not be re-taken per operation"
3462 );
3463 }
3464
3465 #[test]
3470 fn enospc_throws_away_the_cached_free_space() {
3471 let dir = TempDir::new("enospc");
3472 let calls = Arc::new(AtomicUsize::new(0));
3473 let probe: FreeSpaceProbe = {
3474 let calls = Arc::clone(&calls);
3475 Arc::new(move |_: &Path| {
3476 calls.fetch_add(1, Ordering::Relaxed);
3477 Some(1 << 40)
3478 })
3479 };
3480 let store = budgeted_store(
3481 &dir,
3482 1 << 20,
3483 0,
3484 std::time::Duration::from_secs(3600),
3486 probe,
3487 );
3488 let h = hash(930);
3489 put_now(&store, h, block("model-a", 1, 4, 1.0));
3490 store.effective_capacity();
3491 assert_eq!(calls.load(Ordering::Relaxed), 1);
3492
3493 store
3494 .shared
3495 .hooks
3496 .fail_with_enospc
3497 .store(true, Ordering::Relaxed);
3498 let full = hash(931);
3499 let err = store
3500 .put_blocking(full, block("model-a", 1, 4, 2.0))
3501 .expect_err("a full filesystem must be reported, not swallowed");
3502 assert!(matches!(err, StoreError::Io { .. }), "{err}");
3503
3504 let stats = store.stats();
3505 assert_eq!(stats.enospc, 1);
3506 assert!(
3507 calls.load(Ordering::Relaxed) > 1,
3508 "ENOSPC must invalidate the cached reading immediately"
3509 );
3510 assert!(
3511 !store.contains(&full),
3512 "a block that could not be written must not be indexed"
3513 );
3514 assert_eq!(stats.write_failures, 1);
3515 let leftovers: Vec<_> = fs::read_dir(dir.path().join(TMP_DIR))
3516 .expect("tmp dir")
3517 .flatten()
3518 .collect();
3519 assert!(
3520 leftovers.is_empty(),
3521 "a failed write must clean up after itself: {leftovers:?}"
3522 );
3523
3524 assert!(store
3526 .get(&h, &expected("model-a", 1, 4))
3527 .expect("get")
3528 .is_some());
3529
3530 store
3532 .shared
3533 .hooks
3534 .fail_with_enospc
3535 .store(false, Ordering::Relaxed);
3536 put_now(&store, full, block("model-a", 1, 4, 2.0));
3537 assert!(store.contains(&full));
3538 }
3539
3540 #[test]
3544 fn an_unmeasurable_filesystem_falls_back_to_the_configured_budget() {
3545 let dir = TempDir::new("unknowable");
3546 let store = budgeted_store(
3547 &dir,
3548 1 << 20,
3549 1 << 30,
3550 std::time::Duration::ZERO,
3551 Arc::new(|_: &Path| None),
3552 );
3553 assert_eq!(store.effective_capacity(), 1 << 20);
3554 for i in 0..3 {
3555 put_now(&store, hash(940 + i), block("model-a", 1, 4, i as f32));
3556 }
3557 assert_eq!(store.stats().blocks, 3);
3558 assert_eq!(store.stats().evictions, 0);
3559 }
3560
3561 #[test]
3565 #[cfg(unix)]
3566 fn the_platform_probe_measures_a_real_filesystem() {
3567 let dir = TempDir::new("statvfs");
3568 let free = platform_free_bytes(dir.path()).expect("statvfs on a directory that exists");
3569 assert!(free > 0, "a writable temp dir with zero bytes free?");
3570 assert!(
3571 platform_free_bytes(&dir.path().join("no-such-dir")).is_none(),
3572 "a path that does not exist cannot report free space"
3573 );
3574 }
3575}