hyperfile 0.11.2

//! Planning a read apart from doing it.
//!
//! [`HyperFile::plan_reads`](crate::file::file::HyperFile::plan_reads) decides,
//! under the file, where every byte of several ranges comes from and which object
//! requests fetch them; [`ReadPlan::execute`] makes those requests with only a
//! shared reference, so a caller can plan under its lock and fetch outside it, as
//! many plans at once as it likes.
//!
//! What makes fetching outside the lock correct is that an object, once written,
//! never changes: see
//! [Reading outside the lock](https://github.com/daiyy/hyperfile/blob/main/docs/read-planning.md).

use std::io::{Error, ErrorKind, Result};
use std::marker::PhantomData;
use std::sync::Arc;
#[cfg(feature = "wal")]
use std::pin::Pin;
#[cfg(feature = "wal")]
use std::sync::Weak;

use futures::stream::{FuturesUnordered, StreamExt};
use tokio::sync::Semaphore;

use crate::block_source::BlockSource;
use crate::staging::Staging;
use crate::SegmentId;

/// How [`HyperFile::plan_reads`](crate::file::file::HyperFile::plan_reads) turns
/// the pieces it needs from storage into requests.
///
/// Pieces of one object are taken in object order, whatever order the file has
/// them in, and requests never span two objects. The default is the read path's
/// own rule: pieces that touch or overlap are one request, up to `max_get`.
#[derive(Debug, Clone, PartialEq)]
pub struct ReadPolicy {
    /// Join two requests to one object across a gap of at most this many bytes,
    /// fetching the gap and dropping it. 0 joins only pieces that touch.
    ///
    /// A gap of `g` bytes costs `g / B` at a connection's throughput `B`; a request
    /// costs its time to first byte and its charge. Which is cheaper depends on the
    /// store and the network, so it is the caller's to say.
    pub max_gap: usize,
    /// The most bytes one request may fetch, gap included. `None` is the file's
    /// `read_get_max_bytes`. Smaller splits long runs into requests that can be in
    /// flight together, for a caller with more concurrency than requests.
    pub max_get: Option<usize>,
    /// The most gap bytes the whole plan may fetch, as a share of the bytes it
    /// needs: 1.0 allows as many again. `None` is no limit. Gaps are joined
    /// smallest first until this or `max_gap` stops them.
    pub max_overread: Option<f64>,
}

impl Default for ReadPolicy {
    fn default() -> Self {
        Self { max_gap: 0, max_get: None, max_overread: None }
    }
}

/// Where part of one request's bytes go: `len` bytes from `at` within the request
/// to `out` in the caller's buffer.
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub struct GetTarget {
    pub at: u64,
    pub out: u64,
    pub len: u64,
}

/// One object request of a [`ReadPlan`]: `len` bytes at `at` in object `segid`,
/// and where they go. Bytes no target names are a gap, fetched and dropped.
#[derive(Debug, Clone, PartialEq, Eq)]
pub struct PlannedGet {
    pub segid: SegmentId,
    pub at: u64,
    pub len: u64,
    pub targets: Vec<GetTarget>,
}

/// A block written by reference: `len` bytes at `src` in the block source.
#[derive(Debug, Clone, Copy)]
struct RefPiece {
    out: usize,
    src: u64,
    len: usize,
}

/// A block in a segment a WAL-protected flush is still uploading.
#[cfg(feature = "wal")]
#[derive(Debug, Clone)]
struct InmemPiece {
    out: usize,
    segid: SegmentId,
    at: usize,
    len: usize,
    data: Weak<Pin<Box<Vec<u8>>>>,
}

/// Where the bytes of several ranges come from, decided under the file and
/// fetched with [`Self::execute`] without it.
///
/// The ranges are laid out in `out` one after another in the order given, each
/// at its full requested length. Bytes past the end of the file are zero; see
/// [`Self::lens`] for how many of each range are inside it.
pub struct ReadPlan<T, L> {
    staging: T,
    len: usize,
    lens: Vec<u64>,
    /// Zeros: holes, zero blocks, past the end.
    zeros: Vec<(usize, usize)>,
    /// Bytes copied out of the data cache when the plan was made.
    local: Vec<(usize, Vec<u8>)>,
    refs: Vec<RefPiece>,
    source: Option<Arc<dyn BlockSource>>,
    #[cfg(feature = "wal")]
    inmem: Vec<InmemPiece>,
    gets: Vec<PlannedGet>,
    needed: u64,
    _loader: PhantomData<fn() -> L>,
}

impl<T, L> std::fmt::Debug for ReadPlan<T, L> {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        f.debug_struct("ReadPlan").field("len", &self.len).field("lens", &self.lens)
            .field("gets", &self.gets).field("local_bytes", &self.local_bytes()).finish()
    }
}

