This repository has no description
0

Configure Feed

Select the types of activity you want to include in your feed.

core / knot2 / crates / knot-resource / src / disk.rs
4.4 kB 146 lines
1use std::io; 2use std::path::Path; 3use std::sync::Arc; 4use std::sync::atomic::{AtomicU64, Ordering}; 5 6// 3 kinds of u64 that denote "bytes" in their own way 7// & look identical at a callsite. 8// `reserve(path, floor)` *used to* compile perfectly happily. 9knot_types::scalar_newtype! { 10 pub struct DiskFloorBytes(u64); 11 pub struct ReserveBytes(u64); 12 pub struct FreeBytes(u64); 13} 14 15pub fn free_bytes(path: &Path) -> io::Result<FreeBytes> { 16 rustix::fs::statvfs(path) 17 .map(|stat| FreeBytes::new(stat.f_bavail.saturating_mul(stat.f_frsize))) 18 .map_err(io::Error::from) 19} 20 21#[derive(Debug)] 22pub enum ReserveError { 23 BelowFloor { 24 free: FreeBytes, 25 floor: DiskFloorBytes, 26 }, 27 Probe(io::Error), 28} 29 30struct Ledger { 31 floor: DiskFloorBytes, 32 reserved: AtomicU64, 33} 34 35#[derive(Clone)] 36pub struct DiskGovernor(Arc<Ledger>); 37 38impl DiskGovernor { 39 pub fn new(floor: DiskFloorBytes) -> Self { 40 Self(Arc::new(Ledger { 41 floor, 42 reserved: AtomicU64::new(0), 43 })) 44 } 45 46 pub fn reserved_bytes(&self) -> u64 { 47 self.0.reserved.load(Ordering::SeqCst) 48 } 49 50 pub fn reserve( 51 &self, 52 path: &Path, 53 bytes: ReserveBytes, 54 ) -> Result<DiskReservation, ReserveError> { 55 let amount = bytes.get(); 56 let projected = self.0.reserved.fetch_add(amount, Ordering::SeqCst) + amount; 57 let free = match free_bytes(path) { 58 Ok(free) => free, 59 Err(source) => { 60 self.0.reserved.fetch_sub(amount, Ordering::SeqCst); 61 return Err(ReserveError::Probe(source)); 62 } 63 }; 64 if free.get() < self.0.floor.get().saturating_add(projected) { 65 self.0.reserved.fetch_sub(amount, Ordering::SeqCst); 66 return Err(ReserveError::BelowFloor { 67 free, 68 floor: self.0.floor, 69 }); 70 } 71 Ok(DiskReservation { 72 ledger: Arc::clone(&self.0), 73 bytes: amount, 74 }) 75 } 76} 77 78pub struct DiskReservation { 79 ledger: Arc<Ledger>, 80 bytes: u64, 81} 82 83impl Drop for DiskReservation { 84 fn drop(&mut self) { 85 self.ledger.reserved.fetch_sub(self.bytes, Ordering::SeqCst); 86 } 87} 88 89#[cfg(test)] 90mod tests { 91 use super::*; 92 93 #[test] 94 fn a_real_filesystem_reports_headroom() { 95 let dir = std::env::temp_dir(); 96 assert!(free_bytes(&dir).unwrap().get() > 0); 97 } 98 99 #[test] 100 fn a_missing_path_reports_the_fault() { 101 assert!(free_bytes(Path::new("/definitely/not/a/mounted/path")).is_err()); 102 } 103 104 #[test] 105 fn a_reservation_holds_bytes_until_it_drops() { 106 let dir = std::env::temp_dir(); 107 let governor = DiskGovernor::new(DiskFloorBytes::new(0)); 108 assert_eq!(governor.reserved_bytes(), 0); 109 { 110 let _held = governor.reserve(&dir, ReserveBytes::new(4_096)).unwrap(); 111 assert_eq!(governor.reserved_bytes(), 4_096); 112 let _also = governor.reserve(&dir, ReserveBytes::new(1_024)).unwrap(); 113 assert_eq!(governor.reserved_bytes(), 5_120); 114 } 115 assert_eq!(governor.reserved_bytes(), 0); 116 } 117 118 #[test] 119 fn concurrent_reservations_cannot_jointly_punch_through_the_floor() { 120 let dir = std::env::temp_dir(); 121 let free = free_bytes(&dir).unwrap(); 122 let floor = DiskFloorBytes::new(free.get().saturating_sub(6_144)); 123 let governor = DiskGovernor::new(floor); 124 let first = governor.reserve(&dir, ReserveBytes::new(4_096)).unwrap(); 125 let denied = governor.reserve(&dir, ReserveBytes::new(4_096)); 126 assert!( 127 matches!(denied, Err(ReserveError::BelowFloor { .. })), 128 "the second reservation must see the first still held" 129 ); 130 assert_eq!(governor.reserved_bytes(), 4_096); 131 drop(first); 132 assert_eq!(governor.reserved_bytes(), 0); 133 assert!(governor.reserve(&dir, ReserveBytes::new(4_096)).is_ok()); 134 } 135 136 #[test] 137 fn a_probe_fault_leaves_the_ledger_untouched() { 138 let governor = DiskGovernor::new(DiskFloorBytes::new(0)); 139 let fault = governor.reserve( 140 Path::new("/definitely/not/a/mounted/path"), 141 ReserveBytes::new(4_096), 142 ); 143 assert!(matches!(fault, Err(ReserveError::Probe(_)))); 144 assert_eq!(governor.reserved_bytes(), 0); 145 } 146}