clawhdf5/crates/clawhdf5-format/tests/raw_fetch_bounds.rs

//! What a raw-data read fetches from a [`Storage`] without the file in
//! memory is bounded: by batch (at most `RAW_BATCH_BYTES` per
//! `read_ranges` call) and by chunk (never more of a chunk's stored bytes
//! than its decoded size can need), on every path that reads chunks — full,
//! cached, indexed, sweep, selection and the `parallel_read` decoders.
//!
//! A crafted chunk index can point every chunk at one huge extent. Slicing
//! an in-memory file costs nothing there, but a backend that fetches would
//! hold `chunks x extent` bytes before the first chunk failed to decode.

use std::borrow::Cow;
use std::ops::Range;
use std::path::Path;
use std::sync::atomic::{AtomicU64, Ordering::Relaxed};

use clawhdf5_format::chunk_cache::ChunkCache;
use clawhdf5_format::chunked_read::{
    ChunkInfo, SweepContext, list_chunks, read_chunked_data_sweep_in,
};
use clawhdf5_format::data_layout::DataLayout;
use clawhdf5_format::data_read::{
    read_raw_data_cached_in, read_raw_data_full_in, read_raw_data_indexed_in,
    read_raw_data_selection_in,
};
use clawhdf5_format::dataspace::Dataspace;
use clawhdf5_format::datatype::Datatype;
use clawhdf5_format::error::FormatError;
use clawhdf5_format::filter_pipeline::FilterPipeline;
use clawhdf5_format::group_v2;
use clawhdf5_format::message_type::MessageType;
use clawhdf5_format::object_header::ObjectHeader;
use clawhdf5_format::selection::Selection;
use clawhdf5_format::storage::{RAW_BATCH_BYTES, Storage};
use clawhdf5_format::superblock::Superblock;

/// A read_at-only storage that records the most bytes one call fetched
/// (a `read_ranges` call counts all its ranges together) and the total.
struct PeakStorage {
    data: Vec<u8>,
    peak: AtomicU64,
    total: AtomicU64,
}

impl PeakStorage {
    fn new(data: Vec<u8>) -> Self {
        PeakStorage {
            data,
            peak: AtomicU64::new(0),
            total: AtomicU64::new(0),
        }
    }

    fn reset(&self) {
        self.peak.store(0, Relaxed);
        self.total.store(0, Relaxed);
    }

    fn served(&self, offset: u64, len: usize) -> Vec<u8> {
        self.data
            .as_slice()
            .read_at(offset, len)
            .unwrap()
            .into_owned()
    }
}

impl Storage for PeakStorage {
    fn read_at(&self, offset: u64, len: usize) -> Result<Cow<'_, [u8]>, FormatError> {
        let got = self.served(offset, len);
        self.peak.fetch_max(got.len() as u64, Relaxed);
        self.total.fetch_add(got.len() as u64, Relaxed);
        Ok(Cow::Owned(got))
    }

    fn len(&self) -> u64 {
        self.data.len() as u64
    }

    fn read_ranges(&self, ranges: &[Range<u64>]) -> Result<Vec<Cow<'_, [u8]>>, FormatError> {
        let got: Vec<Vec<u8>> = ranges
            .iter()
            .map(|r| self.served(r.start, (r.end - r.start) as usize))
            .collect();
        let bytes: u64 = got.iter().map(|g| g.len() as u64).sum();
        self.peak.fetch_max(bytes, Relaxed);
        self.total.fetch_add(bytes, Relaxed);
        Ok(got.into_iter().map(Cow::Owned).collect())
    }
}

struct Chunked {
    layout: DataLayout,
    dataspace: Dataspace,
    datatype: Datatype,
    pipeline: Option<FilterPipeline>,
    os: u8,
    ls: u8,
}

