spate-s3 0.4.0

//! Measurable entry points into work this crate does not otherwise expose.
//!
//! Benches in this workspace drive their crate's public API. This module is
//! an exception.
//!
//! This crate's public entry point is [`S3Source`](crate::S3Source), which is
//! asynchronous and driven by object-store I/O. Both properties stop an
//! instruction count being deterministic: under a runtime the number becomes
//! a function of how the scheduler interleaved polls. Record framing is pure,
//! synchronous, on the per-object path and so countable, but reaching it means
//! reaching past the async surface.
//!
//! The rules a future seam has to honor:
//!
//! - It is behind the off-by-default `testing` feature and `#[doc(hidden)]`,
//!   so it is not part of this crate's semver surface and no consumer of the
//!   `spate` facade can see it.
//! - It exports **functions and the aliases they need, never this crate's own
//!   types**. `ObjectFramer`, `ObjectEntry` and the rest stay private and free
//!   to change; only the shape of the work is fixed. `MakeFramer` is an alias
//!   over `std` and `spate-core` types, which a caller must be able to name to
//!   supply a framer at all.
//! - Each function is one whole unit of the work a stage does (an object
//!   run) rather than one internal call. Attribution below that level
//!   comes from the callgrind profile the bench already writes.

use crate::config::Compression;
use crate::framer::{Codec, ObjectFramer};
use crate::split::SplitRange;
use crate::trim::RangeTrim;
use bytes::Bytes;
use spate_core::framing::RecordFramer;
use std::io;
use std::sync::Arc;

/// Builds a fresh record framer, as the source does per object.
pub type MakeFramer = Arc<dyn Fn() -> Box<dyn RecordFramer> + Send + Sync>;

/// Frame a run of objects into records through one framer, returning how
/// many records they produced in total.
///
/// Mirrors the lane, which builds a single `ObjectFramer` and cycles
/// `begin_object`/`finish_object` over the members of its split, so a run of
/// several objects here measures the per-object codec resolution and state
/// reset a multi-member split pays, not the cost of constructing a framer
/// per object. One object is a one-element slice.
///
/// Each object is `(key, chunks)`: the key because the codec is resolved from
/// it per object, and the chunks already split the way a fetcher would
/// deliver them, already compressed if `compression` says so, since
/// compressing here would count the compressor rather than the decompressor.
///
/// The framer's record sequence is a pure function of the bytes it is given,
/// so a chunk list starting part-way through a record frames that partial
/// line as a record. [`frame_range`] is the entry a byte-range split takes.
///
/// # Errors
///
/// Whatever the decompressor or the framer reports: a truncated stream, a
/// corrupt frame, or a record over the framer's cap.
pub fn frame_objects(
    compression: Compression,
    make_framer: MakeFramer,
    objects: &[(String, Vec<Bytes>)],
) -> io::Result<usize> {
    let mut framer = ObjectFramer::new(make_framer);
    let mut records = 0;
    for (key, chunks) in objects {
        framer.begin_object(Codec::resolve(compression, key))?;
        for chunk in chunks {
            // A refcount bump, as the lane's own hand-off is.
            framer.push_chunk(chunk.clone())?;
            // Drained per decode step, because the lane drains at the top of
            // every poll-loop iteration rather than at end of object. Draining
            // only at the end would hold a whole object's records live at
            // once and report a DHAT peak the source never reaches. DHAT is
            // attached to this bench, so that shape is part of what it
            // publishes.
            loop {
                while framer.pop_record().is_some() {
                    records += 1;
                }
                if !framer.has_pending_input() {
                    break;
                }
                framer.pump()?;
            }
        }
        framer.finish_object()?;
        // The tail: `finish_object` completes an unterminated final record.
        // A record never spans an object boundary, so the queue is empty
        // before the next `begin_object`.
        while framer.pop_record().is_some() {
            records += 1;
        }
    }
    Ok(records)
}

/// Read the byte range `[start, end)` of one uncompressed object as a
/// byte-range split does, returning how many records it owns.
///
/// The object is fetched in the windows the fetcher schedules for
/// `range_bytes` and `chunk_bytes`, each copied out of `object` as a GET would
/// deliver it. The owned bytes are trimmed out of each window and framed in
/// `chunk_bytes` pieces through one framer.
///
/// # Errors
///
/// Whatever the framer reports, such as a record over its cap.
///
/// # Panics
///
/// If `start >= end` or `end` is past the object.
pub fn frame_range(
    make_framer: MakeFramer,
    object: &[u8],
    start: u64,
    end: u64,
    delimiter: u8,
    range_bytes: u64,
    chunk_bytes: usize,
) -> io::Result<usize> {
    let size = object.len() as u64;
    assert!(
        start < end && end <= size,
        "[{start}, {end}) is not a range of the object"
    );
    let mut trim = RangeTrim::new(SplitRange::new(start, end, delimiter), size);
    let mut framer = ObjectFramer::new(make_framer);
    let mut records = 0;
    framer.begin_object(Codec::Plain)?;
    let mut pos = trim.first_byte();
    while let Some(window_end) = trim.next_window(pos, range_bytes, chunk_bytes as u64) {
        let window = Bytes::copy_from_slice(&object[pos as usize..window_end as usize]);
        let mut owned = trim.feed(pos, window);
        while !owned.is_empty() {
            let chunk = owned.split_to(owned.len().min(chunk_bytes));
            framer.push_chunk(chunk)?;
            while framer.pop_record().is_some() {
                records += 1;
            }
            debug_assert!(
                !framer.has_pending_input(),
                "a plain chunk decodes in one step"
            );
        }
        pos = window_end;
    }
    framer.finish_object()?;
    while framer.pop_record().is_some() {
        records += 1;
    }
    Ok(records)
}