runsync_transfer/io/

writer.rs

//! Destination-side writes.
//!
//! Files are written to a `.part` sidecar and renamed into place only after the
//! last chunk lands and the hash verifies. A crash therefore leaves an obvious
//! partial file that resume can pick up, never a truncated file that looks
//! complete.

use crate::error::Result;
use parking_lot::Mutex;
use std::collections::HashMap;
use std::fs::{File, OpenOptions};
use std::io;
use std::path::{Path, PathBuf};
use std::sync::Arc;

/// Suffix for in-progress files.
pub const PART_SUFFIX: &str = ".rst-part";

/// One destination file, open for concurrent positional writes.
pub struct WriteHandle {
    file: File,
    part_path: PathBuf,
    final_path: PathBuf,
    size: u64,
    /// True when this session created the partial file. Only then is the file
    /// known to be all zeros where nothing has been written, which is what
    /// makes it safe to skip writing an all-zero chunk.
    created_new: bool,
}

impl WriteHandle {
    /// Open the partial file, seeding it from an existing copy where the
    /// filesystem can do that for free.
    ///
    /// Delta sync reuses most blocks of an unchanged file, and the receiver
    /// would otherwise copy every one of them out of the old file and into the
    /// new one — gigabytes of local I/O to reproduce bytes that are already on
    /// the disk. A copy-on-write clone shares the extents instead, so the
    /// seeding is O(1) and only the chunks that genuinely differ are written.
    /// Where the filesystem cannot clone, this falls back to the ordinary path
    /// and the per-chunk copies happen as before.
    pub fn open_cloned(final_path: &Path, size: u64, preallocate: bool) -> Result<Self> {
        let part_path = part_path_for(final_path);
        if !part_path.exists() && final_path.exists() && clone_file(final_path, &part_path) {
            if let Ok(file) = OpenOptions::new().read(true).write(true).open(&part_path) {
                if file.metadata()?.len() != size {
                    file.set_len(size)?;
                }
                return Ok(Self {
                    file,
                    part_path,
                    final_path: final_path.to_path_buf(),
                    size,
                    // A clone is not a file we created empty: it starts full of
                    // the old contents, so nothing may assume unwritten regions
                    // read as zeros.
                    created_new: false,
                });
            }
        }
        Self::open(final_path, size, preallocate)
    }

    /// Open (or reopen) the partial file for `final_path`, sized to `size`.
    pub fn open(final_path: &Path, size: u64, preallocate: bool) -> Result<Self> {
        if let Some(parent) = final_path.parent() {
            std::fs::create_dir_all(parent)
                .map_err(|e| io::Error::new(e.kind(), format!("{}: {e}", parent.display())))?;
        }
        let part_path = part_path_for(final_path);
        // Try to create exclusively first: succeeding tells us the file is
        // brand new, and therefore reads as zeros everywhere we have not
        // written. That fact is what the sparse path depends on.
        let (file, created_new) = match OpenOptions::new()
            .create_new(true)
            .read(true)
            .write(true)
            .open(&part_path)
        {
            Ok(f) => (f, true),
            Err(e) if e.kind() == io::ErrorKind::AlreadyExists => {
                let f = OpenOptions::new()
                    .read(true)
                    .write(true)
                    .truncate(false)
                    .open(&part_path)
                    .map_err(|e| {
                        io::Error::new(e.kind(), format!("{}: {e}", part_path.display()))
                    })?;
                (f, false)
            }
            Err(e) => {
                return Err(
                    io::Error::new(e.kind(), format!("{}: {e}", part_path.display())).into(),
                )
            }
        };

        // Reserve before sizing, and only on a file we just created.
        //
        // macOS `F_PREALLOCATE` appends to the end of the file rather than
        // filling up to a target, so running it after `set_len` reserves a
        // second copy of the whole file. Doing it while the file is still empty
        // is the documented idiom and reserves exactly `size`. A resumed
        // transfer skips it: those blocks are already backed, and re-reserving
        // is precisely what over-allocated.
        if preallocate && size > 0 && created_new {
            reserve(&file, size)?;
        }
        // Then size it: positional writes into a region past EOF would
        // otherwise each extend the file, serialising on the inode.
        if file.metadata()?.len() != size {
            file.set_len(size)?;
        }

        Ok(Self {
            file,
            part_path,
            final_path: final_path.to_path_buf(),
            size,
            created_new,
        })
    }

