1use super::TempStoragePolicy;
4use std::fmt;
5use std::fs::{self, File, OpenOptions};
6use std::io::{self, Read, Seek, SeekFrom, Write};
7use std::path::{Path, PathBuf};
8use std::sync::Arc;
9use std::sync::atomic::{AtomicU64, Ordering};
10
11const FILE_MAGIC: [u8; 8] = *b"PDBXSP01";
12pub(crate) const FILE_HEADER_BYTES: u64 = FILE_MAGIC.len() as u64;
13const RECORD_HEADER_BYTES: usize = 16;
14pub(crate) const RECORD_HEADER_DISK_BYTES: u64 = 16;
15const RECORD_HEADER_REWIND: i64 = -16;
16const HASH_OFFSET: u64 = 0xcbf2_9ce4_8422_2325;
17const HASH_PRIME: u64 = 0x0000_0100_0000_01b3;
18
19#[derive(Debug)]
20pub(crate) struct DiskAccount {
21 root: PathBuf,
22 maximum: u64,
23 used: AtomicU64,
24 spilled: AtomicU64,
25}
26
27impl DiskAccount {
28 pub(crate) fn from_policy(policy: &TempStoragePolicy) -> Option<Arc<Self>> {
29 match policy {
30 TempStoragePolicy::Disabled => None,
31 TempStoragePolicy::Directory { root, max_bytes } => Some(Arc::new(Self {
32 root: root.clone(),
33 maximum: *max_bytes,
34 used: AtomicU64::new(0),
35 spilled: AtomicU64::new(0),
36 })),
37 }
38 }
39
40 fn try_charge(&self, bytes: u64) -> Result<(), SpillError> {
41 let mut current = self.used.load(Ordering::Acquire);
42 loop {
43 let Some(next) = current.checked_add(bytes) else {
44 return Err(self.exhausted(bytes, current));
45 };
46 if next > self.maximum {
47 return Err(self.exhausted(bytes, current));
48 }
49 match self.used.compare_exchange_weak(
50 current,
51 next,
52 Ordering::AcqRel,
53 Ordering::Acquire,
54 ) {
55 Ok(_) => return Ok(()),
56 Err(observed) => current = observed,
57 }
58 }
59 }
60
61 fn exhausted(&self, requested: u64, used: u64) -> SpillError {
62 SpillError::BudgetExceeded {
63 requested,
64 available: self.maximum.saturating_sub(used),
65 }
66 }
67
68 fn release(&self, bytes: u64) {
69 self.used.fetch_sub(bytes, Ordering::AcqRel);
70 }
71
72 pub(crate) fn used_bytes(&self) -> u64 {
73 self.used.load(Ordering::Acquire)
74 }
75
76 pub(crate) fn spilled_bytes(&self) -> u64 {
77 self.spilled.load(Ordering::Acquire)
78 }
79}
80
81#[derive(Debug)]
83pub enum SpillError {
84 Disabled,
86 InvalidKey,
88 BudgetExceeded {
90 requested: u64,
92 available: u64,
94 },
95 RecordTooLarge(u64),
97 RecordExceedsLimit {
99 record: u64,
101 limit: usize,
103 },
104 CorruptHeader,
106 CorruptRecord,
108 Io(io::Error),
110}
111
112impl fmt::Display for SpillError {
113 fn fmt(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result {
114 match self {
115 Self::Disabled => formatter.write_str("temporary disk storage is disabled"),
116 Self::InvalidKey => formatter.write_str("spill key must be one portable file name"),
117 Self::BudgetExceeded {
118 requested,
119 available,
120 } => write!(
121 formatter,
122 "spill requested {requested} bytes with {available} bytes available"
123 ),
124 Self::RecordTooLarge(bytes) => {
125 write!(
126 formatter,
127 "spill record of {bytes} bytes is not addressable"
128 )
129 }
130 Self::RecordExceedsLimit { record, limit } => {
131 write!(
132 formatter,
133 "spill record has {record} bytes but limit is {limit}"
134 )
135 }
136 Self::CorruptHeader => formatter.write_str("spill file header is invalid"),
137 Self::CorruptRecord => formatter.write_str("spill record is truncated or corrupt"),
138 Self::Io(error) => error.fmt(formatter),
139 }
140 }
141}
142
143impl std::error::Error for SpillError {
144 fn source(&self) -> Option<&(dyn std::error::Error + 'static)> {
145 match self {
146 Self::Io(error) => Some(error),
147 _ => None,
148 }
149 }
150}
151
152impl From<io::Error> for SpillError {
153 fn from(error: io::Error) -> Self {
154 Self::Io(error)
155 }
156}
157
158#[derive(Debug)]
160pub struct SpillFile {
161 file: Option<File>,
162 path: Option<PathBuf>,
163 account: Arc<DiskAccount>,
164 charged: u64,
165 records: u64,
166}
167
168impl SpillFile {
169 pub(crate) fn create(account: Arc<DiskAccount>, key: &str) -> Result<Self, SpillError> {
170 if !portable_key(key) {
171 return Err(SpillError::InvalidKey);
172 }
173 fs::create_dir_all(&account.root)?;
174 account.try_charge(FILE_HEADER_BYTES)?;
175 let path = account.root.join(format!("{key}.molframe-spill"));
176 let opened = OpenOptions::new().write(true).create_new(true).open(&path);
177 let mut file = match opened {
178 Ok(file) => file,
179 Err(error) => {
180 account.release(FILE_HEADER_BYTES);
181 return Err(SpillError::Io(error));
182 }
183 };
184 if let Err(error) = file.write_all(&FILE_MAGIC) {
185 account.release(FILE_HEADER_BYTES);
186 let _ignored = fs::remove_file(&path);
187 return Err(SpillError::Io(error));
188 }
189 account
190 .spilled
191 .fetch_add(FILE_HEADER_BYTES, Ordering::AcqRel);
192 Ok(Self {
193 file: Some(file),
194 path: Some(path),
195 account,
196 charged: FILE_HEADER_BYTES,
197 records: 0,
198 })
199 }
200
201 pub fn append(&mut self, payload: &[u8]) -> Result<(), SpillError> {
207 let payload_bytes =
208 u64::try_from(payload.len()).map_err(|_error| SpillError::RecordTooLarge(u64::MAX))?;
209 let Some(frame_bytes) = payload_bytes.checked_add(RECORD_HEADER_DISK_BYTES) else {
210 return Err(SpillError::RecordTooLarge(payload_bytes));
211 };
212 self.account.try_charge(frame_bytes)?;
213 let Some(file) = &mut self.file else {
214 self.account.release(frame_bytes);
215 return Err(SpillError::CorruptHeader);
216 };
217 let original_len = self.charged;
218 let mut header = [0_u8; RECORD_HEADER_BYTES];
219 header[..8].copy_from_slice(&payload_bytes.to_le_bytes());
220 header[8..].copy_from_slice(&checksum(payload).to_le_bytes());
221 if let Err(error) = file
222 .write_all(&header)
223 .and_then(|()| file.write_all(payload))
224 {
225 let _ignored = file.set_len(original_len);
226 let _ignored = file.seek(SeekFrom::End(0));
227 self.account.release(frame_bytes);
228 return Err(SpillError::Io(error));
229 }
230 self.charged += frame_bytes;
231 self.records += 1;
232 self.account
233 .spilled
234 .fetch_add(frame_bytes, Ordering::AcqRel);
235 Ok(())
236 }
237
238 pub fn finish(mut self) -> Result<SpillArtifact, SpillError> {
244 let Some(mut file) = self.file.take() else {
245 return Err(SpillError::CorruptHeader);
246 };
247 file.flush()?;
248 file.sync_data()?;
249 drop(file);
250 let Some(path) = self.path.take() else {
251 return Err(SpillError::CorruptHeader);
252 };
253 let artifact = SpillArtifact {
254 path,
255 account: Arc::clone(&self.account),
256 charged: self.charged,
257 records: self.records,
258 };
259 self.charged = 0;
260 Ok(artifact)
261 }
262}
263
264impl Drop for SpillFile {
265 fn drop(&mut self) {
266 if let Some(path) = self.path.take() {
267 let _ignored = fs::remove_file(path);
268 }
269 self.account.release(self.charged);
270 }
271}
272
273#[derive(Debug)]
275pub struct SpillArtifact {
276 path: PathBuf,
277 account: Arc<DiskAccount>,
278 charged: u64,
279 records: u64,
280}
281
282impl SpillArtifact {
283 pub fn reader(&self) -> Result<SpillReader, SpillError> {
289 SpillReader::open(&self.path)
290 }
291
292 #[must_use]
294 pub const fn bytes(&self) -> u64 {
295 self.charged
296 }
297
298 #[must_use]
300 pub const fn records(&self) -> u64 {
301 self.records
302 }
303
304 #[must_use]
306 pub fn path(&self) -> &Path {
307 &self.path
308 }
309}
310
311impl Drop for SpillArtifact {
312 fn drop(&mut self) {
313 let _ignored = fs::remove_file(&self.path);
314 self.account.release(self.charged);
315 }
316}
317
318#[derive(Debug)]
320pub struct SpillReader {
321 file: File,
322}
323
324impl SpillReader {
325 fn open(path: &Path) -> Result<Self, SpillError> {
326 let mut file = File::open(path)?;
327 let mut magic = [0_u8; FILE_MAGIC.len()];
328 file.read_exact(&mut magic)
329 .map_err(|_error| SpillError::CorruptHeader)?;
330 if magic != FILE_MAGIC {
331 return Err(SpillError::CorruptHeader);
332 }
333 Ok(Self { file })
334 }
335
336 pub fn read_next(
342 &mut self,
343 payload: &mut Vec<u8>,
344 max_bytes: usize,
345 ) -> Result<Option<usize>, SpillError> {
346 let mut header = [0_u8; RECORD_HEADER_BYTES];
347 let mut first = [0_u8; 1];
348 if self.file.read(&mut first)? == 0 {
349 return Ok(None);
350 }
351 header[0] = first[0];
352 self.file
353 .read_exact(&mut header[1..])
354 .map_err(|_error| SpillError::CorruptRecord)?;
355 let length_u64 = u64::from_le_bytes(copy_u64(&header[..8]));
356 let expected = u64::from_le_bytes(copy_u64(&header[8..]));
357 let payload_offset = self.file.stream_position()?;
358 let file_bytes = self.file.metadata()?.len();
359 if length_u64 > file_bytes.saturating_sub(payload_offset) {
360 return Err(SpillError::CorruptRecord);
361 }
362 let Ok(limit_u64) = u64::try_from(max_bytes) else {
363 return Err(SpillError::RecordTooLarge(length_u64));
364 };
365 if length_u64 > limit_u64 {
366 self.file.seek(SeekFrom::Current(RECORD_HEADER_REWIND))?;
367 return Err(SpillError::RecordExceedsLimit {
368 record: length_u64,
369 limit: max_bytes,
370 });
371 }
372 let length =
373 usize::try_from(length_u64).map_err(|_error| SpillError::RecordTooLarge(length_u64))?;
374 payload.resize(length, 0);
375 self.file
376 .read_exact(payload)
377 .map_err(|_error| SpillError::CorruptRecord)?;
378 if checksum(payload) != expected {
379 return Err(SpillError::CorruptRecord);
380 }
381 Ok(Some(length))
382 }
383}
384
385fn portable_key(key: &str) -> bool {
386 !key.is_empty()
387 && key.len() <= 96
388 && key
389 .bytes()
390 .all(|byte| byte.is_ascii_alphanumeric() || matches!(byte, b'-' | b'_'))
391}
392
393pub(crate) fn checksum(bytes: &[u8]) -> u64 {
394 bytes.iter().fold(HASH_OFFSET, |hash, byte| {
395 (hash ^ u64::from(*byte)).wrapping_mul(HASH_PRIME)
396 })
397}
398
399fn copy_u64(bytes: &[u8]) -> [u8; 8] {
400 let mut value = [0_u8; 8];
401 value.copy_from_slice(bytes);
402 value
403}
404
405#[cfg(test)]
406#[path = "spill_tests.rs"]
407mod tests;