gwseq_io/source/

cache.rs

//! Aligned LRU block cache in front of any [`ByteSource`].
//!
//! Reads are aligned to `block_size`, blocks are kept in an LRU capped at
//! `max_blocks`, and a block already being fetched is **waited on, not fetched
//! twice**.
//!
//! Each slot carries its own mutex, and the two locks are taken in one order
//! and never together:
//!
//! - lock the map, find or insert the slot, clone the `Arc`, **unlock**;
//! - lock that one slot and fill it if it is empty.
//!
//! So the map lock covers a hash lookup and never an I/O, while a second
//! caller wanting the same block blocks on the slot's mutex and finds the
//! bytes there when it wakes. There are no workers here at all: the rayon
//! threads the reader already owns are the ones fetching.
//!
//! A fill that fails leaves the slot empty rather than caching the error, so
//! the next request retries. Errors are not `Clone` and a cached one could not
//! be handed to a second waiter anyway, but the behaviour is the one worth
//! having on its own: a transient network failure must not stick to a reader
//! for its lifetime.

use std::ops::Range;
use std::sync::Arc;

use bytes::{Bytes, BytesMut};
use indexmap::IndexMap;
use parking_lot::Mutex;

use crate::error::Result;
use crate::source::ByteSource;

/// The most a single `read_at` reserves up front, whatever it was asked for.
///
/// Comfortably above every read this crate makes in one call — the largest are
/// an R-tree leaf and a BAM chunk span — so the growth path is dead code in
/// practice and live only for a file that lies.
const MAX_RESERVE: usize = 8 << 20;

/// One cached block, empty until somebody fills it.
#[derive(Debug, Default)]
struct Slot {
    data: Mutex<Option<Bytes>>,
}

#[derive(Debug)]
pub struct CachedSource<S: ByteSource> {
    inner: S,
    block_size: u64,
    /// Resident blocks, each with the tick of its last use. See [`Self::slot`].
    blocks: Mutex<IndexMap<u64, (u64, Arc<Slot>)>>,
    /// Monotonic counter the ticks come from. Wrapping it would take 2^64
    /// cache lookups.
    tick: std::sync::atomic::AtomicU64,
    max_blocks: usize,
}

impl<S: ByteSource> CachedSource<S> {
    pub fn new(inner: S, block_size: u64, max_blocks: usize) -> Self {
        assert!(block_size > 0, "block_size must be positive");
        assert!(max_blocks > 0, "max_blocks must be positive");
        Self {
            inner,
            block_size,
            blocks: Mutex::new(IndexMap::with_capacity(max_blocks.min(1024))),
            tick: std::sync::atomic::AtomicU64::new(0),
            max_blocks,
        }
    }

    pub fn into_inner(self) -> S {
        self.inner
    }

    pub fn block_size(&self) -> u64 {
        self.block_size
    }

    /// How many blocks are resident. Tests and benchmarks read this; nothing
    /// else should.
    pub fn cached_blocks(&self) -> usize {
        self.blocks.lock().len()
    }

    /// The blocks a `[offset, offset + len)` request touches.
    ///
    /// `saturating_add` because `len` may be a number a file named: a wrapped
    /// end would put the last block *before* the first and the range would
    /// quietly become one block.
    fn block_range(&self, offset: u64, len: usize) -> Range<u64> {
        let first = offset / self.block_size;
        let last = offset.saturating_add(len as u64).div_ceil(self.block_size);
        first..last.max(first + 1)
    }

    /// `len`, cut down to what the source can actually hand over from `offset`.
    ///
    /// Every reader in the crate goes through this cache, and several of them
    /// pass a length that came straight out of the file: an R-tree leaf's
    /// `size`, a chromosome node's `count * item_size`, a BAM chunk's span, a
    /// HiC block's `size`. `LocalSource` and `MemorySource` clamp such a length
    /// where they read it, but this layer sits in front of them and used to
    /// size a buffer from it first — and a failed allocation is an abort, not
    /// an `Err` anyone can catch.
    ///
    /// A source that cannot say how long it is keeps its `len` unclamped rather
    /// than failing a read that would have worked; the reserve is capped
    /// separately, so nothing is sized from the number either way. Both
    /// implementations that reach here memoise their length, so this costs a
    /// load per read after the header.
    fn clamped(&self, offset: u64, len: usize) -> usize {
        match self.inner.len() {
            Ok(total) => {
                let left = total.saturating_sub(offset);
                len.min(usize::try_from(left).unwrap_or(usize::MAX))
            }
            Err(_) => len,
        }
    }