impl<T, L> ReadPlan<T, L> {
    /// The buffer length [`Self::execute`] fills: every range at its requested
    /// length, one after another.
    pub fn len(&self) -> usize {
        self.len
    }

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

    /// How many bytes of each range are inside the file, in the order given —
    /// what a `read` of it would return.
    pub fn lens(&self) -> &[u64] {
        &self.lens
    }

    /// The object requests, in object order.
    pub fn gets(&self) -> &[PlannedGet] {
        &self.gets
    }

    /// Bytes the requests fetch for the caller.
    pub fn needed_bytes(&self) -> u64 {
        self.needed
    }

    /// Bytes the requests fetch and drop: the gaps joined across.
    pub fn overread_bytes(&self) -> u64 {
        self.gets.iter().map(|g| g.len).sum::<u64>().saturating_sub(self.needed)
    }

    /// Bytes served without an object request: the cache, zeros, references, and
    /// segments still in memory.
    pub fn local_bytes(&self) -> u64 {
        self.len as u64 - self.needed
    }
}

impl<T, L> ReadPlan<T, L>
where
    T: Staging<L> + Sync,
{
    /// Fill `out` with the bytes the plan is for. `out` must be [`Self::len`] long.
    ///
    /// Only a shared reference to the plan is needed, and nothing of the file:
    /// what the plan names on storage does not change. Each object request holds a
    /// permit from `permits` while it is in flight, so one semaphore bounds every
    /// plan that shares it. Requests are made as permits allow, in no fixed order;
    /// the first that fails is the error, and `out` is then partly filled.
    ///
    /// Each request counts in the file's read timing as a read's does. A segment
    /// that was still in memory when the plan was made, and has since been
    /// uploaded and released, is read from storage.
    pub async fn execute(&self, out: &mut [u8], permits: &Semaphore) -> Result<()> {
        if out.len() != self.len {
            return Err(Error::new(ErrorKind::InvalidInput, format!(
                "the plan fills {} bytes, the buffer is {}", self.len, out.len())));
        }
        for (at, len) in self.zeros.iter().copied() {
            out[at..at + len].fill(0);
        }
        for (at, bytes) in self.local.iter() {
            out[*at..*at + bytes.len()].copy_from_slice(bytes);
        }

        // What still has to come from storage: the plan's requests, and any
        // in-memory segment that is in memory no longer.
        #[allow(unused_mut)]
        let mut fetch: Vec<(SegmentId, usize, usize, Vec<GetTarget>)> = self.gets.iter()
            .map(|g| (g.segid, g.at as usize, g.len as usize, g.targets.clone())).collect();
        #[cfg(feature = "wal")]
        for p in self.inmem.iter() {
            match p.data.upgrade() {
                Some(data) => out[p.out..p.out + p.len].copy_from_slice(&data[p.at..p.at + p.len]),
                None => fetch.push((p.segid, p.at, p.len,
                    vec![GetTarget { at: 0, out: p.out as u64, len: p.len as u64 }])),
            }
        }

        let refs = if self.refs.is_empty() { None } else {
            let src = self.source.clone().expect("references are planned only with a source");
            let pieces = self.refs.clone();
            Some(tokio::task::spawn_blocking(move || {
                pieces.into_iter().map(|p| {
                    let mut b = vec![0u8; p.len];
                    src.read_at(p.src, &mut b).map(|_| (p.out, b))
                }).collect::<Result<Vec<_>>>()
            }))
        };

        let mut inflight: FuturesUnordered<_> = fetch.into_iter().map(|(segid, at, len, targets)| async move {
            let _permit = permits.acquire().await.map_err(Error::other)?;
            let mut buf = vec![0u8; len];
            self.staging.load_range(segid, at, &mut buf).await?;
            Ok::<_, Error>((buf, targets))
        }).collect();
        while let Some(done) = inflight.next().await {
            let (buf, targets) = done?;
            for t in targets {
                let (at, to, n) = (t.at as usize, t.out as usize, t.len as usize);
                out[to..to + n].copy_from_slice(&buf[at..at + n]);
            }
        }

        if let Some(join) = refs {
            for (at, bytes) in join.await.map_err(Error::other)?? {
                out[at..at + bytes.len()].copy_from_slice(&bytes);
            }
        }
        Ok(())
    }
}

/// Assembles a plan; see `HyperFile::plan_reads`.
pub(crate) struct ReadPlanBuilder {
    len: usize,
    lens: Vec<u64>,
    zeros: Vec<(usize, usize)>,
    local: Vec<(usize, Vec<u8>)>,
    refs: Vec<RefPiece>,
    #[cfg(feature = "wal")]
    inmem: Vec<InmemPiece>,
    /// (segid, at, len, out) for each piece needed from storage.
    staged: Vec<(SegmentId, u64, u64, u64)>,
}

