This repository has no description
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}