quantumclaw/clawhdf5
3
1
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases 1 Wiki Activity

format: bound and batch every chunk fetch over Storage

Browse Source 
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
2026-09-26 18:30:30 -05:00
co-authored by Claude Opus 5.5
parent f191dc09d5
commit 7d629f49e3
10 changed files with 825 additions and 285 deletions
Download Patch File Download Diff File Expand all files Collapse all files
crates/clawhdf5-format/src/chunked_read.rs
+123 -80
View File
@@ -19,7 +19,7 @@ use crate::filters::{DecodeScratch, decompress_chunk_exact_with};
#[cfg(feature = "std")]
use crate::filters::{all_filters_skipped, decompress_chunk_exact};
use crate::fixed_array::{FixedArrayHeader, read_fixed_array_chunks_in};
use crate::storage::{ExtentBytes, Storage, Window, raw_batches, read_exact_at};
use crate::storage::{ExtentBytes, ExtentReq, Storage, Window, for_each_extent_batch, read_extent};
#[cfg(feature = "std")]
use std::sync::Arc;
@@ -256,52 +256,17 @@ fn fill_from_chunks<S: Storage + ?Sized>(
cache: CacheUse<'_>,
output: &mut [u8],
) -> Result<(), FormatError> {
let contiguous = file_data.as_contiguous().is_some();
let out = OutBuf::new(output);
let batches = raw_batches(chunks.len(), contiguous, |i| chunks[i].chunk_size as usize);
for batch in batches {
fill_batch(
file_data,
&chunks[batch],
pipeline,
placer,
chunk_total_bytes,
cache,
&out,
)?;
}
Ok(())
}
/// Whether a full read looks chunk `c` up in the cache (see
/// [`fill_from_chunks`]): a filtered chunk of a dataset the cache keeps.
#[cfg(feature = "std")]
fn uses_cache(cache: CacheUse<'_>, pipeline: Option<&FilterPipeline>, c: &ChunkInfo) -> bool {
matches!(cache, Some((_, _, true)))
&& pipeline.is_some_and(|pl| !all_filters_skipped(pl, c.filter_mask))
}
/// [`fill_from_chunks`] for one batch of chunks.
#[allow(clippy::too_many_arguments)]
fn fill_batch<S: Storage + ?Sized>(
file_data: &S,
chunks: &[ChunkInfo],
pipeline: Option<&FilterPipeline>,
placer: &ChunkPlacer,
chunk_total_bytes: usize,
cache: CacheUse<'_>,
out: &OutBuf<'_>,
) -> Result<(), FormatError> {
let rank = placer.rank;
let elem_size = placer.elem_size as u32;
#[cfg(not(feature = "std"))]
let _ = cache;
#[cfg(feature = "std")]
let rank = placer.rank;
// Over a backend without the whole file in memory, the decoded chunks
// the cache holds are taken now (so a chunk evicted before it is placed
// is not left without bytes), and only the others are fetched.
#[cfg(feature = "std")]
let hits: Vec<Option<Arc<CacheAlignedBuffer>>> = match cache {
let hits: Vec<CacheHit> = match cache {
Some((cache, key, true)) if file_data.as_contiguous().is_none() => chunks
.iter()
.map(|c| {
@@ -314,25 +279,90 @@ fn fill_batch<S: Storage + ?Sized>(
.collect(),
_ => Vec::new(),
};
let extents: Vec<(u64, usize, bool)> = if file_data.as_contiguous().is_some() {
Vec::new()
} else {
chunks
.iter()
.enumerate()
.map(|(i, c)| {
#[cfg(feature = "std")]
let wanted = hits.get(i).is_none_or(Option::is_none);
#[cfg(not(feature = "std"))]
let wanted = {
let _ = i;
true
};
(c.address, c.chunk_size as usize, wanted)
})
.collect()
};
let raw_bytes = ExtentBytes::fetch(file_data, &extents)?;
#[cfg(not(feature = "std"))]
let hits: Vec<CacheHit> = Vec::new();
let reqs: Vec<ExtentReq> = chunks
.iter()
.enumerate()
.map(|(i, c)| {
let wanted = hits.get(i).is_none_or(Option::is_none);
chunk_req(c, pipeline, chunk_total_bytes, wanted)
})
.collect();
for_each_extent_batch(file_data, &reqs, |batch, raw_bytes| {
fill_batch(
&chunks[batch.clone()],
&reqs[batch.clone()],
hits.get(batch.clone()).unwrap_or_default(),
batch.start,
raw_bytes,
pipeline,
placer,
chunk_total_bytes,
cache,
&out,
)
})
}
/// The extent of chunk `c`'s stored bytes to fetch (`wanted`) or only
/// bounds-check: its whole stored size is checked against the file, and at
/// most [`crate::filters::stored_chunk_limit`] of it is read and decoded.
pub(crate) fn chunk_req(
c: &ChunkInfo,
pipeline: Option<&FilterPipeline>,
chunk_bytes: usize,
wanted: bool,
) -> ExtentReq {
let len = c.chunk_size as usize;
ExtentReq {
addr: c.address,
len,
fetch: wanted.then(|| {
len.min(crate::filters::stored_chunk_limit(
pipeline,
c.filter_mask,
chunk_bytes,
))
}),
}
}
/// A decoded chunk a full read took from the cache.
#[cfg(feature = "std")]
type CacheHit = Option<Arc<CacheAlignedBuffer>>;
#[cfg(not(feature = "std"))]
type CacheHit = Option<core::convert::Infallible>;
/// Whether a full read looks chunk `c` up in the cache (see
/// [`fill_from_chunks`]): a filtered chunk of a dataset the cache keeps.
#[cfg(feature = "std")]
fn uses_cache(cache: CacheUse<'_>, pipeline: Option<&FilterPipeline>, c: &ChunkInfo) -> bool {
matches!(cache, Some((_, _, true)))
&& pipeline.is_some_and(|pl| !all_filters_skipped(pl, c.filter_mask))
}
/// [`fill_from_chunks`] for one batch of chunks: `chunks` (with their
/// extents `reqs` and, over a backend without the file in memory, the cache
/// hits taken for them) are chunks `first..` of the read, whose stored bytes
/// `raw_bytes` holds.
#[allow(clippy::too_many_arguments)]
fn fill_batch(
chunks: &[ChunkInfo],
reqs: &[ExtentReq],
hits: &[CacheHit],
first: usize,
raw_bytes: &ExtentBytes<'_>,
pipeline: Option<&FilterPipeline>,
placer: &ChunkPlacer,
chunk_total_bytes: usize,
cache: CacheUse<'_>,
out: &OutBuf<'_>,
) -> Result<(), FormatError> {
let rank = placer.rank;
let elem_size = placer.elem_size as u32;
#[cfg(not(feature = "std"))]
let _ = cache;
// Decode chunk `i` and place it. Its callers below run it either on one
// thread, or on several for chunks whose regions are pairwise disjoint,
@@ -346,15 +376,14 @@ fn fill_batch<S: Storage + ?Sized>(
)));
}
let offsets = &c.offsets[..rank];
let size = c.chunk_size as usize;
#[cfg(feature = "std")]
if let Some(Some(hit)) = hits.get(i) {
raw_bytes.check(i, c.address, size)?;
raw_bytes.check(first + i, &reqs[i])?;
// SAFETY: see above.
unsafe { placer.place(hit, offsets, out) };
return Ok(());
}
let raw = raw_bytes.get(i, c.address, size)?;
let raw = raw_bytes.get(first + i, &reqs[i])?;
let Some(pl) = pipeline else {
// SAFETY: see above.
unsafe { placer.place(raw, offsets, out) };
@@ -1913,9 +1942,9 @@ pub fn read_chunked_data_sweep_in<S: Storage + ?Sized>(
cached
} else {
// Decompress from file
let c_addr = to_usize(chunk_info.address)?;
let size = chunk_info.chunk_size as usize;
let raw_chunk = read_exact_at(file_data, c_addr as u64, size)?;
to_usize(chunk_info.address)?;
let req = chunk_req(chunk_info, pipeline, chunk_total_bytes, true);
let raw_chunk = read_extent(file_data, &req)?;
let raw_chunk = &*raw_chunk;
let dec = if let Some(pl) = pipeline {
decompress_chunk_exact(
@@ -2052,46 +2081,60 @@ pub fn read_chunked_data_indexed_in<S: Storage + ?Sized>(
let chunk_total_bytes = plan.chunk_total_bytes;
// The decoded chunks the cache holds, and the stored bytes of the
// others: fetched in one batch when the file is not in memory.
// others: fetched batch by batch when the file is not in memory.
let hits: Vec<Option<Arc<CacheAlignedBuffer>>> = plan
.mappings
.iter()
.map(|m| cache.get_decompressed_in(addr, &m.coord))
.collect();
let extents: Vec<(u64, usize, bool)> = plan
let reqs: Vec<ExtentReq> = plan
.mappings
.iter()
.zip(&hits)
.map(|(m, hit)| (m.file_offset, m.file_size as usize, hit.is_none()))
.map(|(m, hit)| {
let len = m.file_size as usize;
ExtentReq {
addr: m.file_offset,
len,
fetch: hit.is_none().then(|| {
len.min(crate::filters::stored_chunk_limit(
pipeline,
m.filter_mask,
chunk_total_bytes,
))
}),
}
})
.collect();
let raw_bytes = ExtentBytes::fetch(file_data, &extents)?;
// Decompress chunks (using LRU cache where possible)
let mut chunk_buffers: Vec<Arc<CacheAlignedBuffer>> = Vec::with_capacity(plan.mappings.len());
for (i, (m, hit)) in plan.mappings.iter().zip(hits).enumerate() {
let (coord, file_offset, file_size, filter_mask) =
(&m.coord, &m.file_offset, &m.file_size, &m.filter_mask);
if let Some(cached) = hit {
chunk_buffers.push(cached);
} else {
let raw_chunk = raw_bytes.get(i, *file_offset, *file_size as usize)?;
let mut hits = hits.into_iter();
for_each_extent_batch(file_data, &reqs, |batch, raw_bytes| {
for i in batch {
let m = &plan.mappings[i];
if let Some(cached) = hits.next().flatten() {
chunk_buffers.push(cached);
continue;
}
let raw_chunk = raw_bytes.get(i, &reqs[i])?;
let decompressed = if let Some(pl) = pipeline {
decompress_chunk_exact(
raw_chunk,
pl,
chunk_total_bytes,
elem_size as u32,
*filter_mask,
coord,
m.filter_mask,
&m.coord,
)?
} else {
raw_chunk.to_vec()
};
let aligned = CacheAlignedBuffer::from_vec(decompressed);
let arc = cache.put_decompressed_aligned_in(addr, coord.clone(), aligned);
chunk_buffers.push(arc);
chunk_buffers.push(cache.put_decompressed_aligned_in(addr, m.coord.clone(), aligned));
}
}
Ok(())
})?;
// Assemble using pre-computed layout
let mut output = vec![0u8; plan.output_bytes];