use std::fs::{self, File};
use std::io::{Read, Seek, SeekFrom, Write};
use std::path::{Path, PathBuf};
use std::sync::atomic::{AtomicU64, Ordering};
use spg_engine::{TempRun, TempStoreError};
const RUN_PREFIX: &str = "spg-sort-";
static RUN_SERIAL: AtomicU64 = AtomicU64::new(0);
fn io_err(what: &str, e: &std::io::Error) -> TempStoreError {
TempStoreError::Io(format!("{what}: {e}"))
}
pub fn temp_dir() -> PathBuf {
std::env::var_os("SPG_TEMP_DIR").map_or_else(std::env::temp_dir, PathBuf::from)
}
const RUN_BUF_BYTES: usize = 256 * 1024;
pub struct FileRun {
pub(crate) path: PathBuf,
file: File,
buf: Vec<u8>,
read_pos: usize,
written: u64,
sealed: bool,
}
impl TempRun for FileRun {
fn append(&mut self, bytes: &[u8]) -> Result<(), TempStoreError> {
debug_assert!(!self.sealed, "append after seal");
self.buf.extend_from_slice(bytes);
if self.buf.len() >= RUN_BUF_BYTES {
self.file
.write_all(&self.buf)
.map_err(|e| io_err("writing spill run", &e))?;
self.buf.clear();
}
self.written += bytes.len() as u64;
Ok(())
}
fn seal(&mut self) -> Result<(), TempStoreError> {
if !self.buf.is_empty() {
self.file
.write_all(&self.buf)
.map_err(|e| io_err("writing spill run", &e))?;
self.buf.clear();
}
self.read_pos = 0;
self.file
.flush()
.map_err(|e| io_err("flushing spill run", &e))?;
self.file
.seek(SeekFrom::Start(0))
.map_err(|e| io_err("rewinding spill run", &e))?;
self.sealed = true;
Ok(())
}
fn read(&mut self, buf: &mut [u8]) -> Result<usize, TempStoreError> {
debug_assert!(self.sealed, "read before seal");
if self.read_pos == self.buf.len() {
self.buf.resize(RUN_BUF_BYTES, 0);
let got = self
.file
.read(&mut self.buf)
.map_err(|e| io_err("reading spill run", &e))?;
self.buf.truncate(got);
self.read_pos = 0;
if got == 0 {
return Ok(0);
}
}
let n = core::cmp::min(buf.len(), self.buf.len() - self.read_pos);
buf[..n].copy_from_slice(&self.buf[self.read_pos..self.read_pos + n]);
self.read_pos += n;
Ok(n)
}
fn bytes_written(&self) -> u64 {
self.written
}
}
impl Drop for FileRun {
fn drop(&mut self) {
let _ = fs::remove_file(&self.path);
}
}
pub fn create_run() -> Result<Box<dyn TempRun>, TempStoreError> {
Ok(Box::new(create_run_in(&temp_dir())?))
}
pub fn create_run_in(dir: &Path) -> Result<FileRun, TempStoreError> {
fs::create_dir_all(dir).map_err(|e| io_err("creating temp dir", &e))?;
let serial = RUN_SERIAL.fetch_add(1, Ordering::Relaxed);
let path = dir.join(format!("{RUN_PREFIX}{}-{serial}.run", std::process::id()));
let file = fs::OpenOptions::new()
.create(true)
.truncate(true)
.read(true)
.write(true)
.open(&path)
.map_err(|e| io_err("opening spill run", &e))?;
Ok(FileRun {
path,
file,
buf: Vec::with_capacity(RUN_BUF_BYTES),
read_pos: 0,
written: 0,
sealed: false,
})
}
pub fn sweep_orphans(dir: &Path) -> usize {
let mine = format!("{RUN_PREFIX}{}-", std::process::id());
let Ok(entries) = fs::read_dir(dir) else {
return 0;
};
let mut removed = 0;
for e in entries.flatten() {
let name = e.file_name();
let Some(name) = name.to_str() else { continue };
if name.starts_with(RUN_PREFIX)
&& !name.starts_with(&mine)
&& fs::remove_file(e.path()).is_ok()
{
removed += 1;
}
}
removed
}