/// The one chunked dataset of fixture `name`.
fn chunked(bytes: &[u8]) -> Chunked {
    let sb = Superblock::parse(bytes, 0).unwrap();
    let (os, ls) = (sb.offset_size, sb.length_size);
    for child in group_v2::resolve_group_children(bytes, &sb, sb.root_group_address).unwrap() {
        let header =
            ObjectHeader::parse(bytes, child.object_header_address as usize, os, ls).unwrap();
        let msg = |t: MessageType| {
            header
                .messages
                .iter()
                .find(|m| m.msg_type == t)
                .map(|m| m.data.clone())
        };
        let Some(dl) = msg(MessageType::DataLayout) else {
            continue;
        };
        let layout = DataLayout::parse(&dl, os, ls).unwrap();
        if !matches!(layout, DataLayout::Chunked { .. }) {
            continue;
        }
        return Chunked {
            layout,
            datatype: Datatype::parse(&msg(MessageType::Datatype).unwrap())
                .unwrap()
                .0,
            dataspace: Dataspace::parse(&msg(MessageType::Dataspace).unwrap(), ls).unwrap(),
            pipeline: msg(MessageType::FilterPipeline).map(|p| FilterPipeline::parse(&p).unwrap()),
            os,
            ls,
        };
    }
    panic!("no chunked dataset");
}

/// Claimed stored size of every chunk in the crafted file.
const HUGE: u32 = 20 << 20;

/// `chunked_large.h5` (1000 `i32` in ten gzip chunks, a v1 B-tree index)
/// with `HUGE` bytes of padding appended and every chunk's index entry
/// rewritten to claim `HUGE` stored bytes at the padding: ten chunks, 200
/// MiB of extents, in a 20 MiB file.
fn crafted() -> (Vec<u8>, Chunked, Vec<ChunkInfo>) {
    let dir = Path::new(env!("CARGO_MANIFEST_DIR")).join("tests/fixtures");
    let mut bytes = std::fs::read(dir.join("chunked_large.h5")).unwrap();
    let ds = chunked(&bytes);
    let es = ds.datatype.type_size() as usize;
    let (chunks, _) = list_chunks(&bytes, &ds.layout, &ds.dataspace, es, ds.os, ds.ls).unwrap();
    assert_eq!(chunks.len(), 10);
    let blob = bytes.len() as u64;
    for c in &chunks {
        // v1 B-tree key (size, filter mask, offsets + 0) then the child
        // address.
        let mut pat = Vec::new();
        pat.extend_from_slice(&c.chunk_size.to_le_bytes());
        pat.extend_from_slice(&c.filter_mask.to_le_bytes());
        // The key holds one offset per dimension plus the element offset
        // (0); `offsets` may or may not list that last one.
        for d in 0..=ds.dataspace.dimensions.len() {
            pat.extend_from_slice(&c.offsets.get(d).copied().unwrap_or(0).to_le_bytes());
        }
        pat.extend_from_slice(&c.address.to_le_bytes());
        let at = bytes
            .windows(pat.len())
            .position(|w| w == pat.as_slice())
            .expect("chunk key");
        bytes[at..at + 4].copy_from_slice(&HUGE.to_le_bytes());
        let a = at + pat.len() - 8;
        bytes[a..a + 8].copy_from_slice(&blob.to_le_bytes());
    }
    bytes.resize(bytes.len() + HUGE as usize, 0x5a);
    let ds = chunked(&bytes);
    let (chunks, _) = list_chunks(&bytes, &ds.layout, &ds.dataspace, es, ds.os, ds.ls).unwrap();
    assert!(
        chunks
            .iter()
            .all(|c| c.chunk_size == HUGE && c.address == blob)
    );
    (bytes, ds, chunks)
}

/// Most a crafted chunk of the fixture may fetch: its decoded size (400
/// bytes) grown by one codec, generously.
const CHUNK_LIMIT: u64 = 400 + 100 + 4096;

