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.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#[derive(Clone, Debug, PartialEq, Eq)]
439pub enum StoreError {
440 Io {
441 op: &'static str,
442 path: PathBuf,
443 message: String,
444 },
445 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
482pub type FreeSpaceProbe = Arc<dyn Fn(&Path) -> Option<u64> + Send + Sync>;
489
490#[cfg(unix)]
494#[allow(clippy::unnecessary_cast)] fn 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 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
520struct FreeSpace {
525 checked_at: Option<Instant>,
526 bytes: Option<u64>,
527}
528
529#[derive(Clone)]
532pub struct DiskConfig {
533 pub root: PathBuf,
535 pub max_bytes: u64,
538 pub shard_chars: usize,
542 pub queue_capacity: usize,
546 pub writer_threads: usize,
549 pub reader_threads: usize,
554 pub prefetch_capacity: usize,
559 pub reserve_bytes: u64,
563 pub free_space_ttl: std::time::Duration,
566 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 inline_writes: AtomicU64,
658 write_failures: AtomicU64,
659 write_skipped: AtomicU64,
662 write_nanos: AtomicU64,
663 write_raced_eviction: AtomicU64,
666 hits: AtomicU64,
667 buffer_hits: AtomicU64,
671 misses: AtomicU64,
672 corrupt: AtomicU64,
674 incompatible: AtomicU64,
677 read_nanos: AtomicU64,
678 evictions: AtomicU64,
679 evicted_bytes: AtomicU64,
680 prefetch_issued: AtomicU64,
682 prefetch_dropped: AtomicU64,
685 prefetch_hits: AtomicU64,
689 prefetch_waits: AtomicU64,
692 async_reads: AtomicU64,
694 enospc: AtomicU64,
697 space_clamped: AtomicU64,
700}
701
702#[derive(Clone, Debug, Default, PartialEq, Eq)]
707pub struct DiskStats {
708 pub blocks: usize,
709 pub bytes: u64,
711 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 pub effective_capacity: u64,
740}
741
742struct Entry {
743 bytes: u64,
744 last_used: u64,
745 published: bool,
749 generation: u64,
754}
755
756struct Index {
757 entries: HashMap<BlockHash, Entry>,
758 bytes: u64,
760 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
797enum Source {
799 Disk(PathBuf),
800 Buffer(Arc<KvBlock>),
801}
802
803struct 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
821struct 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 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
904pub type ReadOutcome = Result<Option<Arc<KvBlock>>, StoreError>;
908
909struct ReadSlot {
913 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
974struct 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 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
1036pub struct ReadHandle {
1043 shared: Arc<Shared>,
1044 hash: BlockHash,
1045 slot: Arc<ReadSlot>,
1046 staged: bool,
1049}
1050
1051impl ReadHandle {
1052 pub fn is_ready(&self) -> bool {
1055 self.slot.is_ready()
1056 }
1057
1058 pub fn try_claim(&self) -> Option<ReadOutcome> {
1061 if !self.slot.is_ready() {
1062 return None;
1063 }
1064 Some(self.claim())
1065 }
1066
1067 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#[cfg(test)]
1098#[derive(Clone, Copy, Debug, Default, PartialEq, Eq)]
1099enum WriteOrder {
1100 #[default]
1101 BufferThenIndex,
1102 IndexBeforeBuffer,
1105 DropBufferBeforeMarking,
1108}
1109
1110#[cfg(test)]
1111#[derive(Default)]
1112struct Hooks {
1113 order: Mutex<WriteOrder>,
1114 after_rename: Mutex<Option<Hook>>,
1118 in_put_window: Mutex<Option<Hook>>,
1120 in_publish_window: Mutex<Option<Hook>>,
1122 fail_with_enospc: std::sync::atomic::AtomicBool,
1126}
1127
1128#[cfg(test)]
1129impl Hooks {
1130 fn fire(slot: &Mutex<Option<Hook>>, hash: &BlockHash) {
1131 let hook = slot.lock().expect("kv disk hook poisoned").clone();
1134 if let Some(hook) = hook {
1135 hook(hash);
1136 }
1137 }
1138}
1139
1140struct Shared {
1145 root: PathBuf,
1146 shard_chars: usize,
1147 max_bytes: u64,
1148 index: Mutex<Index>,
1149 buffer: Mutex<HashMap<BlockHash, Buffered>>,
1157 queue: WriteQueue,
1158 reads: ReadQueue,
1159 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 free_space: Mutex<FreeSpace>,
1171 stats: Stats,
1172 seq: AtomicU64,
1173 generation: AtomicU64,
1174 #[cfg(test)]
1175 hooks: Hooks,
1176}
1177
1178pub 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 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 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 pub fn put(&self, hash: BlockHash, block: KvBlock) -> Result<(), StoreError> {
1295 self.shared.put(hash, block, false)
1296 }
1297
1298 pub fn put_blocking(&self, hash: BlockHash, block: KvBlock) -> Result<(), StoreError> {
1300 self.shared.put(hash, block, true)
1301 }
1302
1303 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 let _ = self
1320 .shared
1321 .queue
1322 .idle
1323 .wait_timeout(state, std::time::Duration::from_millis(1));
1324 }
1325 }
1326
1327 pub fn get(&self, hash: &BlockHash, expected: &CacheSignature) -> ReadOutcome {
1334 self.shared.read_async(hash, expected, true).wait()
1335 }
1336
1337 pub fn read_async(&self, hash: &BlockHash, expected: &CacheSignature) -> ReadHandle {
1340 self.shared.read_async(hash, expected, true)
1341 }
1342
1343 pub fn prefetch(&self, hashes: &[BlockHash], expected: &CacheSignature) {
1353 self.shared.prefetch(hashes, expected);
1354 }
1355
1356 pub fn clear_prefetch(&self) {
1360 self.shared
1361 .staging
1362 .lock()
1363 .expect("kv disk staging poisoned")
1364 .clear();
1365 }
1366
1367 pub fn reindex(&self) -> Result<usize, StoreError> {
1369 self.shared.reindex()
1370 }
1371
1372 pub fn remove(&self, hash: &BlockHash) {
1374 self.shared.quarantine(hash);
1375 }
1376
1377 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 pub fn capacity(&self) -> u64 {
1386 self.shared.max_bytes
1387 }
1388
1389 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 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 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 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 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 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 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 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 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 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 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 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 fn read_async(
1696 self: &Arc<Self>,
1697 hash: &BlockHash,
1698 expected: &CacheSignature,
1699 demand: bool,
1700 ) -> ReadHandle {
1701 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 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 self.stats.prefetch_dropped.fetch_add(1, Ordering::Relaxed);
1816 }
1817 drop(handle);
1821 }
1822 }
1823
1824 fn effective_capacity(&self, used: u64) -> u64 {
1840 let Some(free) = self.free_bytes() else {
1841 return self.max_bytes;
1842 };
1843 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 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 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 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 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 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 fn collect_victims(&self, index: &mut Index, protect: Option<&BlockHash>) -> Vec<Victim> {
2027 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 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
2134struct 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 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 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 fn plenty() -> FreeSpaceProbe {
2236 Arc::new(|_: &Path| Some(1 << 40))
2237 }
2238
2239 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 #[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 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 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 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 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 #[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 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 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 #[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 assert!(store
2439 .get(&h, &expected("model-a", 2, 4))
2440 .expect("get")
2441 .is_some());
2442 }
2443
2444 #[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 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 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 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 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 #[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 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 #[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 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 #[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 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 #[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 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 #[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 let window_at = PREFIX_LEN + 32 + 4 * 6;
2690 bytes[window_at..window_at + 4].copy_from_slice(&6u32.to_le_bytes());
2691 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 #[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 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 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 #[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 #[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 #[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 #[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 .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 #[test]
2947 fn a_queued_block_is_readable_before_it_reaches_disk() {
2948 let dir = TempDir::new("buffered");
2949 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 #[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 .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 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 #[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 #[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 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 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 #[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 #[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 #[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 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 #[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 #[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 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 #[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 #[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 #[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 #[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 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 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 #[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 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 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 #[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 #[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 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 assert!(store
3528 .get(&h, &expected("model-a", 1, 4))
3529 .expect("get")
3530 .is_some());
3531
3532 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 #[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 #[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}