    pub fn size(&self) -> u64 {
        self.size
    }

    pub fn part_path(&self) -> &Path {
        &self.part_path
    }

    /// Whether this session created the partial file, and can therefore assume
    /// unwritten regions read as zeros.
    pub fn created_new(&self) -> bool {
        self.created_new
    }

    /// Write `len` zero bytes at `offset` without allocating a buffer that big.
    pub fn write_zeros_at(&self, offset: u64, len: usize) -> io::Result<()> {
        const BLOCK: usize = 256 * 1024;
        let zeros = [0u8; BLOCK];
        let mut done = 0usize;
        while done < len {
            let n = BLOCK.min(len - done);
            self.write_at(offset + done as u64, &zeros[..n])?;
            done += n;
        }
        Ok(())
    }

    /// Write `buf` at `offset`. Safe to call concurrently from many threads.
    pub fn write_at(&self, offset: u64, buf: &[u8]) -> io::Result<()> {
        // A write landing past the declared size means the sender's manifest and
        // its frames disagree. Refuse rather than grow the file.
        let end = offset
            .checked_add(buf.len() as u64)
            .ok_or_else(|| io::Error::new(io::ErrorKind::InvalidInput, "write offset overflow"))?;
        if end > self.size {
            return Err(io::Error::new(
                io::ErrorKind::InvalidInput,
                format!(
                    "write of {} bytes at {offset} exceeds declared size {}",
                    buf.len(),
                    self.size
                ),
            ));
        }
        let mut written = 0usize;
        while written < buf.len() {
            let n = pwrite(&self.file, &buf[written..], offset + written as u64)?;
            if n == 0 {
                return Err(io::Error::new(
                    io::ErrorKind::WriteZero,
                    "pwrite made no progress",
                ));
            }
            written += n;
        }
        Ok(())
    }

    pub fn read_at(&self, offset: u64, buf: &mut [u8]) -> io::Result<usize> {
        pread(&self.file, buf, offset)
    }

    /// Is `want` already what the file holds at `offset`?
    ///
    /// Used to avoid rewriting bytes a copy-on-write clone already shares:
    /// writing them would break the sharing and cost the very I/O the clone
    /// was there to avoid.
    pub fn matches_at(&self, offset: u64, want: &[u8]) -> io::Result<bool> {
        let mut scratch = vec![0u8; want.len()];
        let n = pread(&self.file, &mut scratch, offset)?;
        Ok(n == want.len() && scratch == want)
    }

    /// Flush to stable storage. Called once per file, not per chunk — an fsync
    /// per chunk would dominate the cost of the transfer.
    pub fn sync(&self) -> io::Result<()> {
        self.file.sync_all()
    }

    /// Rename the partial file into place and apply metadata.
    ///
    /// Takes `&self` rather than `self` because several workers may still hold
    /// a clone of the `Arc` around this handle when the last chunk lands.
    /// Renaming a file that is still open is safe on unix, and on Windows too:
    /// Rust's `OpenOptions` opens with `FILE_SHARE_DELETE`. Callers guarantee
    /// this runs once per file.
    pub fn commit(&self, mode: u32, mtime: i64, preserve_metadata: bool) -> Result<()> {
        self.file.sync_all()?;
        if preserve_metadata {
            #[cfg(unix)]
            if mode != 0 {
                use std::os::unix::fs::PermissionsExt;
                let _ = self
                    .file
                    .set_permissions(std::fs::Permissions::from_mode(mode & 0o7777));
            }
        }

        std::fs::rename(&self.part_path, &self.final_path).map_err(|e| {
            io::Error::new(
                e.kind(),
                format!(
                    "rename {} -> {}: {e}",
                    self.part_path.display(),
                    self.final_path.display()
                ),
            )
        })?;

        if preserve_metadata && mtime > 0 {
            set_mtime(&self.final_path, mtime);
        }
        Ok(())
    }

