Skip to main content

hdf5_pure/
chunk_cache.rs

1//! Chunk cache with hash-based index and LRU eviction.
2//!
3//! The [`ChunkCache`] avoids re-traversing B-trees on repeated reads of chunked
4//! datasets.  On first access it scans the B-tree once and builds a
5//! `HashMap<ChunkCoord, ChunkInfo>` (the *chunk index*).  Decompressed chunk
6//! data is cached with LRU eviction controlled by a byte-budget.
7
8#[cfg(not(feature = "std"))]
9extern crate alloc;
10
11#[cfg(not(feature = "std"))]
12use alloc::vec::Vec;
13
14#[cfg(not(feature = "std"))]
15use crate::nosync::Mutex;
16#[cfg(feature = "std")]
17use std::sync::Mutex;
18
19#[cfg(not(feature = "std"))]
20use alloc::collections::BTreeMap;
21#[cfg(feature = "std")]
22use std::collections::HashMap;
23
24use crate::chunked_read::ChunkInfo;
25
26/// Coordinate key for a chunk — the N-dimensional offset vector.
27pub type ChunkCoord = Vec<u64>;
28
29/// Default maximum bytes of decompressed chunk data to cache.
30pub const DEFAULT_CACHE_BYTES: usize = 1024 * 1024; // 1 MiB
31
32/// Default maximum number of cached decompressed chunks.
33pub const DEFAULT_MAX_SLOTS: usize = 16;
34
35/// Configuration for a per-dataset chunk cache.
36///
37/// The byte and slot limits are the `hdf5-pure` counterpart of the
38/// `rdcc_nbytes` and `rdcc_nslots` raw-data chunk-cache settings from HDF5's
39/// `H5Pset_cache`. They apply to decompressed raw chunk data. The optional
40/// chunk-index cache controls whether `hdf5-pure` retains the parsed chunk
41/// address index between reads of the same [`crate::Dataset`]. Disabling the
42/// index cache lowers retained metadata memory at the cost of re-scanning the
43/// on-disk chunk index for repeated reads.
44#[derive(Debug, Clone, Copy, PartialEq, Eq)]
45pub struct ChunkCacheConfig {
46    max_bytes: usize,
47    max_slots: usize,
48    cache_index: bool,
49}
50
51impl ChunkCacheConfig {
52    /// Create a config matching the historical defaults: 1 MiB of decompressed
53    /// chunks, 16 slots, and retained parsed chunk indexes.
54    pub const fn new() -> Self {
55        Self {
56            max_bytes: DEFAULT_CACHE_BYTES,
57            max_slots: DEFAULT_MAX_SLOTS,
58            cache_index: true,
59        }
60    }
61
62    /// Create a config from HDF5 `H5Pset_cache` raw data chunk-cache values.
63    ///
64    /// `rdcc_nslots` maps to the maximum retained chunk slots and
65    /// `rdcc_nbytes` maps to the maximum retained decompressed chunk bytes.
66    /// Modern HDF5 ignores `H5Pset_cache`'s `mdc_nelmts`; use
67    /// [`crate::MetadataCacheConfig`] for the metadata-cache budget.
68    ///
69    /// # `rdcc_w0` has nothing to decide here
70    ///
71    /// C's third field is a head start, in entries, for a rule that preempts only
72    /// chunks the caller consumed *in full* before the unqualified rule joins the
73    /// walk (`H5D__chunk_cache_prune`, discriminating on
74    /// `H5D_rdcc_ent_t::rd_count` / `wr_count`). Nothing here answers to it. This
75    /// cache holds decompressed data read from the file, so there is no dirty
76    /// chunk for `wr_count` to describe; and it records only whether a chunk is
77    /// held, never how much of it a caller took, so the `rd_count` class it would
78    /// sort by does not exist to be sorted.
79    ///
80    /// What `w0` exists to prevent — a scan evicting chunks for chunks it will not
81    /// ask for again — is handled here without a knob: a whole read fills the
82    /// cache and stops offering, so it has no prune for a preemption policy to
83    /// order. The field-by-field argument, including why a row window is the one
84    /// path that does evict and why recency already keeps what it needs, is in
85    /// the [property-support reference].
86    ///
87    /// [property-support reference]: https://github.com/CramBL/hdf5-pure/blob/main/docs/reference/property-support.md
88    pub const fn from_h5p_cache(rdcc_nslots: usize, rdcc_nbytes: usize) -> Self {
89        Self {
90            max_bytes: rdcc_nbytes,
91            max_slots: rdcc_nslots,
92            cache_index: true,
93        }
94    }
95
96    /// Disable retained decompressed chunks and parsed chunk indexes.
97    pub const fn disabled() -> Self {
98        Self {
99            max_bytes: 0,
100            max_slots: 0,
101            cache_index: false,
102        }
103    }
104
105    /// Set the maximum decompressed chunk bytes retained per dataset.
106    pub const fn with_max_bytes(mut self, max_bytes: usize) -> Self {
107        self.max_bytes = max_bytes;
108        self
109    }
110
111    /// Set the maximum number of decompressed chunk slots retained per dataset.
112    pub const fn with_max_slots(mut self, max_slots: usize) -> Self {
113        self.max_slots = max_slots;
114        self
115    }
116
117    /// Enable or disable retaining the parsed chunk index between reads.
118    pub const fn with_index_cache(mut self, enabled: bool) -> Self {
119        self.cache_index = enabled;
120        self
121    }
122
123    /// Return the maximum decompressed chunk bytes retained per dataset.
124    pub const fn max_bytes(&self) -> usize {
125        self.max_bytes
126    }
127
128    /// Return the maximum decompressed chunk slots retained per dataset.
129    pub const fn max_slots(&self) -> usize {
130        self.max_slots
131    }
132
133    /// Return whether parsed chunk indexes are retained between reads.
134    pub const fn index_cache_enabled(&self) -> bool {
135        self.cache_index
136    }
137}
138
139impl Default for ChunkCacheConfig {
140    fn default() -> Self {
141        Self::new()
142    }
143}
144
145/// What a dataset's chunk cache has done, and what it is holding.
146///
147/// Returned by [`crate::Dataset::chunk_cache_stats`].
148/// [`index_loaded`](Self::index_loaded), [`cached_chunks`](Self::cached_chunks)
149/// and [`cached_bytes`](Self::cached_bytes) are a point-in-time view of
150/// occupancy; the counters are cumulative since the handle was opened or since
151/// the last [`reset_chunk_cache_stats`](crate::Dataset::reset_chunk_cache_stats).
152///
153/// The reason to look is that [`ChunkCacheConfig`] is a budget chosen before a
154/// single chunk has been read, and nothing else reports whether it was the right
155/// one:
156///
157/// - [`hit_rate`](Self::hit_rate) says whether the cache is earning its memory.
158/// - [`rejections`](Self::rejections) and [`evictions`](Self::evictions)
159///   together say whether the working set outgrew the budget. **Read both**:
160///   which of the two moves depends on how the dataset is being read, and each
161///   is structurally zero on the path the other reports. A whole read fills the
162///   cache and then stops offering, so it rejects and never evicts; a row window
163///   ([`Dataset::read_raw_rows`](crate::Dataset::read_raw_rows) and the typed
164///   `read_*_rows`) evicts by plain LRU, so it evicts and never rejects. Either
165///   one climbing while [`hit_rate`](Self::hit_rate) stays low is the same
166///   finding.
167/// - [`oversize_chunks`](Self::oversize_chunks) says whether a single chunk is
168///   larger than the whole byte budget, in which case no slot count helps.
169/// - [`invalidations`](Self::invalidations) says how many retained chunks a
170///   commit in this session threw away.
171#[derive(Debug, Clone, Copy, PartialEq, Eq, Default)]
172pub struct ChunkCacheStats {
173    index_loaded: bool,
174    cached_chunks: usize,
175    cached_bytes: usize,
176    counters: Counters,
177}
178
179impl ChunkCacheStats {
180    /// Whether the parsed chunk index is currently held in memory.
181    pub const fn index_loaded(&self) -> bool {
182        self.index_loaded
183    }
184
185    /// Number of decompressed chunks currently retained.
186    pub const fn cached_chunks(&self) -> usize {
187        self.cached_chunks
188    }
189
190    /// Total bytes of decompressed chunk data currently retained.
191    pub const fn cached_bytes(&self) -> usize {
192        self.cached_bytes
193    }
194
195    /// Chunk lookups served from retained decompressed data.
196    pub const fn hits(&self) -> u64 {
197        self.counters.hits
198    }
199
200    /// Chunk lookups that had to be fetched and decoded.
201    pub const fn misses(&self) -> u64 {
202        self.counters.misses
203    }
204
205    /// Retained chunks dropped to make room for another chunk.
206    ///
207    /// A *whole* read never evicts a chunk it placed or was served, so on that
208    /// path this counts only what one read took from an earlier read's chunks
209    /// that it did not itself use — zero for a dataset read the same way twice,
210    /// and no evidence either way about the budget. See
211    /// [`rejections`](Self::rejections) for that. A *row window* asks for the
212    /// plain LRU rule instead, since the chunk its successor needs is the one it
213    /// finished on; there this is the budget signal and `rejections` is the
214    /// figure that stays at zero.
215    pub const fn evictions(&self) -> u64 {
216        self.counters.evictions
217    }
218
219    /// Chunks offered to the cache and not retained, because it was already full
220    /// of chunks the same read had placed or been served.
221    ///
222    /// A whole read visits each of its chunks exactly once, so once it has filled
223    /// the cache there is nothing to gain by evicting one of its own chunks for
224    /// another; it stops offering instead. That makes this the budget signal on
225    /// that path and [`evictions`](Self::evictions) useless there: a whole read of
226    /// a dataset eight times the budget reports seven eighths of its chunks here
227    /// and no evictions at all. A row window reverses the two — it evicts by plain
228    /// LRU and never reaches this.
229    ///
230    /// Counted apart from [`oversize_chunks`](Self::oversize_chunks), which more
231    /// slots would not fix.
232    pub const fn rejections(&self) -> u64 {
233        self.counters.rejections
234    }
235
236    /// Chunks offered to the cache whose decompressed size exceeds the whole
237    /// byte budget, [`ChunkCacheConfig::max_bytes`].
238    ///
239    /// Raising [`max_slots`](ChunkCacheConfig::max_slots) does nothing for these;
240    /// only a byte budget above one chunk will admit them. Like
241    /// [`rejections`](Self::rejections) this counts offers, so a chunk too large
242    /// for the budget is counted again on every read that reaches it.
243    pub const fn oversize_chunks(&self) -> u64 {
244        self.counters.oversize_chunks
245    }
246
247    /// Retained chunks dropped because a commit in this session may have made
248    /// them stale.
249    ///
250    /// Session-wide, not handle-wide: a commit through *any* handle advances the
251    /// file's content revision, and every dataset handle drops its chunks the
252    /// next time it resolves. Only a read-write session invalidates; this stays
253    /// zero on a read-only open. Invalidations approaching
254    /// [`misses`](Self::misses) mean the session is rewriting the chunks it is
255    /// caching, and a larger budget will not change that.
256    pub const fn invalidations(&self) -> u64 {
257        self.counters.invalidations
258    }
259
260    /// Every chunk lookup: hits plus misses.
261    pub const fn lookups(&self) -> u64 {
262        self.counters.hits.saturating_add(self.counters.misses)
263    }
264
265    /// The fraction of chunk lookups served from the cache, or `None` before any
266    /// lookup has happened.
267    ///
268    /// `None` rather than `0.0`, which is also what a cache that missed every
269    /// lookup reports: the two mean opposite things to a caller deciding whether
270    /// to raise the budget, and only one of them is a reason to. A *disabled*
271    /// cache is the third case reporting `Some(0.0)` — every lookup is counted a
272    /// miss whether or not the cache was ever allowed to answer — so read it
273    /// beside [`ChunkCacheConfig`], which says whether there was a cache at all.
274    pub fn hit_rate(&self) -> Option<f64> {
275        let lookups = self.lookups();
276        if lookups == 0 {
277            return None;
278        }
279        #[expect(
280            clippy::cast_precision_loss,
281            reason = "a hit rate is a ratio; f64 holds these counts exactly far past any \
282                      chunk count a process will reach"
283        )]
284        Some(self.counters.hits as f64 / lookups as f64)
285    }
286}
287
288// ---------------------------------------------------------------------------
289// LRU entry
290// ---------------------------------------------------------------------------
291
292struct CachedChunk {
293    coord: ChunkCoord,
294    data: Vec<u8>,
295    /// Monotonically increasing access counter for LRU ordering.
296    last_access: u64,
297    /// The read pass that stored it, so that pass cannot evict it again.
298    stored_by: u64,
299}
300
301/// One read's pass over a set of chunks, as far as cache admission is concerned.
302///
303/// Every read path in this crate visits each of its chunks exactly once. Within
304/// one such pass, evicting a chunk to make room for another is work with no
305/// upside *to that pass*: the evicted chunk has already been placed and will not
306/// be asked for again, and neither will the one that displaced it. A whole read
307/// of a dataset larger than the cache did exactly that — 2,048 chunks offered to
308/// 16 slots, 2,032 of them evicted by the same read that stored them, an
309/// allocator round trip each and, on the unfiltered path, a copy of the chunk as
310/// well (issue #228).
311///
312/// So a pass fills the cache and then stops offering, and what it leaves behind
313/// is the chunks it reached first rather than the ones it reached last.
314///
315/// # Which half is worth keeping is the caller's question, not this type's
316///
317/// That last sentence is the whole trade, and it does not go the same way for
318/// every read. Keeping the *tail* is only possible by offering every chunk and
319/// evicting, which is the cost this exists to remove — so a read that wants the
320/// tail asks for [`CachePass::LRU`] and pays for it.
321///
322/// A read of a whole dataset does not want it: a caller who reads it again
323/// starts at the beginning, so a retained prefix is worth at least as much as a
324/// retained suffix, and it costs a fraction as much to keep. A *lone* row window
325/// does want it, because its successor is the adjacent window and the chunk they
326/// share is the one this read finished on — which is why
327/// [`Dataset::read_raw_rows`](crate::Dataset::read_raw_rows) asks for `LRU` while
328/// the two whole-dataset loops open a real pass.
329///
330/// The windowed reader itself takes the pass from its caller rather than
331/// choosing, because the same window means different things to different
332/// callers: a sweep of a whole dataset in windows opens one real pass for all of
333/// them, since it asks for each chunk exactly once and has no more use for the
334/// last window's chunks than for the first's.
335///
336/// # Being served a chunk claims it, too
337///
338/// The same reasoning reaches one step further than placing a chunk does. A pass
339/// owns any slot it has been *served* from as well as any it filled, because a
340/// chunk this read has already used is worth at least as much as one it has not
341/// reached yet — and giving it back buys the same nothing.
342///
343/// Left out, that cost a repeat read everything the previous read had saved for
344/// it. A second read of a dataset larger than the cache hit all `max_slots`
345/// retained chunks, then handed every one of them back on its next `max_slots`
346/// misses, so the third read hit nothing and the count alternated `max_slots`,
347/// 0, `max_slots`, 0 forever — half the available hit rate, plus a placement and
348/// an eviction per chunk to arrive at a set the next read would destroy.
349///
350/// This is narrower than refusing to touch an earlier pass's chunks at all,
351/// which would strand the cache on whatever filled it first. A read claims only
352/// what it is actually served, so a read that wants chunks the cache does not
353/// hold claims nothing and takes the slots it needs; one that half overlaps
354/// keeps the half it used and replaces the half it did not. The cache still
355/// tracks the most recent access pattern, and now stops paying to re-track the
356/// same one.
357#[derive(Clone, Copy, Debug, PartialEq, Eq)]
358pub struct CachePass(u64);
359
360impl CachePass {
361    /// Admission by the plain LRU rule this cache had before passes existed:
362    /// every slot is evictable, including one this same identity stored.
363    ///
364    /// It is recognized by *being* this value rather than by its number. Zero is
365    /// outside the range [`ChunkCache::begin_pass`] hands out, which keeps a real
366    /// pass from ever being mistaken for it — but that alone would not be enough
367    /// in the other direction, since every `LRU` insert records the same
368    /// `stored_by` and would then look like its own pass's work.
369    pub const LRU: CachePass = CachePass(0);
370}
371
372// ---------------------------------------------------------------------------
373// ChunkCache
374// ---------------------------------------------------------------------------
375
376/// A per-dataset chunk cache with hash-based index and LRU eviction.
377///
378/// # Usage
379///
380/// ```ignore
381/// let cache = ChunkCache::new();
382/// // Pass &cache to read_chunked_data — it will populate the index lazily.
383/// ```
384///
385/// The cache is wrapped in `Mutex` internally so it can be mutated through
386/// shared references (thread-safe).
387pub struct ChunkCache {
388    inner: Mutex<CacheInner>,
389}
390
391struct CacheInner {
392    /// Hash index: chunk coordinate → ChunkInfo (offset + size in file).
393    /// Populated once per dataset on first access.
394    #[cfg(feature = "std")]
395    index: Option<HashMap<ChunkCoord, ChunkInfo>>,
396    #[cfg(not(feature = "std"))]
397    index: Option<BTreeMap<ChunkCoord, ChunkInfo>>,
398
399    /// LRU cache of decompressed chunk data.
400    slots: Vec<CachedChunk>,
401
402    /// Current total bytes of cached decompressed data.
403    current_bytes: usize,
404
405    /// Maximum bytes of decompressed data to cache.
406    max_bytes: usize,
407
408    /// Maximum number of slots.
409    max_slots: usize,
410
411    /// Monotonic counter for LRU ordering.
412    tick: u64,
413
414    /// Monotonic counter handing out [`CachePass`] identities.
415    pass: u64,
416
417    /// Whether the parsed chunk index should be retained between reads.
418    cache_index: bool,
419
420    /// Cumulative counters reported by [`ChunkCache::stats`].
421    counters: Counters,
422}
423
424/// The cumulative half of [`ChunkCacheStats`], kept beside the slots it
425/// describes so a snapshot is taken under one lock.
426#[derive(Debug, Clone, Copy, PartialEq, Eq, Default)]
427struct Counters {
428    hits: u64,
429    misses: u64,
430    evictions: u64,
431    rejections: u64,
432    oversize_chunks: u64,
433    invalidations: u64,
434}
435
436impl ChunkCache {
437    /// Create a new chunk cache with default limits (1 MiB, 16 slots).
438    pub fn new() -> Self {
439        Self::with_capacity(DEFAULT_CACHE_BYTES, DEFAULT_MAX_SLOTS)
440    }
441
442    /// Create a new chunk cache with custom byte budget and slot count.
443    pub fn with_capacity(max_bytes: usize, max_slots: usize) -> Self {
444        Self::with_config(
445            ChunkCacheConfig::new()
446                .with_max_bytes(max_bytes)
447                .with_max_slots(max_slots),
448        )
449    }
450
451    /// Create a new chunk cache from a full configuration.
452    pub fn with_config(config: ChunkCacheConfig) -> Self {
453        Self {
454            inner: Mutex::new(CacheInner {
455                index: None,
456                slots: Vec::with_capacity(config.max_slots.min(64)),
457                current_bytes: 0,
458                max_bytes: config.max_bytes,
459                max_slots: config.max_slots,
460                tick: 0,
461                pass: 0,
462                cache_index: config.cache_index,
463                counters: Counters::default(),
464            }),
465        }
466    }
467
468    /// Snapshot what this cache is holding and what it has done.
469    ///
470    /// This is the public, read-only way to observe whether a chunk-cache
471    /// configuration is taking effect. It locks the cache briefly to read a
472    /// consistent snapshot.
473    pub fn stats(&self) -> ChunkCacheStats {
474        let inner = self.inner.lock().unwrap();
475        ChunkCacheStats {
476            index_loaded: inner.index.is_some(),
477            cached_chunks: inner.slots.len(),
478            cached_bytes: inner.current_bytes,
479            counters: inner.counters,
480        }
481    }
482
483    /// Zero the cumulative counters, leaving the retained index and chunks
484    /// alone.
485    ///
486    /// Occupancy is unaffected: this resets what the cache *has done*, not what
487    /// it is holding, so a caller can measure one read without the reads that
488    /// warmed the cache for it.
489    pub fn reset_stats(&self) {
490        let mut inner = self.inner.lock().unwrap();
491        inner.counters = Counters::default();
492    }
493
494    // ----- Index operations -----
495
496    /// Build the chunk index from a pre-collected list of `ChunkInfo`.
497    ///
498    /// The `rank` parameter is used to truncate offsets to spatial dims only
499    /// (B-tree v1 stores rank+1 offsets).
500    pub fn populate_index(&self, chunks: &[ChunkInfo], rank: usize) {
501        let mut inner = self.inner.lock().unwrap();
502        if !inner.cache_index {
503            return;
504        }
505        if inner.index.is_some() {
506            return; // already populated
507        }
508        #[cfg(feature = "std")]
509        let mut map = HashMap::with_capacity(chunks.len());
510        #[cfg(not(feature = "std"))]
511        let mut map = BTreeMap::new();
512
513        for ci in chunks {
514            let coord: ChunkCoord = ci.offsets.iter().take(rank).copied().collect();
515            map.insert(coord, ci.clone());
516        }
517        inner.index = Some(map);
518    }
519
520    /// Return all indexed chunks as a `Vec<ChunkInfo>` (order unspecified).
521    pub fn all_indexed_chunks(&self) -> Option<Vec<ChunkInfo>> {
522        self.indexed_chunks_matching(|_| true)
523    }
524
525    /// Return the indexed chunks `keep` accepts, as a `Vec<ChunkInfo>` (order
526    /// unspecified).
527    ///
528    /// A row window wants the chunks its rows overlap and nothing else, and the
529    /// difference is not a nicety: each [`ChunkInfo`] owns a coordinate `Vec`, so
530    /// taking the whole index and discarding most of it costs an allocation per
531    /// chunk *of the dataset* per window. A sweep of a dataset in windows paid
532    /// that as a product — 8 windows over 2,048 chunks, 16,384 allocations to
533    /// visit 2,048 chunks (issue #289). The filter runs while the lock is held,
534    /// over borrowed entries, so a rejected chunk costs no allocation at all.
535    pub fn indexed_chunks_matching(
536        &self,
537        keep: impl Fn(&ChunkInfo) -> bool,
538    ) -> Option<Vec<ChunkInfo>> {
539        let inner = self.inner.lock().unwrap();
540        inner
541            .index
542            .as_ref()
543            .map(|m| m.values().filter(|ci| keep(ci)).cloned().collect())
544    }
545
546    // ----- Decompressed data cache (LRU) -----
547
548    /// Run `f` over a borrowed view of a cached chunk's decompressed bytes, if
549    /// present, returning its result.
550    ///
551    /// The closure runs while the cache lock is held, which lets the caller copy
552    /// the chunk straight into its output buffer with no intermediate `Vec`
553    /// allocation or clone. The closure must not touch this cache (it would
554    /// deadlock); the chunk-assembly scatter it is used for does not.
555    pub fn with_decompressed<R>(
556        &self,
557        pass: CachePass,
558        coord: &[u64],
559        f: impl FnOnce(&[u8]) -> R,
560    ) -> Option<R> {
561        let mut guard = self.inner.lock().unwrap();
562        // Reborrowed once so the loop below borrows `slots` and `counters`
563        // separately; through the guard itself each is a borrow of the whole.
564        let inner = &mut *guard;
565        inner.tick += 1;
566        let tick = inner.tick;
567        for slot in inner.slots.iter_mut() {
568            if slot.coord.as_slice() == coord {
569                slot.last_access = tick;
570                // A pass that has been served a chunk owns it for the rest of
571                // that pass, exactly as if it had placed it. Without this a
572                // repeat read gives back every chunk it was just served, one per
573                // miss, and the read after it hits nothing; see [`CachePass`].
574                //
575                // Not for [`CachePass::LRU`], which is entitled to evict its own
576                // and would only be overwriting the provenance of a real pass
577                // running beside it on another thread.
578                if pass != CachePass::LRU {
579                    slot.stored_by = pass.0;
580                }
581                inner.counters.hits += 1;
582                return Some(f(&slot.data));
583            }
584        }
585        inner.counters.misses += 1;
586        None
587    }
588
589    /// Opens a read pass. See [`CachePass`] for what one is and why it exists.
590    pub fn begin_pass(&self) -> CachePass {
591        let mut inner = self.inner.lock().unwrap();
592        inner.pass += 1;
593        CachePass(inner.pass)
594    }
595
596    /// Makes room for a `data_len`-byte chunk at `coord`, reporting whether the
597    /// caller should go on to store it.
598    ///
599    /// The single place the admission policy lives, so the owned and borrowed
600    /// entry points below cannot drift — and so the borrowed one learns it has
601    /// nowhere to put the chunk *before* copying it rather than after.
602    ///
603    /// **On `true` the caller must push a slot**: the byte total has already been
604    /// charged for it, and a caller that returned instead would leave the cache
605    /// believing it holds bytes nothing occupies. On `false` nothing was changed
606    /// beyond the LRU tick.
607    fn reserve(inner: &mut CacheInner, pass: CachePass, coord: &[u64], data_len: usize) -> bool {
608        // A disabled cache counts nothing: it was never offered the chunk, and a
609        // caller who turned it off is not looking for a budget signal.
610        if inner.max_bytes == 0 || inner.max_slots == 0 {
611            return false;
612        }
613        // A chunk larger than the whole budget can never be retained, at any slot
614        // count, so it is counted apart from the chunks that merely did not fit.
615        if data_len > inner.max_bytes {
616            inner.counters.oversize_chunks += 1;
617            return false;
618        }
619
620        // Check if already present
621        inner.tick += 1;
622        let tick = inner.tick;
623        for slot in inner.slots.iter_mut() {
624            if slot.coord == coord {
625                slot.last_access = tick;
626                return false; // already cached
627            }
628        }
629
630        // A chunk this same pass stored is not taken back: that trades a chunk
631        // nobody will ask for again for another one nobody will ask for again.
632        //
633        // Unless this is [`CachePass::LRU`], which asks for the plain rule and
634        // must therefore be allowed to evict what it stored itself. Testing that
635        // by identity rather than leaning on `LRU`'s number is the whole of it:
636        // every `LRU` insert records the same `stored_by`, so an identity
637        // comparison alone would make the second one see the first as its own and
638        // refuse — turning the plain rule into fill-once for the life of the
639        // cache. `a_pass_marked_lru_evicts_its_own_chunks` is that bug's test.
640        let evicts_its_own = pass == CachePass::LRU;
641        let reclaimable = |slot: &&CachedChunk| evicts_its_own || slot.stored_by != pass.0;
642
643        // Whether the reclaimable slots can make room *at all*, decided before
644        // anything is removed. Evicting some and then finding the rest untouchable
645        // would leave the cache holding less and storing nothing — a chunk given
646        // up for no one. Removing every reclaimable slot is the most room there is
647        // to be had, so the test is the loop's own exit condition evaluated
648        // against that state.
649        let (freed_slots, freed_bytes) = inner
650            .slots
651            .iter()
652            .filter(reclaimable)
653            .fold((0usize, 0usize), |(n, b), slot| {
654                (n + 1, b + slot.data.len())
655            });
656        let (least_slots, least_bytes) = (
657            inner.slots.len() - freed_slots,
658            inner.current_bytes - freed_bytes,
659        );
660        if least_slots >= inner.max_slots
661            || (least_bytes + data_len > inner.max_bytes && least_slots > 0)
662        {
663            inner.counters.rejections += 1;
664            return false;
665        }
666
667        // Evict in LRU order until there is room. The check above proves a
668        // reclaimable slot exists for as long as this condition holds, so the
669        // `else` below cannot be reached; it returns rather than storing over
670        // budget in case that reasoning is ever made false.
671        while inner.slots.len() >= inner.max_slots
672            || (inner.current_bytes + data_len > inner.max_bytes && !inner.slots.is_empty())
673        {
674            // The LRU slot among those an earlier pass stored.
675            let lru_idx = inner
676                .slots
677                .iter()
678                .enumerate()
679                .filter(|(_, s)| reclaimable(s))
680                .min_by_key(|(_, s)| s.last_access)
681                .map(|(i, _)| i);
682            let Some(lru_idx) = lru_idx else {
683                debug_assert!(
684                    false,
685                    "the feasibility check above admitted a chunk this pass cannot make room for"
686                );
687                return false;
688            };
689            let removed = inner.slots.swap_remove(lru_idx);
690            inner.current_bytes -= removed.data.len();
691            inner.counters.evictions += 1;
692        }
693
694        inner.current_bytes += data_len;
695        true
696    }
697
698    /// Insert decompressed chunk data into the LRU cache, taking ownership of the
699    /// buffer (no copy). A chunk too large for the budget, a disabled cache, or a
700    /// pass that has already filled the cache drops the buffer instead of storing
701    /// it.
702    ///
703    /// `coord` is borrowed and copied only on the path that stores it, so a
704    /// caller in a loop needs no owned coordinate per chunk.
705    pub fn put_decompressed(&self, pass: CachePass, coord: &[u64], data: Vec<u8>) {
706        let mut inner = self.inner.lock().unwrap();
707        if !Self::reserve(&mut inner, pass, coord, data.len()) {
708            return;
709        }
710        let last_access = inner.tick;
711        inner.slots.push(CachedChunk {
712            coord: coord.to_vec(),
713            data,
714            last_access,
715            stored_by: pass.0,
716        });
717    }
718
719    /// The coordinates of every chunk whose decompressed bytes this cache
720    /// currently holds, in no particular order.
721    ///
722    /// A reader planning coalesced reads (see [`crate::chunk_span`]) uses this
723    /// to leave the chunks it already has out of the plan: a span covering a
724    /// chunk the read will skip fetches those bytes for nothing. It answers in
725    /// one lock over at most [`ChunkCacheConfig::max_slots`] entries, where
726    /// probing per chunk would take a lock apiece.
727    ///
728    /// The answer is a snapshot. A chunk named here can be evicted before the
729    /// read reaches it — by this very read, admitting later chunks — which
730    /// costs the coalescing for that chunk and nothing else: a chunk in no span
731    /// is read directly.
732    pub fn decompressed_coords(&self) -> Vec<ChunkCoord> {
733        let inner = self.inner.lock().unwrap();
734        inner.slots.iter().map(|s| s.coord.clone()).collect()
735    }
736
737    /// Insert a copy of `data` into the LRU cache, but only if it will actually
738    /// be kept. This lets the unfiltered read path scatter directly from the file
739    /// buffer and copy into the cache only when the chunk is going to stay there
740    /// — no copy at all when caching is off, when the chunk is over the budget,
741    /// or when this pass has already filled the cache.
742    pub fn put_decompressed_slice(&self, pass: CachePass, coord: &[u64], data: &[u8]) {
743        let mut inner = self.inner.lock().unwrap();
744        if !Self::reserve(&mut inner, pass, coord, data.len()) {
745            return;
746        }
747        let last_access = inner.tick;
748        inner.slots.push(CachedChunk {
749            coord: coord.to_vec(),
750            data: data.to_vec(),
751            last_access,
752            stored_by: pass.0,
753        });
754    }
755
756    /// Clear the entire cache (index + decompressed data).
757    ///
758    /// Called after a mutation through the owning [`Dataset`](crate::Dataset)
759    /// handle: an append relocates the trailing chunk and adds new index
760    /// entries, so both the cached chunk index and any retained decompressed
761    /// chunks may be stale.
762    pub fn clear(&self) {
763        let mut inner = self.inner.lock().unwrap();
764        inner.index = None;
765        // Counted before the drop, and as chunks rather than as one event: what a
766        // caller wants to compare against its miss count is how much retained
767        // data a write threw away, not how many times a write happened.
768        inner.counters.invalidations += inner.slots.len() as u64;
769        inner.slots.clear();
770        inner.current_bytes = 0;
771        inner.tick = 0;
772    }
773}
774
775impl Default for ChunkCache {
776    fn default() -> Self {
777        Self::new()
778    }
779}
780
781// ---------------------------------------------------------------------------
782// Tests
783// ---------------------------------------------------------------------------
784
785#[cfg(test)]
786mod tests {
787    use super::*;
788
789    fn make_chunk(offsets: Vec<u64>, address: u64, size: u32) -> ChunkInfo {
790        ChunkInfo {
791            chunk_size: size,
792            filter_mask: 0,
793            offsets,
794            address,
795        }
796    }
797
798    #[test]
799    fn index_populate_and_lookup() {
800        let cache = ChunkCache::new();
801        let chunks = vec![
802            make_chunk(vec![0, 0, 0], 0x1000, 80),
803            make_chunk(vec![10, 0, 0], 0x2000, 80),
804        ];
805        cache.populate_index(&chunks, 2); // rank=2, truncate to [0,0] and [10,0]
806        assert!(cache.stats().index_loaded());
807
808        let mut addrs: Vec<u64> = cache
809            .all_indexed_chunks()
810            .unwrap()
811            .iter()
812            .map(|c| c.address)
813            .collect();
814        addrs.sort_unstable();
815        assert_eq!(addrs, vec![0x1000, 0x2000]);
816    }
817
818    /// Test helper: clone a cached chunk's bytes if present (the production
819    /// path uses `with_decompressed` to avoid this copy).
820    fn get_decompressed(cache: &ChunkCache, coord: &[u64]) -> Option<Vec<u8>> {
821        cache.with_decompressed(CachePass::LRU, coord, <[u8]>::to_vec)
822    }
823
824    #[test]
825    fn decompressed_cache_hit() {
826        let cache = ChunkCache::new();
827        cache.put_decompressed(cache.begin_pass(), &[0, 0], vec![1, 2, 3, 4]);
828        let got = get_decompressed(&cache, &[0, 0]).unwrap();
829        assert_eq!(got, vec![1, 2, 3, 4]);
830    }
831
832    #[test]
833    fn lru_eviction_by_slots() {
834        let cache = ChunkCache::with_capacity(1024 * 1024, 2); // max 2 slots
835
836        cache.put_decompressed(cache.begin_pass(), &[0], vec![1; 10]);
837        cache.put_decompressed(cache.begin_pass(), &[1], vec![2; 10]);
838        assert_eq!(cache.stats().cached_chunks(), 2);
839
840        // Access slot 0 to make it more recent
841        get_decompressed(&cache, &[0]);
842
843        // Insert slot 2 — should evict slot 1 (LRU)
844        cache.put_decompressed(cache.begin_pass(), &[2], vec![3; 10]);
845        assert_eq!(cache.stats().cached_chunks(), 2);
846
847        assert!(get_decompressed(&cache, &[0]).is_some());
848        assert!(get_decompressed(&cache, &[1]).is_none()); // evicted
849        assert!(get_decompressed(&cache, &[2]).is_some());
850    }
851
852    #[test]
853    fn lru_eviction_by_bytes() {
854        let cache = ChunkCache::with_capacity(50, 100); // 50 bytes max
855
856        cache.put_decompressed(cache.begin_pass(), &[0], vec![0; 20]);
857        cache.put_decompressed(cache.begin_pass(), &[1], vec![0; 20]);
858        assert_eq!(cache.stats().cached_bytes(), 40);
859
860        // This needs 20 bytes but only 10 free — evict LRU
861        cache.put_decompressed(cache.begin_pass(), &[2], vec![0; 20]);
862        assert!(cache.stats().cached_bytes() <= 50);
863        assert!(get_decompressed(&cache, &[0]).is_none()); // evicted (LRU)
864    }
865
866    #[test]
867    fn put_decompressed_slice_only_copies_when_admitted() {
868        // Disabled cache: the slice is not copied or stored.
869        let cache = ChunkCache::with_config(ChunkCacheConfig::disabled());
870        cache.put_decompressed_slice(cache.begin_pass(), &[0], &[1, 2, 3]);
871        assert_eq!(cache.stats().cached_chunks(), 0);
872
873        // Enabled cache within budget: stored.
874        let cache = ChunkCache::with_capacity(1024, 16);
875        cache.put_decompressed_slice(cache.begin_pass(), &[0], &[1, 2, 3, 4]);
876        assert_eq!(get_decompressed(&cache, &[0]).unwrap(), vec![1, 2, 3, 4]);
877
878        // Over the per-chunk budget: not stored.
879        let cache = ChunkCache::with_capacity(2, 16);
880        cache.put_decompressed_slice(cache.begin_pass(), &[0], &[1, 2, 3, 4]);
881        assert_eq!(cache.stats().cached_chunks(), 0);
882    }
883
884    /// The rule [`CachePass`] exists for: one pass fills the cache and then
885    /// stops, rather than spending a copy per chunk to evict what it just
886    /// stored. A later pass is free to replace all of it.
887    #[test]
888    fn a_pass_fills_the_cache_and_then_stops_evicting_itself() {
889        let cache = ChunkCache::with_capacity(1024 * 1024, 2);
890
891        // One pass over four chunks, as a read of a four-chunk dataset makes.
892        let pass = cache.begin_pass();
893        for c in 0..4u64 {
894            cache.put_decompressed(pass, &[c], vec![c as u8; 10]);
895        }
896
897        // The two it reached first are the two it kept: chunks 2 and 3 were
898        // never copied, and chunks 0 and 1 were not evicted to make room for
899        // them.
900        assert_eq!(cache.stats().cached_chunks(), 2);
901        assert!(get_decompressed(&cache, &[0]).is_some());
902        assert!(get_decompressed(&cache, &[1]).is_some());
903        assert!(get_decompressed(&cache, &[2]).is_none());
904        assert!(get_decompressed(&cache, &[3]).is_none());
905
906        // A second read is a second pass, and it may take both slots back.
907        let next = cache.begin_pass();
908        cache.put_decompressed(next, &[9], vec![9; 10]);
909        cache.put_decompressed(next, &[8], vec![8; 10]);
910        assert_eq!(cache.stats().cached_chunks(), 2);
911        assert!(get_decompressed(&cache, &[9]).is_some());
912        assert!(get_decompressed(&cache, &[8]).is_some());
913    }
914
915    /// One eviction is counted per chunk dropped, not per admission that had to
916    /// drop something.
917    ///
918    /// Only the byte budget can force an admission to drop more than one chunk,
919    /// and only when chunks differ in size — which a real dataset's do not, since
920    /// every chunk decompresses to the same length. So this exercises the cache
921    /// directly: two small chunks, then one that needs the room of both.
922    #[test]
923    fn evictions_count_chunks_dropped_not_admissions_that_dropped_them() {
924        let cache = ChunkCache::with_capacity(250, 8);
925        let first = cache.begin_pass();
926        cache.put_decompressed(first, &[0], vec![0u8; 100]);
927        cache.put_decompressed(first, &[1], vec![0u8; 100]);
928        assert_eq!(cache.stats().cached_chunks(), 2);
929        assert_eq!(cache.stats().evictions(), 0);
930
931        // 200 held, 250 allowed: a 200-byte chunk needs both slots gone.
932        let second = cache.begin_pass();
933        cache.put_decompressed(second, &[2], vec![0u8; 200]);
934
935        let stats = cache.stats();
936        assert_eq!(stats.cached_chunks(), 1);
937        assert_eq!(stats.evictions(), 2);
938    }
939
940    /// A pass keeps what it was served even if a windowed read is served the same
941    /// chunk in between.
942    ///
943    /// [`CachePass::LRU`] is entitled to give up its own chunks, so it must not
944    /// stamp its identity onto a slot a real pass owns — that would hand the
945    /// chunk back to the very pass that is relying on keeping it. Two live passes
946    /// are what it takes to see this, which no single read produces; the cache
947    /// API hands them out directly, so this needs no threads.
948    #[test]
949    fn an_lru_hit_does_not_release_a_chunk_of_another_pass_to_it() {
950        let cache = ChunkCache::with_capacity(1024 * 1024, 2);
951        let reader = cache.begin_pass();
952        cache.put_decompressed(reader, &[0], vec![0u8; 100]);
953        cache.put_decompressed(reader, &[1], vec![0u8; 100]);
954
955        // A row window, elsewhere, is served the chunk `reader` placed.
956        assert!(
957            cache
958                .with_decompressed(CachePass::LRU, &[0], |_| ())
959                .is_some()
960        );
961
962        // `reader` offers a third chunk. Both slots are still its own, so there
963        // is nothing it may take, and the offer is refused.
964        cache.put_decompressed(reader, &[2], vec![0u8; 100]);
965
966        assert_eq!(cache.stats().cached_chunks(), 2);
967        assert!(get_decompressed(&cache, &[0]).is_some());
968        assert!(get_decompressed(&cache, &[1]).is_some());
969    }
970
971    /// [`CachePass::LRU`] is a sentinel: it works only because a real pass is
972    /// never numbered zero. A `begin_pass` that started counting at zero would
973    /// silently turn the windowed reader's plain-LRU admission into fill-once and
974    /// lose it the boundary chunk its successor window needs.
975    #[test]
976    fn a_pass_marked_lru_evicts_its_own_chunks() {
977        let cache = ChunkCache::with_capacity(1024 * 1024, 2);
978
979        for c in 0..4u64 {
980            cache.put_decompressed(CachePass::LRU, &[c], vec![c as u8; 10]);
981        }
982
983        // The last two, where a fill-once pass would have kept the first two.
984        assert_eq!(cache.stats().cached_chunks(), 2);
985        assert!(get_decompressed(&cache, &[2]).is_some());
986        assert!(get_decompressed(&cache, &[3]).is_some());
987        assert!(get_decompressed(&cache, &[0]).is_none());
988
989        // The property that makes the sentinel sound, asserted rather than
990        // assumed: no real pass can collide with it.
991        assert_ne!(cache.begin_pass(), CachePass::LRU);
992    }
993
994    /// A pass that gives up must not have taken anything with it. Reclaiming some
995    /// slots and then finding the rest untouchable would leave the cache holding
996    /// less and storing nothing — a chunk dropped for no one.
997    #[test]
998    fn a_pass_that_cannot_make_room_evicts_nothing() {
999        // 100 bytes, plenty of slots: only the byte budget can bite.
1000        let cache = ChunkCache::with_capacity(100, 16);
1001
1002        let first = cache.begin_pass();
1003        cache.put_decompressed(first, &[0], vec![0; 10]);
1004
1005        let second = cache.begin_pass();
1006        cache.put_decompressed(second, &[1], vec![1; 80]);
1007        assert_eq!(cache.stats().cached_chunks(), 2);
1008
1009        // 80 bytes more will not fit even with the 10-byte chunk from `first`
1010        // reclaimed, and the 80-byte one belongs to this pass. The old code
1011        // evicted the reclaimable chunk first and gave up afterwards.
1012        cache.put_decompressed(second, &[2], vec![2; 80]);
1013        assert_eq!(cache.stats().cached_chunks(), 2);
1014        assert_eq!(cache.stats().cached_bytes(), 90);
1015        assert!(get_decompressed(&cache, &[0]).is_some());
1016    }
1017
1018    /// The same rule on the borrowed entry point, where it also decides whether
1019    /// the chunk is copied at all.
1020    #[test]
1021    fn a_full_pass_does_not_copy_the_chunk_it_cannot_store() {
1022        let cache = ChunkCache::with_capacity(1024 * 1024, 1);
1023        let pass = cache.begin_pass();
1024
1025        cache.put_decompressed_slice(pass, &[0], &[1; 10]);
1026        cache.put_decompressed_slice(pass, &[1], &[2; 10]);
1027
1028        assert_eq!(cache.stats().cached_chunks(), 1);
1029        assert_eq!(cache.stats().cached_bytes(), 10);
1030        assert!(get_decompressed(&cache, &[0]).is_some());
1031    }
1032
1033    #[test]
1034    fn oversized_chunk_not_cached() {
1035        let cache = ChunkCache::with_capacity(10, 16);
1036        cache.put_decompressed(cache.begin_pass(), &[0], vec![0; 100]); // too big
1037        assert_eq!(cache.stats().cached_chunks(), 0);
1038    }
1039
1040    #[test]
1041    fn disabled_cache_retains_no_index_or_chunks() {
1042        let cache = ChunkCache::with_config(ChunkCacheConfig::disabled());
1043        let chunks = vec![make_chunk(vec![0, 0], 0x1000, 80)];
1044        cache.populate_index(&chunks, 1);
1045        assert!(!cache.stats().index_loaded());
1046
1047        cache.put_decompressed(cache.begin_pass(), &[0], vec![1, 2, 3]);
1048        assert_eq!(cache.stats().cached_chunks(), 0);
1049        assert_eq!(cache.stats().cached_bytes(), 0);
1050    }
1051
1052    #[test]
1053    fn h5p_cache_constructor_maps_raw_data_chunk_settings() {
1054        let config = ChunkCacheConfig::from_h5p_cache(521, 2 * 1024 * 1024);
1055        assert_eq!(config.max_slots(), 521);
1056        assert_eq!(config.max_bytes(), 2 * 1024 * 1024);
1057        assert!(config.index_cache_enabled());
1058    }
1059
1060    #[test]
1061    fn clear_resets_everything() {
1062        let cache = ChunkCache::new();
1063        let chunks = vec![make_chunk(vec![0, 0], 0x1000, 80)];
1064        cache.populate_index(&chunks, 1);
1065        cache.put_decompressed(cache.begin_pass(), &[0], vec![1, 2, 3]);
1066
1067        cache.clear();
1068        assert!(!cache.stats().index_loaded());
1069        assert_eq!(cache.stats().cached_chunks(), 0);
1070        assert_eq!(cache.stats().cached_bytes(), 0);
1071    }
1072
1073    #[test]
1074    fn duplicate_insert_is_noop() {
1075        let cache = ChunkCache::new();
1076        cache.put_decompressed(cache.begin_pass(), &[0], vec![1, 2, 3]);
1077        cache.put_decompressed(cache.begin_pass(), &[0], vec![1, 2, 3]); // duplicate
1078        assert_eq!(cache.stats().cached_chunks(), 1);
1079        assert_eq!(cache.stats().cached_bytes(), 3);
1080    }
1081}