clawhdf5/crates/clawhdf5-format/src/parallel_read.rs

//! Parallel chunk decompression using rayon with lane partitioning.
//!
//! When reading a chunked+compressed dataset with many chunks, this module
//! uses lane-partitioned parallel decompression: each thread receives a
//! deterministic, disjoint subset of chunks — no overlap, no coordination.
//!
//! The lane assignment is seeded by dataset metadata so repeated reads of
//! the same region produce identical partitions (cache-friendly, reproducible).

use crate::chunked_read::ChunkInfo;
use crate::error::FormatError;
use crate::filter_pipeline::FilterPipeline;
use crate::filters::decompress_chunk_exact;
use crate::lane_partition::{self, LaneStats, PartitionStats};
use crate::storage::{ExtentReq, Storage, for_each_extent_batch};

/// The extents of `chunks`' stored bytes (see
/// [`crate::chunked_read::chunk_req`]), fetched batch by batch with
/// [`for_each_extent_batch`] when the file is not in memory (each chunk's
/// bounds error is reported when that chunk is decoded, as before).
fn chunk_reqs(
    chunks: &[ChunkInfo],
    pipeline: Option<&FilterPipeline>,
    chunk_total_bytes: usize,
) -> Vec<ExtentReq> {
    chunks
        .iter()
        .map(|c| crate::chunked_read::chunk_req(c, pipeline, chunk_total_bytes, true))
        .collect()
}

/// Threshold: only use parallel decompression when chunk count exceeds this.
const PARALLEL_THRESHOLD: usize = 4;

/// Result of decompressing a single chunk, tagged with its index for ordering.
struct DecompressedChunk {
    index: usize,
    data: Vec<u8>,
}

/// Returns `true` if the parallel path should be used for the given chunk count.
pub fn should_use_parallel(chunk_count: usize) -> bool {
    chunk_count > PARALLEL_THRESHOLD
}

/// Whether handing a read's chunks to rayon can decode them faster than the
/// calling thread would alone.
///
/// `false` when the pool the work would go to (the current pool inside a
/// rayon worker, else the global one) has a single thread. Handing work to
/// that pool is then worse than useless: the caller blocks while the one
/// worker decodes, and every other thread reading at the same time queues
/// behind the same worker, so N reader threads decode on one core. (That is
/// how full reads with `--decode-threads 1` stopped scaling at about 2x in
/// the `concurrent_read` benchmark.)
pub fn pool_can_parallelise() -> bool {
    rayon::current_num_threads() > 1
}

/// How many rayon workers [`run_with_helpers`] should ask to help with
/// `items` work items, given that the calling thread works too: the pool's
/// other threads (all of them when the caller is not one), at most one per
/// item beyond the caller's first.
pub(crate) fn helper_count(items: usize) -> usize {
    let pool = rayon::current_num_threads();
    // A one-thread pool means "decode on the calling thread" (the setting
    // benchmarks use to compare with h5py, where each call decodes on its
    // caller): no helper, so one read never uses two cores.
    if pool <= 1 {
        return 0;
    }
    let others = if rayon::current_thread_index().is_some() {
        pool.saturating_sub(1)
    } else {
        pool
    };
    others.min(items.saturating_sub(1))
}

