diff --git a/src/data/mnist.rs b/src/data/mnist.rs index aef3fdb..54ddf08 100644 --- a/src/data/mnist.rs +++ b/src/data/mnist.rs @@ -4,6 +4,7 @@ use std::fs; use std::io::{Read, Write}; use std::path::{Path, PathBuf}; +use std::sync::atomic::{AtomicU32, Ordering}; use crate::error::{Error, Result}; use crate::tensor::Tensor; @@ -59,10 +60,8 @@ impl Mnist { /// Load a split from an explicit directory. pub fn load_from(cache: &Path, split: Split) -> Result { let (image_file, label_file) = split.files(); - download_missing(cache, &[image_file, label_file])?; - - let (images, image_count) = parse_images(&fs::read(cache.join(image_file))?)?; - let (labels, label_count) = parse_labels(&fs::read(cache.join(label_file))?)?; + let (images, image_count) = load_cached(cache, image_file, parse_images, fetch)?; + let (labels, label_count) = load_cached(cache, label_file, parse_labels, fetch)?; if image_count != label_count { return Err(Error::Dataset(format!( "mnist: {image_count} images but {label_count} labels" @@ -119,32 +118,78 @@ pub fn default_cache_dir() -> PathBuf { .join("mnist") } -/// Download and decompress any of `files` not already cached. -fn download_missing(cache: &Path, files: &[&str]) -> Result<()> { - fs::create_dir_all(cache)?; - for name in files { - let target = cache.join(name); - if target.exists() { - continue; +/// Parse a cached file, downloading it first when it is missing. +/// +/// A cached file that does not parse is deleted and fetched once more. Without +/// that, a file truncated by an interrupted write would be trusted forever, +/// since the cache check only asks whether the file exists. +fn load_cached( + cache: &Path, + name: &str, + parse: fn(&[u8]) -> Result, + fetch: impl Fn(&Path, &str) -> Result>, +) -> Result { + let target = cache.join(name); + if target.exists() { + match parse(&fs::read(&target)?) { + Ok(parsed) => return Ok(parsed), + Err(e) => { + eprintln!( + "mnist: {} is not usable ({e}), downloading it again", + target.display() + ); + let _ = fs::remove_file(&target); + } } + } + fs::create_dir_all(cache)?; + let bytes = fetch(cache, name)?; + parse(&bytes).map_err(|e| Error::Dataset(format!("mnist: {}: {e}", target.display()))) +} - let url = format!("{MIRROR}/{name}.gz"); - eprintln!("mnist: downloading {url}"); - let response = ureq::get(&url) - .call() - .map_err(|e| Error::Dataset(format!("could not fetch {url}: {e}")))?; +/// Download `name` from the mirror, decompress it, and write it into the cache. +fn fetch(cache: &Path, name: &str) -> Result> { + let url = format!("{MIRROR}/{name}.gz"); + eprintln!("mnist: downloading {url}"); + let response = ureq::get(&url) + .call() + .map_err(|e| Error::Dataset(format!("could not fetch {url}: {e}")))?; - let mut compressed = Vec::new(); - response - .into_body() - .into_reader() - .read_to_end(&mut compressed)?; - let mut raw = Vec::new(); - flate2::read::GzDecoder::new(&compressed[..]).read_to_end(&mut raw)?; + let mut compressed = Vec::new(); + response + .into_body() + .into_reader() + .read_to_end(&mut compressed)?; + let mut raw = Vec::new(); + flate2::read::GzDecoder::new(&compressed[..]).read_to_end(&mut raw)?; - fs::File::create(&target)?.write_all(&raw)?; + write_atomic(&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); } - Ok(()) + written.map_err(Error::from) } /// IDX images: magic, count, rows, cols, then one byte per pixel. @@ -204,3 +249,134 @@ fn read_header(bytes: &[u8], words: usize, magic: u32, what: &str) -> Result TempDir { + let dir = + std::env::temp_dir().join(format!("fastnn-mnist-{}-{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); + } + } + + /// A valid IDX image file holding `count` one-pixel images. + fn image_file(count: u32) -> Vec { + let mut bytes = Vec::new(); + for word in [IMAGE_MAGIC, count, 1, 1] { + bytes.extend_from_slice(&word.to_be_bytes()); + } + bytes.resize(IMAGE_HEADER + count as usize, 7); + 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"); + fs::write(dir.0.join("images"), image_file(3)).unwrap(); + let (_, count) = load_cached(&dir.0, "images", parse_images, |_, _| { + panic!("a valid cached file must not be fetched") + }) + .unwrap(); + assert_eq!(count, 3); + } + + #[test] + fn a_missing_file_is_fetched() { + 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)?; + Ok(bytes) + }) + .unwrap(); + assert_eq!(count, 2); + assert!(parse_images(&fs::read(dir.0.join("images")).unwrap()).is_ok()); + } + + #[test] + fn a_truncated_cached_file_is_fetched_again() { + let dir = TempDir::new("truncated"); + // Valid magic and header, body cut short: what an interrupted write leaves. + let mut truncated = image_file(5); + truncated.truncate(IMAGE_HEADER + 2); + fs::write(dir.0.join("images"), &truncated).unwrap(); + + let fetches = std::cell::Cell::new(0); + 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)?; + Ok(bytes) + }) + .unwrap(); + assert_eq!(count, 5); + assert_eq!(fetches.get(), 1); + } + + #[test] + fn a_bad_download_names_the_file() { + let dir = TempDir::new("bad-download"); + fs::write(dir.0.join("images"), b"junk").unwrap(); + let err = load_cached(&dir.0, "images", parse_images, |_, _| { + Ok(b"still junk".to_vec()) + }) + .unwrap_err(); + assert!(err.to_string().contains("images"), "{err}"); + } +}