use std::fs::{self, File, OpenOptions};
use std::io::{self, BufWriter, Write};
use std::path::{Path, PathBuf};
pub struct AtomicFileWriter {
writer: Option<BufWriter<File>>,
tmp_path: PathBuf,
dest_path: PathBuf,
finished: bool,
}
impl AtomicFileWriter {
pub fn new(dest: impl AsRef<Path>) -> io::Result<Self> {
Self::open(dest, false)
}
pub fn new_private(dest: impl AsRef<Path>) -> io::Result<Self> {
Self::open(dest, true)
}
#[cfg_attr(not(unix), allow(unused_variables))]
fn open(dest: impl AsRef<Path>, private: bool) -> 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)?;
#[cfg(unix)]
if private {
use std::os::unix::fs::PermissionsExt;
file.set_permissions(fs::Permissions::from_mode(0o600))?;
}
Ok(Self {
writer: Some(BufWriter::new(file)),
tmp_path,
dest_path,
finished: false,
})
}
pub fn finish(mut self) -> io::Result<()> {
let mut writer = self
.writer
.take()
.expect("AtomicFileWriter::finish called after writer already consumed");
writer.flush()?;
let file = writer.into_inner().map_err(|e| e.into_error())?;
file.sync_all()?;
drop(file);
if let Err(e) = rename_with_retry(&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
}
}
fn rename_with_retry(from: &Path, to: &Path) -> io::Result<()> {
#[cfg(not(windows))]
{
fs::rename(from, to)
}
#[cfg(windows)]
{
use std::thread::sleep;
use std::time::{Duration, Instant};
let deadline = Instant::now() + Duration::from_secs(1);
loop {
match fs::rename(from, to) {
Ok(()) => return Ok(()),
Err(e) => {
let transient = matches!(
e.raw_os_error(),
Some(5)
| Some(32)
);
if !transient || Instant::now() >= deadline {
return Err(e);
}
sleep(Duration::from_millis(50));
}
}
}
}
}
impl Write for AtomicFileWriter {
fn write(&mut self, buf: &[u8]) -> io::Result<usize> {
self.writer
.as_mut()
.expect("write after AtomicFileWriter::finish")
.write(buf)
}
fn flush(&mut self) -> io::Result<()> {
self.writer
.as_mut()
.expect("flush after AtomicFileWriter::finish")
.flush()
}
}
impl io::Seek for AtomicFileWriter {
fn seek(&mut self, pos: io::SeekFrom) -> io::Result<u64> {
let writer = self
.writer
.as_mut()
.expect("seek after AtomicFileWriter::finish");
writer.flush()?;
writer.get_mut().seek(pos)
}
}
impl Drop for AtomicFileWriter {
fn drop(&mut self) {
if !self.finished {
drop(self.writer.take());
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()
}
pub fn atomic_write_private(dest: impl AsRef<Path>, data: &[u8]) -> io::Result<()> {
let mut writer = AtomicFileWriter::new_private(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);
}
}