/// Run `body` on the calling thread and on up to `helpers` rayon workers at
/// once, returning when the caller's call has finished and every worker that
/// started one has too. `body` shares its work out itself (typically by
/// claiming items from an atomic counter until none are left).
///
/// The caller never waits for a worker to *become* free: helpers are queued
/// on the pool, and one that only gets to run after the caller has finished
/// returns without calling `body`. So a busy or small pool can only fail to
/// speed a read up, never hold it back — with `par_iter`, the calling thread
/// (not a pool worker) handed all the work to the pool and slept, and N
/// threads reading through a 2-worker pool decoded on 2 cores.
///
/// A panic in `body`, on any thread, is resumed on the caller once every
/// helper that started has stopped.
pub(crate) fn run_with_helpers(helpers: usize, body: &(dyn Fn() + Sync)) {
    use std::panic::{AssertUnwindSafe, catch_unwind, resume_unwind};
    use std::sync::{Arc, Condvar, Mutex, PoisonError};

    if helpers == 0 {
        body();
        return;
    }

    type Body = dyn Fn() + Sync + 'static;
    struct Shared {
        /// `body`, its lifetime erased. Only dereferenced by a helper that
        /// registered in `state` while it was open (see below).
        body: *const Body,
        /// (closed, helpers inside `body`).
        state: Mutex<(bool, usize)>,
        idle: Condvar,
        panic: Mutex<Option<Box<dyn core::any::Any + Send>>>,
    }
    // SAFETY: `body` points to a `Sync` closure, so calling it from other
    // threads is allowed; the pointer is only used under the protocol below,
    // which keeps it from outliving the closure.
    unsafe impl Send for Shared {}
    unsafe impl Sync for Shared {}

    fn help(shared: &Shared) {
        {
            let mut state = shared.state.lock().unwrap_or_else(PoisonError::into_inner);
            if state.0 {
                return;
            }
            state.1 += 1;
        }
        // SAFETY: registered while open, so the caller of `run_with_helpers`
        // is still inside it (it closes, then waits until no helper is
        // registered, before returning), and `body` is alive.
        let body = unsafe { &*shared.body };
        if let Err(payload) = catch_unwind(AssertUnwindSafe(body)) {
            shared
                .panic
                .lock()
                .unwrap_or_else(PoisonError::into_inner)
                .get_or_insert(payload);
        }
        let mut state = shared.state.lock().unwrap_or_else(PoisonError::into_inner);
        state.1 -= 1;
        if state.1 == 0 {
            shared.idle.notify_all();
        }
    }

    let body_ptr: *const (dyn Fn() + Sync + '_) = body;
    // SAFETY: only the lifetime changes (same fat-pointer layout). The
    // pointer is dereferenced only while this function is running: see
    // `help` and the wait below.
    let body_ptr: *const Body = unsafe { core::mem::transmute(body_ptr) };
    let shared = Arc::new(Shared {
        body: body_ptr,
        state: Mutex::new((false, 0)),
        idle: Condvar::new(),
        panic: Mutex::new(None),
    });
    for _ in 0..helpers {
        let shared = Arc::clone(&shared);
        rayon::spawn(move || help(&shared));
    }
    let caller = catch_unwind(AssertUnwindSafe(body));
    {
        // Close, then wait for the helpers inside `body`; later ones return
        // at once. This must happen even if `body` panicked on this thread.
        let mut state = shared.state.lock().unwrap_or_else(PoisonError::into_inner);
        state.0 = true;
        while state.1 > 0 {
            state = shared
                .idle
                .wait(state)
                .unwrap_or_else(PoisonError::into_inner);
        }
    }
    if let Err(payload) = caller {
        resume_unwind(payload);
    }
    let helper_panic = shared
        .panic
        .lock()
        .unwrap_or_else(PoisonError::into_inner)
        .take();
    if let Some(payload) = helper_panic {
        resume_unwind(payload);
    }
}

/// Decompress chunks in parallel using lane-partitioned assignment.
///
/// Instead of naive `par_iter`, chunks are deterministically assigned to lanes
/// (threads) using a seeded pseudorandom permutation.  Each lane processes
/// only its assigned chunks — no redundant work, no coordination.
///
/// # Arguments
///
/// * `seed` - Seed for the partition permutation (e.g. dataset address + chunk range hash).
/// * `num_lanes` - Number of parallel lanes.  Pass `None` to auto-detect from available cores.
///
/// # Errors
///
/// Returns the first error encountered by any worker thread.
pub fn decompress_chunks_lane_partitioned(
    file_data: &[u8],
    chunks: &[ChunkInfo],
    pipeline: &FilterPipeline,
    chunk_total_bytes: usize,
    element_size: u32,
    seed: u64,
    num_lanes: Option<usize>,
) -> Result<(Vec<Vec<u8>>, PartitionStats), FormatError> {
    decompress_chunks_lane_partitioned_in(
        file_data,
        chunks,
        pipeline,
        chunk_total_bytes,
        element_size,
        seed,
        num_lanes,
    )
}

/// [`decompress_chunks_lane_partitioned`] over any [`Storage`].
pub fn decompress_chunks_lane_partitioned_in<S: Storage + ?Sized>(
    file_data: &S,
    chunks: &[ChunkInfo],
    pipeline: &FilterPipeline,
    chunk_total_bytes: usize,
    element_size: u32,
    seed: u64,
    num_lanes: Option<usize>,
) -> Result<(Vec<Vec<u8>>, PartitionStats), FormatError> {
    use rayon::prelude::*;

    let lanes = num_lanes.unwrap_or_else(|| {
        std::thread::available_parallelism()
            .map(|n| n.get())
            .unwrap_or(1)
    });

    let reqs = chunk_reqs(chunks, Some(pipeline), chunk_total_bytes);
    let mut ordered: Vec<Vec<u8>> = Vec::with_capacity(chunks.len());
    let mut partition_stats = PartitionStats::new(0);
    partition_stats.total_chunks = chunks.len();
    // Each batch of fetched chunks is partitioned into lanes and decoded
    // before the next batch is fetched (with the file in memory there is
    // one batch: all the chunks).
    for_each_extent_batch(file_data, &reqs, |batch, raw_bytes| {
        let assignments = lane_partition::partition_chunks(batch.len(), lanes, seed);
        // Each lane processes its assigned chunks and returns results + stats.
        let lane_results: Result<Vec<(Vec<DecompressedChunk>, LaneStats)>, FormatError> =
            assignments
                .into_par_iter()
                .map(|indices| {
                    let mut results = Vec::with_capacity(indices.len());
                    let mut stats = LaneStats::default();

                    for &local in &indices {
                        let index = batch.start + local;
                        let chunk_info = &chunks[index];
                        let size = chunk_info.chunk_size as usize;
                        let raw_chunk = raw_bytes.get(index, &reqs[index])?;

                        let decompressed = decompress_chunk_exact(
                            raw_chunk,
                            pipeline,
                            chunk_total_bytes,
                            element_size,
                            chunk_info.filter_mask,
                            &chunk_info.offsets,
                        )?;

                        stats.chunks_processed += 1;
                        stats.compressed_bytes += size as u64;
                        stats.decompressed_bytes += decompressed.len() as u64;

                        results.push(DecompressedChunk {
                            index,
                            data: decompressed,
                        });
                    }

                    Ok((results, stats))
                })
                .collect();
        let lane_results = lane_results?;

        // Aggregate stats
        if partition_stats.per_lane.len() < lane_results.len() {
            partition_stats
                .per_lane
                .resize_with(lane_results.len(), LaneStats::default);
            partition_stats.num_lanes = lane_results.len();
        }
        for (lane, (_, stats)) in partition_stats.per_lane.iter_mut().zip(&lane_results) {
            lane.chunks_processed += stats.chunks_processed;
            lane.compressed_bytes += stats.compressed_bytes;
            lane.decompressed_bytes += stats.decompressed_bytes;
        }

        // Flatten and sort by original index to restore order
        let mut all_chunks: Vec<DecompressedChunk> = lane_results
            .into_iter()
            .flat_map(|(chunks, _)| chunks)
            .collect();
        all_chunks.sort_by_key(|dc| dc.index);
        ordered.extend(all_chunks.into_iter().map(|dc| dc.data));
        Ok(())
    })?;
    Ok((ordered, partition_stats))
}

/// Decompress chunks in parallel using rayon (legacy par_iter path).
///
/// Each chunk is read from `file_data` at the address in the corresponding
/// `ChunkInfo`, decompressed through `pipeline`, and collected in order.
///
/// # Errors
///
/// Returns the first error encountered by any worker thread.
pub fn decompress_chunks_parallel(
    file_data: &[u8],
    chunks: &[ChunkInfo],
    pipeline: &FilterPipeline,
    chunk_total_bytes: usize,
    element_size: u32,
) -> Result<Vec<Vec<u8>>, FormatError> {
    decompress_chunks_parallel_in(file_data, chunks, pipeline, chunk_total_bytes, element_size)
}

/// [`decompress_chunks_parallel`] over any [`Storage`].
pub fn decompress_chunks_parallel_in<S: Storage + ?Sized>(
    file_data: &S,
    chunks: &[ChunkInfo],
    pipeline: &FilterPipeline,
    chunk_total_bytes: usize,
    element_size: u32,
) -> Result<Vec<Vec<u8>>, FormatError> {
    use rayon::prelude::*;

    let reqs = chunk_reqs(chunks, Some(pipeline), chunk_total_bytes);
    let mut ordered: Vec<Vec<u8>> = Vec::with_capacity(chunks.len());
    for_each_extent_batch(file_data, &reqs, |batch, raw_bytes| {
        let results: Result<Vec<DecompressedChunk>, FormatError> = batch
            .clone()
            .into_par_iter()
            .map(|index| {
                let chunk_info = &chunks[index];
                let raw_chunk = raw_bytes.get(index, &reqs[index])?;

                let decompressed = decompress_chunk_exact(
                    raw_chunk,
                    pipeline,
                    chunk_total_bytes,
                    element_size,
                    chunk_info.filter_mask,
                    &chunk_info.offsets,
                )?;

                Ok(DecompressedChunk {
                    index,
                    data: decompressed,
                })
            })
            .collect();

        let mut result_vec = results?;
        result_vec.sort_by_key(|dc| dc.index);
        ordered.extend(result_vec.into_iter().map(|dc| dc.data));
        Ok(())
    })?;
    Ok(ordered)
}

/// Decompress chunks sequentially (fallback when parallel is not warranted).
pub fn decompress_chunks_sequential(
    file_data: &[u8],
    chunks: &[ChunkInfo],
    pipeline: Option<&FilterPipeline>,
    chunk_total_bytes: usize,
    element_size: u32,
) -> Result<Vec<Vec<u8>>, FormatError> {
    decompress_chunks_sequential_in(file_data, chunks, pipeline, chunk_total_bytes, element_size)
}

/// [`decompress_chunks_sequential`] over any [`Storage`].
pub fn decompress_chunks_sequential_in<S: Storage + ?Sized>(
    file_data: &S,
    chunks: &[ChunkInfo],
    pipeline: Option<&FilterPipeline>,
    chunk_total_bytes: usize,
    element_size: u32,
) -> Result<Vec<Vec<u8>>, FormatError> {
    let reqs = chunk_reqs(chunks, pipeline, chunk_total_bytes);
    let mut result = Vec::with_capacity(chunks.len());
    for_each_extent_batch(file_data, &reqs, |batch, raw_bytes| {
        for i in batch {
            let chunk_info = &chunks[i];
            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,
                    element_size,
                    chunk_info.filter_mask,
                    &chunk_info.offsets,
                )?
            } else {
                raw_chunk.to_vec()
            };
            result.push(decompressed);
        }
        Ok(())
    })?;
    Ok(result)
}

