use std::io::Write;
use std::path::Path;
use std::sync::atomic::{AtomicU64, Ordering};
use crate::error::{Error, Result};
static TMP_COUNTER: AtomicU64 = AtomicU64::new(0);
#[doc(hidden)]
pub mod faults {
use std::sync::atomic::{AtomicI64, AtomicU64, Ordering};
use crate::error::{Error, Result};
const DISABLED: i64 = -1;
static FAIL_AT: AtomicI64 = AtomicI64::new(DISABLED);
static COUNT: AtomicU64 = AtomicU64::new(0);
pub fn arm(n: u64) {
COUNT.store(0, Ordering::SeqCst);
FAIL_AT.store(n as i64, Ordering::SeqCst);
}
pub fn disarm() {
FAIL_AT.store(DISABLED, Ordering::SeqCst);
COUNT.store(0, Ordering::SeqCst);
}
pub fn writes_seen() -> u64 {
COUNT.load(Ordering::SeqCst)
}
pub(super) fn guard() -> Result<()> {
if FAIL_AT.load(Ordering::Relaxed) == DISABLED {
return Ok(());
}
let n = COUNT.fetch_add(1, Ordering::SeqCst) as i64;
let target = FAIL_AT.load(Ordering::SeqCst);
if target != DISABLED && n >= target {
return Err(Error::other(format!("injected crash at atomic write #{n}")));
}
Ok(())
}
}
fn tmp_path(path: &Path) -> std::path::PathBuf {
let n = TMP_COUNTER.fetch_add(1, Ordering::Relaxed);
let pid = std::process::id();
let mut name = path
.file_name()
.map(|s| s.to_os_string())
.unwrap_or_default();
name.push(format!(".tmp.{pid}.{n}"));
match path.parent() {
Some(dir) => dir.join(name),
None => std::path::PathBuf::from(name),
}
}
pub fn write_atomic(path: &Path, bytes: &[u8]) -> Result<()> {
faults::guard()?;
let tmp = tmp_path(path);
let res = (|| {
let mut f = std::fs::File::create(&tmp).map_err(|e| Error::io(&tmp, e))?;
f.write_all(bytes).map_err(|e| Error::io(&tmp, e))?;
f.sync_all().map_err(|e| Error::io(&tmp, e))?;
std::fs::rename(&tmp, path).map_err(|e| Error::io(path, e))
})();
match res {
Ok(()) => {
sync_parent_dir(path);
Ok(())
}
Err(e) => {
let _ = std::fs::remove_file(&tmp);
Err(e)
}
}
}
pub struct AtomicFile {
final_path: std::path::PathBuf,
tmp_path: std::path::PathBuf,
file: Option<std::fs::File>,
}
impl AtomicFile {
pub fn create(path: &Path) -> Result<AtomicFile> {
let tmp = tmp_path(path);
let file = std::fs::File::create(&tmp).map_err(|e| Error::io(&tmp, e))?;
Ok(AtomicFile {
final_path: path.to_path_buf(),
tmp_path: tmp,
file: Some(file),
})
}
pub fn file(&mut self) -> &mut std::fs::File {
self.file.as_mut().expect("AtomicFile already committed")
}
pub fn commit(mut self) -> Result<()> {
faults::guard()?;
let f = self.file.take().expect("AtomicFile already committed");
let res = (|| {
f.sync_all().map_err(|e| Error::io(&self.tmp_path, e))?;
std::fs::rename(&self.tmp_path, &self.final_path)
.map_err(|e| Error::io(&self.final_path, e))
})();
match res {
Ok(()) => {
sync_parent_dir(&self.final_path);
Ok(())
}
Err(e) => {
let _ = std::fs::remove_file(&self.tmp_path);
Err(e)
}
}
}
}
impl Drop for AtomicFile {
fn drop(&mut self) {
if self.file.is_some() {
let _ = std::fs::remove_file(&self.tmp_path);
}
}
}
fn sync_parent_dir(path: &Path) {
if let Some(dir) = path.parent() {
if let Ok(d) = std::fs::File::open(dir) {
let _ = d.sync_all();
}
}
}