gwseq_io/source/

sink.rs

//! The write side.
//!
//! Mirrors [`ByteSource`](crate::source::ByteSource) for the bbi writer, which
//! needs positioned writes: its header carries counts and offsets whose values
//! are only known once the sections, the R-tree and the zoom levels have been
//! built, so they are reserved at open and patched at close.
//!
//! One handle serves the writer's read-back pass too. Building zoom levels
//! means reading the data section back while the file is still open for
//! writing; a second read-only handle would see only what had been flushed. A
//! `std::fs::File` opened `.read(true).write(true)` does both, and `read_at`
//! on it is the same positioned read as everywhere else — one handle, one
//! flush, no second view of the file that can lag.

use std::fs::File;
use std::path::Path;

use bytes::Bytes;

use crate::error::{Error, Result};
use crate::source::ByteSource;

/// Bytes staged before a write reaches the file. A block is a few kilobytes, so
/// this batches a few hundred of them into one call.
pub const WRITE_BUFFER_SIZE: usize = 1 << 20;

pub trait ByteSink: Send {
    fn path(&self) -> &str;

    /// Append, returning the offset the data landed at.
    fn append(&mut self, data: &[u8]) -> Result<u64>;

    /// Overwrite bytes already written — a reserved header field, an item count.
    fn write_all_at(&mut self, offset: u64, data: &[u8]) -> Result<()>;

    /// Bytes written so far, which is also the offset the next append lands at.
    fn position(&self) -> u64;

    fn flush(&mut self) -> Result<()>;

    /// Finish. Idempotent, and also called from `Drop` — but a `Drop` that
    /// fails cannot report, so the writer's own `close()` is what callers use
    /// and what returns the error.
    fn close(&mut self) -> Result<()>;
}

/// A local file opened for both directions, so the zoom pass reads back through
/// the handle that wrote.
#[derive(Debug)]
pub struct LocalSink {
    /// `None` once closed.
    file: Option<File>,
    path: String,
    /// Staged bytes and the offset the first of them belongs at. `position` is
    /// the end of the staged bytes, i.e. what a caller sees as the file's
    /// length — the writer hands out offsets from it before the bytes have
    /// reached the disk.
    stage: Vec<u8>,
    stage_offset: u64,
    position: u64,
}

impl LocalSink {
    pub fn create(path: impl AsRef<Path>) -> Result<Self> {
        let path = path.as_ref();
        let display = path.to_string_lossy().into_owned();
        let file = File::options()
            // Truncating rather than appending: a writer starts a file, it does
            // not continue one, and a stale tail past the trailing magic would
            // be the one thing a bbi reader cannot notice.
            .read(true)
            .write(true)
            .create(true)
            .truncate(true)
            .open(path)
            .map_err(|e| Error::io(&display, e))?;
        Ok(Self {
            file: Some(file),
            path: display,
            stage: Vec::with_capacity(WRITE_BUFFER_SIZE),
            stage_offset: 0,
            position: 0,
        })
    }

    /// Drop the handle and remove the file, without flushing what is staged.
    ///
    /// What a writer that has given up wants: `create` truncated the path at
    /// open, so leaving quietly still leaves a zero-magic stub behind that the
    /// caller has to clean up. Nothing staged is worth writing to a file about
    /// to be unlinked.
    ///
    /// The handle goes first because Windows refuses to remove an open file.
    /// A removal that fails — the path already gone, a directory that is not
    /// writable — is not reported: this is the cleanup path of an error that is
    /// already on its way to the caller, and a second error would only hide the
    /// first.
    pub fn discard(&mut self) {
        self.file = None;
        self.stage.clear();
        let _ = std::fs::remove_file(&self.path);
    }

    /// A read view over what has been written, for the zoom pass. Flushes
    /// first, so the view never lags the writer.
    pub fn as_source(&mut self) -> Result<SinkSource<'_>> {
        self.flush()?;
        let path = self.path.clone();
        let file = self.file.as_ref().ok_or(Error::Closed { path })?;
        Ok(SinkSource {
            file,
            path: &self.path,
            len: self.position,
        })
    }

    fn handle(&self) -> Result<&File> {
        self.file.as_ref().ok_or_else(|| Error::Closed {
            path: self.path.clone(),
        })
    }

    /// Push the staged bytes out.
    ///
    /// `what` describes the moment for the error message: a full device is
    /// usually only reported here, the appends before it having gone into the
    /// buffer, so the wording says which of the three moments it was.
    fn spill(&mut self, what: &str) -> Result<()> {
        if self.stage.is_empty() {
            return Ok(());
        }
        let offset = self.stage_offset;
        {
            let file = self.handle()?;
            // The `io::Error` says whether the device is full, read-only or
            // gone, which is the difference between "make room" and "this was
            // never going to work". Dropping it left every failed write saying
            // the same sentence.
            pwrite(file, offset, &self.stage)
                .map_err(|e| Error::write(&self.path, format!("{what}: {e}")))?;
        }
        self.stage_offset += self.stage.len() as u64;
        self.stage.clear();
        Ok(())
    }
}

