osobhandClaude Fable 5.1
0addf328bc
Same-moment A/B on a 64 MB f64 dataset: chunked+deflate 110 -> 69 ms, chunked 72 -> 60 ms, contiguous 56 -> 30 ms. - read_chunked_data_cached — the path the facade uses — decompressed chunks one at a time; only the uncached reader was parallel. Cache misses are now decoded in bounded batches (128), in parallel with the `parallel` feature. - Every chunk was pushed into the 16 MiB chunk cache, which a larger dataset just churns (insert, evict moments later). Chunks are cached only when the whole dataset fits (new ChunkCache::max_bytes). - Unfiltered chunks went file -> Vec -> aligned cache buffer -> output. They are copied straight from the file bytes. - The facade's typed reads convert a contiguous dataset straight from the borrowed file bytes instead of copying it into a Vec first. - The native little-endian fast paths allocated vec![0; n] and then overwrote it; they now fill an uninitialised buffer in one copy (native_le_to_vec). alloc_output requests zeroed memory from the allocator instead of reserving and filling. The unit test that expected unfiltered chunks to land in the decompressed cache now asserts the new design (index reused, cache not involved). Co-Authored-By: Claude Fable 5.1 <[email protected]>
945 lines
32 KiB
Rust
945 lines
32 KiB
Rust
//! Chunk cache with hash-based index and LRU eviction.
|
|
//!
|
|
//! The [`ChunkCache`] avoids re-traversing B-trees on repeated reads of chunked
|
|
//! datasets. On first access it scans the B-tree once and builds a
|
|
//! `HashMap<ChunkCoord, ChunkInfo>` (the *chunk index*). Decompressed chunk
|
|
//! data is cached with LRU eviction controlled by a byte-budget.
|
|
|
|
extern crate alloc;
|
|
|
|
#[cfg(not(feature = "std"))]
|
|
use alloc::{vec, vec::Vec};
|
|
|
|
use core::ops::{Deref, DerefMut};
|
|
|
|
#[cfg(not(feature = "std"))]
|
|
use alloc::collections::BTreeMap;
|
|
#[cfg(feature = "std")]
|
|
use std::collections::HashMap;
|
|
#[cfg(feature = "std")]
|
|
use std::sync::Arc;
|
|
|
|
use crate::chunk_index::{ChunkIndex, ChunkLayout};
|
|
use crate::chunked_read::ChunkInfo;
|
|
|
|
// ---------------------------------------------------------------------------
|
|
// Cache-line alignment constants (TVL — Tensor Virtualization Layout)
|
|
// ---------------------------------------------------------------------------
|
|
|
|
/// Cache line size in bytes for the target architecture.
|
|
///
|
|
/// ARM64 uses 128-byte cache lines; x86_64 uses 64-byte. We align all chunk
|
|
/// buffers to this boundary so SIMD operations can assume aligned input.
|
|
#[cfg(target_arch = "aarch64")]
|
|
pub const CACHE_LINE_SIZE: usize = 128;
|
|
|
|
#[cfg(target_arch = "x86_64")]
|
|
pub const CACHE_LINE_SIZE: usize = 64;
|
|
|
|
#[cfg(not(any(target_arch = "aarch64", target_arch = "x86_64")))]
|
|
pub const CACHE_LINE_SIZE: usize = 64;
|
|
|
|
/// Round `size` up to the next multiple of [`CACHE_LINE_SIZE`].
|
|
#[inline]
|
|
pub fn align_to_cache_line(size: usize) -> usize {
|
|
(size + CACHE_LINE_SIZE - 1) & !(CACHE_LINE_SIZE - 1)
|
|
}
|
|
|
|
// ---------------------------------------------------------------------------
|
|
// CacheAlignedBuffer
|
|
// ---------------------------------------------------------------------------
|
|
|
|
/// A byte buffer whose data pointer is aligned to [`CACHE_LINE_SIZE`].
|
|
///
|
|
/// This enables SIMD operations to use aligned loads/stores when processing
|
|
/// chunk data, avoiding the penalty of misaligned memory accesses.
|
|
///
|
|
/// The buffer is backed by `core::alloc::Layout`-controlled allocation and
|
|
/// uses `alloc::alloc` for the actual allocation, making it compatible with
|
|
/// `no_std` (requires the `alloc` crate). It dereferences to `&[u8]` /
|
|
/// `&mut [u8]` for seamless use.
|
|
pub struct CacheAlignedBuffer {
|
|
ptr: *mut u8,
|
|
len: usize,
|
|
capacity: usize,
|
|
}
|
|
|
|
// SAFETY: The raw pointer is exclusively owned — no aliasing.
|
|
unsafe impl Send for CacheAlignedBuffer {}
|
|
// SAFETY: `CacheAlignedBuffer` exposes its contents only via `&[u8]`/`&mut
|
|
// [u8]` through the ordinary borrow-checked `Deref`/`DerefMut` impls below —
|
|
// the same access pattern as `Vec<u8>`, which is `Sync`. Needed so
|
|
// `Arc<CacheAlignedBuffer>` (used by the chunk cache) is itself `Send`.
|
|
unsafe impl Sync for CacheAlignedBuffer {}
|
|
|
|
impl CacheAlignedBuffer {
|
|
/// Allocate a new cache-line-aligned buffer of exactly `len` bytes,
|
|
/// initialized to zero.
|
|
pub fn zeroed(len: usize) -> Self {
|
|
if len == 0 {
|
|
return Self {
|
|
ptr: core::ptr::NonNull::dangling().as_ptr(),
|
|
len: 0,
|
|
capacity: 0,
|
|
};
|
|
}
|
|
let capacity = align_to_cache_line(len);
|
|
let layout = core::alloc::Layout::from_size_align(capacity, CACHE_LINE_SIZE)
|
|
.expect("invalid layout");
|
|
// SAFETY: layout has non-zero size.
|
|
let ptr = unsafe { alloc::alloc::alloc_zeroed(layout) };
|
|
if ptr.is_null() {
|
|
alloc::alloc::handle_alloc_error(layout);
|
|
}
|
|
Self { ptr, len, capacity }
|
|
}
|
|
|
|
/// Create a cache-line-aligned copy of an existing byte slice.
|
|
pub fn from_slice(data: &[u8]) -> Self {
|
|
let mut buf = Self::zeroed(data.len());
|
|
buf.as_mut_slice()[..data.len()].copy_from_slice(data);
|
|
buf
|
|
}
|
|
|
|
/// Create from an existing `Vec<u8>`, copying into an aligned allocation.
|
|
pub fn from_vec(v: Vec<u8>) -> Self {
|
|
Self::from_slice(&v)
|
|
}
|
|
|
|
/// The length of the valid data (may be less than capacity).
|
|
#[inline]
|
|
pub fn len(&self) -> usize {
|
|
self.len
|
|
}
|
|
|
|
/// Whether the buffer is empty.
|
|
#[inline]
|
|
pub fn is_empty(&self) -> bool {
|
|
self.len == 0
|
|
}
|
|
|
|
/// The underlying aligned pointer.
|
|
#[inline]
|
|
pub fn as_ptr(&self) -> *const u8 {
|
|
self.ptr
|
|
}
|
|
|
|
/// Mutable pointer to the data.
|
|
#[inline]
|
|
pub fn as_mut_ptr(&mut self) -> *mut u8 {
|
|
self.ptr
|
|
}
|
|
|
|
/// Borrow as a byte slice.
|
|
#[inline]
|
|
pub fn as_slice(&self) -> &[u8] {
|
|
if self.len == 0 {
|
|
return &[];
|
|
}
|
|
// SAFETY: ptr is valid for `len` bytes and properly aligned.
|
|
unsafe { core::slice::from_raw_parts(self.ptr, self.len) }
|
|
}
|
|
|
|
/// Borrow as a mutable byte slice.
|
|
#[inline]
|
|
pub fn as_mut_slice(&mut self) -> &mut [u8] {
|
|
if self.len == 0 {
|
|
return &mut [];
|
|
}
|
|
// SAFETY: ptr is valid for `len` bytes and properly aligned.
|
|
unsafe { core::slice::from_raw_parts_mut(self.ptr, self.len) }
|
|
}
|
|
|
|
/// Convert to a `Vec<u8>` (copies data into a standard allocation).
|
|
pub fn to_vec(&self) -> Vec<u8> {
|
|
self.as_slice().to_vec()
|
|
}
|
|
|
|
/// Returns `true` if the data pointer is aligned to `CACHE_LINE_SIZE`.
|
|
#[inline]
|
|
pub fn is_aligned(&self) -> bool {
|
|
self.len == 0 || (self.ptr as usize).is_multiple_of(CACHE_LINE_SIZE)
|
|
}
|
|
}
|
|
|
|
impl Drop for CacheAlignedBuffer {
|
|
fn drop(&mut self) {
|
|
if self.capacity > 0 {
|
|
let layout = core::alloc::Layout::from_size_align(self.capacity, CACHE_LINE_SIZE)
|
|
.expect("invalid layout");
|
|
// SAFETY: ptr was allocated with this layout.
|
|
unsafe { alloc::alloc::dealloc(self.ptr, layout) };
|
|
}
|
|
}
|
|
}
|
|
|
|
impl Clone for CacheAlignedBuffer {
|
|
fn clone(&self) -> Self {
|
|
Self::from_slice(self.as_slice())
|
|
}
|
|
}
|
|
|
|
impl Deref for CacheAlignedBuffer {
|
|
type Target = [u8];
|
|
#[inline]
|
|
fn deref(&self) -> &[u8] {
|
|
self.as_slice()
|
|
}
|
|
}
|
|
|
|
impl DerefMut for CacheAlignedBuffer {
|
|
#[inline]
|
|
fn deref_mut(&mut self) -> &mut [u8] {
|
|
self.as_mut_slice()
|
|
}
|
|
}
|
|
|
|
impl core::fmt::Debug for CacheAlignedBuffer {
|
|
fn fmt(&self, f: &mut core::fmt::Formatter<'_>) -> core::fmt::Result {
|
|
f.debug_struct("CacheAlignedBuffer")
|
|
.field("len", &self.len)
|
|
.field("capacity", &self.capacity)
|
|
.field("aligned", &self.is_aligned())
|
|
.finish()
|
|
}
|
|
}
|
|
|
|
/// Coordinate key for a chunk -- the N-dimensional offset vector.
|
|
pub type ChunkCoord = Vec<u64>;
|
|
|
|
// ---------------------------------------------------------------------------
|
|
// ChunkCache and supporting types require std::sync::Mutex.
|
|
// ---------------------------------------------------------------------------
|
|
|
|
/// Default maximum bytes of decompressed chunk data to cache (16 MiB).
|
|
///
|
|
/// Increased from the HDF5 C library default of 1 MiB to better accommodate
|
|
/// modern workloads with large embedding datasets and high-dimensional chunks.
|
|
pub const DEFAULT_CACHE_BYTES: usize = 16 * 1024 * 1024; // 16 MiB
|
|
|
|
/// Default maximum number of cached decompressed chunks.
|
|
///
|
|
/// 521 is prime, which provides better hash distribution for the chunk
|
|
/// coordinate map and reduces collision chains compared to power-of-two sizes.
|
|
pub const DEFAULT_MAX_SLOTS: usize = 521;
|
|
|
|
// ---------------------------------------------------------------------------
|
|
// LRU entry
|
|
// ---------------------------------------------------------------------------
|
|
|
|
#[cfg(feature = "std")]
|
|
struct CachedChunk {
|
|
coord: ChunkCoord,
|
|
/// Shared so a cache hit is a refcount bump, not a copy of the whole
|
|
/// (potentially large) decompressed chunk.
|
|
data: Arc<CacheAlignedBuffer>,
|
|
/// Monotonically increasing access counter for LRU ordering.
|
|
last_access: u64,
|
|
}
|
|
|
|
// ---------------------------------------------------------------------------
|
|
// ChunkCache
|
|
// ---------------------------------------------------------------------------
|
|
|
|
/// A per-dataset chunk cache with hash-based index and LRU eviction.
|
|
///
|
|
/// # Usage
|
|
///
|
|
/// ```ignore
|
|
/// let cache = ChunkCache::new();
|
|
/// // Pass &cache to read_chunked_data — it will populate the index lazily.
|
|
/// ```
|
|
///
|
|
/// The cache is wrapped in `Mutex` internally so it can be mutated through
|
|
/// shared references (thread-safe).
|
|
///
|
|
/// Only available with the `std` feature because it requires `std::sync::Mutex`.
|
|
#[cfg(feature = "std")]
|
|
pub struct ChunkCache {
|
|
inner: std::sync::Mutex<CacheInner>,
|
|
}
|
|
|
|
#[cfg(feature = "std")]
|
|
struct CacheInner {
|
|
/// Hash index: chunk coordinate -> ChunkInfo (offset + size in file).
|
|
/// Populated once per dataset on first access.
|
|
index: Option<HashMap<ChunkCoord, ChunkInfo>>,
|
|
|
|
/// Address of the dataset (its chunk-index base address) that the cached
|
|
/// index, chunk index, layout, and decompressed slots currently belong to.
|
|
/// The cache is shared per file across datasets, so every cached-read entry
|
|
/// checks this and resets the per-dataset state when the dataset changes —
|
|
/// otherwise one dataset's chunk index (with its own rank) would be reused
|
|
/// for another, corrupting reads.
|
|
index_addr: Option<u64>,
|
|
|
|
/// LRU cache of decompressed chunk data.
|
|
slots: Vec<CachedChunk>,
|
|
|
|
/// Coordinate -> index into `slots`, for O(1) lookup instead of a linear
|
|
/// scan. Kept in sync with `slots` on every insert/evict/clear — in
|
|
/// particular, `slots.swap_remove(i)` moves the last element into slot
|
|
/// `i`, so the moved element's index entry must be updated too.
|
|
slot_index: HashMap<ChunkCoord, usize>,
|
|
|
|
/// Current total bytes of cached decompressed data.
|
|
current_bytes: usize,
|
|
|
|
/// Maximum bytes of decompressed data to cache.
|
|
max_bytes: usize,
|
|
|
|
/// Maximum number of slots.
|
|
max_slots: usize,
|
|
|
|
/// Monotonic counter for LRU ordering.
|
|
tick: u64,
|
|
|
|
/// Last accessed chunk coordinate (for sequential detection).
|
|
last_coord: Option<ChunkCoord>,
|
|
|
|
/// Access pattern statistics.
|
|
stats: AccessStats,
|
|
|
|
/// Pre-built chunk index for O(1) coordinate lookups.
|
|
chunk_index: Option<ChunkIndex>,
|
|
|
|
/// Pre-computed chunk layout for fast assembly.
|
|
chunk_layout: Option<ChunkLayout>,
|
|
}
|
|
|
|
/// Access pattern statistics tracked by the chunk cache.
|
|
///
|
|
/// Updated on each `get_decompressed` / `put_decompressed` call to help
|
|
/// the sweep detector understand the workload.
|
|
#[cfg(feature = "std")]
|
|
#[derive(Debug, Clone, Default)]
|
|
pub struct AccessStats {
|
|
/// Number of accesses that followed a sequential pattern.
|
|
pub sequential_count: u64,
|
|
/// Number of accesses that appeared random (non-sequential).
|
|
pub random_count: u64,
|
|
/// Last detected sweep direction description (informational).
|
|
pub sweep_direction: Option<&'static str>,
|
|
/// Number of cache hits (decompressed data found in cache).
|
|
pub hits: u64,
|
|
/// Number of cache misses (decompressed data not in cache).
|
|
pub misses: u64,
|
|
/// Number of LRU evictions performed.
|
|
pub evictions: u64,
|
|
/// Total bytes read from cached data.
|
|
pub bytes_read: u64,
|
|
}
|
|
|
|
#[cfg(feature = "std")]
|
|
impl AccessStats {
|
|
/// Cache hit rate as a fraction in [0.0, 1.0].
|
|
///
|
|
/// Returns 0.0 if no accesses have been recorded.
|
|
pub fn hit_rate(&self) -> f64 {
|
|
let total = self.hits + self.misses;
|
|
if total == 0 {
|
|
0.0
|
|
} else {
|
|
self.hits as f64 / total as f64
|
|
}
|
|
}
|
|
}
|
|
|
|
#[cfg(feature = "std")]
|
|
impl ChunkCache {
|
|
/// Create a new chunk cache with default limits (16 MiB, 521 slots).
|
|
pub fn new() -> Self {
|
|
Self::with_capacity(DEFAULT_CACHE_BYTES, DEFAULT_MAX_SLOTS)
|
|
}
|
|
|
|
/// Create a new chunk cache with custom byte budget and slot count.
|
|
pub fn with_capacity(max_bytes: usize, max_slots: usize) -> Self {
|
|
Self {
|
|
inner: std::sync::Mutex::new(CacheInner {
|
|
index: None,
|
|
index_addr: None,
|
|
slots: Vec::with_capacity(max_slots.min(64)),
|
|
slot_index: HashMap::with_capacity(max_slots.min(64)),
|
|
current_bytes: 0,
|
|
max_bytes,
|
|
max_slots,
|
|
tick: 0,
|
|
last_coord: None,
|
|
stats: AccessStats::default(),
|
|
chunk_index: None,
|
|
chunk_layout: None,
|
|
}),
|
|
}
|
|
}
|
|
|
|
// ----- Index operations -----
|
|
|
|
/// The most decompressed bytes this cache will hold.
|
|
pub fn max_bytes(&self) -> usize {
|
|
self.inner.lock().map(|g| g.max_bytes).unwrap_or(0)
|
|
}
|
|
|
|
/// Bind the cache to the dataset at chunk-index address `addr`.
|
|
///
|
|
/// The cache is shared per file across all of its datasets. If the cache
|
|
/// currently holds state for a different dataset, all per-dataset state
|
|
/// (chunk index, chunk-index map, layout, and decompressed slots) is
|
|
/// dropped so the next access rebuilds it for this dataset. Reading the
|
|
/// same dataset again is a no-op, preserving the cache's benefit for
|
|
/// repeated/sequential access. Returns `true` if a reset occurred.
|
|
pub fn ensure_dataset(&self, addr: u64) -> bool {
|
|
let mut inner = self.inner.lock().unwrap_or_else(|e| e.into_inner());
|
|
if inner.index_addr == Some(addr) {
|
|
return false;
|
|
}
|
|
inner.index = None;
|
|
inner.chunk_index = None;
|
|
inner.chunk_layout = None;
|
|
inner.slots.clear();
|
|
inner.slot_index.clear();
|
|
inner.current_bytes = 0;
|
|
inner.last_coord = None;
|
|
inner.index_addr = Some(addr);
|
|
true
|
|
}
|
|
|
|
/// Returns `true` if the chunk index has been built.
|
|
pub fn has_index(&self) -> bool {
|
|
self.inner
|
|
.lock()
|
|
.unwrap_or_else(|e| e.into_inner())
|
|
.index
|
|
.is_some()
|
|
}
|
|
|
|
/// Build the chunk index from a pre-collected list of `ChunkInfo`.
|
|
///
|
|
/// The `rank` parameter is used to truncate offsets to spatial dims only
|
|
/// (B-tree v1 stores rank+1 offsets).
|
|
pub fn populate_index(&self, chunks: &[ChunkInfo], rank: usize) {
|
|
let mut inner = self.inner.lock().unwrap_or_else(|e| e.into_inner());
|
|
if inner.index.is_some() {
|
|
return; // already populated
|
|
}
|
|
let mut map = HashMap::with_capacity(chunks.len());
|
|
|
|
for ci in chunks {
|
|
let coord: ChunkCoord = ci.offsets.iter().take(rank).copied().collect();
|
|
map.insert(coord, ci.clone());
|
|
}
|
|
inner.index = Some(map);
|
|
}
|
|
|
|
/// Look up a chunk by its spatial coordinate in the index.
|
|
pub fn lookup_index(&self, coord: &[u64]) -> Option<ChunkInfo> {
|
|
let inner = self.inner.lock().unwrap_or_else(|e| e.into_inner());
|
|
inner.index.as_ref()?.get(coord).cloned()
|
|
}
|
|
|
|
/// Return all indexed chunks as a `Vec<ChunkInfo>` (order unspecified).
|
|
pub fn all_indexed_chunks(&self) -> Option<Vec<ChunkInfo>> {
|
|
let inner = self.inner.lock().unwrap_or_else(|e| e.into_inner());
|
|
inner.index.as_ref().map(|m| m.values().cloned().collect())
|
|
}
|
|
|
|
// ----- Chunk index (pre-built coordinate → ChunkInfo map) -----
|
|
|
|
/// Returns `true` if the chunk B-tree index has been built.
|
|
pub fn has_chunk_index(&self) -> bool {
|
|
self.inner
|
|
.lock()
|
|
.unwrap_or_else(|e| e.into_inner())
|
|
.chunk_index
|
|
.is_some()
|
|
}
|
|
|
|
/// Build and store the chunk B-tree index from a pre-collected list of `ChunkInfo`.
|
|
pub fn populate_chunk_index(&self, chunks: &[ChunkInfo], rank: usize) {
|
|
let mut inner = self.inner.lock().unwrap_or_else(|e| e.into_inner());
|
|
if inner.chunk_index.is_some() {
|
|
return;
|
|
}
|
|
inner.chunk_index = Some(ChunkIndex::build(chunks, rank));
|
|
}
|
|
|
|
// ----- Chunk layout (pre-computed assembly plan) -----
|
|
|
|
/// Returns `true` if the chunk layout has been computed.
|
|
pub fn has_chunk_layout(&self) -> bool {
|
|
self.inner
|
|
.lock()
|
|
.unwrap_or_else(|e| e.into_inner())
|
|
.chunk_layout
|
|
.is_some()
|
|
}
|
|
|
|
/// Build and store the pre-computed chunk layout for fast assembly.
|
|
pub fn populate_chunk_layout(&self, ds_dims: &[usize], chunk_dims: &[usize], elem_size: usize) {
|
|
let mut inner = self.inner.lock().unwrap_or_else(|e| e.into_inner());
|
|
if inner.chunk_layout.is_some() {
|
|
return;
|
|
}
|
|
if let Some(ref idx) = inner.chunk_index {
|
|
inner.chunk_layout = Some(ChunkLayout::build(idx, ds_dims, chunk_dims, elem_size));
|
|
}
|
|
}
|
|
|
|
/// Execute a function with a reference to the chunk layout.
|
|
///
|
|
/// Returns `None` if the layout hasn't been computed yet.
|
|
pub fn with_chunk_layout<F, R>(&self, f: F) -> Option<R>
|
|
where
|
|
F: FnOnce(&ChunkLayout) -> R,
|
|
{
|
|
let inner = self.inner.lock().unwrap_or_else(|e| e.into_inner());
|
|
inner.chunk_layout.as_ref().map(f)
|
|
}
|
|
|
|
// ----- Decompressed data cache (LRU) -----
|
|
|
|
/// Try to get cached decompressed data for a chunk coordinate.
|
|
///
|
|
/// O(1) lookup. Returns an owned copy for API compatibility with callers
|
|
/// that need a `Vec<u8>`; prefer [`Self::get_decompressed_aligned`] when
|
|
/// an `Arc`-shared buffer works for the caller, since that avoids the
|
|
/// copy entirely.
|
|
pub fn get_decompressed(&self, coord: &[u64]) -> Option<Vec<u8>> {
|
|
self.get_decompressed_aligned(coord)
|
|
.map(|arc| arc.as_slice().to_vec())
|
|
}
|
|
|
|
/// Try to get a reference-counted clone of the aligned buffer for a chunk.
|
|
///
|
|
/// O(1) index lookup; the clone is an `Arc` refcount bump, not a copy of
|
|
/// the underlying decompressed data.
|
|
pub fn get_decompressed_aligned(&self, coord: &[u64]) -> Option<Arc<CacheAlignedBuffer>> {
|
|
let mut inner = self.inner.lock().unwrap_or_else(|e| e.into_inner());
|
|
inner.tick += 1;
|
|
let tick = inner.tick;
|
|
|
|
// Track sequential vs random access
|
|
let is_sequential = inner.last_coord.as_ref().is_some_and(|prev| {
|
|
// Sequential if exactly one dimension changed
|
|
let changes: usize = prev
|
|
.iter()
|
|
.zip(coord.iter())
|
|
.filter(|(a, b)| a != b)
|
|
.count();
|
|
changes <= 1
|
|
});
|
|
if is_sequential {
|
|
inner.stats.sequential_count += 1;
|
|
} else if inner.last_coord.is_some() {
|
|
inner.stats.random_count += 1;
|
|
}
|
|
inner.last_coord = Some(coord.to_vec());
|
|
|
|
let found = if let Some(&idx) = inner.slot_index.get(coord) {
|
|
inner.slots[idx].last_access = tick;
|
|
Some(Arc::clone(&inner.slots[idx].data))
|
|
} else {
|
|
None
|
|
};
|
|
if let Some(ref data) = found {
|
|
inner.stats.hits += 1;
|
|
inner.stats.bytes_read += data.len() as u64;
|
|
} else {
|
|
inner.stats.misses += 1;
|
|
}
|
|
found
|
|
}
|
|
|
|
/// Insert decompressed chunk data into the LRU cache.
|
|
///
|
|
/// The data is stored in a [`CacheAlignedBuffer`] so subsequent reads
|
|
/// return cache-line-aligned memory. Returns the `Arc`-shared buffer that
|
|
/// is now cached (or already was), so the caller can reuse it directly
|
|
/// instead of holding a separate copy of the same data.
|
|
pub fn put_decompressed(&self, coord: ChunkCoord, data: Vec<u8>) -> Arc<CacheAlignedBuffer> {
|
|
let aligned = CacheAlignedBuffer::from_vec(data);
|
|
self.put_decompressed_aligned(coord, aligned)
|
|
}
|
|
|
|
/// Insert an already-aligned buffer into the LRU cache.
|
|
///
|
|
/// Returns the `Arc`-shared buffer now held by the cache (the one just
|
|
/// inserted, or the existing cached copy if `coord` was already present).
|
|
pub fn put_decompressed_aligned(
|
|
&self,
|
|
coord: ChunkCoord,
|
|
data: CacheAlignedBuffer,
|
|
) -> Arc<CacheAlignedBuffer> {
|
|
let data = Arc::new(data);
|
|
let mut inner = self.inner.lock().unwrap_or_else(|e| e.into_inner());
|
|
let data_len = data.len();
|
|
|
|
// Don't cache if single chunk exceeds budget — still return the data
|
|
// to the caller, just don't retain it.
|
|
if data_len > inner.max_bytes {
|
|
return data;
|
|
}
|
|
|
|
// Check if already present
|
|
inner.tick += 1;
|
|
let tick = inner.tick;
|
|
if let Some(&idx) = inner.slot_index.get(&coord) {
|
|
inner.slots[idx].last_access = tick;
|
|
return Arc::clone(&inner.slots[idx].data); // already cached
|
|
}
|
|
|
|
// Evict until we have room
|
|
while inner.slots.len() >= inner.max_slots
|
|
|| (inner.current_bytes + data_len > inner.max_bytes && !inner.slots.is_empty())
|
|
{
|
|
// Find LRU slot
|
|
let lru_idx = inner
|
|
.slots
|
|
.iter()
|
|
.enumerate()
|
|
.min_by_key(|(_, s)| s.last_access)
|
|
.map(|(i, _)| i)
|
|
.unwrap();
|
|
let removed = inner.slots.swap_remove(lru_idx);
|
|
inner.slot_index.remove(&removed.coord);
|
|
// swap_remove moved the former last element into `lru_idx` (unless
|
|
// it *was* the last element) — fix up that element's index entry.
|
|
if lru_idx < inner.slots.len() {
|
|
let moved_coord = inner.slots[lru_idx].coord.clone();
|
|
inner.slot_index.insert(moved_coord, lru_idx);
|
|
}
|
|
inner.current_bytes -= removed.data.len();
|
|
inner.stats.evictions += 1;
|
|
}
|
|
|
|
inner.current_bytes += data_len;
|
|
let new_idx = inner.slots.len();
|
|
inner.slot_index.insert(coord.clone(), new_idx);
|
|
inner.slots.push(CachedChunk {
|
|
coord,
|
|
data: Arc::clone(&data),
|
|
last_access: tick,
|
|
});
|
|
data
|
|
}
|
|
|
|
/// Clear the entire cache (index + decompressed data).
|
|
pub fn clear(&self) {
|
|
let mut inner = self.inner.lock().unwrap_or_else(|e| e.into_inner());
|
|
inner.index = None;
|
|
inner.index_addr = None;
|
|
inner.slots.clear();
|
|
inner.slot_index.clear();
|
|
inner.current_bytes = 0;
|
|
inner.tick = 0;
|
|
inner.last_coord = None;
|
|
inner.stats = AccessStats::default();
|
|
inner.chunk_index = None;
|
|
inner.chunk_layout = None;
|
|
}
|
|
|
|
/// Record that the given chunk coordinates are predicted to be accessed
|
|
/// soon (bookkeeping only).
|
|
///
|
|
/// This does **not** prefetch or pre-decompress anything — it only
|
|
/// checks whether each coordinate is already in the chunk index and
|
|
/// updates access-pattern stats accordingly. Real prefetching (e.g.
|
|
/// background pre-decompression) is not implemented.
|
|
pub fn prefetch_hint(&self, next_coords: &[ChunkCoord]) {
|
|
let inner = self.inner.lock().unwrap_or_else(|e| e.into_inner());
|
|
if inner.index.is_none() {
|
|
return;
|
|
}
|
|
drop(inner);
|
|
// For each predicted coordinate, verify it exists in the index.
|
|
// The index is already populated, so this is a no-op for known chunks.
|
|
// The purpose is to signal intent — callers can pre-decompress if needed.
|
|
// We touch the stats to record that prefetch hints were issued.
|
|
let mut inner = self.inner.lock().unwrap_or_else(|e| e.into_inner());
|
|
for coord in next_coords {
|
|
let exists = inner
|
|
.index
|
|
.as_ref()
|
|
.map(|idx| idx.contains_key(coord))
|
|
.unwrap_or(false);
|
|
if exists {
|
|
inner.stats.sequential_count += 1;
|
|
}
|
|
}
|
|
}
|
|
|
|
/// Return the current access pattern statistics.
|
|
pub fn access_stats(&self) -> AccessStats {
|
|
self.inner
|
|
.lock()
|
|
.unwrap_or_else(|e| e.into_inner())
|
|
.stats
|
|
.clone()
|
|
}
|
|
|
|
/// Update the sweep direction label in the access stats.
|
|
pub fn set_sweep_direction(&self, direction: &'static str) {
|
|
self.inner
|
|
.lock()
|
|
.unwrap_or_else(|e| e.into_inner())
|
|
.stats
|
|
.sweep_direction = Some(direction);
|
|
}
|
|
|
|
/// Number of decompressed chunks currently cached.
|
|
pub fn cached_chunk_count(&self) -> usize {
|
|
self.inner
|
|
.lock()
|
|
.unwrap_or_else(|e| e.into_inner())
|
|
.slots
|
|
.len()
|
|
}
|
|
|
|
/// Total bytes of decompressed data currently cached.
|
|
pub fn cached_bytes(&self) -> usize {
|
|
self.inner
|
|
.lock()
|
|
.unwrap_or_else(|e| e.into_inner())
|
|
.current_bytes
|
|
}
|
|
}
|
|
|
|
#[cfg(feature = "std")]
|
|
impl Default for ChunkCache {
|
|
fn default() -> Self {
|
|
Self::new()
|
|
}
|
|
}
|
|
|
|
// ---------------------------------------------------------------------------
|
|
// Tests
|
|
// ---------------------------------------------------------------------------
|
|
|
|
#[cfg(test)]
|
|
mod tests {
|
|
use super::*;
|
|
|
|
fn make_chunk(offsets: Vec<u64>, address: u64, size: u32) -> ChunkInfo {
|
|
ChunkInfo {
|
|
chunk_size: size,
|
|
filter_mask: 0,
|
|
offsets,
|
|
address,
|
|
}
|
|
}
|
|
|
|
#[test]
|
|
fn index_populate_and_lookup() {
|
|
let cache = ChunkCache::new();
|
|
let chunks = vec![
|
|
make_chunk(vec![0, 0, 0], 0x1000, 80),
|
|
make_chunk(vec![10, 0, 0], 0x2000, 80),
|
|
];
|
|
cache.populate_index(&chunks, 2); // rank=2, truncate to [0,0] and [10,0]
|
|
assert!(cache.has_index());
|
|
|
|
let c0 = cache.lookup_index(&[0, 0]).unwrap();
|
|
assert_eq!(c0.address, 0x1000);
|
|
|
|
let c1 = cache.lookup_index(&[10, 0]).unwrap();
|
|
assert_eq!(c1.address, 0x2000);
|
|
|
|
assert!(cache.lookup_index(&[5, 0]).is_none());
|
|
}
|
|
|
|
#[test]
|
|
fn decompressed_cache_hit() {
|
|
let cache = ChunkCache::new();
|
|
cache.put_decompressed(vec![0, 0], vec![1, 2, 3, 4]);
|
|
let got = cache.get_decompressed(&[0, 0]).unwrap();
|
|
assert_eq!(got, vec![1, 2, 3, 4]);
|
|
}
|
|
|
|
#[test]
|
|
fn lru_eviction_by_slots() {
|
|
let cache = ChunkCache::with_capacity(1024 * 1024, 2); // max 2 slots
|
|
|
|
cache.put_decompressed(vec![0], vec![1; 10]);
|
|
cache.put_decompressed(vec![1], vec![2; 10]);
|
|
assert_eq!(cache.cached_chunk_count(), 2);
|
|
|
|
// Access slot 0 to make it more recent
|
|
cache.get_decompressed(&[0]);
|
|
|
|
// Insert slot 2 — should evict slot 1 (LRU)
|
|
cache.put_decompressed(vec![2], vec![3; 10]);
|
|
assert_eq!(cache.cached_chunk_count(), 2);
|
|
|
|
assert!(cache.get_decompressed(&[0]).is_some());
|
|
assert!(cache.get_decompressed(&[1]).is_none()); // evicted
|
|
assert!(cache.get_decompressed(&[2]).is_some());
|
|
}
|
|
|
|
#[test]
|
|
fn lru_eviction_by_bytes() {
|
|
let cache = ChunkCache::with_capacity(50, 100); // 50 bytes max
|
|
|
|
cache.put_decompressed(vec![0], vec![0; 20]);
|
|
cache.put_decompressed(vec![1], vec![0; 20]);
|
|
assert_eq!(cache.cached_bytes(), 40);
|
|
|
|
// This needs 20 bytes but only 10 free — evict LRU
|
|
cache.put_decompressed(vec![2], vec![0; 20]);
|
|
assert!(cache.cached_bytes() <= 50);
|
|
assert!(cache.get_decompressed(&[0]).is_none()); // evicted (LRU)
|
|
}
|
|
|
|
#[test]
|
|
fn oversized_chunk_not_cached() {
|
|
let cache = ChunkCache::with_capacity(10, 16);
|
|
cache.put_decompressed(vec![0], vec![0; 100]); // too big
|
|
assert_eq!(cache.cached_chunk_count(), 0);
|
|
}
|
|
|
|
#[test]
|
|
fn clear_resets_everything() {
|
|
let cache = ChunkCache::new();
|
|
let chunks = vec![make_chunk(vec![0, 0], 0x1000, 80)];
|
|
cache.populate_index(&chunks, 1);
|
|
cache.put_decompressed(vec![0], vec![1, 2, 3]);
|
|
|
|
cache.clear();
|
|
assert!(!cache.has_index());
|
|
assert_eq!(cache.cached_chunk_count(), 0);
|
|
assert_eq!(cache.cached_bytes(), 0);
|
|
}
|
|
|
|
#[test]
|
|
fn duplicate_insert_is_noop() {
|
|
let cache = ChunkCache::new();
|
|
cache.put_decompressed(vec![0], vec![1, 2, 3]);
|
|
cache.put_decompressed(vec![0], vec![1, 2, 3]); // duplicate
|
|
assert_eq!(cache.cached_chunk_count(), 1);
|
|
assert_eq!(cache.cached_bytes(), 3);
|
|
}
|
|
|
|
#[test]
|
|
fn slot_index_consistent_after_many_evictions() {
|
|
// Force repeated swap_remove evictions (small slot budget, many
|
|
// inserts) and confirm the coord -> slot index stays correct: every
|
|
// remaining coord must still resolve to its own data, not another
|
|
// slot's (which would happen if swap_remove's index fixup were wrong).
|
|
let cache = ChunkCache::with_capacity(1024 * 1024, 4); // max 4 slots
|
|
|
|
for i in 0..50u64 {
|
|
cache.put_decompressed(vec![i], vec![(i % 256) as u8; 8]);
|
|
// Interleave reads of a couple of earlier coords to churn LRU
|
|
// order (and thus which slot gets swap_remove'd) beyond simple
|
|
// FIFO eviction.
|
|
if i >= 2 {
|
|
let _ = cache.get_decompressed(&[i - 2]);
|
|
}
|
|
}
|
|
|
|
// Whatever remains in the cache (at most 4 slots) must return its
|
|
// own correct data.
|
|
for i in 0..50u64 {
|
|
if let Some(data) = cache.get_decompressed(&[i]) {
|
|
assert_eq!(
|
|
data,
|
|
vec![(i % 256) as u8; 8],
|
|
"coord {i} returned wrong data after eviction churn"
|
|
);
|
|
}
|
|
}
|
|
assert!(cache.cached_chunk_count() <= 4);
|
|
}
|
|
|
|
#[test]
|
|
fn get_decompressed_aligned_shares_arc_on_hit() {
|
|
let cache = ChunkCache::new();
|
|
cache.put_decompressed(vec![0, 0], vec![9, 9, 9, 9]);
|
|
let a = cache.get_decompressed_aligned(&[0, 0]).unwrap();
|
|
let b = cache.get_decompressed_aligned(&[0, 0]).unwrap();
|
|
// A cache hit clones the Arc (refcount bump), not the underlying
|
|
// buffer — both handles point at the same allocation.
|
|
assert!(Arc::ptr_eq(&a, &b));
|
|
assert_eq!(a.as_slice(), &[9, 9, 9, 9]);
|
|
}
|
|
|
|
// --- CacheAlignedBuffer tests ---
|
|
|
|
#[test]
|
|
fn aligned_buffer_basic() {
|
|
let buf = CacheAlignedBuffer::zeroed(256);
|
|
assert_eq!(buf.len(), 256);
|
|
assert!(buf.is_aligned());
|
|
assert_eq!(&buf[..4], &[0, 0, 0, 0]);
|
|
}
|
|
|
|
#[test]
|
|
fn aligned_buffer_from_slice() {
|
|
let data = vec![1u8, 2, 3, 4, 5];
|
|
let buf = CacheAlignedBuffer::from_slice(&data);
|
|
assert_eq!(buf.len(), 5);
|
|
assert!(buf.is_aligned());
|
|
assert_eq!(buf.to_vec(), data);
|
|
}
|
|
|
|
#[test]
|
|
fn aligned_buffer_from_vec() {
|
|
let data = vec![42u8; 1024];
|
|
let buf = CacheAlignedBuffer::from_vec(data.clone());
|
|
assert!(buf.is_aligned());
|
|
assert_eq!(buf.to_vec(), data);
|
|
}
|
|
|
|
#[test]
|
|
fn aligned_buffer_empty() {
|
|
let buf = CacheAlignedBuffer::zeroed(0);
|
|
assert!(buf.is_empty());
|
|
assert!(buf.is_aligned());
|
|
assert_eq!(buf.to_vec(), Vec::<u8>::new());
|
|
}
|
|
|
|
#[test]
|
|
fn aligned_buffer_clone_is_aligned() {
|
|
let buf = CacheAlignedBuffer::from_slice(&[1, 2, 3, 4]);
|
|
let cloned = buf.clone();
|
|
assert!(cloned.is_aligned());
|
|
assert_eq!(buf.to_vec(), cloned.to_vec());
|
|
}
|
|
|
|
#[test]
|
|
fn aligned_buffer_deref_works() {
|
|
let buf = CacheAlignedBuffer::from_slice(&[10, 20, 30]);
|
|
assert_eq!(buf[0], 10);
|
|
assert_eq!(buf[1], 20);
|
|
assert_eq!(buf[2], 30);
|
|
}
|
|
|
|
#[test]
|
|
fn aligned_buffer_various_sizes() {
|
|
// Test alignment for various sizes including non-power-of-two
|
|
for size in [1, 7, 63, 64, 65, 127, 128, 129, 255, 256, 1000, 4096] {
|
|
let buf = CacheAlignedBuffer::zeroed(size);
|
|
assert!(buf.is_aligned(), "not aligned for size {size}");
|
|
assert_eq!(buf.len(), size);
|
|
}
|
|
}
|
|
|
|
#[test]
|
|
fn cached_data_is_aligned() {
|
|
let cache = ChunkCache::new();
|
|
cache.put_decompressed(vec![0, 0], vec![1, 2, 3, 4, 5, 6, 7, 8]);
|
|
let aligned = cache.get_decompressed_aligned(&[0, 0]).unwrap();
|
|
assert!(aligned.is_aligned());
|
|
assert_eq!(aligned.to_vec(), vec![1, 2, 3, 4, 5, 6, 7, 8]);
|
|
}
|
|
|
|
#[test]
|
|
fn align_to_cache_line_values() {
|
|
assert_eq!(align_to_cache_line(0), 0);
|
|
assert_eq!(align_to_cache_line(1), CACHE_LINE_SIZE);
|
|
assert_eq!(align_to_cache_line(CACHE_LINE_SIZE), CACHE_LINE_SIZE);
|
|
assert_eq!(
|
|
align_to_cache_line(CACHE_LINE_SIZE + 1),
|
|
CACHE_LINE_SIZE * 2
|
|
);
|
|
}
|
|
}
|