impl ReadPlanBuilder {
    pub(crate) fn new() -> Self {
        Self {
            len: 0, lens: Vec::new(), zeros: Vec::new(), local: Vec::new(), refs: Vec::new(),
            #[cfg(feature = "wal")]
            inmem: Vec::new(),
            staged: Vec::new(),
        }
    }

    /// Start a range of `len` bytes, `within` of them inside the file. Returns
    /// where it starts in `out`.
    pub(crate) fn range(&mut self, len: usize, within: usize) -> usize {
        let at = self.len;
        self.len += len;
        self.lens.push(within as u64);
        if within < len {
            self.zero(at + within, len - within);
        }
        at
    }

    pub(crate) fn zero(&mut self, out: usize, len: usize) {
        match self.zeros.last_mut() {
            Some((at, n)) if *at + *n == out => *n += len,
            _ => self.zeros.push((out, len)),
        }
    }

    /// `len` bytes for `out`, to be filled by `fill` now.
    pub(crate) fn local(&mut self, out: usize, len: usize, fill: impl FnOnce(&mut [u8])) {
        match self.local.last_mut() {
            Some((at, bytes)) if *at + bytes.len() == out => {
                let from = bytes.len();
                bytes.resize(from + len, 0);
                fill(&mut bytes[from..]);
            },
            _ => {
                let mut bytes = vec![0u8; len];
                fill(&mut bytes);
                self.local.push((out, bytes));
            },
        }
    }

    pub(crate) fn reference(&mut self, out: usize, src: u64, len: usize) {
        self.refs.push(RefPiece { out, src, len });
    }

    #[cfg(feature = "wal")]
    pub(crate) fn inmem(&mut self, out: usize, segid: SegmentId, at: usize, len: usize, data: Weak<Pin<Box<Vec<u8>>>>) {
        self.inmem.push(InmemPiece { out, segid, at, len, data });
    }

    pub(crate) fn staged(&mut self, out: usize, segid: SegmentId, at: usize, len: usize) {
        self.staged.push((segid, at as u64, len as u64, out as u64));
    }

    pub(crate) fn build<T, L>(self, staging: T, source: Option<Arc<dyn BlockSource>>, policy: &ReadPolicy, max_get: usize)
        -> ReadPlan<T, L>
    {
        let needed = self.staged.iter().map(|p| p.2).sum();
        let gets = join(self.staged, policy, max_get);
        ReadPlan {
            staging,
            len: self.len,
            lens: self.lens,
            zeros: self.zeros,
            local: self.local,
            source: if self.refs.is_empty() { None } else { source },
            refs: self.refs,
            #[cfg(feature = "wal")]
            inmem: self.inmem,
            gets,
            needed,
            _loader: PhantomData,
        }
    }
}

