hyperfile 0.9.0

//! Scatter-gather body for segment uploads.
//!
//! `SegmentBody` is the data-plane analogue of `Vec<u8>` for a flush
//! segment, except the bytes never live in one contiguous allocation.
//! It accumulates `Bytes` pieces — one for the segment header /
//! summary, one per dirty meta block, one per dirty data block — and
//! produces an `aws_smithy_types::body::SdkBody` that streams those
//! pieces directly to the S3 SDK without a final memcpy into a
//! single buffer.
//!
//! Why this matters: with the previous `Vec<u8>` segment buffer,
//! every flush of a *D*-byte payload allocated a contiguous
//! `segment_buffer_size` (typically 256 MiB or 1 GiB) buffer up
//! front, regardless of *D*. The actual data was then memcpy'd in,
//! and the AWS SDK's `SdkBody::from(&[u8])` did a second copy via
//! `Bytes::copy_from_slice`. Memory peak during flush was
//! `segment_buffer_size + payload`, dominated by the upfront
//! allocation.
//!
//! With `SegmentBody`, the only allocation is the sum of the
//! pieces themselves (≈ payload + a few KiB of header). The SDK
//! pulls one frame at a time via `http_body::Body::poll_data`, so
//! it does not need to materialize the full body in its own
//! buffers either.

use std::convert::Infallible;
use std::sync::Arc;

use crate::block_source::BlockSource;
use std::pin::Pin;
use std::task::{Context, Poll};

use bytes::Bytes;
use http::HeaderMap;
use http_body::{Body, SizeHint};

/// One piece of a segment body: bytes in hand, or a range to be read from a
/// [`BlockSource`](crate::block_source::BlockSource) when the part holding it is
/// uploaded (see [`SegmentBody::materialize`]).
// The WAL build keeps whole segments in one buffer and never references.
#[cfg_attr(feature = "wal", allow(dead_code))]
#[derive(Clone)]
pub(crate) enum Piece {
    Bytes(Bytes),
    Lazy { src: Arc<dyn BlockSource>, off: u64, len: usize },
}

impl Piece {
    fn len(&self) -> usize {
        match self {
            Piece::Bytes(b) => b.len(),
            Piece::Lazy { len, .. } => *len,
        }
    }

    /// `[at, at + take)` of this piece, without reading anything.
    fn sub(&self, at: usize, take: usize) -> Piece {
        match self {
            Piece::Bytes(b) => Piece::Bytes(b.slice(at..at + take)),
            Piece::Lazy { src, off, .. } =>
                Piece::Lazy { src: src.clone(), off: off + at as u64, len: take },
        }
    }
}

/// A growable list of pieces that together form the body of one S3 segment.
/// Push pieces in upload order; cheap to clone (bytes are refcounted, a lazy
/// piece is a source handle and a range).
#[derive(Clone, Default)]
pub struct SegmentBody {
    pieces: Vec<Piece>,
    total_len: usize,
}

impl SegmentBody {
    pub(crate) fn new() -> Self {
        Self {
            pieces: Vec::new(),
            total_len: 0,
        }
    }

    pub(crate) fn with_capacity(n_pieces: usize) -> Self {
        Self {
            pieces: Vec::with_capacity(n_pieces),
            total_len: 0,
        }
    }

    pub(crate) fn push(&mut self, b: Bytes) {
        self.push_piece(Piece::Bytes(b));
    }

    /// Append `len` bytes to be read from `src` at `off` when uploaded. Adjacent
    /// lazy ranges of the same source are merged, so a run of referenced blocks is
    /// one read.
    // The WAL build keeps whole segments in one buffer and never references.
    #[cfg_attr(feature = "wal", allow(dead_code))]
    pub(crate) fn push_lazy(&mut self, src: Arc<dyn BlockSource>, off: u64, len: usize) {
        if let Some(Piece::Lazy { src: s, off: o, len: l }) = self.pieces.last_mut() {
            if Arc::ptr_eq(s, &src) && *o + *l as u64 == off {
                *l += len;
                self.total_len += len;
                return;
            }
        }
        self.push_piece(Piece::Lazy { src, off, len });
    }

    pub(crate) fn push_piece(&mut self, p: Piece) {
        self.total_len += p.len();
        self.pieces.push(p);
    }

    /// Total payload length across all pieces. Equivalent to the
    /// final segment file size on S3.
    pub fn len(&self) -> usize {
        self.total_len
    }

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