    /// Find or create the slot for a block, evicting the least recently used if
    /// the map is full. Never fetches — the caller does that outside the lock.
    ///
    /// Recency is a tick per entry rather than the map's order. Keeping the
    /// order meant `move_index` on every *hit*, which shifts every entry
    /// between — fine at the default 128 blocks, quadratic for a caller who
    /// asks for ten thousand. A tick makes a hit two stores, and moves the cost
    /// to eviction: a scan for the smallest tick, which happens only on a miss
    /// that overflows, and a miss is about to do I/O anyway.
    fn slot(&self, block: u64) -> Arc<Slot> {
        let mut map = self.blocks.lock();
        let tick = self.tick.fetch_add(1, std::sync::atomic::Ordering::Relaxed);
        if let Some(entry) = map.get_mut(&block) {
            entry.0 = tick;
            return entry.1.clone();
        }
        let slot = Arc::new(Slot::default());
        map.insert(block, (tick, slot.clone()));
        while map.len() > self.max_blocks {
            let oldest = map
                .iter()
                .enumerate()
                .min_by_key(|(_, (_, (tick, _)))| *tick)
                .map(|(index, _)| index)
                .expect("the map is over its limit, so it is not empty");
            // `swap_remove` rather than `shift_remove`: the map's order carries
            // no meaning now that recency does not live in it, and shifting
            // would put back the cost this removed.
            map.swap_remove_index(oldest);
        }
        slot
    }

    /// Fill a slot if it is empty, and hand back its bytes.
    ///
    /// The slot's lock is held across the read, which is the point: a second
    /// caller wanting this block waits here instead of issuing its own request.
    /// The map lock is not held, so the rest of the cache stays live.
    fn fill(&self, block: u64, slot: &Slot) -> Result<Bytes> {
        let mut guard = slot.data.lock();
        if let Some(data) = guard.as_ref() {
            return Ok(data.clone());
        }
        let data = self
            .inner
            .read_at(block * self.block_size, self.block_size as usize)?;
        *guard = Some(data.clone());
        Ok(data)
    }

    fn block(&self, block: u64) -> Result<Bytes> {
        let slot = self.slot(block);
        self.fill(block, &slot)
    }
}

impl<S: ByteSource> ByteSource for CachedSource<S> {
    fn path(&self) -> &str {
        self.inner.path()
    }

    fn len(&self) -> Result<u64> {
        self.inner.len()
    }

    fn read_at(&self, offset: u64, len: usize) -> Result<Bytes> {
        let len = self.clamped(offset, len);
        if len == 0 {
            return Ok(Bytes::new());
        }
        let range = self.block_range(offset, len);

        // One block: hand back a slice of it. `Bytes::slice` is a refcount
        // bump, so a request served from cache copies nothing at all.
        if range.end - range.start == 1 {
            let data = self.block(range.start)?;
            let start = (offset - range.start * self.block_size) as usize;
            if start >= data.len() {
                return Ok(Bytes::new());
            }
            let end = (start + len).min(data.len());
            return Ok(data.slice(start..end));
        }

        // Reserved rather than allocated outright: `len` is clamped to the
        // file above, but a file may honestly be enormous and the loop below
        // appends only what the blocks really hold. Growth is amortised
        // doubling, which on the few-hundred-KB reads this path actually sees
        // never happens at all.
        let mut out = BytesMut::with_capacity(len.min(MAX_RESERVE));
        let mut wanted = len;
        let mut position = offset;
        for index in range {
            if wanted == 0 {
                break;
            }
            let data = self.block(index)?;
            let start = (position - index * self.block_size) as usize;
            if start >= data.len() {
                break; // end of file inside this block
            }
            let take = wanted.min(data.len() - start);
            out.extend_from_slice(&data[start..start + take]);
            position += take as u64;
            wanted -= take;
            if data.len() < self.block_size as usize {
                break; // a short block is the last one
            }
        }
        Ok(out.freeze())
    }

    fn prefetch(&self, ranges: &[Range<u64>]) {
        self.inner.prefetch(ranges)
    }

