Something went wrong. Try again.
Monorepo for Tangled
Something went wrong. Try again.
123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147use std::io;use std::path::Path;use std::sync::Arc;use std::sync::atomic::{AtomicU64, Ordering};
// 3 kinds of u64 that denote "bytes" in their own way// & look identical at a callsite.// `reserve(path, floor)` *used to* compile perfectly happily.knot_types::scalar_newtype! { pub struct DiskFloorBytes(u64); pub struct ReserveBytes(u64); pub struct FreeBytes(u64);}
pub fn free_bytes(path: &Path) -> io::Result<FreeBytes> { rustix::fs::statvfs(path) .map(|stat| FreeBytes::new(stat.f_bavail.saturating_mul(stat.f_frsize))) .map_err(io::Error::from)}
#[derive(Debug)]pub enum ReserveError { BelowFloor { free: FreeBytes, floor: DiskFloorBytes, }, Probe(io::Error),}
struct Ledger { floor: DiskFloorBytes, reserved: AtomicU64,}
#[derive(Clone)]pub struct DiskGovernor(Arc<Ledger>);
impl DiskGovernor { pub fn new(floor: DiskFloorBytes) -> Self { Self(Arc::new(Ledger { floor, reserved: AtomicU64::new(0), })) }
pub fn reserved_bytes(&self) -> u64 { self.0.reserved.load(Ordering::SeqCst) }
pub fn reserve( &self, path: &Path, bytes: ReserveBytes, ) -> Result<DiskReservation, ReserveError> { let amount = bytes.get(); let projected = self.0.reserved.fetch_add(amount, Ordering::SeqCst) + amount; let free = match free_bytes(path) { Ok(free) => free, Err(source) => { self.0.reserved.fetch_sub(amount, Ordering::SeqCst); return Err(ReserveError::Probe(source)); } }; if free.get() < self.0.floor.get().saturating_add(projected) { self.0.reserved.fetch_sub(amount, Ordering::SeqCst); return Err(ReserveError::BelowFloor { free, floor: self.0.floor, }); } Ok(DiskReservation { ledger: Arc::clone(&self.0), bytes: amount, }) }}
pub struct DiskReservation { ledger: Arc<Ledger>, bytes: u64,}
impl Drop for DiskReservation { fn drop(&mut self) { self.ledger.reserved.fetch_sub(self.bytes, Ordering::SeqCst); }}
#[cfg(test)]mod tests { use super::*;
#[test] fn a_real_filesystem_reports_headroom() { let dir = std::env::temp_dir(); assert!(free_bytes(&dir).unwrap().get() > 0); }
#[test] fn a_missing_path_reports_the_fault() { assert!(free_bytes(Path::new("/definitely/not/a/mounted/path")).is_err()); }
#[test] fn a_reservation_holds_bytes_until_it_drops() { let dir = std::env::temp_dir(); let governor = DiskGovernor::new(DiskFloorBytes::new(0)); assert_eq!(governor.reserved_bytes(), 0); { let _held = governor.reserve(&dir, ReserveBytes::new(4_096)).unwrap(); assert_eq!(governor.reserved_bytes(), 4_096); let _also = governor.reserve(&dir, ReserveBytes::new(1_024)).unwrap(); assert_eq!(governor.reserved_bytes(), 5_120); } assert_eq!(governor.reserved_bytes(), 0); }
#[test] fn concurrent_reservations_cannot_jointly_punch_through_the_floor() { let dir = std::env::temp_dir(); let free = free_bytes(&dir).unwrap(); let floor = DiskFloorBytes::new(free.get().saturating_sub(6_144)); let governor = DiskGovernor::new(floor); let first = governor.reserve(&dir, ReserveBytes::new(4_096)).unwrap(); let denied = governor.reserve(&dir, ReserveBytes::new(4_096)); assert!( matches!(denied, Err(ReserveError::BelowFloor { .. })), "the second reservation must see the first still held" ); assert_eq!(governor.reserved_bytes(), 4_096); drop(first); assert_eq!(governor.reserved_bytes(), 0); assert!(governor.reserve(&dir, ReserveBytes::new(4_096)).is_ok()); }
#[test] fn a_probe_fault_leaves_the_ledger_untouched() { let governor = DiskGovernor::new(DiskFloorBytes::new(0)); let fault = governor.reserve( Path::new("/definitely/not/a/mounted/path"), ReserveBytes::new(4_096), ); assert!(matches!(fault, Err(ReserveError::Probe(_)))); assert_eq!(governor.reserved_bytes(), 0); }}