    // The WAL build keeps whole segments in one buffer and never references.
    #[cfg_attr(feature = "wal", allow(dead_code))]
    pub(crate) fn pieces(&self) -> &[Piece] {
        &self.pieces
    }

    /// Whether any piece still has to be read from its source.
    pub(crate) fn has_lazy(&self) -> bool {
        self.pieces.iter().any(|p| matches!(p, Piece::Lazy { .. }))
    }

    /// The same body with every lazy piece read into memory.
    ///
    /// Called per part of a multipart upload, just before that part is sent, so
    /// what is resident at once is bounded by the part size times the upload
    /// concurrency, not by the segment. The reads run on the blocking pool: a
    /// source does ordinary file I/O and must not stall the runtime that drives
    /// the other parts.
    pub(crate) async fn materialize(self) -> std::io::Result<SegmentBody> {
        if !self.has_lazy() {
            return Ok(self);
        }
        tokio::task::spawn_blocking(move || self.materialize_blocking()).await
            .map_err(std::io::Error::other)?
    }

    pub(crate) fn materialize_blocking(self) -> std::io::Result<SegmentBody> {
        let mut out = SegmentBody::with_capacity(self.pieces.len());
        for p in self.pieces {
            match p {
                Piece::Bytes(b) => out.push(b),
                Piece::Lazy { src, off, len } => {
                    let mut buf = vec![0u8; len];
                    src.read_at(off, &mut buf)?;
                    out.push(Bytes::from(buf));
                }
            }
        }
        Ok(out)
    }

    /// The bytes, for a body with no lazy piece left. Panics otherwise: every
    /// caller either materialized first or never had one.
    pub(crate) fn byte_pieces(&self) -> impl Iterator<Item = &Bytes> {
        self.pieces.iter().map(|p| match p {
            Piece::Bytes(b) => b,
            Piece::Lazy { .. } => panic!("a lazy piece reached a path that needs bytes; materialize first"),
        })
    }

    /// Take a sub-body covering byte range `[start, start + len)`
    /// of the concatenation of pieces, walking them and slicing
    /// individual pieces at the boundaries (a byte slice is a
    /// refcounted view and a lazy slice a narrower range; nothing is read).
    ///
    /// Used by the multipart-upload path to slice the full body
    /// into part-sized sub-bodies.
    pub(crate) fn slice(&self, start: usize, len: usize) -> Self {
        assert!(
            start + len <= self.total_len,
            "SegmentBody::slice out of range: start={} len={} total={}",
            start, len, self.total_len,
        );
        if len == 0 {
            return Self::new();
        }
        // Sized for the range, not the whole body: each part of a multipart upload is
        // a slice, and room for every piece of the segment in each of them came to
        // 64 lists of 8 MiB for 1 GiB in 16 MiB parts. In a segment every piece is at
        // least 4 KiB -- a data block, a meta block, or the summary, rounded up to
        // 4 KiB -- so a range spans at most one piece per 4 KiB and one more at each
        // end. A body of smaller pieces is still sliced correctly; the list grows.
        let mut out = Self::with_capacity((len / 4096 + 2).min(self.pieces.len()));
        let mut cursor = 0usize;
        let mut remaining = len;
        for piece in &self.pieces {
            if remaining == 0 {
                break;
            }
            let piece_start = cursor;
            let piece_end = cursor + piece.len();
            cursor = piece_end;

            // Skip pieces entirely before `start`.
            if piece_end <= start {
                continue;
            }
            // Trim the head of the first overlapping piece.
            let piece_off = if piece_start < start {
                start - piece_start
            } else {
                0
            };
            let avail = piece.len() - piece_off;
            let take = avail.min(remaining);
            out.push_piece(piece.sub(piece_off, take));
            remaining -= take;
        }
        out
    }

    /// Convert into an SdkBody that streams the pieces one frame
    /// at a time. The body is wrapped in `SdkBody::retryable` so
    /// the SDK can rebuild a fresh body for retries (cloning the
    /// `Vec<Bytes>` is cheap — each piece is refcounted).
    ///
    /// The body must have been [materialized](Self::materialize).
    pub(crate) fn into_sdk_body(self) -> aws_smithy_types::body::SdkBody {
        let pieces: Vec<Bytes> = self.byte_pieces().cloned().collect();
        aws_smithy_types::body::SdkBody::retryable(move || {
            aws_smithy_types::body::SdkBody::from_body_0_4(ScatterBody::new(pieces.clone()))
        })
    }
}