#[cfg(test)]
mod tests {
    use super::*;
    use crate::filter_pipeline::{FILTER_SHUFFLE, FilterDescription};

    /// Eight shuffled 32-byte chunks; chunk 5 is stored short when `short`.
    fn chunks(short: bool) -> (Vec<u8>, Vec<ChunkInfo>) {
        let mut file = Vec::new();
        let mut infos = Vec::new();
        for i in 0..8u64 {
            let len = if short && i == 5 { 16 } else { 32 };
            infos.push(ChunkInfo {
                chunk_size: len as u32,
                filter_mask: 0,
                offsets: vec![i * 8],
                address: file.len() as u64,
            });
            file.extend(core::iter::repeat_n(i as u8, len));
        }
        (file, infos)
    }

    /// Every item is processed exactly once, whatever mix of caller and
    /// helpers ends up doing it.
    #[test]
    fn run_with_helpers_shares_all_work() {
        use core::sync::atomic::{AtomicUsize, Ordering};
        for helpers in [0, 1, 3, 16] {
            let n = 1000;
            let next = AtomicUsize::new(0);
            let done: Vec<AtomicUsize> = (0..n).map(|_| AtomicUsize::new(0)).collect();
            run_with_helpers(helpers, &|| {
                loop {
                    let i = next.fetch_add(1, Ordering::Relaxed);
                    if i >= n {
                        break;
                    }
                    done[i].fetch_add(1, Ordering::Relaxed);
                }
            });
            assert!(done.iter().all(|d| d.load(Ordering::Relaxed) == 1));
        }
    }