    /// Delete the partial file. Used when a transfer aborts unrecoverably.
    pub fn discard(&self) {
        let _ = std::fs::remove_file(&self.part_path);
    }
}

/// Copy-on-write clone, where the platform offers one. Returns whether it
/// worked; a failure is never fatal, only slower.
// The `return`s below are load-bearing: each arm is behind a `cfg`, so a tail
// expression would be parsed as the result on every platform at once.
#[allow(clippy::needless_return)]
fn clone_file(src: &Path, dst: &Path) -> bool {
    #[cfg(target_os = "macos")]
    {
        use std::ffi::CString;
        let (Ok(s), Ok(d)) = (
            CString::new(src.as_os_str().as_encoded_bytes()),
            CString::new(dst.as_os_str().as_encoded_bytes()),
        ) else {
            return false;
        };
        // APFS shares the extents; the copy costs metadata only.
        return unsafe { libc::clonefile(s.as_ptr(), d.as_ptr(), 0) } == 0;
    }
    #[cfg(target_os = "linux")]
    {
        use std::os::unix::io::AsRawFd;
        // FICLONE: btrfs, XFS with reflinks, bcachefs.
        const FICLONE: libc::c_ulong = 0x4004_9409;
        let (Ok(from), Ok(to)) = (
            std::fs::File::open(src),
            OpenOptions::new().create_new(true).write(true).open(dst),
        ) else {
            return false;
        };
        let ok = unsafe { libc::ioctl(to.as_raw_fd(), FICLONE, from.as_raw_fd()) } == 0;
        if !ok {
            drop(to);
            let _ = std::fs::remove_file(dst);
        }
        ok
    }
    #[cfg(not(any(target_os = "macos", target_os = "linux")))]
    {
        let _ = (src, dst);
        false
    }
}

pub fn part_path_for(final_path: &Path) -> PathBuf {
    let mut s = final_path.as_os_str().to_os_string();
    s.push(PART_SUFFIX);
    PathBuf::from(s)
}

/// Ask the filesystem to actually reserve blocks.
///
/// `set_len` only records a size; on a sparse-capable filesystem the blocks are
/// not committed until written. For a 100 GB transfer that is the difference
/// between failing on the first byte and failing eight hours in, so this asks
/// for real reservation where the platform supports it.
// The explicit `return`s below are load-bearing: each arm is behind a `cfg`,
// so a tail expression would be parsed as the function's result on every
// platform at once.
#[allow(clippy::needless_return)]
fn reserve(file: &File, size: u64) -> Result<()> {
    #[cfg(target_os = "linux")]
    {
        use std::os::unix::io::AsRawFd;
        // Keep the size set by set_len; only allocate the blocks.
        let rc = unsafe {
            libc::fallocate(
                file.as_raw_fd(),
                libc::FALLOC_FL_KEEP_SIZE,
                0,
                size as libc::off_t,
            )
        };
        if rc != 0 {
            let e = io::Error::last_os_error();
            match e.raw_os_error() {
                // Filesystem does not implement it (tmpfs, some network mounts).
                // set_len already gave us the size we need; carry on.
                Some(libc::EOPNOTSUPP) | Some(libc::ENOSYS) | Some(libc::EINVAL) => {}
                // Out of space is exactly what this call exists to surface early.
                _ => return Err(crate::error::Error::Io(e)),
            }
        }
        return Ok(());
    }
    #[cfg(target_os = "macos")]
    {
        use std::os::unix::fs::MetadataExt;
        use std::os::unix::io::AsRawFd;

        // `F_PEOFPOSMODE` allocates `fst_length` bytes *beyond what the file
        // already has*, not up to a target size. Asking for `size` on a file
        // that `set_len` already sized therefore reserves a second copy — and
        // repeating it reserves another each time. Ask only for the shortfall,
        // which also makes this a no-op on a file that is already backed.
        let allocated = file.metadata()?.blocks().saturating_mul(512);
        if allocated >= size {
            return Ok(());
        }
        let want = (size - allocated) as libc::off_t;

        let mut store = libc::fstore_t {
            fst_flags: libc::F_ALLOCATECONTIG,
            fst_posmode: libc::F_PEOFPOSMODE,
            fst_offset: 0,
            fst_length: want,
            fst_bytesalloc: 0,
        };
        unsafe {
            // Contiguous first; fall back to any available extents.
            if libc::fcntl(file.as_raw_fd(), libc::F_PREALLOCATE, &mut store) == -1 {
                store.fst_flags = libc::F_ALLOCATEALL;
                let _ = libc::fcntl(file.as_raw_fd(), libc::F_PREALLOCATE, &mut store);
            }
        }
        return Ok(());
    }
    #[cfg(not(any(target_os = "linux", target_os = "macos")))]
    {
        let _ = (file, size);
        Ok(())
    }
}

