use std::fs::{self, File, OpenOptions};
use std::io::{self, BufWriter, Write};
use std::path::{Path, PathBuf};
pub struct AtomicFileWriter {
writer: BufWriter<File>,
tmp_path: PathBuf,
dest_path: PathBuf,
finished: bool,
}
impl AtomicFileWriter {
pub fn new(dest: impl AsRef<Path>) -> io::Result<Self> {
let dest_path = dest.as_ref().to_path_buf();
let dir = dest_path.parent().unwrap_or(Path::new("."));
let base_name = dest_path
.file_name()
.and_then(|n| n.to_str())
.unwrap_or("out");
let random_suffix: u64 = rand::random();
let tmp_name = format!(".{}.{:016x}.tmp", base_name, random_suffix);
let tmp_path = dir.join(tmp_name);
let file = OpenOptions::new()
.write(true)
.create_new(true)
.open(&tmp_path)?;
Ok(Self {
writer: BufWriter::new(file),
tmp_path,
dest_path,
finished: false,
})
}
pub fn finish(mut self) -> io::Result<()> {
self.writer.flush()?;
self.writer.get_ref().sync_all()?;
if let Err(e) = fs::rename(&self.tmp_path, &self.dest_path) {
let _ = fs::remove_file(&self.tmp_path);
return Err(e);
}
self.finished = true;
Ok(())
}
#[must_use]
pub fn tmp_path(&self) -> &Path {
&self.tmp_path
}
#[must_use]
pub fn dest_path(&self) -> &Path {
&self.dest_path
}
}
impl Write for AtomicFileWriter {
fn write(&mut self, buf: &[u8]) -> io::Result<usize> {
self.writer.write(buf)
}
fn flush(&mut self) -> io::Result<()> {
self.writer.flush()
}
}
impl Drop for AtomicFileWriter {
fn drop(&mut self) {
if !self.finished {
let _ = fs::remove_file(&self.tmp_path);
}
}
}
pub fn atomic_write(dest: impl AsRef<Path>, data: &[u8]) -> io::Result<()> {
let mut writer = AtomicFileWriter::new(dest)?;
writer.write_all(data)?;
writer.finish()
}
#[cfg(test)]
mod tests {
use super::*;
use std::fs;
#[test]
fn atomic_write_creates_file() {
let dir = tempfile::tempdir().unwrap();
let dest = dir.path().join("output.txt");
atomic_write(&dest, b"hello world").unwrap();
assert_eq!(fs::read_to_string(&dest).unwrap(), "hello world");
let tmp = dir.path().join("output.txt.tmp");
assert!(!tmp.exists());
}
#[test]
fn atomic_writer_drop_cleans_up() {
let dir = tempfile::tempdir().unwrap();
let dest = dir.path().join("output.txt");
{
let mut w = AtomicFileWriter::new(&dest).unwrap();
w.write_all(b"partial").unwrap();
}
assert!(!dest.exists(), "dest should not exist after aborted write");
let tmp = dir.path().join("output.txt.tmp");
assert!(!tmp.exists(), "temp file should be cleaned up");
}
#[test]
fn atomic_writer_streaming() {
let dir = tempfile::tempdir().unwrap();
let dest = dir.path().join("streamed.txt");
let mut w = AtomicFileWriter::new(&dest).unwrap();
for i in 0..100 {
writeln!(w, "line {}", i).unwrap();
}
w.finish().unwrap();
let content = fs::read_to_string(&dest).unwrap();
assert_eq!(content.lines().count(), 100);
}
}