diff --git a/src/data/mnist.rs b/src/data/mnist.rs index aef3fdb..fde1fcf 100644 --- a/src/data/mnist.rs +++ b/src/data/mnist.rs @@ -61,8 +61,8 @@ impl Mnist { 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) = parse_images(&read_or_redownload(cache, image_file)?)?; + let (labels, label_count) = parse_labels(&read_or_redownload(cache, label_file)?)?; if image_count != label_count { return Err(Error::Dataset(format!( "mnist: {image_count} images but {label_count} labels" @@ -119,7 +119,49 @@ pub fn default_cache_dir() -> PathBuf { .join("mnist") } +/// Write `bytes` to `target` atomically: write a per-process temporary file, +/// `sync_all`, then rename into place. A crash at any point leaves the +/// target either absent or complete. +fn write_atomic(target: &Path, bytes: &[u8]) -> Result<()> { + let name = target + .file_name() + .unwrap_or_default() + .to_string_lossy(); + let partial = target.with_file_name(format!("{name}.{}.partial", std::process::id())); + + let mut file = fs::File::create(&partial)?; + file.write_all(bytes)?; + file.sync_all()?; + fs::rename(&partial, target)?; + Ok(()) +} + +/// Fetch a single file from the mirror, decompress, and write atomically. +fn fetch_and_write(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)?; + + write_atomic(&cache.join(name), &raw) +} + /// Download and decompress any of `files` not already cached. +/// +/// Writes to a per-process `.partial` temporary file, calls `sync_all`, and +/// renames on success so that an interrupted or concurrent download never +/// leaves a truncated file in the cache. If a cached file exists but is +/// corrupt (truncated by an earlier crash), it is deleted and re-downloaded +/// once. fn download_missing(cache: &Path, files: &[&str]) -> Result<()> { fs::create_dir_all(cache)?; for name in files { @@ -127,26 +169,30 @@ fn download_missing(cache: &Path, files: &[&str]) -> Result<()> { if target.exists() { continue; } - - 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)?; - - fs::File::create(&target)?.write_all(&raw)?; + fetch_and_write(cache, name)?; } Ok(()) } +/// Try to parse a cached file; if it fails, delete it and re-download once. +fn read_or_redownload(cache: &Path, name: &str) -> Result> { + let target = cache.join(name); + let bytes = fs::read(&target)?; + let header_ok = if bytes.len() >= 4 { + let magic = u32::from_be_bytes([bytes[0], bytes[1], bytes[2], bytes[3]]); + magic == IMAGE_MAGIC || magic == LABEL_MAGIC + } else { + false + }; + if header_ok { + return Ok(bytes); + } + eprintln!("mnist: cached {name} is corrupt, re-downloading"); + let _ = fs::remove_file(&target); + fetch_and_write(cache, name)?; + Ok(fs::read(&target)?) +} + /// IDX images: magic, count, rows, cols, then one byte per pixel. fn parse_images(bytes: &[u8]) -> Result<(Vec, usize)> { let header = read_header(bytes, 4, IMAGE_MAGIC, "images")?; @@ -186,6 +232,78 @@ fn parse_labels(bytes: &[u8]) -> Result<(Vec, usize)> { )) } +#[cfg(test)] +mod tests { + use super::*; + + fn test_dir(name: &str) -> PathBuf { + let dir = std::env::temp_dir() + .join("fastnn_test") + .join(name); + let _ = fs::remove_dir_all(&dir); + fs::create_dir_all(&dir).unwrap(); + dir + } + + #[test] + fn truncated_file_rejected_by_parser() { + let cache = test_dir("truncated_file"); + let target = cache.join("train-images-idx3-ubyte"); + + fs::write(&target, b"truncated").unwrap(); + assert!(target.exists()); + + let bytes = fs::read(&target).unwrap(); + assert!(parse_images(&bytes).is_err()); + } + + #[test] + fn partial_suffix_not_mistaken_for_final() { + let cache = test_dir("partial_suffix"); + let name = "train-images-idx3-ubyte"; + let partial = cache.join(format!("{name}.partial")); + + fs::write(&partial, b"junk").unwrap(); + + let target = cache.join(name); + assert!(!target.exists(), ".partial must not satisfy the exists() check"); + } + + #[test] + fn write_atomic_creates_file_with_correct_content() { + let dir = test_dir("write_atomic_ok"); + let target = dir.join("output.bin"); + let payload = b"hello, atomic world"; + + write_atomic(&target, payload).unwrap(); + assert_eq!(fs::read(&target).unwrap(), payload); + } + + #[test] + fn write_atomic_leaves_no_partial_on_success() { + let dir = test_dir("write_atomic_no_partial"); + let target = dir.join("data.bin"); + write_atomic(&target, b"ok").unwrap(); + + let partials: Vec<_> = fs::read_dir(&dir) + .unwrap() + .filter_map(|e| e.ok()) + .filter(|e| e.file_name().to_string_lossy().contains(".partial")) + .collect(); + assert!(partials.is_empty(), "no .partial files should remain"); + } + + #[test] + fn write_atomic_overwrites_existing() { + let dir = test_dir("write_atomic_overwrite"); + let target = dir.join("file.bin"); + fs::write(&target, b"old").unwrap(); + + write_atomic(&target, b"new").unwrap(); + assert_eq!(fs::read(&target).unwrap(), b"new"); + } +} + /// Read `words` big-endian u32s and check the magic number. fn read_header(bytes: &[u8], words: usize, magic: u32, what: &str) -> Result> { if bytes.len() < words * 4 {