1use crate::block_io::BlockDevice;
40use crate::error::Result;
41use std::collections::HashMap;
42use std::sync::Mutex;
43
44struct CacheState {
55 capacity: usize,
56 entries: HashMap<u64, (Vec<u8>, u64)>,
57 pinned: HashMap<u64, Vec<u8>>,
58 next_seq: u64,
59 hits: u64,
60 misses: u64,
61}
62
63impl CacheState {
64 fn new(capacity: usize) -> Self {
65 Self {
66 capacity,
67 entries: HashMap::with_capacity(capacity.min(1024)),
68 pinned: HashMap::new(),
69 next_seq: 0,
70 hits: 0,
71 misses: 0,
72 }
73 }
74
75 fn next_seq(&mut self) -> u64 {
76 let s = self.next_seq;
77 self.next_seq = s.wrapping_add(1);
78 s
79 }
80
81 fn get(&mut self, block: u64) -> Option<Vec<u8>> {
83 if let Some(bytes) = self.pinned.get(&block) {
84 self.hits += 1;
85 return Some(bytes.clone());
86 }
87 let seq = self.next_seq();
88 if let Some(slot) = self.entries.get_mut(&block) {
89 slot.1 = seq;
90 self.hits += 1;
91 Some(slot.0.clone())
92 } else {
93 self.misses += 1;
94 None
95 }
96 }
97
98 fn put(&mut self, block: u64, bytes: Vec<u8>) {
101 if let std::collections::hash_map::Entry::Occupied(mut e) = self.pinned.entry(block) {
105 e.insert(bytes);
106 return;
107 }
108 if self.capacity == 0 {
112 return;
113 }
114 if self.entries.len() >= self.capacity {
115 if let Some((&victim, _)) = self.entries.iter().min_by_key(|(_, (_, seq))| *seq) {
116 self.entries.remove(&victim);
117 }
118 }
119 let seq = self.next_seq();
120 self.entries.insert(block, (bytes, seq));
121 }
122
123 fn pin(&mut self, block: u64, bytes: Vec<u8>) {
127 self.entries.remove(&block);
128 self.pinned.insert(block, bytes);
129 }
130
131 fn unpin_all(&mut self) {
133 let drained: Vec<(u64, Vec<u8>)> = self.pinned.drain().collect();
136 for (block, bytes) in drained {
137 if self.capacity == 0 {
138 continue;
139 }
140 if self.entries.len() >= self.capacity {
141 if let Some((&victim, _)) = self.entries.iter().min_by_key(|(_, (_, seq))| *seq) {
142 self.entries.remove(&victim);
143 }
144 }
145 let seq = self.next_seq();
146 self.entries.insert(block, (bytes, seq));
147 }
148 }
149}
150
151pub struct CachedDevice {
154 inner: std::sync::Arc<dyn BlockDevice>,
155 block_size: u32,
156 state: Mutex<CacheState>,
157}
158
159impl CachedDevice {
160 pub fn new(inner: std::sync::Arc<dyn BlockDevice>, block_size: u32, capacity: usize) -> Self {
164 Self {
165 inner,
166 block_size,
167 state: Mutex::new(CacheState::new(capacity)),
168 }
169 }
170
171 pub fn stats(&self) -> (u64, u64) {
173 let s = self.state.lock().expect("cache mutex poisoned");
174 (s.hits, s.misses)
175 }
176
177 pub fn pinned_blocks(&self) -> usize {
180 let s = self.state.lock().expect("cache mutex poisoned");
181 s.pinned.len()
182 }
183}
184
185impl BlockDevice for CachedDevice {
186 fn read_at(&self, offset: u64, buf: &mut [u8]) -> Result<()> {
187 let bs = self.block_size as u64;
188 let block = offset / bs;
189 let off_in_block = (offset % bs) as usize;
190 let len = buf.len();
191
192 if off_in_block + len <= bs as usize {
194 {
196 let mut state = self.state.lock().expect("cache mutex poisoned");
197 if let Some(blk) = state.get(block) {
198 buf.copy_from_slice(&blk[off_in_block..off_in_block + len]);
199 return Ok(());
200 }
201 }
202 let mut blk = vec![0u8; bs as usize];
204 self.inner.read_at(block * bs, &mut blk)?;
205 buf.copy_from_slice(&blk[off_in_block..off_in_block + len]);
206 let mut state = self.state.lock().expect("cache mutex poisoned");
207 state.put(block, blk);
208 return Ok(());
209 }
210
211 self.inner.read_at(offset, buf)?;
216 let state = self.state.lock().expect("cache mutex poisoned");
217 if state.pinned.is_empty() {
218 return Ok(());
219 }
220 let end = offset + len as u64;
221 for block in offset / bs..end.div_ceil(bs) {
222 let Some(img) = state.pinned.get(&block) else {
223 continue;
224 };
225 let block_start = block * bs;
226 let from = offset.max(block_start);
227 let to = end.min(block_start + bs);
228 buf[(from - offset) as usize..(to - offset) as usize]
229 .copy_from_slice(&img[(from - block_start) as usize..(to - block_start) as usize]);
230 }
231 Ok(())
232 }
233
234 fn size_bytes(&self) -> u64 {
235 self.inner.size_bytes()
236 }
237
238 fn write_at(&self, offset: u64, buf: &[u8]) -> Result<()> {
239 let bs = self.block_size as u64;
240 self.inner.write_at(offset, buf)?;
243
244 let off_in_block = (offset % bs) as usize;
249 let len = buf.len();
250 if off_in_block == 0 && len == bs as usize {
251 let block = offset / bs;
252 let mut state = self.state.lock().expect("cache mutex poisoned");
253 state.put(block, buf.to_vec());
254 return Ok(());
255 }
256
257 let first_block = offset / bs;
279 let last_block = (offset + buf.len() as u64).saturating_sub(1) / bs;
280 let buf_end_byte = offset + buf.len() as u64;
281 let mut state = self.state.lock().expect("cache mutex poisoned");
282 for b in first_block..=last_block {
283 let block_start = b * bs;
284 let block_end = block_start + bs;
285 let write_start = offset.max(block_start);
286 let write_end = buf_end_byte.min(block_end);
287 let in_block_off = (write_start - block_start) as usize;
288 let in_block_end = (write_end - block_start) as usize;
289 let buf_start = (write_start - offset) as usize;
290 let buf_end = (write_end - offset) as usize;
291
292 if let Some(img) = state.pinned.get_mut(&b) {
293 img[in_block_off..in_block_end].copy_from_slice(&buf[buf_start..buf_end]);
294 }
295 state.entries.remove(&b);
296 }
297 Ok(())
298 }
299
300 fn flush(&self) -> Result<()> {
301 self.inner.flush()
302 }
303
304 fn is_writable(&self) -> bool {
305 self.inner.is_writable()
306 }
307
308 fn populate_cache(&self, block: u64, bytes: Vec<u8>) {
309 let mut state = self.state.lock().expect("cache mutex poisoned");
310 state.pin(block, bytes);
311 }
312
313 fn unpin_all(&self) {
314 let mut state = self.state.lock().expect("cache mutex poisoned");
315 state.unpin_all();
316 }
317
318 fn invalidate_cache(&self) -> Result<()> {
319 let mut state = self.state.lock().expect("cache mutex poisoned");
320 if !state.pinned.is_empty() {
321 return Err(crate::error::Error::Corrupt(
322 "cache contains uncheckpointed journal metadata",
323 ));
324 }
325 self.inner.invalidate_cache()?;
326 state.entries.clear();
327 Ok(())
328 }
329}
330
331#[cfg(test)]
332mod tests {
333 use super::*;
334 use std::sync::Arc;
335
336 #[test]
337 fn invalidation_requires_checkpoint_and_then_reads_physical_bytes() {
338 let raw = CountingDevice::new(8192, true);
339 let cache = CachedDevice::new(raw.clone(), 4096, 8);
340 let mut bytes = vec![0; 4096];
341 cache.read_at(0, &mut bytes).unwrap();
342 raw.write_at(0, &vec![7; 4096]).unwrap();
343 cache.read_at(0, &mut bytes).unwrap();
344 assert_eq!(bytes[0], 0);
345 cache.invalidate_cache().unwrap();
346 cache.read_at(0, &mut bytes).unwrap();
347 assert_eq!(bytes[0], 7);
348 cache.populate_cache(0, vec![9; 4096]);
349 assert!(
350 cache.invalidate_cache().is_err(),
351 "uncheckpointed metadata cannot be discarded"
352 );
353 cache.read_at(0, &mut bytes).unwrap();
354 assert_eq!(bytes[0], 9);
355 raw.write_at(0, &vec![9; 4096]).unwrap();
356 cache.unpin_all();
357 cache.invalidate_cache().unwrap();
358 cache.read_at(0, &mut bytes).unwrap();
359 assert_eq!(bytes[0], 9);
360 assert_eq!(raw.reads(), 3);
361 }
362
363 struct CountingDevice {
365 bytes: Mutex<Vec<u8>>,
366 reads: std::sync::atomic::AtomicU64,
367 writes: std::sync::atomic::AtomicU64,
368 writable: bool,
369 }
370
371 impl CountingDevice {
372 fn new(size: usize, writable: bool) -> Arc<Self> {
373 Arc::new(Self {
374 bytes: Mutex::new(vec![0u8; size]),
375 reads: std::sync::atomic::AtomicU64::new(0),
376 writes: std::sync::atomic::AtomicU64::new(0),
377 writable,
378 })
379 }
380 fn reads(&self) -> u64 {
381 self.reads.load(std::sync::atomic::Ordering::SeqCst)
382 }
383 fn writes(&self) -> u64 {
384 self.writes.load(std::sync::atomic::Ordering::SeqCst)
385 }
386 }
387
388 impl BlockDevice for CountingDevice {
389 fn read_at(&self, offset: u64, buf: &mut [u8]) -> Result<()> {
390 self.reads.fetch_add(1, std::sync::atomic::Ordering::SeqCst);
391 let b = self.bytes.lock().unwrap();
392 let off = offset as usize;
393 buf.copy_from_slice(&b[off..off + buf.len()]);
394 Ok(())
395 }
396 fn size_bytes(&self) -> u64 {
397 self.bytes.lock().unwrap().len() as u64
398 }
399 fn write_at(&self, offset: u64, buf: &[u8]) -> Result<()> {
400 self.writes
401 .fetch_add(1, std::sync::atomic::Ordering::SeqCst);
402 let mut b = self.bytes.lock().unwrap();
403 let off = offset as usize;
404 b[off..off + buf.len()].copy_from_slice(buf);
405 Ok(())
406 }
407 fn is_writable(&self) -> bool {
408 self.writable
409 }
410 }
411
412 #[test]
416 fn a_multi_block_read_sees_pinned_blocks() {
417 let inner = CountingDevice::new(4096 * 16, false);
418 let cached = CachedDevice::new(inner.clone(), 4096, 8);
419 cached.populate_cache(3, vec![0xBBu8; 4096]);
420 let mut buf = vec![0u8; 8192];
422 cached.read_at(2 * 4096 + 2048, &mut buf).unwrap();
423 assert_eq!(&buf[..2048], &[0u8; 2048][..], "block 2 is the device's");
424 assert_eq!(&buf[2048..6144], &[0xBBu8; 4096][..], "block 3 is pinned");
425 assert_eq!(&buf[6144..], &[0u8; 2048][..], "block 4 is the device's");
426 }
427
428 #[test]
429 fn second_read_of_same_block_is_a_cache_hit() {
430 let inner = CountingDevice::new(4096 * 16, false);
431 let cached = CachedDevice::new(inner.clone(), 4096, 8);
432 let mut buf = vec![0u8; 100];
433 cached.read_at(0, &mut buf).unwrap();
434 cached.read_at(0, &mut buf).unwrap();
435 cached.read_at(50, &mut buf).unwrap(); assert_eq!(inner.reads(), 1, "cache should serve all 3 from one read");
437 let (hits, misses) = cached.stats();
438 assert_eq!(hits, 2);
439 assert_eq!(misses, 1);
440 }
441
442 #[test]
443 fn unaligned_write_merges_into_pinned_block() {
444 let inner = CountingDevice::new(4096 * 16, true);
451 let cached = CachedDevice::new(inner.clone(), 4096, 8);
452 let mut journaled = vec![0u8; 4096];
455 journaled[0] = 0xAA;
456 journaled[3000] = 0xBB;
457 cached.populate_cache(5, journaled);
458 cached.write_at(5 * 4096 + 100, &[0xCCu8; 100]).unwrap();
460 let mut buf = vec![0u8; 4096];
462 for i in 0..4096 {
463 cached
464 .read_at(5 * 4096 + i as u64, &mut buf[i..i + 1])
465 .unwrap();
466 }
467 assert_eq!(buf[0], 0xAA, "pinned journaled byte 0 must survive");
468 assert_eq!(buf[100], 0xCC, "new sub-block bytes must be visible");
469 assert_eq!(buf[199], 0xCC);
470 assert_eq!(buf[200], 0x00, "untouched portion stays at journaled value");
471 assert_eq!(
472 buf[3000], 0xBB,
473 "pinned bytes outside the write window survive"
474 );
475 }
476
477 #[test]
478 fn unaligned_write_invalidates_cache_entry() {
479 let inner = CountingDevice::new(4096 * 16, true);
483 let cached = CachedDevice::new(inner.clone(), 4096, 8);
484 let mut buf = vec![0u8; 100];
485 cached.read_at(0, &mut buf).unwrap();
486 cached.write_at(0, &[42u8; 100]).unwrap(); cached.read_at(0, &mut buf).unwrap();
488 assert_eq!(inner.reads(), 2);
489 assert_eq!(buf[0], 42, "post-write read should see the new bytes");
490 assert_eq!(
495 inner.writes(),
496 1,
497 "an unaligned write must reach the device once, not be held in cache"
498 );
499 }
500
501 #[test]
502 fn aligned_write_updates_cache_entry() {
503 let inner = CountingDevice::new(4096 * 16, true);
508 let cached = CachedDevice::new(inner.clone(), 4096, 8);
509 let mut buf = vec![0u8; 100];
511 cached.read_at(0, &mut buf).unwrap();
512 assert_eq!(inner.reads(), 1);
513 let new_block = vec![0xCDu8; 4096];
515 cached.write_at(0, &new_block).unwrap();
516 let mut readback = vec![0u8; 100];
518 cached.read_at(0, &mut readback).unwrap();
519 assert_eq!(inner.reads(), 1, "aligned write-through skips disk read");
520 assert_eq!(readback[0], 0xCD);
521 let mut from_disk = vec![0u8; 100];
523 inner.read_at(0, &mut from_disk).unwrap();
524 assert_eq!(from_disk[0], 0xCD);
525 }
526
527 #[test]
528 fn populate_cache_pins_entry_against_lru() {
529 let inner = CountingDevice::new(4096 * 16, false);
533 let cached = CachedDevice::new(inner.clone(), 4096, 2); cached.populate_cache(7, vec![0xAAu8; 4096]);
536 let mut throwaway = vec![0u8; 8];
538 for blk in 0..5u64 {
539 cached.read_at(blk * 4096, &mut throwaway).unwrap();
540 }
541 let mut buf = vec![0u8; 8];
544 cached.read_at(7 * 4096, &mut buf).unwrap();
545 assert_eq!(buf, vec![0xAA; 8],
546 "pinned entry must survive LRU pressure — otherwise journaled writes vanish before checkpoint");
547 }
548
549 #[test]
550 fn unpin_all_lets_pinned_entries_evict_normally() {
551 let inner = CountingDevice::new(4096 * 16, false);
555 let cached = CachedDevice::new(inner.clone(), 4096, 1); cached.populate_cache(3, vec![0xBBu8; 4096]);
557 cached.unpin_all();
558 let mut throwaway = vec![0u8; 8];
561 cached.read_at(0, &mut throwaway).unwrap();
562 let mut buf3 = vec![0u8; 8];
572 cached.read_at(3 * 4096, &mut buf3).unwrap();
573 assert_eq!(
574 buf3,
575 vec![0; 8],
576 "after unpin + LRU eviction, the inner device's bytes win"
577 );
578 }
579
580 #[test]
581 fn multi_block_read_bypasses_cache() {
582 let inner = CountingDevice::new(4096 * 16, false);
583 let cached = CachedDevice::new(inner.clone(), 4096, 8);
584 let mut buf = vec![0u8; 8000]; cached.read_at(0, &mut buf).unwrap();
586 cached.read_at(0, &mut buf).unwrap();
588 assert_eq!(inner.reads(), 2);
589 let (hits, misses) = cached.stats();
590 assert_eq!(hits, 0);
591 assert_eq!(misses, 0, "multi-block reads bypass entirely");
592 }
593
594 #[test]
595 fn lru_evicts_oldest_when_capacity_exceeded() {
596 let inner = CountingDevice::new(4096 * 16, false);
597 let cached = CachedDevice::new(inner.clone(), 4096, 2); let mut buf = vec![0u8; 8];
599 for blk in 0..3u64 {
601 cached.read_at(blk * 4096, &mut buf).unwrap();
602 }
603 cached.read_at(0, &mut buf).unwrap();
605 cached.read_at(2 * 4096, &mut buf).unwrap();
607 let (hits, misses) = cached.stats();
608 assert_eq!(misses, 4, "blocks 0,1,2 + re-read of 0 (evicted)");
609 assert_eq!(hits, 1, "re-read of 2 still cached");
610 }
611
612 #[test]
613 fn lru_keeps_recently_touched_block_alive() {
614 let inner = CountingDevice::new(4096 * 16, false);
618 let cached = CachedDevice::new(inner.clone(), 4096, 2);
619 let mut buf = vec![0u8; 8];
620 cached.read_at(0, &mut buf).unwrap(); cached.read_at(4096, &mut buf).unwrap(); cached.read_at(0, &mut buf).unwrap(); cached.read_at(2 * 4096, &mut buf).unwrap(); cached.read_at(0, &mut buf).unwrap(); let (hits, misses) = cached.stats();
626 assert_eq!(hits, 2);
627 assert_eq!(misses, 3);
628 cached.read_at(4096, &mut buf).unwrap(); let (_, m_after) = cached.stats();
631 assert_eq!(m_after, 4, "block 1 was evicted, re-read is a miss");
632 }
633}