commonware_runtime/utils/buffer/

write.rs

use crate::{
    Blob, Buf, Error, Handle, IoBufs, ReadOptions, WriteOptions,
    buffer::{SyncState, tip::Buffer},
};
use std::num::NonZeroUsize;

/// A writer that buffers the raw content of a [Blob] to optimize the performance of appending or
/// updating data.
///
/// # Allocation Semantics
///
/// - [Self::new] starts with a detached tip buffer and allocates backing on first buffered write.
/// - Subsequent writes reuse that backing, copy-on-write allocation only occurs when buffered data
///   is shared (for example, after handing out immutable views) or a merge needs more capacity.
/// - Sparse writes merged into tip extend logical length and zero-fill any gap in-buffer.
/// - Flush paths ([Self::sync], [Self::resize], overlap flushes in [Self::write_at]) hand drained
///   bytes to the blob and leave the tip detached until the next buffered write.
///
/// # Access
///
/// [Write] is a single-owner buffered handle that owns mutation ordering and durability
/// bookkeeping for the wrapped [Blob]. Raw [Blob] handles shared before wrapping observe only
/// flushed data and may not see the latest buffered writes until [Self::sync], [Self::resize], or
/// an overlapping [Self::write_at] flushes them. Those raw handles must not be used to write,
/// resize, or otherwise mutate the blob while a [Write] exists. External mutations bypass the
/// buffer state and [Self::sync] may use [Blob::write_at] with [WriteOptions::SYNC], which is
/// not a durability barrier for those external mutations.
///
/// Storage-mutating functions consume the writer and return it only on success: an error (or a
/// dropped future) destroys the handle.
///
/// # Example
///
/// ```
/// use commonware_runtime::{Runner, buffer::{Write, Read}, Blob, Error, Storage, deterministic};
/// use commonware_utils::NZUsize;
///
/// let executor = deterministic::Runner::default();
/// executor.start(|context| async move {
///     // Open a blob for writing
///     let (blob, size) = context.open("my_partition", b"my_data").await.expect("unable to open blob");
///     assert_eq!(size, 0);
///
///     // Create a buffered writer with 16-byte buffer
///     let blob = Write::new(blob, 0, NZUsize!(16));
///     let blob = blob.write_at(0, b"hello").await.expect("write failed");
///     let blob = blob.sync().await.expect("sync failed");
///
///     // Write more data in multiple flushes
///     let blob = blob.write_at(5, b" world").await.expect("write failed");
///     let blob = blob.write_at(11, b"!").await.expect("write failed");
///     let blob = blob.sync().await.expect("sync failed");
///
///     // Release the writer and read back the persisted data through a new open.
///     drop(blob);
///     let (blob, size) = context.open("my_partition", b"my_data").await.expect("unable to reopen blob");
///     let mut reader = Read::new(blob, size, NZUsize!(8));
///     let buf = reader.read(size as usize).await.expect("read failed");
///     assert_eq!(buf.coalesce().as_ref(), b"hello world!");
/// });
/// ```
pub struct Write<B: Blob> {
    /// The underlying blob to write to.
    blob: B,

    /// Buffered bytes at the logical tip of the blob.
    buffer: Buffer,

    /// Durability state for plain writes and range-sync writes.
    sync_state: SyncState,
}

impl<B: Blob> Write<B> {
    /// Creates a new [Write] that buffers up to `capacity` bytes of data to be appended to the tip
    /// of `blob` with the provided `size`.
    pub fn new(blob: B, size: u64, capacity: NonZeroUsize) -> Self {
        Self {
            blob,
            buffer: Buffer::new(size, capacity.get()),
            // Existing blob contents may not be durable yet.
            sync_state: SyncState::Dirty,
        }
    }

    /// Returns the current logical size of the blob including any buffered data.
    ///
    /// This represents the total size of data that would be present after flushing.
    pub const fn size(&self) -> u64 {
        self.buffer.size()
    }

    /// Whether [Self::sync] would write buffered bytes, sync the blob, or observe a started sync.
    pub const fn needs_sync(&self) -> bool {
        !self.buffer.is_empty() || !self.sync_state.is_clean()
    }

    /// Read exactly `len` immutable bytes starting at `offset`.
    pub async fn read_at(&self, offset: u64, len: usize) -> Result<IoBufs, Error> {
        // Ensure the read doesn't overflow.
        let end_offset = offset
            .checked_add(len as u64)
            .ok_or(Error::OffsetOverflow)?;

        // If the data required is beyond the size of the blob, return an error.
        if end_offset > self.buffer.size() {
            return Err(Error::BlobInsufficientLength);
        }

        // Keep the zero-length fast path after the bounds check so offset > size still preserves
        // the BlobInsufficientLength contract.
        if len == 0 {
            return Ok(IoBufs::default());
        }

        // Entirely in buffered tip.
        if offset >= self.buffer.offset {
            let start = (offset - self.buffer.offset) as usize;
            let end = start + len;
            return Ok(self.buffer.slice(start..end).into());
        }

        // Entirely in blob.
        if end_offset <= self.buffer.offset {
            return self.read_blob(offset, len).await;
        }

        // Overlaps blob and buffered tip.
        let blob_len = (self.buffer.offset - offset) as usize;
        let tip_len = len - blob_len;
        let tip = self.buffer.slice(..tip_len);

        let mut blob = self.read_blob(offset, blob_len).await?;
        blob.append(tip);
        Ok(blob)
    }

    /// Read bytes from the underlying blob.
    async fn read_blob(&self, offset: u64, len: usize) -> Result<IoBufs, Error> {
        Ok(self
            .blob
            .read_at(offset, len, ReadOptions::default())
            .await?
            .freeze())
    }