fn set_mtime(path: &Path, mtime: i64) {
    #[cfg(unix)]
    {
        use std::ffi::CString;
        if let Ok(c) = CString::new(path.as_os_str().as_encoded_bytes()) {
            let times = [
                libc::timeval {
                    tv_sec: mtime as libc::time_t,
                    tv_usec: 0,
                },
                libc::timeval {
                    tv_sec: mtime as libc::time_t,
                    tv_usec: 0,
                },
            ];
            unsafe {
                libc::utimes(c.as_ptr(), times.as_ptr());
            }
        }
    }
    #[cfg(not(unix))]
    let _ = (path, mtime);
}

#[cfg(unix)]
fn pwrite(file: &File, buf: &[u8], offset: u64) -> io::Result<usize> {
    use std::os::unix::fs::FileExt;
    file.write_at(buf, offset)
}

#[cfg(windows)]
fn pwrite(file: &File, buf: &[u8], offset: u64) -> io::Result<usize> {
    use std::os::windows::fs::FileExt;
    file.seek_write(buf, offset)
}

#[cfg(unix)]
fn pread(file: &File, buf: &mut [u8], offset: u64) -> io::Result<usize> {
    use std::os::unix::fs::FileExt;
    file.read_at(buf, offset)
}

#[cfg(windows)]
fn pread(file: &File, buf: &mut [u8], offset: u64) -> io::Result<usize> {
    use std::os::windows::fs::FileExt;
    file.seek_read(buf, offset)
}

// ---------------------------------------------------------------------------
// Registry
// ---------------------------------------------------------------------------

/// One file's slot. The inner lock makes opening happen exactly once.
#[derive(Default)]
struct Slot(Mutex<Option<Arc<WriteHandle>>>);

/// Open write handles keyed by file id.
///
/// Handles are opened on the first chunk that arrives for a file rather than
/// up front: a manifest may name a million files, and a million open
/// descriptors is not a thing we can ask the OS for.
#[derive(Default)]
pub struct FileWriters {
    map: Mutex<HashMap<u32, Arc<Slot>>>,
}

impl FileWriters {
    pub fn new() -> Self {
        Self::default()
    }

    /// Fetch the handle for `file_id`, opening it with `make` if absent.
    ///
    /// `make` runs **at most once per file**, even when every worker hits the
    /// same file's first chunk simultaneously. That is not just an efficiency
    /// point: opening a destination preallocates it, so letting several racing
    /// callers each open the same file reserved the file's full size several
    /// times over and could exhaust the disk on a large transfer.
    ///
    /// The lock taken during `make` is per file, so opening one file never
    /// blocks writers on another.
    pub fn get_or_open<F>(&self, file_id: u32, make: F) -> Result<Arc<WriteHandle>>
    where
        F: FnOnce() -> Result<WriteHandle>,
    {
        let slot = {
            let mut map = self.map.lock();
            map.entry(file_id).or_default().clone()
        };
        let mut guard = slot.0.lock();
        if let Some(h) = guard.as_ref() {
            return Ok(h.clone());
        }
        let handle = Arc::new(make()?);
        *guard = Some(handle.clone());
        Ok(handle)
    }

