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

fix(format): key the shared chunk cache by dataset

Browse Source 
A File is Send + Sync and keeps one ChunkCache for all its datasets.
The cached readers bound that cache to "the current dataset" with
ensure_dataset(addr), then checked, built and read its index and its
decompressed chunks in separate lock acquisitions. Two threads reading
two chunked datasets interleaved those steps, so one could store its
chunk index under the other's binding, or get the other's decompressed
chunk for the same coordinate: wrong data, or an index-out-of-bounds
panic when the ranks differed (16 threads x 40 reads over 24 datasets
panicked on every run).

The cache now keeps per-dataset state keyed by chunk-index address:
the chunk index, ChunkIndex and ChunkLayout per dataset (held as Arcs,
built outside the lock, first writer wins), and decompressed chunks
keyed by (address, coordinate). The chunked readers use the new
addr-taking methods (chunks_for, chunk_layout_for, get/put_decompressed_in,
prefetch_hint_in) exclusively. Memory stays bounded: decompressed data by
the existing byte/slot budget across datasets, indexes by at most 64
datasets and 2^20 index entries in total, dropping the least recently
used dataset's index first. Switching datasets no longer throws away the
other datasets' cached chunks.

The address-less methods remain and act on the dataset last bound with
ensure_dataset; they are documented as not for concurrent readers.

Regression: threads_reading_different_datasets_get_their_own_chunks
(crates/clawhdf5/tests/concurrent_chunk_cache.rs), plus cache unit tests
datasets_sharing_coordinates_stay_separate, dataset_indexes_are_bounded
and concurrent_readers_of_different_datasets_see_their_own_chunks.

