format: bound and batch every chunk fetch over Storage
Only the full read split its chunk fetches into 64 MiB batches. The selection path, the indexed read and the parallel_read decoders fetched every chunk's stored bytes in one read_ranges call, each extent bounded only by the file length, so a crafted chunk index pointing many chunks at one large extent made File::open_storage hold chunks x extent bytes (3.3 GB from a 16.8 MB file) before the first decode error. - storage::for_each_extent_batch is now the one way raw-data reads fetch chunk bytes: batches of at most RAW_BATCH_BYTES (now pub), each decoded before the next is fetched. Used by the full, cached, indexed, selection and parallel_read paths; the sweep read uses read_extent per chunk. - ExtentReq carries each chunk's claimed extent (bounds-checked as before, same errors) and the prefix actually fetched: filters::stored_chunk_limit — the chunk size if unfiltered, else each applied filter's worst-case growth (n + n/4 + 4096 per codec; unbounded only for an application-registered codec). The in-memory path cuts the slice it decodes the same way, so both paths still agree. - tests/raw_fetch_bounds.rs: a crafted chunked_large.h5 (ten chunks all claiming 20 MiB at one padding blob) read through every path over a storage that records the largest single fetch; and 160 MiB of legitimate unfiltered chunks fetched batch by batch. Before: one 80 MiB fetch (selection) and one 160 MiB fetch; after: within the budget. Co-Authored-By: Claude Opus 5.5 (1M context) <[email protected]>
This commit is contained in:
osobh
@@ -0,0 +1,332 @@
|
||||
//! 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);
|
||||
}
|
||||
Reference in New Issue
Block a user