impl ByteSink for LocalSink {
    fn path(&self) -> &str {
        &self.path
    }

    fn append(&mut self, data: &[u8]) -> Result<u64> {
        let offset = self.position;
        self.stage.extend_from_slice(data);
        self.position += data.len() as u64;
        if self.stage.len() >= WRITE_BUFFER_SIZE {
            self.spill("the device reported an error while writing to it")?;
        }
        Ok(offset)
    }

    fn write_all_at(&mut self, offset: u64, data: &[u8]) -> Result<()> {
        // Staged bytes go out first: a patch may land inside them — the header
        // is patched at close, and by then nothing of it is staged — and
        // writing around them would let the spill overwrite the patch.
        self.spill("the device reported an error while writing to it")?;
        let file = self.handle()?;
        pwrite(file, offset, data).map_err(|e| {
            Error::write(
                &self.path,
                format!("the device reported an error while writing to it: {e}"),
            )
        })?;
        // A patch may reach past what has been appended: the writer reserves
        // its prefix by appending zeroes, so in practice it never does, but a
        // sink that let `position` lag the file would hand out an offset that
        // overwrites.
        self.position = self.position.max(offset + data.len() as u64);
        self.stage_offset = self.stage_offset.max(self.position);
        Ok(())
    }

    fn position(&self) -> u64 {
        self.position
    }

    fn flush(&mut self) -> Result<()> {
        self.spill("the device reported an error while flushing it")?;
        // No `sync_all`: the read view is the same handle, so it sees the
        // written bytes without the file reaching the disk, and forcing it to
        // would cost a fsync per zoom level for nothing.
        Ok(())
    }

    fn close(&mut self) -> Result<()> {
        if self.file.is_none() {
            return Ok(());
        }
        self.spill("the device reported an error while closing it, so the file is incomplete")?;
        // Dropped rather than `sync_all`ed, for the same reason `flush` does
        // not: durability across a power cut is not what this file promises,
        // and the caller can fsync the path if it is.
        self.file = None;
        Ok(())
    }
}

impl Drop for LocalSink {
    fn drop(&mut self) {
        let _ = self.close();
    }
}

/// The read half of a [`LocalSink`], borrowed for as long as the zoom pass
/// needs it.
///
/// Not `LocalSource`: that snapshots the length at open, which is exactly wrong
/// for a file still being written. This one carries the sink's own position,
/// which is the only length that is true.
#[derive(Debug)]
pub struct SinkSource<'a> {
    file: &'a File,
    path: &'a str,
    len: u64,
}

impl ByteSource for SinkSource<'_> {
    fn path(&self) -> &str {
        self.path
    }

    fn len(&self) -> Result<u64> {
        Ok(self.len)
    }

    fn read_at(&self, offset: u64, len: usize) -> Result<Bytes> {
        if len == 0 || offset >= self.len {
            return Ok(Bytes::new());
        }
        let len = len.min((self.len - offset) as usize);
        let mut buf = vec![0u8; len];
        let read = pread(self.file, offset, &mut buf).map_err(|e| Error::io(self.path, e))?;
        buf.truncate(read);
        Ok(Bytes::from(buf))
    }

    fn close(&self) {
        // The sink owns the handle; a borrowed view does not get to close it.
    }
}

/// Positioned write, looping until the whole buffer is out.
///
/// One `write_at` may take less than it was given, so the loop is not optional.
#[cfg(unix)]
fn pwrite(file: &File, offset: u64, buf: &[u8]) -> std::io::Result<()> {
    use std::os::unix::fs::FileExt;
    let mut written = 0;
    while written < buf.len() {
        match file.write_at(&buf[written..], offset + written as u64) {
            Ok(0) => {
                return Err(std::io::Error::new(
                    std::io::ErrorKind::WriteZero,
                    "write returned zero bytes",
                ))
            }
            Ok(n) => written += n,
            Err(e) if e.kind() == std::io::ErrorKind::Interrupted => {}
            Err(e) => return Err(e),
        }
    }
    Ok(())
}

#[cfg(windows)]
fn pwrite(file: &File, offset: u64, buf: &[u8]) -> std::io::Result<()> {
    use std::os::windows::fs::FileExt;
    let mut written = 0;
    while written < buf.len() {
        match file.seek_write(&buf[written..], offset + written as u64) {
            Ok(0) => {
                return Err(std::io::Error::new(
                    std::io::ErrorKind::WriteZero,
                    "write returned zero bytes",
                ))
            }
            Ok(n) => written += n,
            Err(e) if e.kind() == std::io::ErrorKind::Interrupted => {}
            Err(e) => return Err(e),
        }
    }
    Ok(())
}