    /// Merge `buf` into the in-memory tip buffer at `offset`.
    ///
    /// Returns whether `buf` fit. Returns `false` without changing the writer when it does not
    /// fit. Performs no I/O and does not make the write durable. Callers can fall back to
    /// [`Self::write_at`] when it does not fit.
    pub fn try_write_at(&mut self, offset: u64, buf: &[u8]) -> bool {
        buf.is_empty()
            || (offset.checked_add(buf.len() as u64).is_some() && self.buffer.merge(buf, offset))
    }

    /// Write bytes from `buf` at `offset`.
    ///
    /// Data is merged into the in-memory tip buffer when possible, otherwise buffered data may be
    /// flushed and chunks are written directly to the underlying blob.
    ///
    /// Returns [Error::OffsetOverflow] when `offset + bufs.len()` overflows.
    pub async fn write_at(
        mut self,
        offset: u64,
        bufs: impl Into<IoBufs> + Send,
    ) -> Result<Self, Error> {
        let mut bufs = bufs.into();

        // Ensure the write doesn't overflow.
        offset
            .checked_add(bufs.remaining() as u64)
            .ok_or(Error::OffsetOverflow)?;

        // Process each chunk of the input buffer, attempting to merge into the tip buffer
        // or writing directly to the underlying blob.
        let mut current_offset = offset;
        while bufs.has_remaining() {
            let chunk = bufs.chunk();
            let chunk_len = chunk.len();

            // Chunk falls entirely within the buffer's current range and can be merged.
            if self.buffer.merge(chunk, current_offset) {
                bufs.advance(chunk_len);
                current_offset += chunk_len as u64;
                continue;
            }

            // Chunk cannot be merged, so flush the buffer if the range overlaps, and check
            // if merge is possible after.
            let chunk_end = current_offset + chunk_len as u64;
            if self.buffer.offset < chunk_end
                && let Some((old_buf, old_offset)) = self.buffer.take()
            {
                self.sync_state
                    .write_at(&self.blob, old_offset, old_buf, WriteOptions::default())
                    .await?;
                if self.buffer.merge(chunk, current_offset) {
                    bufs.advance(chunk_len);
                    current_offset += chunk_len as u64;
                    continue;
                }
            }

            // Chunk could not be merged (exceeds buffer capacity or outside its range), so
            // write directly. Note that we may end up writing an intersecting range twice:
            // once when the buffer is flushed above, then again when we write the chunk
            // below. Removing this inefficiency may not be worth the additional complexity.
            let direct = bufs.split_to(chunk_len);
            self.sync_state
                .write_at(&self.blob, current_offset, direct, WriteOptions::default())
                .await?;
            current_offset += chunk_len as u64;

            // Maintain the "buffer at tip" invariant by advancing offset to the end of this
            // write if it extended the underlying blob.
            self.buffer.offset = self.buffer.offset.max(current_offset);
        }

        Ok(self)
    }

    /// Resize the logical blob to `len`.
    ///
    /// If buffered data exists and the resize extends beyond current size, buffered data is flushed
    /// before resizing the underlying blob.
    pub async fn resize(mut self, len: u64) -> Result<Self, Error> {
        // Flush buffered data to the underlying blob.
        //
        // This can only happen if the new size is greater than the current size.
        if let Some((buf, offset)) = self.buffer.resize(len) {
            self.sync_state
                .write_at(&self.blob, offset, buf, WriteOptions::default())
                .await?;
        }

        self.sync_state.resize(&self.blob, len).await?;

        Ok(self)
    }

    /// Flush buffered bytes and durably sync mutations tracked by this writer.
    pub async fn sync(mut self) -> Result<Self, Error> {
        match self.buffer.take() {
            Some((buf, offset)) => self.write_blob_sync(offset, buf).await?,
            None => self.sync_blob().await?,
        }
        Ok(self)
    }

    /// Flush buffered bytes and begin durably syncing mutations tracked by this writer.
    ///
    /// Awaiting the returned [`Handle`] waits for the same durability guarantee as [`Self::sync`]
    /// for the state flushed by this call. Later calls to [`Self::sync`] and writer methods that
    /// mutate the blob wait before issuing blob operations. A failure of the started sync is
    /// reported by the handle and by the next such call. Flush errors are returned directly.
    pub async fn start_sync(mut self) -> Result<(Self, Handle<()>), Error> {
        // Complete buffered writes before requesting their durability barrier.
        if let Some((buf, offset)) = self.buffer.take() {
            self.sync_state
                .write_at(&self.blob, offset, buf, WriteOptions::default())
                .await?;
        }

        // The runtime runs the sync. The returned handle observes its completion.
        let handle = self.sync_state.start_sync(&self.blob).await;
        Ok((self, handle))
    }

    /// Wait for any started sync to complete without starting a new sync.
    pub async fn wait_for_sync(mut self) -> Result<Self, Error> {
        self.sync_state.wait_for_pending().await?;
        Ok(self)
    }

    /// Write bytes to the underlying blob and make them durable.
    ///
    /// Uses [`Blob::write_at`] with [`WriteOptions::SYNC`] when there are no earlier unsynced
    /// mutations. Otherwise, writes the bytes and then syncs the blob.
    async fn write_blob_sync(
        &mut self,
        offset: u64,
        bufs: impl Into<IoBufs> + Send,
    ) -> Result<(), Error> {
        self.sync_state
            .write_at(&self.blob, offset, bufs, WriteOptions::SYNC)
            .await
    }

    /// Sync the underlying blob if there are unsynced mutations.
    async fn sync_blob(&mut self) -> Result<(), Error> {
        self.sync_state.sync(&self.blob).await
    }
}