pub mod global_key_offset;
pub(crate) use global_key_offset::GlobalKeyCacheSnapshot;
pub use global_key_offset::{
GenerationIdentity, GlobalKeyOffsetCache, PartitionLoc, DEFAULT_GLOBAL_KEY_CACHE_BYTES,
DEFAULT_GLOBAL_KEY_CACHE_SHARDS,
};
use bytes::Bytes;
use lru::LruCache;
use std::hash::{Hash, Hasher};
use std::sync::atomic::{AtomicU64, Ordering};
use std::sync::Mutex;
pub const DEFAULT_SHARDS: usize = 16;
pub const DEFAULT_BUDGET_BYTES: usize = 256 * 1024 * 1024;
#[derive(Clone, Copy, Debug, PartialEq, Eq, Hash)]
pub struct ChunkKey {
pub sstable: u64,
pub chunk_index: u64,
pub aux: u64,
}
impl ChunkKey {
#[inline]
pub fn new(sstable: u64, chunk_index: u64) -> Self {
Self {
sstable,
chunk_index,
aux: 0,
}
}
#[inline]
pub fn with_aux(sstable: u64, chunk_index: u64, aux: u64) -> Self {
Self {
sstable,
chunk_index,
aux,
}
}
}
struct Shard {
lru: LruCache<ChunkKey, Bytes>,
current_bytes: usize,
}
impl Shard {
fn new() -> Self {
Self {
lru: LruCache::unbounded(),
current_bytes: 0,
}
}
}
pub struct DecompressedChunkCache {
shards: Box<[Mutex<Shard>]>,
mask: usize,
budget_per_shard: usize,
disabled: bool,
hits: AtomicU64,
misses: AtomicU64,
evictions: AtomicU64,
}
impl std::fmt::Debug for DecompressedChunkCache {
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
f.debug_struct("DecompressedChunkCache")
.field("shards", &self.shards.len())
.field("budget_per_shard", &self.budget_per_shard)
.field("resident_bytes", &self.resident_bytes())
.field("hits", &self.hits.load(Ordering::Relaxed))
.field("misses", &self.misses.load(Ordering::Relaxed))
.field("evictions", &self.evictions.load(Ordering::Relaxed))
.finish()
}
}
impl DecompressedChunkCache {
pub fn with_budget_bytes(total_budget_bytes: usize) -> Self {
Self::with_budget_and_shards(total_budget_bytes, DEFAULT_SHARDS)
}
pub fn with_budget_and_shards(total_budget_bytes: usize, shard_count: usize) -> Self {
let shard_count = shard_count.max(1).next_power_of_two();
let budget_per_shard = (total_budget_bytes / shard_count).max(1);
let mut shards = Vec::with_capacity(shard_count);
for _ in 0..shard_count {
shards.push(Mutex::new(Shard::new()));
}
Self {
shards: shards.into_boxed_slice(),
mask: shard_count - 1,
budget_per_shard,
disabled: false,
hits: AtomicU64::new(0),
misses: AtomicU64::new(0),
evictions: AtomicU64::new(0),
}
}
pub fn disabled() -> Self {
Self {
shards: vec![Mutex::new(Shard::new())].into_boxed_slice(),
mask: 0,
budget_per_shard: 0,
disabled: true,
hits: AtomicU64::new(0),
misses: AtomicU64::new(0),
evictions: AtomicU64::new(0),
}
}
#[inline]
fn lock(m: &Mutex<Shard>) -> std::sync::MutexGuard<'_, Shard> {
m.lock().unwrap_or_else(|e| e.into_inner())
}
#[inline]
fn shard_for(&self, key: &ChunkKey) -> &Mutex<Shard> {
let mut h = std::collections::hash_map::DefaultHasher::new();
key.hash(&mut h);
let idx = (h.finish() as usize) & self.mask;
&self.shards[idx]
}
pub fn get(&self, key: &ChunkKey) -> Option<Bytes> {
if self.disabled {
return None;
}
let mut guard = Self::lock(self.shard_for(key));
match guard.lru.get(key) {
Some(v) => {
let v = v.clone();
drop(guard);
self.hits.fetch_add(1, Ordering::Relaxed);
Some(v)
}
None => {
drop(guard);
self.misses.fetch_add(1, Ordering::Relaxed);
None
}
}
}
pub fn insert(&self, key: ChunkKey, data: Vec<u8>) -> Bytes {
let bytes = Bytes::from(data);
if self.disabled {
return bytes;
}
let len = bytes.len();
let mut guard = Self::lock(self.shard_for(&key));
if let Some(old) = guard.lru.put(key, bytes.clone()) {
guard.current_bytes = guard.current_bytes.saturating_sub(old.len());
}
guard.current_bytes = guard.current_bytes.saturating_add(len);
let mut evicted_here: u64 = 0;
while guard.current_bytes > self.budget_per_shard && guard.lru.len() > 1 {
match guard.lru.pop_lru() {
Some((_, evicted)) => {
guard.current_bytes = guard.current_bytes.saturating_sub(evicted.len());
evicted_here += 1;
}
None => break,
}
}
drop(guard);
if evicted_here > 0 {
self.evictions.fetch_add(evicted_here, Ordering::Relaxed);
}
bytes
}
pub fn resident_bytes(&self) -> usize {
self.shards
.iter()
.map(|m| Self::lock(m).current_bytes)
.sum()
}
pub fn len(&self) -> usize {
self.shards.iter().map(|m| Self::lock(m).lru.len()).sum()
}
pub fn is_empty(&self) -> bool {
self.len() == 0
}
pub fn budget_bytes(&self) -> usize {
self.budget_per_shard * self.shards.len()
}
pub fn hit_count(&self) -> u64 {
self.hits.load(Ordering::Relaxed)
}
pub fn miss_count(&self) -> u64 {
self.misses.load(Ordering::Relaxed)
}
pub fn eviction_count(&self) -> u64 {
self.evictions.load(Ordering::Relaxed)
}
}
impl Default for DecompressedChunkCache {
fn default() -> Self {
Self::with_budget_bytes(DEFAULT_BUDGET_BYTES)
}
}
#[cfg(test)]
mod tests {
use super::*;
use std::sync::Arc;
fn chunk(byte: u8, len: usize) -> Vec<u8> {
vec![byte; len]
}
#[test]
fn eviction_order_lru() {
let cache = DecompressedChunkCache::with_budget_and_shards(200, 1);
let a = ChunkKey::new(1, 0);
let b = ChunkKey::new(1, 1);
let c = ChunkKey::new(1, 2);
cache.insert(a, chunk(0xAA, 100));
cache.insert(b, chunk(0xBB, 100));
assert!(cache.get(&a).is_some());
cache.insert(c, chunk(0xCC, 100));
assert!(cache.get(&a).is_some(), "A (recently used) must survive");
assert!(cache.get(&c).is_some(), "C (just inserted) must survive");
assert!(
cache.get(&b).is_none(),
"B (least recently used) must be evicted"
);
assert!(cache.resident_bytes() <= 200);
}
#[test]
fn byte_budget_bounded() {
let budget = 500usize;
let cache = DecompressedChunkCache::with_budget_and_shards(budget, 1);
for i in 0..50u64 {
cache.insert(ChunkKey::new(7, i), chunk(i as u8, 100));
assert!(
cache.resident_bytes() <= budget,
"resident {} exceeded budget {} after insert {}",
cache.resident_bytes(),
budget,
i
);
}
assert!(
cache.len() <= 5,
"entry count must stay bounded (got {})",
cache.len()
);
}
#[test]
fn eviction_count_tracks_real_evictions() {
let cache = DecompressedChunkCache::with_budget_and_shards(200, 1);
assert_eq!(cache.eviction_count(), 0, "fresh cache has evicted nothing");
cache.insert(ChunkKey::new(1, 0), chunk(0x00, 100));
cache.insert(ChunkKey::new(1, 1), chunk(0x01, 100));
assert_eq!(cache.eviction_count(), 0);
cache.insert(ChunkKey::new(1, 2), chunk(0x02, 100));
assert_eq!(cache.eviction_count(), 1);
cache.insert(ChunkKey::new(1, 3), chunk(0x03, 100));
assert_eq!(cache.eviction_count(), 2);
}
#[test]
fn disabled_cache_reports_zero_evictions() {
let cache = DecompressedChunkCache::disabled();
for i in 0..10u64 {
cache.insert(ChunkKey::new(1, i), chunk(i as u8, 4096));
}
assert_eq!(cache.eviction_count(), 0);
}
#[test]
fn single_oversized_entry_retained() {
let cache = DecompressedChunkCache::with_budget_and_shards(100, 1);
let k = ChunkKey::new(1, 0);
cache.insert(k, chunk(0xEE, 4096));
assert!(
cache.get(&k).is_some(),
"oversized entry must remain resident"
);
assert_eq!(cache.len(), 1);
}
#[test]
fn zero_copy_hit_pointer_identity() {
let cache = DecompressedChunkCache::with_budget_and_shards(1 << 20, 1);
let k = ChunkKey::new(3, 9);
let inserted = cache.insert(k, chunk(0x42, 1024));
let h1 = cache.get(&k).expect("first hit");
let h2 = cache.get(&k).expect("second hit");
assert_eq!(h1.as_ptr(), inserted.as_ptr());
assert_eq!(h2.as_ptr(), inserted.as_ptr());
assert_eq!(&*h1, &chunk(0x42, 1024)[..]);
}
#[test]
fn concurrency_soundness() {
use std::thread;
let cache = Arc::new(DecompressedChunkCache::with_budget_bytes(64 * 1024));
let chunk_len = 256usize;
let n_keys = 512u64;
let mut handles = Vec::new();
for t in 0..8u64 {
let cache = Arc::clone(&cache);
handles.push(thread::spawn(move || {
for round in 0..2000u64 {
let idx = (t.wrapping_mul(31).wrapping_add(round)) % n_keys;
let key = ChunkKey::new(1, idx);
let expect_byte = idx as u8;
let got = match cache.get(&key) {
Some(v) => v,
None => cache.insert(key, vec![expect_byte; chunk_len]),
};
assert_eq!(got.len(), chunk_len);
assert!(got.iter().all(|&b| b == expect_byte));
}
}));
}
for h in handles {
h.join().expect("worker thread must not panic");
}
assert!(
cache.resident_bytes() <= cache.budget_bytes(),
"resident {} exceeded budget {}",
cache.resident_bytes(),
cache.budget_bytes()
);
}
#[test]
fn poisoned_lock_recovers() {
use std::panic::{catch_unwind, AssertUnwindSafe};
use std::thread;
let cache = Arc::new(DecompressedChunkCache::with_budget_and_shards(1 << 20, 1));
let k = ChunkKey::new(1, 0);
cache.insert(k, chunk(0x55, 16));
let poison_cache = Arc::clone(&cache);
let _ = thread::spawn(move || {
let _guard = DecompressedChunkCache::lock(&poison_cache.shards[0]);
panic!("intentional poison");
})
.join();
let res = catch_unwind(AssertUnwindSafe(|| {
let hit = cache.get(&k);
cache.insert(ChunkKey::new(1, 1), chunk(0x66, 16));
hit.is_some()
}));
assert_eq!(
res.ok(),
Some(true),
"cache must recover from a poisoned lock"
);
}
#[test]
fn shard_count_rounds_to_power_of_two() {
let cache = DecompressedChunkCache::with_budget_and_shards(1 << 20, 10);
assert_eq!(cache.shards.len(), 16);
assert_eq!(cache.mask, 15);
}
#[test]
fn ranged_key_does_not_alias_same_offset_different_size() {
let cache = DecompressedChunkCache::with_budget_bytes(1 << 20);
let sstable = 0xABCD_u64;
let offset = 4096_u64;
let k16 = ChunkKey::with_aux(sstable, offset, 16);
let v16 = cache.insert(k16, chunk(0x11, 16));
let k32 = ChunkKey::with_aux(sstable, offset, 32);
assert!(
cache.get(&k32).is_none(),
"same offset with a different size must not alias the cached range"
);
let again = cache.get(&k16).expect("original ranged key still resident");
assert_eq!(
v16.as_ptr(),
again.as_ptr(),
"same (offset,size) key returns the same buffer"
);
let k0 = ChunkKey::new(sstable, offset);
assert!(
cache.get(&k0).is_none(),
"whole-chunk key (aux=0) must not alias a sized range read"
);
}
#[test]
fn disabled_cache_is_a_genuine_no_op() {
let cache = DecompressedChunkCache::disabled();
assert_eq!(cache.budget_bytes(), 0, "disabled cache has no budget");
let key = ChunkKey::new(1, 0);
let arc = cache.insert(key, chunk(0xEE, 4096));
assert_eq!(
arc.len(),
4096,
"insert still hands back the produced buffer"
);
assert!(cache.get(&key).is_none(), "disabled cache never retains");
assert_eq!(cache.resident_bytes(), 0);
assert_eq!(cache.len(), 0);
assert!(cache.is_empty());
assert_eq!(cache.hit_count(), 0);
assert_eq!(cache.miss_count(), 0);
}
}