/// `http_body::Body` implementation backing `SegmentBody`. One
/// `Bytes` is yielded per `poll_data` call.
struct ScatterBody {
    iter: std::vec::IntoIter<Bytes>,
    remaining: u64,
}

impl ScatterBody {
    fn new(pieces: Vec<Bytes>) -> Self {
        let remaining: u64 = pieces.iter().map(|b| b.len() as u64).sum();
        Self {
            iter: pieces.into_iter(),
            remaining,
        }
    }
}

impl Body for ScatterBody {
    type Data = Bytes;
    type Error = Infallible;

    fn poll_data(
        self: Pin<&mut Self>,
        _cx: &mut Context<'_>,
    ) -> Poll<Option<Result<Bytes, Self::Error>>> {
        let me = self.get_mut();
        match me.iter.next() {
            Some(b) => {
                me.remaining = me.remaining.saturating_sub(b.len() as u64);
                Poll::Ready(Some(Ok(b)))
            }
            None => Poll::Ready(None),
        }
    }

    fn poll_trailers(
        self: Pin<&mut Self>,
        _cx: &mut Context<'_>,
    ) -> Poll<Result<Option<HeaderMap>, Self::Error>> {
        Poll::Ready(Ok(None))
    }

    fn size_hint(&self) -> SizeHint {
        SizeHint::with_exact(self.remaining)
    }

    fn is_end_stream(&self) -> bool {
        self.remaining == 0
    }
}

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

    #[test]
    fn segment_body_len_tracks_pieces() {
        let mut b = SegmentBody::new();
        assert_eq!(b.len(), 0);
        assert!(b.is_empty());
        b.push(Bytes::from_static(b"hello "));
        b.push(Bytes::from_static(b"world"));
        assert_eq!(b.len(), 11);
        assert!(!b.is_empty());
        assert_eq!(b.pieces().len(), 2);
    }

    #[test]
    fn segment_body_slice_within_one_piece() {
        let mut b = SegmentBody::new();
        b.push(Bytes::from_static(b"abcdefgh"));
        let s = b.slice(2, 4);
        assert_eq!(s.len(), 4);
        assert_eq!(s.byte_pieces().next().unwrap().as_ref(), b"cdef");
    }

    #[test]
    fn segment_body_slice_across_pieces() {
        let mut b = SegmentBody::new();
        b.push(Bytes::from_static(b"abc"));
        b.push(Bytes::from_static(b"defghi"));
        b.push(Bytes::from_static(b"jkl"));
        // request bytes 1..10 = "bcdefghij"
        let s = b.slice(1, 9);
        let collected: Vec<u8> = s.byte_pieces().flat_map(|p| p.iter().copied()).collect();
        assert_eq!(collected.as_slice(), b"bcdefghij");
        assert_eq!(s.len(), 9);
    }

    #[test]
    fn segment_body_slice_full() {
        let mut b = SegmentBody::new();
        b.push(Bytes::from_static(b"abc"));
        b.push(Bytes::from_static(b"def"));
        let s = b.slice(0, 6);
        assert_eq!(s.len(), 6);
    }

    #[test]
    fn segment_body_slice_empty() {
        let mut b = SegmentBody::new();
        b.push(Bytes::from_static(b"abc"));
        let s = b.slice(2, 0);
        assert_eq!(s.len(), 0);
        assert!(s.is_empty());
    }

    #[test]
    #[should_panic(expected = "SegmentBody::slice out of range")]
    fn segment_body_slice_panics_out_of_range() {
        let mut b = SegmentBody::new();
        b.push(Bytes::from_static(b"abc"));
        let _ = b.slice(2, 5);
    }

    #[tokio::test]
    async fn scatter_body_polls_pieces_in_order() {
        let pieces = vec![
            Bytes::from_static(b"part-1"),
            Bytes::from_static(b"part-2"),
            Bytes::from_static(b"part-3"),
        ];
        let body = ScatterBody::new(pieces);
        let mut body = Box::pin(body);
        let mut collected = Vec::new();
        while let Some(frame) = std::future::poll_fn(|cx| body.as_mut().poll_data(cx)).await {
            let b = frame.expect("infallible");
            collected.extend_from_slice(&b);
        }
        assert_eq!(&collected, b"part-1part-2part-3");
    }
}