/// Turn pieces into requests by `policy`.
///
/// First, within each object in object order, pieces that touch or overlap are
/// one run, while the run stays within `max_get`. Then the gaps between runs, at
/// most `max_gap`, are joined smallest first — ties in object and offset order, so
/// the plan is the same every time — while the joined request stays within
/// `max_get` and the gaps joined stay within the `max_overread` budget.
fn join(mut pieces: Vec<(SegmentId, u64, u64, u64)>, policy: &ReadPolicy, max_get: usize) -> Vec<PlannedGet> {
    let max_get = max_get.max(1) as u64;
    pieces.sort_by_key(|p| (p.0, p.1, p.3));
    let needed: u64 = pieces.iter().map(|p| p.2).sum();

    // Runs: (segid, start, end, pieces).
    let mut runs: Vec<(SegmentId, u64, u64, Vec<(u64, u64, u64)>)> = Vec::new();
    for (segid, at, len, out) in pieces {
        let end = at + len;
        if let Some(r) = runs.last_mut() {
            if r.0 == segid && at <= r.2 && end.max(r.2) - r.1 <= max_get {
                r.2 = r.2.max(end);
                r.3.push((at, len, out));
                continue;
            }
        }
        runs.push((segid, at, end, vec![(at, len, out)]));
    }

    // Gaps between neighbouring runs of one object, as candidates.
    let mut gaps: Vec<(u64, usize)> = (0..runs.len().saturating_sub(1))
        .filter(|&i| runs[i].0 == runs[i + 1].0)
        .map(|i| (runs[i + 1].1.saturating_sub(runs[i].2), i))
        .filter(|&(g, _)| g as usize <= policy.max_gap)
        .collect();
    gaps.sort();
    let mut budget = policy.max_overread.map(|r| (needed as f64 * r.max(0.0)) as u64);

    // Joined runs as intervals of run indices: `lo[root]..=hi[root]`.
    let n = runs.len();
    let mut parent: Vec<usize> = (0..n).collect();
    let lo: Vec<usize> = (0..n).collect();
    let mut hi: Vec<usize> = (0..n).collect();
    fn root(parent: &mut [usize], mut i: usize) -> usize {
        while parent[i] != i {
            parent[i] = parent[parent[i]];
            i = parent[i];
        }
        i
    }
    for (gap, i) in gaps {
        if let Some(b) = budget {
            if gap > b {
                continue;
            }
        }
        let (a, c) = (root(&mut parent, i), root(&mut parent, i + 1));
        let (l, h) = (lo[a], hi[c]);
        if runs[h].2 - runs[l].1 > max_get {
            continue;
        }
        parent[c] = a;
        hi[a] = h;
        if let Some(b) = budget.as_mut() {
            *b -= gap;
        }
    }

    let mut gets = Vec::new();
    let mut i = 0;
    while i < n {
        let r = root(&mut parent, i);
        let (l, h) = (lo[r], hi[r]);
        let (segid, start, end) = (runs[l].0, runs[l].1, runs[h].2);
        let targets = runs[l..=h].iter().flat_map(|run| run.3.iter())
            .map(|&(at, len, out)| GetTarget { at: at - start, out, len }).collect();
        gets.push(PlannedGet { segid, at: start, len: end - start, targets });
        i = h + 1;
    }
    gets
}

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

    fn seg(n: u64) -> SegmentId {
        SegmentId::new_from_cno(n)
    }

    fn shape(gets: &[PlannedGet]) -> Vec<(u64, u64, u64)> {
        gets.iter().map(|g| (g.segid.as_cno(), g.at, g.len)).collect()
    }

    #[test]
    fn touching_pieces_are_one_request_in_object_order() {
        // File order 2, 0, 1; object order 0, 1, 2.
        let p = vec![(seg(1), 8, 4, 0), (seg(1), 0, 4, 4), (seg(1), 4, 4, 8)];
        let g = join(p, &ReadPolicy::default(), 1 << 20);
        assert_eq!(shape(&g), vec![(1, 0, 12)]);
        assert_eq!(g[0].targets, vec![
            GetTarget { at: 0, out: 4, len: 4 }, GetTarget { at: 4, out: 8, len: 4 },
            GetTarget { at: 8, out: 0, len: 4 }]);
    }

    #[test]
    fn requests_never_span_objects() {
        let p = vec![(seg(1), 0, 4, 0), (seg(2), 4, 4, 4)];
        assert_eq!(shape(&join(p, &ReadPolicy { max_gap: 100, ..Default::default() }, 1 << 20)),
            vec![(1, 0, 4), (2, 4, 4)]);
    }

    #[test]
    fn a_gap_is_joined_only_within_max_gap() {
        let p = || vec![(seg(1), 0, 4, 0), (seg(1), 10, 4, 4), (seg(1), 30, 4, 8)];
        assert_eq!(shape(&join(p(), &ReadPolicy::default(), 1 << 20)), vec![(1, 0, 4), (1, 10, 4), (1, 30, 4)]);
        assert_eq!(shape(&join(p(), &ReadPolicy { max_gap: 6, ..Default::default() }, 1 << 20)),
            vec![(1, 0, 14), (1, 30, 4)]);
        assert_eq!(shape(&join(p(), &ReadPolicy { max_gap: 16, ..Default::default() }, 1 << 20)),
            vec![(1, 0, 34)]);
    }

    #[test]
    fn the_overread_budget_takes_the_smallest_gaps_first() {
        // Gaps 16 then 6; needed 12, so a budget of 0.5 allows 6 and not 16.
        let p = vec![(seg(1), 0, 4, 0), (seg(1), 20, 4, 4), (seg(1), 30, 4, 8)];
        let g = join(p, &ReadPolicy { max_gap: 100, max_overread: Some(0.5), ..Default::default() }, 1 << 20);
        assert_eq!(shape(&g), vec![(1, 0, 4), (1, 20, 14)]);
    }

    #[test]
    fn max_get_bounds_runs_and_joins() {
        let p = (0..8u64).map(|k| (seg(1), k * 4, 4, k * 4)).collect::<Vec<_>>();
        assert_eq!(shape(&join(p, &ReadPolicy::default(), 12)), vec![(1, 0, 12), (1, 12, 12), (1, 24, 8)]);
        let p = vec![(seg(1), 0, 4, 0), (seg(1), 6, 4, 4)];
        assert_eq!(shape(&join(p, &ReadPolicy { max_gap: 4, ..Default::default() }, 8)), vec![(1, 0, 4), (1, 6, 4)]);
    }

    #[test]
    fn overlapping_pieces_share_a_request() {
        // The same block asked for by two ranges.
        let p = vec![(seg(1), 0, 8, 0), (seg(1), 4, 8, 8)];
        let g = join(p, &ReadPolicy::default(), 1 << 20);
        assert_eq!(shape(&g), vec![(1, 0, 12)]);
        assert_eq!(g[0].targets.len(), 2);
    }
}