    fn close(&self) {
        self.blocks.lock().clear();
        self.inner.close();
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use crate::error::Error;
    use std::sync::atomic::{AtomicUsize, Ordering};

    /// An in-memory source that counts the reads reaching it, so the tests can
    /// tell a cache hit from a miss.
    #[derive(Debug)]
    struct CountingSource {
        data: Bytes,
        reads: AtomicUsize,
        fail_until: AtomicUsize,
    }

    impl CountingSource {
        fn new(data: &[u8]) -> Self {
            Self {
                data: Bytes::copy_from_slice(data),
                reads: AtomicUsize::new(0),
                fail_until: AtomicUsize::new(0),
            }
        }
    }

    impl ByteSource for CountingSource {
        fn path(&self) -> &str {
            "memory"
        }
        fn len(&self) -> Result<u64> {
            Ok(self.data.len() as u64)
        }
        fn read_at(&self, offset: u64, len: usize) -> Result<Bytes> {
            self.reads.fetch_add(1, Ordering::SeqCst);
            if self.fail_until.load(Ordering::SeqCst) > 0 {
                self.fail_until.fetch_sub(1, Ordering::SeqCst);
                return Err(Error::invalid("transient"));
            }
            let start = (offset as usize).min(self.data.len());
            let end = (start + len).min(self.data.len());
            Ok(self.data.slice(start..end))
        }
    }

    fn data(n: usize) -> Vec<u8> {
        (0..n).map(|i| (i % 251) as u8).collect()
    }

    #[test]
    fn cached_reads_match_uncached_ones_for_every_range() {
        let raw = data(1000);
        for block_size in [1u64, 7, 64, 512, 4096] {
            let cache = CachedSource::new(CountingSource::new(&raw), block_size, 64);
            let plain = CountingSource::new(&raw);
            for offset in [0u64, 1, 63, 64, 65, 511, 999, 1000, 1500] {
                for len in [0usize, 1, 5, 64, 200, 1000] {
                    let a = cache.read_at(offset, len).unwrap();
                    let b = plain.read_at(offset, len).unwrap();
                    assert_eq!(a, b, "block_size={block_size} offset={offset} len={len}");
                }
            }
        }
    }

    /// The one allocation in this file sized from a caller's number, and the
    /// caller is often a file: a bbi leaf `size`, a chromosome node's
    /// `count * item_size`, a BAM chunk span, a HiC block `size`. Before the
    /// clamp, `1 << 60` here aborted the process — no `Err` can carry a failed
    /// allocation, so this is the only place it can be stopped.
    #[test]
    fn a_length_larger_than_the_file_is_not_allocated() {
        let raw = data(100);
        for block_size in [4u64, 32, 4096] {
            let cache = CachedSource::new(CountingSource::new(&raw), block_size, 4);
            assert_eq!(cache.read_at(0, 1 << 60).unwrap(), &raw[..]);
            assert_eq!(cache.read_at(50, usize::MAX).unwrap(), &raw[50..]);
            assert!(cache.read_at(100, 1 << 60).unwrap().is_empty());
            // An offset past the end plus a length that would wrap `u64`.
            assert!(cache.read_at(u64::MAX - 1, 1 << 60).unwrap().is_empty());
        }
    }

    #[test]
    fn a_repeated_read_does_not_reach_the_source_again() {
        let raw = data(1000);
        let cache = CachedSource::new(CountingSource::new(&raw), 128, 64);
        cache.read_at(0, 100).unwrap();
        let after_first = cache.inner.reads.load(Ordering::SeqCst);
        assert_eq!(after_first, 1);
        for _ in 0..10 {
            cache.read_at(0, 100).unwrap();
            cache.read_at(20, 50).unwrap();
        }
        assert_eq!(cache.inner.reads.load(Ordering::SeqCst), after_first);
    }

    #[test]
    fn the_lru_evicts_the_oldest_and_a_hit_renews_it() {
        let raw = data(10_000);
        let cache = CachedSource::new(CountingSource::new(&raw), 100, 3);
        for block in 0..3u64 {
            cache.read_at(block * 100, 10).unwrap();
        }
        assert_eq!(cache.cached_blocks(), 3);
        // Touch block 0 so block 1 becomes the oldest, then force an eviction.
        cache.read_at(0, 10).unwrap();
        cache.read_at(300, 10).unwrap();
        assert_eq!(cache.cached_blocks(), 3);
        let before = cache.inner.reads.load(Ordering::SeqCst);
        cache.read_at(0, 10).unwrap(); // still resident
        assert_eq!(cache.inner.reads.load(Ordering::SeqCst), before);
        cache.read_at(100, 10).unwrap(); // was evicted
        assert_eq!(cache.inner.reads.load(Ordering::SeqCst), before + 1);
    }

    #[test]
    fn a_failed_fill_is_retried_rather_than_cached() {
        let raw = data(500);
        let cache = CachedSource::new(CountingSource::new(&raw), 128, 8);
        cache.inner.fail_until.store(1, Ordering::SeqCst);
        assert!(cache.read_at(0, 10).is_err());
        // The next request must reach the source again and succeed.
        assert_eq!(&cache.read_at(0, 10).unwrap()[..], &raw[0..10]);
    }

    #[test]
    fn concurrent_readers_of_one_block_fetch_it_once() {
        let raw = data(1 << 16);
        let cache = Arc::new(CachedSource::new(CountingSource::new(&raw), 1 << 12, 64));
        let threads: Vec<_> = (0..16)
            .map(|_| {
                let cache = cache.clone();
                let raw = raw.clone();
                std::thread::spawn(move || {
                    for _ in 0..200 {
                        let got = cache.read_at(4096, 4096).unwrap();
                        assert_eq!(&got[..], &raw[4096..8192]);
                    }
                })
            })
            .collect();
        for t in threads {
            t.join().unwrap();
        }
        // Sixteen threads, one block, one read of it.
        assert_eq!(cache.inner.reads.load(Ordering::SeqCst), 1);
    }

    #[test]
    fn close_drops_the_blocks_and_the_handle() {
        let cache = CachedSource::new(CountingSource::new(&data(500)), 128, 8);
        cache.read_at(0, 10).unwrap();
        assert_eq!(cache.cached_blocks(), 1);
        cache.close();
        assert_eq!(cache.cached_blocks(), 0);
    }
}