This repository has no description
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 struct BelowFloor {
23 pub free: FreeBytes,
24 pub floor: DiskFloorBytes,
25}
26
27struct Ledger {
28 floor: DiskFloorBytes,
29 reserved: AtomicU64,
30}
31
32#[derive(Clone)]
33pub struct DiskGovernor(Arc<Ledger>);
34
35impl DiskGovernor {
36 pub fn new(floor: DiskFloorBytes) -> Self {
37 Self(Arc::new(Ledger {
38 floor,
39 reserved: AtomicU64::new(0),
40 }))
41 }
42
43 pub fn reserved_bytes(&self) -> u64 {
44 self.0.reserved.load(Ordering::SeqCst)
45 }
46
47 pub fn reserve_against(
48 &self,
49 free: FreeBytes,
50 bytes: ReserveBytes,
51 ) -> Result<DiskReservation, BelowFloor> {
52 let amount = bytes.get();
53 let projected = self.0.reserved.fetch_add(amount, Ordering::SeqCst) + amount;
54 if free.get() < self.0.floor.get().saturating_add(projected) {
55 self.0.reserved.fetch_sub(amount, Ordering::SeqCst);
56 return Err(BelowFloor {
57 free,
58 floor: self.0.floor,
59 });
60 }
61 Ok(DiskReservation {
62 ledger: Arc::clone(&self.0),
63 bytes: amount,
64 })
65 }
66}
67
68pub struct DiskReservation {
69 ledger: Arc<Ledger>,
70 bytes: u64,
71}
72
73impl Drop for DiskReservation {
74 fn drop(&mut self) {
75 self.ledger.reserved.fetch_sub(self.bytes, Ordering::SeqCst);
76 }
77}
78
79#[cfg(test)]
80mod tests {
81 use super::*;
82
83 #[test]
84 fn a_real_filesystem_reports_headroom() {
85 let dir = std::env::temp_dir();
86 assert!(free_bytes(&dir).unwrap().get() > 0);
87 }
88
89 #[test]
90 fn a_missing_path_reports_the_fault() {
91 assert!(free_bytes(Path::new("/definitely/not/a/mounted/path")).is_err());
92 }
93
94 #[test]
95 fn a_reservation_holds_bytes_until_it_drops() {
96 let free = FreeBytes::new(1 << 20);
97 let governor = DiskGovernor::new(DiskFloorBytes::new(0));
98 assert_eq!(governor.reserved_bytes(), 0);
99 {
100 let _held = governor
101 .reserve_against(free, ReserveBytes::new(4_096))
102 .unwrap();
103 assert_eq!(governor.reserved_bytes(), 4_096);
104 let _also = governor
105 .reserve_against(free, ReserveBytes::new(1_024))
106 .unwrap();
107 assert_eq!(governor.reserved_bytes(), 5_120);
108 }
109 assert_eq!(governor.reserved_bytes(), 0);
110 }
111
112 #[test]
113 fn concurrent_reservations_cannot_jointly_punch_through_the_floor() {
114 let free = FreeBytes::new(10_240);
115 let governor = DiskGovernor::new(DiskFloorBytes::new(4_096));
116 let first = governor
117 .reserve_against(free, ReserveBytes::new(4_096))
118 .unwrap();
119 let denied = governor.reserve_against(free, ReserveBytes::new(4_096));
120 assert!(
121 denied.is_err(),
122 "the second reservation must see the first still held"
123 );
124 assert_eq!(governor.reserved_bytes(), 4_096);
125 drop(first);
126 assert_eq!(governor.reserved_bytes(), 0);
127 assert!(
128 governor
129 .reserve_against(free, ReserveBytes::new(4_096))
130 .is_ok(),
131 "a free-space reading the floor leaves room in must admit a lone reservation, or \
132 the refusal above was the floor rather than the reservation still being held"
133 );
134 }
135
136 #[test]
137 fn a_refused_reservation_leaves_the_ledger_untouched() {
138 let governor = DiskGovernor::new(DiskFloorBytes::new(4_096));
139 let refused = governor.reserve_against(FreeBytes::new(4_096), ReserveBytes::new(1));
140 assert!(refused.is_err());
141 assert_eq!(
142 governor.reserved_bytes(),
143 0,
144 "a refusal that left its bytes on the ledger would deny every later upload too"
145 );
146 }
147}