#[cfg(unix)]
fn pread(file: &File, offset: u64, buf: &mut [u8]) -> std::io::Result<usize> {
    use std::os::unix::fs::FileExt;
    let mut filled = 0;
    while filled < buf.len() {
        match file.read_at(&mut buf[filled..], offset + filled as u64) {
            Ok(0) => break,
            Ok(n) => filled += n,
            Err(e) if e.kind() == std::io::ErrorKind::Interrupted => {}
            Err(e) => return Err(e),
        }
    }
    Ok(filled)
}

#[cfg(windows)]
fn pread(file: &File, offset: u64, buf: &mut [u8]) -> std::io::Result<usize> {
    use std::os::windows::fs::FileExt;
    let mut filled = 0;
    while filled < buf.len() {
        match file.seek_read(&mut buf[filled..], offset + filled as u64) {
            Ok(0) => break,
            Ok(n) => filled += n,
            Err(e) if e.kind() == std::io::ErrorKind::Interrupted => {}
            Err(e) => return Err(e),
        }
    }
    Ok(filled)
}

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

    fn temp(name: &str) -> std::path::PathBuf {
        let dir = std::env::temp_dir().join("gwseq_sink_tests");
        std::fs::create_dir_all(&dir).unwrap();
        dir.join(name)
    }

    #[test]
    fn appends_land_end_to_end_and_report_their_offsets() {
        let path = temp("append.bin");
        let mut sink = LocalSink::create(&path).unwrap();
        assert_eq!(sink.append(b"hello").unwrap(), 0);
        assert_eq!(sink.append(b" world").unwrap(), 5);
        assert_eq!(sink.position(), 11);
        sink.close().unwrap();
        assert_eq!(std::fs::read(&path).unwrap(), b"hello world");
        std::fs::remove_file(&path).ok();
    }

    #[test]
    fn a_patch_overwrites_staged_bytes_rather_than_racing_them() {
        // The staged bytes have not reached the file when the patch is made,
        // so a sink that wrote the patch and then spilled the stage would lose
        // it. This is the header case, where every reserved field is patched.
        let path = temp("patch.bin");
        let mut sink = LocalSink::create(&path).unwrap();
        sink.append(&[0u8; 16]).unwrap();
        sink.append(b"tail").unwrap();
        sink.write_all_at(4, b"ABCD").unwrap();
        sink.close().unwrap();
        let bytes = std::fs::read(&path).unwrap();
        assert_eq!(&bytes[0..4], &[0, 0, 0, 0]);
        assert_eq!(&bytes[4..8], b"ABCD");
        assert_eq!(&bytes[16..20], b"tail");
        std::fs::remove_file(&path).ok();
    }

    #[test]
    fn the_read_view_sees_what_has_only_been_staged() {
        // What the zoom pass depends on: it reads back sections it wrote a
        // moment ago, most of which are still in the buffer.
        let path = temp("readback.bin");
        let mut sink = LocalSink::create(&path).unwrap();
        sink.append(b"0123456789").unwrap();
        {
            let source = sink.as_source().unwrap();
            assert_eq!(source.len().unwrap(), 10);
            assert_eq!(&source.read_at(3, 4).unwrap()[..], b"3456");
        }
        // And writing goes on afterwards, at the right offset.
        assert_eq!(sink.append(b"abc").unwrap(), 10);
        sink.close().unwrap();
        assert_eq!(std::fs::read(&path).unwrap(), b"0123456789abc");
        std::fs::remove_file(&path).ok();
    }

    #[test]
    fn a_write_larger_than_the_buffer_still_lands_whole() {
        let path = temp("large.bin");
        let mut sink = LocalSink::create(&path).unwrap();
        let big = vec![7u8; WRITE_BUFFER_SIZE + 1234];
        sink.append(&big).unwrap();
        sink.append(b"end").unwrap();
        sink.close().unwrap();
        let bytes = std::fs::read(&path).unwrap();
        assert_eq!(bytes.len(), big.len() + 3);
        assert!(bytes[..big.len()].iter().all(|b| *b == 7));
        assert_eq!(&bytes[big.len()..], b"end");
        std::fs::remove_file(&path).ok();
    }

    #[test]
    fn closing_twice_is_harmless_and_writing_after_it_is_not() {
        let path = temp("closed.bin");
        let mut sink = LocalSink::create(&path).unwrap();
        sink.append(b"x").unwrap();
        sink.close().unwrap();
        sink.close().unwrap();
        // The stage is empty, so the failure comes from the handle being gone.
        let err = sink.write_all_at(0, b"y").unwrap_err();
        assert!(matches!(err, Error::Closed { .. }), "{err}");
        std::fs::remove_file(&path).ok();
    }

    #[test]
    fn a_truncating_create_leaves_no_tail_of_an_older_file() {
        let path = temp("truncate.bin");
        std::fs::write(&path, vec![9u8; 4096]).unwrap();
        let mut sink = LocalSink::create(&path).unwrap();
        sink.append(b"new").unwrap();
        sink.close().unwrap();
        assert_eq!(std::fs::read(&path).unwrap(), b"new");
        std::fs::remove_file(&path).ok();
    }
}