    pub fn take(&self, file_id: u32) -> Option<Arc<WriteHandle>> {
        let slot = self.map.lock().remove(&file_id)?;
        let mut guard = slot.0.lock();
        guard.take()
    }

    pub fn len(&self) -> usize {
        self.map
            .lock()
            .values()
            .filter(|s| s.0.lock().is_some())
            .count()
    }

    pub fn is_empty(&self) -> bool {
        self.len() == 0
    }

    /// Drop every handle, deleting partial files. Used on unrecoverable abort.
    pub fn discard_all(&self) {
        let slots: Vec<_> = self.map.lock().drain().map(|(_, s)| s).collect();
        for s in slots {
            if let Some(h) = s.0.lock().take() {
                h.discard();
            }
        }
    }
}

#[cfg(test)]
mod tests {
    use super::*;

    #[test]
    fn writes_out_of_order_then_commits() {
        let tmp = tempfile::tempdir().unwrap();
        let dest = tmp.path().join("nested/dir/out.bin");
        let h = WriteHandle::open(&dest, 3000, true).unwrap();

        // Deliberately out of order, as chunks arrive across parallel streams.
        h.write_at(2000, &[3u8; 1000]).unwrap();
        h.write_at(0, &[1u8; 1000]).unwrap();
        h.write_at(1000, &[2u8; 1000]).unwrap();

        assert!(h.part_path().exists());
        assert!(!dest.exists(), "must not appear until committed");

        h.commit(0o640, 1_600_000_000, true).unwrap();
        assert!(dest.exists());
        assert!(!part_path_for(&dest).exists());

        let got = std::fs::read(&dest).unwrap();
        assert_eq!(got.len(), 3000);
        assert_eq!(&got[0..1000], &[1u8; 1000]);
        assert_eq!(&got[1000..2000], &[2u8; 1000]);
        assert_eq!(&got[2000..3000], &[3u8; 1000]);

        #[cfg(unix)]
        {
            use std::os::unix::fs::PermissionsExt;
            let m = std::fs::metadata(&dest).unwrap();
            assert_eq!(m.permissions().mode() & 0o777, 0o640);
        }
    }

    #[test]
    fn refuses_writes_past_declared_size() {
        let tmp = tempfile::tempdir().unwrap();
        let h = WriteHandle::open(&tmp.path().join("a.bin"), 100, false).unwrap();
        assert!(h.write_at(90, &[0u8; 20]).is_err());
        assert!(h.write_at(u64::MAX, &[0u8; 1]).is_err());
        h.write_at(90, &[0u8; 10]).unwrap();
    }

    #[test]
    fn concurrent_writers_share_one_handle() {
        let tmp = tempfile::tempdir().unwrap();
        let dest = tmp.path().join("big.bin");
        let h = Arc::new(WriteHandle::open(&dest, 64 * 4096, true).unwrap());

        let threads: Vec<_> = (0..8u64)
            .map(|t| {
                let h = h.clone();
                std::thread::spawn(move || {
                    for i in 0..8u64 {
                        let idx = t * 8 + i;
                        h.write_at(idx * 4096, &vec![idx as u8; 4096]).unwrap();
                    }
                })
            })
            .collect();
        for t in threads {
            t.join().unwrap();
        }

        h.commit(0o644, 0, true).unwrap();
        let got = std::fs::read(&dest).unwrap();
        for idx in 0..64u64 {
            let s = (idx * 4096) as usize;
            assert!(
                got[s..s + 4096].iter().all(|&b| b == idx as u8),
                "chunk {idx}"
            );
        }
    }