Co-Authored-By: Claude Opus 5.5 (1M context) <[email protected]>
This commit is contained in:
osobh
2026-09-25 21:12:52 -05:00
co-authored by Claude Opus 5.5
parent d074385944
commit 9066d34eaa
3 changed files with 660 additions and 351 deletions
Download Patch File Download Diff File Expand all files Collapse all files
crates/clawhdf5-format/src/chunked_read.rs
+54 -77
View File
@@ -836,24 +836,20 @@ pub fn read_chunked_data_cached(
)));
}
// The per-file cache is shared across datasets; bind it to this one so a
// different dataset's chunk index is never reused for this read.
cache.ensure_dataset(addr);
// Populate chunk index on first access
if !cache.has_index() {
let (chunks, _) = list_chunks(
// The per-file cache is shared across datasets (and threads); every
// lookup is keyed by this dataset's chunk-index address, so another
// dataset's index or chunks are never used for this read.
let chunks = cache.chunks_for(addr, rank, || {
list_chunks(
file_data,
layout,
dataspace,
elem_size,
offset_size,
length_size,
)?;
cache.populate_index(&chunks, rank);
}
let chunks = cache.all_indexed_chunks().unwrap_or_default();
)
.map(|(chunks, _)| chunks)
})?;
// Assemble output
let total_bytes = checked_byte_len(dataspace.checked_num_elements()?, elem_size)?;
@@ -920,7 +916,7 @@ pub fn read_chunked_data_cached(
continue;
}
let coord: Vec<u64> = chunk_info.offsets.iter().take(rank).copied().collect();
match cache.get_decompressed_aligned(&coord) {
match cache.get_decompressed_in(addr, &coord) {
Some(cached) => place(&cached, chunk_info),
None => misses.push(chunk_info),
}
@@ -957,7 +953,7 @@ pub fn read_chunked_data_cached(
let data = data?;
if cache_them {
let coord: Vec<u64> = chunk_info.offsets.iter().take(rank).copied().collect();
let cached = cache.put_decompressed(coord, data);
let cached = cache.put_decompressed_in(addr, coord, data);
place(&cached, chunk_info);
} else {
place(&data, chunk_info);
@@ -1161,24 +1157,20 @@ pub fn read_chunked_data_sweep(
)));
}
// The per-file cache is shared across datasets; bind it to this one so a
// different dataset's chunk index is never reused for this read.
cache.ensure_dataset(addr);
// Populate chunk index on first access
if !cache.has_index() {
let (chunks, _) = list_chunks(
// The per-file cache is shared across datasets (and threads); every
// lookup is keyed by this dataset's chunk-index address, so another
// dataset's index or chunks are never used for this read.
let chunks = cache.chunks_for(addr, rank, || {
list_chunks(
file_data,
layout,
dataspace,
elem_size,
offset_size,
length_size,
)?;
cache.populate_index(&chunks, rank);
}
let chunks = cache.all_indexed_chunks().unwrap_or_default();
)
.map(|(chunks, _)| chunks)
})?;
// Assemble output
let total_bytes = checked_byte_len(dataspace.checked_num_elements()?, elem_size)?;
@@ -1209,12 +1201,12 @@ pub fn read_chunked_data_sweep(
// Issue prefetch hint for predicted next chunks
if !sweep.predicted_next.is_empty() {
cache.prefetch_hint(&sweep.predicted_next);
cache.prefetch_hint_in(addr, &sweep.predicted_next);
cache.set_sweep_direction(sweep.direction);
}
// Try decompressed cache first
let decompressed = if let Some(cached) = cache.get_decompressed_aligned(&coord) {
let decompressed = if let Some(cached) = cache.get_decompressed_in(addr, &coord) {
cached
} else {
// Decompress from file
@@ -1233,7 +1225,7 @@ pub fn read_chunked_data_sweep(
} else {
raw_chunk.to_vec()
};
cache.put_decompressed(coord, dec)
cache.put_decompressed_in(addr, coord, dec)
};
let chunk_offsets: Vec<usize> = chunk_info
@@ -1317,48 +1309,34 @@ pub fn read_chunked_data_indexed(
)));
}
// The per-file cache is shared across datasets; bind it to this one so a
// different dataset's chunk index is never reused for this read.
cache.ensure_dataset(addr);
// Build chunk index on first access
if !cache.has_chunk_index() {
let (chunks, _) = list_chunks(
file_data,
layout,
dataspace,
elem_size,
offset_size,
length_size,
)?;
cache.populate_chunk_index(&chunks, rank);
// Also populate the legacy index for compatibility
if !cache.has_index() {
cache.populate_index(&chunks, rank);
}
}
// Build chunk layout on first access
if !cache.has_chunk_layout() {
cache.populate_chunk_layout(&ds_dims, &chunk_dims, elem_size);
}
// Get the layout info (mappings, output size, chunk total bytes)
let (mappings_info, output_bytes, chunk_total_bytes) = cache
.with_chunk_layout(|layout| {
let info: Vec<_> = layout
.mappings
.iter()
.map(|m| (m.coord.clone(), m.file_offset, m.file_size, m.filter_mask))
.collect();
(info, layout.output_bytes, layout.chunk_total_bytes)
})
.ok_or_else(|| FormatError::ChunkedReadError("chunk layout not available".into()))?;
// Chunk index and assembly plan for this dataset, built on first access
// and kept per dataset (keyed by chunk-index address) in the shared cache.
let plan = cache.chunk_layout_for(
addr,
rank,
|| {
list_chunks(
file_data,
layout,
dataspace,
elem_size,
offset_size,
length_size,
)
.map(|(chunks, _)| chunks)
},
&ds_dims,
&chunk_dims,
elem_size,
)?;
let chunk_total_bytes = plan.chunk_total_bytes;
// Decompress chunks (using LRU cache where possible)
let mut chunk_buffers: Vec<Arc<CacheAlignedBuffer>> = Vec::with_capacity(mappings_info.len());
for (coord, file_offset, file_size, filter_mask) in &mappings_info {
if let Some(cached) = cache.get_decompressed_aligned(coord) {
let mut chunk_buffers: Vec<Arc<CacheAlignedBuffer>> = Vec::with_capacity(plan.mappings.len());
for m in &plan.mappings {
let (coord, file_offset, file_size, filter_mask) =
(&m.coord, &m.file_offset, &m.file_size, &m.filter_mask);
if let Some(cached) = cache.get_decompressed_in(addr, coord) {
chunk_buffers.push(cached);
} else {
let c_addr = *file_offset as usize;
@@ -1377,17 +1355,15 @@ pub fn read_chunked_data_indexed(
raw_chunk.to_vec()
};
let aligned = CacheAlignedBuffer::from_vec(decompressed);
let arc = cache.put_decompressed_aligned(coord.clone(), aligned);
let arc = cache.put_decompressed_aligned_in(addr, coord.clone(), aligned);
chunk_buffers.push(arc);
}
}
// Assemble using pre-computed layout
let mut output = vec![0u8; output_bytes];
let mut output = vec![0u8; plan.output_bytes];
let data_refs: Vec<&[u8]> = chunk_buffers.iter().map(|b| b.as_slice()).collect();
cache.with_chunk_layout(|layout| {
layout.assemble(&data_refs, &mut output);
});
plan.assemble(&data_refs, &mut output);
Ok(output)
}
@@ -2307,12 +2283,12 @@ mod tests {
let datatype = make_f64_type();
let cache = ChunkCache::new();
assert!(!cache.has_index());
assert_eq!(cache.indexed_dataset_count(), 0);
let raw = read_chunked_data_cached(
&file_data, &layout, &dataspace, &datatype, None, 8, 8, &cache,
)
.unwrap();
assert!(cache.has_index());
assert_eq!(cache.indexed_dataset_count(), 1);
assert_eq!(raw.len(), 20 * 8);
for i in 0..20 {
let val = f64::from_le_bytes(raw[i * 8..(i + 1) * 8].try_into().unwrap());
@@ -2334,7 +2310,7 @@ mod tests {
&file_data, &layout, &dataspace, &datatype, None, 8, 8, &cache,
)
.unwrap();
assert!(cache.has_index());
assert_eq!(cache.indexed_dataset_count(), 1);
assert_eq!(cache.cached_chunk_count(), 0);
// Second read — reuses the cached index
@@ -2343,6 +2319,7 @@ mod tests {
)
.unwrap();
assert_eq!(raw1, raw2);
assert_eq!(cache.indexed_dataset_count(), 1);
}
#[test]