One atomic write helper for checkpoints and the dataset cache

src/data/mnist.rs

@@ -2,11 +2,11 @@
 //! `~/.fastnn/datasets/mnist`.
 
 use std::fs;
-use std::io::{Read, Write};
+use std::io::Read;
 use std::path::{Path, PathBuf};
-use std::sync::atomic::{AtomicU32, Ordering};
 
 use crate::error::{Error, Result};
+use crate::serialize::atomic::write_atomic_bytes;
 use crate::tensor::Tensor;
 
 use super::dataset::Dataset;
@@ -163,35 +163,10 @@ fn fetch(cache: &Path, name: &str) -> Result<Vec<u8>> {
     let mut raw = Vec::new();
     flate2::read::GzDecoder::new(&compressed[..]).read_to_end(&mut raw)?;
 
-    write_atomic(&cache.join(name), &raw)?;
+    write_atomic_bytes(&cache.join(name), &raw)?;
     Ok(raw)
 }
 
-/// Write `bytes` to `target` so a crash leaves it either absent or complete:
-/// a temporary file next to it, `sync_all`, then a rename into place.
-fn write_atomic(target: &Path, bytes: &[u8]) -> Result<()> {
-    // The pid and a counter keep concurrent writers, in and across processes,
-    // off each other's temporary file.
-    static SERIAL: AtomicU32 = AtomicU32::new(0);
-    let name = target.file_name().unwrap_or_default().to_string_lossy();
-    let partial = target.with_file_name(format!(
-        "{name}.{}.{}.partial",
-        std::process::id(),
-        SERIAL.fetch_add(1, Ordering::Relaxed)
-    ));
-
-    let written = (|| {
-        let mut file = fs::File::create(&partial)?;
-        file.write_all(bytes)?;
-        file.sync_all()?;
-        fs::rename(&partial, target)
-    })();
-    if written.is_err() {
-        let _ = fs::remove_file(&partial);
-    }
-    written.map_err(Error::from)
-}
-
 /// IDX images: magic, count, rows, cols, then one byte per pixel.
 fn parse_images(bytes: &[u8]) -> Result<(Vec<f32>, usize)> {
     let header = read_header(bytes, 4, IMAGE_MAGIC, "images")?;
@@ -283,48 +258,6 @@ mod tests {
         bytes
     }
 
-    fn partials_in(dir: &Path) -> usize {
-        fs::read_dir(dir)
-            .unwrap()
-            .filter(|entry| {
-                entry
-                    .as_ref()
-                    .unwrap()
-                    .file_name()
-                    .to_string_lossy()
-                    .ends_with(".partial")
-            })
-            .count()
-    }
-
-    #[test]
-    fn write_atomic_writes_the_bytes_and_leaves_no_partial() {
-        let dir = TempDir::new("write");
-        let target = dir.0.join("data.bin");
-        write_atomic(&target, b"hello").unwrap();
-        assert_eq!(fs::read(&target).unwrap(), b"hello");
-        assert_eq!(partials_in(&dir.0), 0);
-    }
-
-    #[test]
-    fn write_atomic_replaces_an_existing_file() {
-        let dir = TempDir::new("replace");
-        let target = dir.0.join("data.bin");
-        fs::write(&target, b"old").unwrap();
-        write_atomic(&target, b"new").unwrap();
-        assert_eq!(fs::read(&target).unwrap(), b"new");
-    }
-
-    #[test]
-    fn write_atomic_cleans_up_when_the_rename_fails() {
-        let dir = TempDir::new("rename-fails");
-        // A file cannot replace a directory, so the rename fails after the write.
-        let target = dir.0.join("taken");
-        fs::create_dir(&target).unwrap();
-        assert!(write_atomic(&target, b"x").is_err());
-        assert_eq!(partials_in(&dir.0), 0);
-    }
-
     #[test]
     fn a_valid_cached_file_is_used_without_fetching() {
         let dir = TempDir::new("valid");
@@ -341,7 +274,7 @@ mod tests {
         let dir = TempDir::new("missing");
         let (_, count) = load_cached(&dir.0, "images", parse_images, |cache, name| {
             let bytes = image_file(2);
-            write_atomic(&cache.join(name), &bytes)?;
+            fs::write(cache.join(name), &bytes)?;
             Ok(bytes)
         })
         .unwrap();
@@ -361,7 +294,7 @@ mod tests {
         let (_, count) = load_cached(&dir.0, "images", parse_images, |cache, name| {
             fetches.set(fetches.get() + 1);
             let bytes = image_file(5);
-            write_atomic(&cache.join(name), &bytes)?;
+            fs::write(cache.join(name), &bytes)?;
             Ok(bytes)
         })
         .unwrap();

src/serialize/atomic.rs

@@ -1,39 +1,131 @@
-//! Atomic file writes: write to a per-process temp, `sync_all`, then rename.
+//! One atomic file write for checkpoints and the dataset cache: a temporary
+//! file next to the target, `sync_all`, then a rename into place.
 
 use std::ffi::OsString;
-use std::fs::File;
+use std::fs::{self, File};
+use std::io::Write;
 use std::path::{Path, PathBuf};
+use std::sync::atomic::{AtomicU32, Ordering};
 
 use crate::error::Result;
 
-/// Write to `path` atomically.
+/// Write to `path` so a crash leaves it either absent or complete.
 ///
-/// Writes into `{path}.partial.{pid}`, calls [`File::sync_all`], then renames
-/// the temp to `path`. A crash mid-write leaves the previous file intact. On
-/// error the temp is removed so it never accumulates.
-pub(crate) fn write_atomic(
-    path: &Path,
-    write_fn: impl FnOnce(&mut File) -> Result<()>,
-) -> Result<()> {
-    let tmp = tmp_path(path);
-    let result = write_and_sync(&tmp, write_fn);
-    if result.is_err() {
-        let _ = std::fs::remove_file(&tmp);
-        return result;
-    }
-    std::fs::rename(&tmp, path)?;
-    Ok(())
+/// `write` fills a temporary file next to `path`, which is synced and renamed
+/// into place only when `write` succeeds, and removed when anything fails. The
+/// previous file is never touched until the rename.
+pub(crate) fn write_atomic(path: &Path, write: impl FnOnce(&mut File) -> Result<()>) -> Result<()> {
+    let partial = partial_path(path);
+    let written: Result<()> = (|| {
+        let mut file = File::create(&partial)?;
+        write(&mut file)?;
+        file.sync_all()?;
+        fs::rename(&partial, path)?;
+        Ok(())
+    })();
+    if written.is_err() {
+        let _ = fs::remove_file(&partial);
+    }
+    written
 }
 
-fn write_and_sync(path: &Path, write_fn: impl FnOnce(&mut File) -> Result<()>) -> Result<()> {
-    let mut file = File::create(path)?;
-    write_fn(&mut file)?;
-    file.sync_all()?;
-    Ok(())
+/// [`write_atomic`] for bytes already in memory.
+pub(crate) fn write_atomic_bytes(path: &Path, bytes: &[u8]) -> Result<()> {
+    write_atomic(path, |file| Ok(file.write_all(bytes)?))
 }
 
-fn tmp_path(path: &Path) -> PathBuf {
+/// `{path}.{pid}.{serial}.partial`. The pid and the counter keep concurrent
+/// writers, in and across processes, off each other's temporary file.
+fn partial_path(path: &Path) -> PathBuf {
+    static SERIAL: AtomicU32 = AtomicU32::new(0);
     let mut name = OsString::from(path.as_os_str());
-    name.push(format!(".partial.{}", std::process::id()));
+    name.push(format!(
+        ".{}.{}.partial",
+        std::process::id(),
+        SERIAL.fetch_add(1, Ordering::Relaxed)
+    ));
     PathBuf::from(name)
 }
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+    use crate::error::Error;
+
+    /// A fresh directory under the system temp dir, removed on drop.
+    struct TempDir(PathBuf);
+
+    impl TempDir {
+        fn new(name: &str) -> TempDir {
+            let dir =
+                std::env::temp_dir().join(format!("fastnn-atomic-{}-{name}", std::process::id()));
+            let _ = fs::remove_dir_all(&dir);
+            fs::create_dir_all(&dir).unwrap();
+            TempDir(dir)
+        }
+    }
+
+    impl Drop for TempDir {
+        fn drop(&mut self) {
+            let _ = fs::remove_dir_all(&self.0);
+        }
+    }
+
+    fn partials_in(dir: &Path) -> usize {
+        fs::read_dir(dir)
+            .unwrap()
+            .filter(|entry| {
+                entry
+                    .as_ref()
+                    .unwrap()
+                    .file_name()
+                    .to_string_lossy()
+                    .ends_with(".partial")
+            })
+            .count()
+    }
+
+    #[test]
+    fn writes_the_bytes_and_leaves_no_partial() {
+        let dir = TempDir::new("write");
+        let target = dir.0.join("data.bin");
+        write_atomic_bytes(&target, b"hello").unwrap();
+        assert_eq!(fs::read(&target).unwrap(), b"hello");
+        assert_eq!(partials_in(&dir.0), 0);
+    }
+
+    #[test]
+    fn replaces_an_existing_file() {
+        let dir = TempDir::new("replace");
+        let target = dir.0.join("data.bin");
+        fs::write(&target, b"old").unwrap();
+        write_atomic_bytes(&target, b"new").unwrap();
+        assert_eq!(fs::read(&target).unwrap(), b"new");
+    }
+
+    #[test]
+    fn a_failed_write_leaves_the_old_file_and_no_partial() {
+        let dir = TempDir::new("write-fails");
+        let target = dir.0.join("data.bin");
+        fs::write(&target, b"old").unwrap();
+
+        let result = write_atomic(&target, |file| {
+            file.write_all(b"half")?;
+            Err(Error::Checkpoint("disk full".into()))
+        });
+
+        assert!(result.is_err());
+        assert_eq!(fs::read(&target).unwrap(), b"old");
+        assert_eq!(partials_in(&dir.0), 0);
+    }
+
+    #[test]
+    fn a_failed_rename_leaves_no_partial() {
+        let dir = TempDir::new("rename-fails");
+        // A file cannot replace a directory, so the rename fails after the write.
+        let target = dir.0.join("taken");
+        fs::create_dir(&target).unwrap();
+        assert!(write_atomic_bytes(&target, b"x").is_err());
+        assert_eq!(partials_in(&dir.0), 0);
+    }
+}

src/serialize/checkpoint.rs

@@ -69,8 +69,7 @@ pub fn load(module: &dyn Module, path: impl AsRef<Path>) -> Result<()> {
 
 /// Write named tensors to `path`.
 ///
-/// The write is atomic: data lands in a per-process temp file, is fsynced, then
-/// renamed over `path`, so a crash mid-write never destroys an existing file.
+/// Written atomically: a crash mid-save leaves the previous file intact.
 pub fn save_tensors(tensors: &Checkpoint, path: impl AsRef<Path>) -> Result<()> {
     write_atomic(path.as_ref(), |file| {
         let mut out = BufWriter::new(&mut *file);

src/serialize/safetensors.rs

@@ -31,8 +31,7 @@ pub struct LoadReport {
 
 /// Write a module's parameters and buffers as `F32` safetensors.
 ///
-/// The write is atomic: data lands in a per-process temp file, is fsynced, then
-/// renamed over `path`, so a crash mid-write never destroys an existing file.
+/// Written atomically: a crash mid-save leaves the previous file intact.
 pub fn save_safetensors(module: &dyn Module, path: impl AsRef<Path>) -> Result<()> {
     let mut tensors: BTreeMap<String, Tensor> = BTreeMap::new();
     for (name, param) in module.named_parameters() {

tests/robustness.rs

@@ -408,10 +408,8 @@ fn fallible_ops_still_build_the_graph() {
     assert!(w.grad().is_some(), "try_matmul dropped the graph");
 }
 
-// ---- Atomic checkpoint writes --------------------------------------------------
-
 #[test]
-fn truncated_fdl_checkpoint_does_not_destroy_the_good_copy() {
+fn a_corrupt_checkpoint_can_be_saved_over() {
     let file = TempFile::new("atomic-fdl");
     let model = small_model();
     save(&model, file.path()).unwrap();
@@ -432,8 +430,6 @@ fn truncated_fdl_checkpoint_does_not_destroy_the_good_copy() {
         "a truncated fdl checkpoint should not parse"
     );
 
-    // Overwrite the truncated file with a fresh save -- the atomic write must
-    // succeed even when the target already exists and is corrupt.
     save(&model, file.path()).unwrap();
     let fresh = small_model();
     load(&fresh, file.path()).unwrap();
@@ -445,59 +441,39 @@ fn truncated_fdl_checkpoint_does_not_destroy_the_good_copy() {
 }
 
 #[test]
-fn truncated_safetensors_does_not_destroy_the_good_copy() {
+fn a_corrupt_safetensors_file_can_be_saved_over() {
     let file = TempFile::new("atomic-st");
-    let path = file.path().with_extension("safetensors");
+    let path = file.path();
 
     manual_seed(42);
     let model = Sequential::new().add(Linear::new(4, 2));
-    save_safetensors(&model, &path).unwrap();
+    save_safetensors(&model, path).unwrap();
 
     let good = model.parameters()[0].value().to_vec();
 
     // Simulate a crash: truncate the file.
-    std::fs::write(&path, [0u8; 4]).unwrap();
+    std::fs::write(path, [0u8; 4]).unwrap();
 
     // The truncated file must fail to load.
     let fresh = Sequential::new().add(Linear::new(4, 2));
     assert!(
-        load_safetensors(&fresh, &path).is_err(),
+        load_safetensors(&fresh, path).is_err(),
         "a truncated safetensors file should not parse"
     );
 
-    // Re-save over the corrupt file; the atomic write must succeed.
-    save_safetensors(&model, &path).unwrap();
+    save_safetensors(&model, path).unwrap();
     let reloaded = Sequential::new().add(Linear::new(4, 2));
-    let report = load_safetensors(&reloaded, &path).unwrap();
+    let report = load_safetensors(&reloaded, path).unwrap();
     assert!(report.missing.is_empty());
     assert_eq!(
         reloaded.parameters()[0].value().to_vec(),
         good,
         "re-saved safetensors should match the original weights"
     );
-    let _ = std::fs::remove_file(&path);
-}
-
-#[test]
-fn atomic_write_cleans_up_temp_on_fdl_save() {
-    // After a successful save the .partial.{pid} file must not linger.
-    let file = TempFile::new("atomic-cleanup");
-    let model = small_model();
-    save(&model, file.path()).unwrap();
-
-    let tmp = file.path().with_file_name(format!(
-        "{}.partial.{}",
-        file.path().file_name().unwrap().to_str().unwrap(),
-        std::process::id()
-    ));
-    assert!(
-        !tmp.exists(),
-        "temp file should be renamed away after a successful save"
-    );
 }
 
 #[test]
-fn save_tensors_round_trips_through_atomic_write() {
+fn save_tensors_round_trips() {
     let file = TempFile::new("atomic-rt");
     let mut tensors = Checkpoint::new();
     tensors.insert("a".into(), Tensor::from_vec(vec![1.0, 2.0, 3.0], &[3]));