    #[test]
    fn reopening_a_partial_preserves_bytes() {
        let tmp = tempfile::tempdir().unwrap();
        let dest = tmp.path().join("resume.bin");
        {
            let h = WriteHandle::open(&dest, 8192, true).unwrap();
            h.write_at(0, &[9u8; 4096]).unwrap();
            h.sync().unwrap();
        }
        // A second open must not truncate what the first one wrote.
        let h = WriteHandle::open(&dest, 8192, true).unwrap();
        let mut buf = [0u8; 4096];
        h.read_at(0, &mut buf).unwrap();
        assert!(buf.iter().all(|&b| b == 9));
    }

    #[test]
    fn registry_opens_once_under_contention() {
        let tmp = tempfile::tempdir().unwrap();
        let reg = Arc::new(FileWriters::new());
        let opens = Arc::new(std::sync::atomic::AtomicUsize::new(0));

        let threads: Vec<_> = (0..16)
            .map(|_| {
                let reg = reg.clone();
                let opens = opens.clone();
                let path = tmp.path().join("shared.bin");
                std::thread::spawn(move || {
                    reg.get_or_open(1, || {
                        opens.fetch_add(1, std::sync::atomic::Ordering::SeqCst);
                        WriteHandle::open(&path, 1024, false)
                    })
                    .unwrap()
                })
            })
            .collect();
        let handles: Vec<_> = threads.into_iter().map(|t| t.join().unwrap()).collect();

        // Every caller must observe the same handle, whatever the race did.
        for h in &handles {
            assert!(Arc::ptr_eq(h, &handles[0]));
        }
        assert_eq!(reg.len(), 1);
        // The part this originally missed: opening is side-effecting — it
        // preallocates — so it must happen exactly once, not merely produce one
        // surviving handle.
        assert_eq!(
            opens.load(std::sync::atomic::Ordering::SeqCst),
            1,
            "the destination was opened more than once"
        );
    }

    #[cfg(unix)]
    fn allocated_bytes(p: &Path) -> u64 {
        use std::os::unix::fs::MetadataExt;
        std::fs::metadata(p)
            .map(|m| m.blocks().saturating_mul(512))
            .unwrap_or(0)
    }

    /// A preallocated file must reserve its own size and no more.
    ///
    /// This is the check that was missing: the reservation used to be issued
    /// after the file was sized, and macOS `F_PREALLOCATE` appends rather than
    /// fills, so an N-byte destination consumed 2N — or N per racing worker.
    #[cfg(unix)]
    #[test]
    fn preallocation_reserves_exactly_the_file_size() {
        let tmp = tempfile::tempdir().unwrap();
        let dest = tmp.path().join("prealloc.bin");
        let size = 64 * 1024 * 1024;

        {
            let h = WriteHandle::open(&dest, size, true).unwrap();
            h.write_at(0, &[1u8; 4096]).unwrap();
            h.sync().unwrap();
        }
        let after_first = allocated_bytes(&part_path_for(&dest));
        assert!(
            after_first <= size + size / 8,
            "a {size} byte file reserved {after_first} bytes on first open"
        );

        // Reopening, as a resumed transfer does, must not reserve again.
        for _ in 0..3 {
            let h = WriteHandle::open(&dest, size, true).unwrap();
            h.write_at(0, &[1u8; 4096]).unwrap();
            h.sync().unwrap();
        }
        let after_reopens = allocated_bytes(&part_path_for(&dest));
        assert!(
            after_reopens <= size + size / 8,
            "reopening grew the reservation to {after_reopens} for a {size} byte file"
        );
    }

    #[test]
    fn zero_length_file_commits() {
        let tmp = tempfile::tempdir().unwrap();
        let dest = tmp.path().join("empty.bin");
        let h = WriteHandle::open(&dest, 0, true).unwrap();
        h.commit(0o644, 0, true).unwrap();
        assert!(dest.exists());
        assert_eq!(std::fs::metadata(&dest).unwrap().len(), 0);
    }
}