osobhandClaude Opus 4.8
4fa7e89a46
clawhdf5-agent: fixed-length string datasets (memory text chunks, session summaries, ids, tags, entity/relation names) were stored uncompressed behind a stale "chunked compound not yet supported" comment. Chunked writes work for fixed-size string/compound datatypes like any other, so write_string_dataset now chunks + deflates once a dataset's payload reaches 4 KiB — large, redundant NullPad content compresses well while tiny metadata stays contiguous (no chunk-overhead bloat). The dead `compress` parameter is removed in favor of this size heuristic. clawhdf5-format: enabling string compression exposed a latent bug — the per-file ChunkCache built its chunk index once and reused it for every chunked dataset in the file, keyed only by chunk coordinate with no dataset discrimination. With one chunked dataset per file this never surfaced; with two of different rank (a 1-D compressed string array and the 2-D embeddings matrix) the first dataset's rank-1 index was reused for the second, panicking with an out-of-bounds chunk coordinate. The cache now binds to a dataset by its chunk-index address and rebinds — dropping the index, chunk-index map, layout, and decompressed slots — whenever the dataset being read changes, while still caching repeated/sequential access to the same dataset. Tests: facade regression reading a 1-D compressed string dataset and a 2-D compressed f32 dataset through one shared File cache (verified to panic without the fix); existing agent e2e tests (large text chunks, migration round-trip) now pass with compression on. Co-Authored-By: Claude Opus 4.8 <[email protected]>
869 lines
28 KiB
Rust
869 lines
28 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;
|
|
|
|
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 {}
|
|
|
|
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,
|
|
data: 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>,
|
|
|
|
/// 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)),
|
|
current_bytes: 0,
|
|
max_bytes,
|
|
max_slots,
|
|
tick: 0,
|
|
last_coord: None,
|
|
stats: AccessStats::default(),
|
|
chunk_index: None,
|
|
chunk_layout: None,
|
|
}),
|
|
}
|
|
}
|
|
|
|
// ----- Index operations -----
|
|
|
|
/// 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.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.
|
|
///
|
|
/// Returns a clone of the cache-line-aligned buffer.
|
|
pub fn get_decompressed(&self, coord: &[u64]) -> Option<Vec<u8>> {
|
|
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 mut found = None;
|
|
for slot in inner.slots.iter_mut() {
|
|
if slot.coord.as_slice() == coord {
|
|
slot.last_access = tick;
|
|
found = Some(slot.data.to_vec());
|
|
break;
|
|
}
|
|
}
|
|
if let Some(ref data) = found {
|
|
inner.stats.hits += 1;
|
|
inner.stats.bytes_read += data.len() as u64;
|
|
} else {
|
|
inner.stats.misses += 1;
|
|
}
|
|
found
|
|
}
|
|
|
|
/// Try to get a reference-counted clone of the aligned buffer for a chunk.
|
|
pub fn get_decompressed_aligned(&self, coord: &[u64]) -> Option<CacheAlignedBuffer> {
|
|
let mut inner = self.inner.lock().unwrap_or_else(|e| e.into_inner());
|
|
inner.tick += 1;
|
|
let tick = inner.tick;
|
|
let mut found = None;
|
|
for slot in inner.slots.iter_mut() {
|
|
if slot.coord.as_slice() == coord {
|
|
slot.last_access = tick;
|
|
found = Some(slot.data.clone());
|
|
break;
|
|
}
|
|
}
|
|
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.
|
|
pub fn put_decompressed(&self, coord: ChunkCoord, data: Vec<u8>) {
|
|
let aligned = CacheAlignedBuffer::from_slice(&data);
|
|
self.put_decompressed_aligned(coord, aligned);
|
|
}
|
|
|
|
/// Insert an already-aligned buffer into the LRU cache.
|
|
pub fn put_decompressed_aligned(&self, coord: ChunkCoord, data: CacheAlignedBuffer) {
|
|
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
|
|
if data_len > inner.max_bytes {
|
|
return;
|
|
}
|
|
|
|
// Check if already present
|
|
inner.tick += 1;
|
|
let tick = inner.tick;
|
|
for slot in inner.slots.iter_mut() {
|
|
if slot.coord == coord {
|
|
slot.last_access = tick;
|
|
return; // 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.current_bytes -= removed.data.len();
|
|
inner.stats.evictions += 1;
|
|
}
|
|
|
|
inner.current_bytes += data_len;
|
|
inner.slots.push(CachedChunk {
|
|
coord,
|
|
data,
|
|
last_access: tick,
|
|
});
|
|
}
|
|
|
|
/// 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.current_bytes = 0;
|
|
inner.tick = 0;
|
|
inner.last_coord = None;
|
|
inner.stats = AccessStats::default();
|
|
inner.chunk_index = None;
|
|
inner.chunk_layout = None;
|
|
}
|
|
|
|
/// Hint that the given chunk coordinates will be accessed soon.
|
|
///
|
|
/// Pre-populates the chunk index for these coordinates so that
|
|
/// subsequent lookups are O(1). This does NOT pre-decompress the
|
|
/// chunks — it only ensures the index entries exist.
|
|
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);
|
|
}
|
|
|
|
// --- 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
|
|
);
|
|
}
|
|
}
|