    /// A panic in the shared body reaches the caller whichever thread it
    /// happened on, and only after the helpers inside the body have left it
    /// (they borrow the caller's stack).
    #[test]
    fn run_with_helpers_propagates_panics() {
        use core::sync::atomic::{AtomicUsize, Ordering};
        use std::panic::{AssertUnwindSafe, catch_unwind};
        let caller = std::thread::current().id();
        for panic_on_caller in [true, false] {
            let inside = AtomicUsize::new(0);
            let calls = AtomicUsize::new(0);
            let result = catch_unwind(AssertUnwindSafe(|| {
                run_with_helpers(4, &|| {
                    inside.fetch_add(1, Ordering::SeqCst);
                    calls.fetch_add(1, Ordering::SeqCst);
                    let on_caller = std::thread::current().id() == caller;
                    std::thread::sleep(std::time::Duration::from_millis(20));
                    inside.fetch_sub(1, Ordering::SeqCst);
                    if on_caller == panic_on_caller {
                        panic!("boom");
                    }
                });
            }));
            // A helper may never have run (the pool was slow to start it),
            // in which case nothing panicked when `panic_on_caller` is false.
            if panic_on_caller || calls.load(Ordering::SeqCst) > 1 {
                assert!(result.is_err());
            }
            assert_eq!(inside.load(Ordering::SeqCst), 0);
        }
    }

    /// Every parallel decoder refuses a chunk that decodes short, naming it.
    #[test]
    fn short_decoded_chunk_is_an_error() {
        let pipeline = FilterPipeline {
            version: 2,
            filters: vec![FilterDescription {
                filter_id: FILTER_SHUFFLE,
                name: None,
                flags: 0,
                client_data: vec![4],
            }],
        };
        let (file, good) = chunks(false);
        assert_eq!(
            decompress_chunks_parallel(&file, &good, &pipeline, 32, 4).unwrap()[5],
            [5u8; 32]
        );
        let (file, bad) = chunks(true);
        let errs = [
            decompress_chunks_lane_partitioned(&file, &bad, &pipeline, 32, 4, 1, Some(3))
                .map(|_| ())
                .unwrap_err(),
            decompress_chunks_parallel(&file, &bad, &pipeline, 32, 4)
                .map(|_| ())
                .unwrap_err(),
            decompress_chunks_sequential(&file, &bad, Some(&pipeline), 32, 4)
                .map(|_| ())
                .unwrap_err(),
        ];
        for e in errs {
            assert!(e.to_string().contains("[40]"), "{e}");
        }
    }
}