#[test]
fn crafted_chunk_index_cannot_amplify_fetches() {
    let (bytes, ds, chunks) = crafted();
    let st = PeakStorage::new(bytes.clone());
    let pl = ds.pipeline.as_ref();
    let (dl, sp, dt, os, ls) = (&ds.layout, &ds.dataspace, &ds.datatype, ds.os, ds.ls);
    let check =
        |what: &str, got: Result<Vec<u8>, FormatError>, want: Result<Vec<u8>, FormatError>| {
            // Same outcome as slicing the whole file.
            assert_eq!(got, want, "{what}");
            let (peak, total) = (st.peak.load(Relaxed), st.total.load(Relaxed));
            assert!(
                peak <= RAW_BATCH_BYTES as u64,
                "{what}: one fetch of {peak} bytes"
            );
            // Every chunk's fetch is bounded by what it can need, whatever its
            // index entry claims (plus the index and header reads).
            assert!(
                total <= chunks.len() as u64 * CHUNK_LIMIT + 64 * 1024,
                "{what}: fetched {total} bytes"
            );
            st.reset();
        };
    let slice: &[u8] = &bytes;

    let sel = Selection::Hyperslab {
        start: vec![100],
        stride: vec![1],
        count: vec![400],
        block: vec![1],
    };
    check(
        "selection",
        read_raw_data_selection_in(&st, dl, sp, dt, pl, os, ls, &sel),
        read_raw_data_selection_in(slice, dl, sp, dt, pl, os, ls, &sel),
    );
    check(
        "full",
        read_raw_data_full_in(&st, dl, sp, dt, pl, os, ls),
        read_raw_data_full_in(slice, dl, sp, dt, pl, os, ls),
    );
    check(
        "cached",
        read_raw_data_cached_in(&st, dl, sp, dt, pl, os, ls, &ChunkCache::new()),
        read_raw_data_cached_in(slice, dl, sp, dt, pl, os, ls, &ChunkCache::new()),
    );
    check(
        "indexed",
        read_raw_data_indexed_in(&st, dl, sp, dt, pl, os, ls, &ChunkCache::new()),
        read_raw_data_indexed_in(slice, dl, sp, dt, pl, os, ls, &ChunkCache::new()),
    );
    check(
        "sweep",
        read_chunked_data_sweep_in(
            &st,
            dl,
            sp,
            dt,
            pl,
            os,
            ls,
            &ChunkCache::new(),
            &mut SweepContext::new(4, 2),
        ),
        read_chunked_data_sweep_in(
            slice,
            dl,
            sp,
            dt,
            pl,
            os,
            ls,
            &ChunkCache::new(),
            &mut SweepContext::new(4, 2),
        ),
    );
    #[cfg(feature = "parallel")]
    {
        use clawhdf5_format::parallel_read::{
            decompress_chunks_lane_partitioned_in, decompress_chunks_parallel_in,
            decompress_chunks_sequential_in,
        };
        let pl = pl.unwrap();
        let flat = |r: Result<Vec<Vec<u8>>, FormatError>| r.map(|v| v.concat());
        check(
            "parallel",
            flat(decompress_chunks_parallel_in(&st, &chunks, pl, 400, 4)),
            flat(decompress_chunks_parallel_in(slice, &chunks, pl, 400, 4)),
        );
        check(
            "sequential",
            flat(decompress_chunks_sequential_in(
                &st,
                &chunks,
                Some(pl),
                400,
                4,
            )),
            flat(decompress_chunks_sequential_in(
                slice,
                &chunks,
                Some(pl),
                400,
                4,
            )),
        );
        check(
            "lane partitioned",
            flat(
                decompress_chunks_lane_partitioned_in(&st, &chunks, pl, 400, 4, 7, Some(3))
                    .map(|(v, _)| v),
            ),
            flat(
                decompress_chunks_lane_partitioned_in(slice, &chunks, pl, 400, 4, 7, Some(3))
                    .map(|(v, _)| v),
            ),
        );
    }
}

/// Legitimately large chunks (unfiltered, 4 MiB each, 160 MiB in all) are
/// fetched batch by batch: no call holds more than the batch budget, and
/// the data is right.
#[cfg(feature = "parallel")]
#[test]
fn large_reads_are_fetched_in_batches() {
    use clawhdf5_format::parallel_read::decompress_chunks_sequential_in;
    const CHUNK: usize = 4 << 20;
    let data: Vec<u8> = (0..2 * CHUNK).map(|i| (i % 251) as u8).collect();
    let chunks: Vec<ChunkInfo> = (0..40u64)
        .map(|i| ChunkInfo {
            chunk_size: CHUNK as u32,
            filter_mask: 0,
            offsets: vec![i * CHUNK as u64],
            address: (i % 2) * CHUNK as u64,
        })
        .collect();
    let st = PeakStorage::new(data.clone());
    let got = decompress_chunks_sequential_in(&st, &chunks, None, CHUNK, 1).unwrap();
    assert_eq!(got.len(), 40);
    for (i, c) in got.iter().enumerate() {
        let at = (i % 2) * CHUNK;
        assert!(c == &data[at..at + CHUNK], "chunk {i}");
    }
    let peak = st.peak.load(Relaxed);
    assert!(peak <= RAW_BATCH_BYTES as u64, "one fetch of {peak} bytes");
    assert_eq!(st.total.load(Relaxed), 40 * CHUNK as u64);
}