use std::sync::atomic::AtomicBool;
use std::sync::{Arc, OnceLock, Weak};
use crate::portability::{AtomicUsize, Ordering};
use kovan_map::HopscotchMap;
use crate::sync::Mutex;
use xxhash_rust::xxh3::xxh3_64;
use super::block::Block;
use super::filter_block::FilterBlock;
use super::index_block::IndexBlock;
use crate::options::MAX_BLOCK_CACHE_SHARD_BITS;
use crate::statistics::{Statistics, Ticker};
#[derive(Hash, Eq, PartialEq, Clone, Copy)]
struct CacheKey {
file_id: u64,
offset: u64,
}
const MAX_SHARD_BITS: u32 = MAX_BLOCK_CACHE_SHARD_BITS;
const MIN_SHARD_CAPACITY: usize = 64 * 1024;
const ENTRY_OVERHEAD: usize = 160;
const ESTIMATED_ENTRY_BYTES: usize = 4 * 1024 + ENTRY_OVERHEAD;
const MIN_MAP_BUCKETS: usize = 64;
const MAX_MAP_BUCKETS: usize = 64 * 1024;
const REMOVALS_PER_RECLAIM: u64 = 2048;
enum CacheEntry {
Data(Arc<Block>),
Index(Arc<IndexBlock>),
Filter(Arc<FilterBlock>),
}
impl CacheEntry {
fn payload_charge(&self) -> usize {
match self {
Self::Data(block) => block.charge(),
Self::Index(block) => block.charge(),
Self::Filter(block) => block.charge(),
}
}
fn clone_ref(&self) -> Self {
match self {
Self::Data(b) => Self::Data(Arc::clone(b)),
Self::Index(b) => Self::Index(Arc::clone(b)),
Self::Filter(b) => Self::Filter(Arc::clone(b)),
}
}
}
fn entry_charge(entry: &CacheEntry) -> usize {
entry.payload_charge() + ENTRY_OVERHEAD
}
struct ClockEntry {
entry: CacheEntry,
key: CacheKey,
charge: usize,
slot: u32,
referenced: AtomicBool,
}
struct ClockRing {
slots: Vec<Option<Arc<ClockEntry>>>,
free: Vec<usize>,
hand: usize,
used: usize,
removals: u64,
}
impl ClockRing {
const fn new() -> Self {
Self {
slots: Vec::new(),
free: Vec::new(),
hand: 0,
used: 0,
removals: 0,
}
}
fn alloc_slot(&mut self) -> usize {
if let Some(slot) = self.free.pop() {
return slot;
}
self.slots.push(None);
self.slots.len() - 1
}
fn record_removal(&mut self) {
self.removals += 1;
if self.removals.is_multiple_of(REMOVALS_PER_RECLAIM) {
kovan::flush();
}
}
fn release(&mut self, slot: usize, charge: usize) {
if let Some(held) = self.slots.get_mut(slot)
&& held.take().is_some()
{
self.free.push(slot);
self.used = self.used.saturating_sub(charge);
}
}
fn reset(&mut self) {
self.slots = Vec::new();
self.free = Vec::new();
self.hand = 0;
self.used = 0;
self.removals = 0;
}
}
struct CacheShard {
map: OnceLock<Box<HopscotchMap<CacheKey, Weak<ClockEntry>>>>,
capacity: usize,
ring: Mutex<ClockRing>,
}
impl CacheShard {
fn new(capacity: usize) -> Self {
Self {
map: OnceLock::new(),
capacity,
ring: Mutex::new(ClockRing::new()),
}
}
fn map(&self) -> &HopscotchMap<CacheKey, Weak<ClockEntry>> {
self.map.get_or_init(|| {
let estimate =
(self.capacity / ESTIMATED_ENTRY_BYTES).clamp(MIN_MAP_BUCKETS, MAX_MAP_BUCKETS);
Box::new(HopscotchMap::with_capacity(estimate))
})
}
fn get(&self, key: &CacheKey) -> Option<CacheEntry> {
let entry = self.map.get()?.get(key)?.upgrade()?;
if entry.key != *key {
return None;
}
entry.referenced.store(true, Ordering::Relaxed);
Some(entry.entry.clone_ref())
}
fn take_existing(&self, ring: &mut ClockRing, key: &CacheKey) {
let Some(map) = self.map.get() else {
return;
};
if let Some(entry) = map.force_remove(key).as_ref().and_then(Weak::upgrade) {
ring.release(entry.slot as usize, entry.charge);
}
}
fn evict_one(&self, ring: &mut ClockRing) -> bool {
let len = ring.slots.len();
if len == 0 {
return false;
}
if ring.hand >= len {
ring.hand = 0;
}
for step in 0..=2 * len {
let hand = ring.hand;
ring.hand = if hand + 1 == len { 0 } else { hand + 1 };
let forced = step >= len;
let Some(entry) = ring.slots[hand]
.take_if(|entry| forced || !entry.referenced.swap(false, Ordering::Relaxed))
else {
continue;
};
ring.free.push(hand);
ring.used = ring.used.saturating_sub(entry.charge);
if let Some(map) = self.map.get() {
map.force_remove(&entry.key);
}
ring.record_removal();
return true;
}
false
}
fn insert_within_budget(
&self,
ring: &mut ClockRing,
key: CacheKey,
entry: &CacheEntry,
size: usize,
) -> bool {
if size > self.capacity {
return false;
}
self.take_existing(ring, &key);
while ring.used + size > self.capacity {
if !self.evict_one(ring) {
break;
}
}
self.store(ring, key, entry.clone_ref(), size);
true
}
fn store(&self, ring: &mut ClockRing, key: CacheKey, entry: CacheEntry, size: usize) {
let slot = ring.alloc_slot();
let slot_entry = Arc::new(ClockEntry {
entry,
key,
charge: size,
slot: slot as u32,
referenced: AtomicBool::new(false),
});
self.map().insert(key, Arc::downgrade(&slot_entry));
ring.slots[slot] = Some(slot_entry);
ring.used += size;
}
fn replace_all_with(
&self,
ring: &mut ClockRing,
key: CacheKey,
entry: CacheEntry,
size: usize,
) {
self.clear(ring);
self.store(ring, key, entry, size);
}
fn evict_file(&self, ring: &mut ClockRing, file_id: u64) {
let Some(map) = self.map.get() else {
return;
};
for slot in 0..ring.slots.len() {
let Some(entry) = ring.slots[slot].as_ref() else {
continue;
};
if entry.key.file_id != file_id {
continue;
}
let (key, charge) = (entry.key, entry.charge);
map.force_remove(&key);
ring.release(slot, charge);
}
}
fn clear(&self, ring: &mut ClockRing) {
if let Some(map) = self.map.get() {
map.clear();
}
ring.reset();
}
}
pub(crate) struct BlockCache {
shards: Box<[CacheShard]>,
capacity: usize,
shard_mask: u64,
#[cfg(test)]
num_shards: usize,
total_used: AtomicUsize,
strict: bool,
stats: Option<Arc<Statistics>>,
}
impl BlockCache {
#[cfg(test)]
pub(crate) fn new(capacity_bytes: usize) -> Self {
Self::with_config(capacity_bytes, 6, false)
}
pub(crate) fn with_config(
capacity_bytes: usize,
shard_bits: u32,
strict_capacity_limit: bool,
) -> Self {
if capacity_bytes == 0 {
return Self {
shards: Vec::new().into_boxed_slice(),
capacity: 0,
shard_mask: 0,
#[cfg(test)]
num_shards: 0,
total_used: AtomicUsize::new(0),
strict: strict_capacity_limit,
stats: None,
};
}
let shard_bits = shard_bits.min(MAX_SHARD_BITS);
let mut num_shards: usize = 1usize << shard_bits;
while num_shards > 1 && capacity_bytes / num_shards < MIN_SHARD_CAPACITY {
num_shards /= 2;
}
let per_shard = capacity_bytes / num_shards;
let shards: Box<[CacheShard]> = (0..num_shards)
.map(|_| CacheShard::new(per_shard))
.collect::<Vec<_>>()
.into_boxed_slice();
Self {
shards,
capacity: per_shard * num_shards,
shard_mask: (num_shards - 1) as u64,
#[cfg(test)]
num_shards,
total_used: AtomicUsize::new(0),
strict: strict_capacity_limit,
stats: None,
}
}
pub(crate) fn with_stats(mut self, stats: Option<Arc<Statistics>>) -> Self {
self.stats = stats;
self
}
fn shard_index(&self, key: &CacheKey) -> usize {
let mut buf = [0u8; 16];
buf[..8].copy_from_slice(&key.file_id.to_le_bytes());
buf[8..].copy_from_slice(&key.offset.to_le_bytes());
(xxh3_64(&buf) & self.shard_mask) as usize
}
fn lookup<T>(
&self,
file_id: u64,
offset: u64,
project: fn(&CacheEntry) -> Option<Arc<T>>,
) -> Option<Arc<T>> {
if self.shards.is_empty() {
return None;
}
let key = CacheKey { file_id, offset };
let idx = self.shard_index(&key);
let hit = self.shards[idx].get(&key).as_ref().and_then(project);
if let Some(s) = self.stats.as_deref() {
if hit.is_some() {
s.add(Ticker::BlockCacheHit, 1);
} else {
s.add(Ticker::BlockCacheMiss, 1);
}
}
crate::perf_context::record_block_cache_lookup(hit.is_some());
hit
}
pub(crate) fn get(&self, file_id: u64, offset: u64) -> Option<Arc<Block>> {
self.lookup(file_id, offset, |entry| match entry {
CacheEntry::Data(block) => Some(Arc::clone(block)),
_ => None,
})
}
pub(crate) fn get_index(&self, file_id: u64, offset: u64) -> Option<Arc<IndexBlock>> {
self.lookup(file_id, offset, |entry| match entry {
CacheEntry::Index(block) => Some(Arc::clone(block)),
_ => None,
})
}
pub(crate) fn get_filter(&self, file_id: u64, offset: u64) -> Option<Arc<FilterBlock>> {
self.lookup(file_id, offset, |entry| match entry {
CacheEntry::Filter(block) => Some(Arc::clone(block)),
_ => None,
})
}
pub(crate) fn insert_index(&self, file_id: u64, offset: u64, block: Arc<IndexBlock>) -> bool {
self.store(file_id, offset, CacheEntry::Index(block))
}
pub(crate) fn insert_filter(&self, file_id: u64, offset: u64, block: Arc<FilterBlock>) -> bool {
self.store(file_id, offset, CacheEntry::Filter(block))
}
pub(crate) fn insert(&self, file_id: u64, offset: u64, block: Arc<Block>) {
self.store(file_id, offset, CacheEntry::Data(block));
}
fn store(&self, file_id: u64, offset: u64, entry: CacheEntry) -> bool {
if self.shards.is_empty() {
return false;
}
let key = CacheKey { file_id, offset };
let size = entry_charge(&entry);
let idx = self.shard_index(&key);
let stored = {
let shard = &self.shards[idx];
let mut ring = shard.ring.lock();
let before = ring.used;
if shard.insert_within_budget(&mut ring, key, &entry, size) {
self.publish(before, ring.used);
true
} else if self.strict || size > self.capacity {
false
} else {
loop {
let current = self.total_used.load(Ordering::Acquire);
let after = current.saturating_sub(before).saturating_add(size);
if after > self.capacity {
return false;
}
if self
.total_used
.compare_exchange_weak(current, after, Ordering::AcqRel, Ordering::Acquire)
.is_ok()
{
break;
}
}
shard.replace_all_with(&mut ring, key, entry, size);
true
}
};
if stored && let Some(s) = self.stats.as_deref() {
s.add(Ticker::BlockCacheAdd, 1);
}
stored
}
fn publish(&self, before: usize, after: usize) {
if after >= before {
self.total_used.fetch_add(after - before, Ordering::Relaxed);
} else {
self.total_used.fetch_sub(before - after, Ordering::Relaxed);
}
}
pub(crate) fn record_bloom_useful(&self) {
if let Some(s) = self.stats.as_deref() {
s.add(Ticker::BloomFilterUseful, 1);
}
crate::perf_context::record_bloom_check(true);
}
pub(crate) fn record_bloom_full_positive(&self) {
if let Some(s) = self.stats.as_deref() {
s.add(Ticker::BloomFilterFullPositive, 1);
}
crate::perf_context::record_bloom_check(false);
}
pub(crate) fn evict_file(&self, file_id: u64) {
for shard in self.shards.iter() {
let (before, after) = {
let mut ring = shard.ring.lock();
let before = ring.used;
shard.evict_file(&mut ring, file_id);
(before, ring.used)
};
if before > after {
self.total_used.fetch_sub(before - after, Ordering::Relaxed);
}
}
}
pub(crate) fn clear(&self) {
for shard in self.shards.iter() {
let freed = {
let mut ring = shard.ring.lock();
let freed = ring.used;
shard.clear(&mut ring);
freed
};
if freed > 0 {
self.total_used.fetch_sub(freed, Ordering::Relaxed);
}
}
}
pub(crate) fn usage(&self) -> usize {
self.total_used.load(Ordering::Relaxed)
}
pub(crate) fn capacity(&self) -> usize {
self.capacity
}
#[cfg(test)]
pub(crate) fn num_shards(&self) -> usize {
self.num_shards
}
#[cfg(test)]
pub(crate) fn populated_shards(&self) -> usize {
self.shards
.iter()
.filter(|s| s.ring.lock().used > 0)
.count()
}
#[cfg(test)]
pub(crate) fn true_usage(&self) -> usize {
let held: Vec<_> = self.shards.iter().map(|s| s.ring.lock()).collect();
held.iter().map(|ring| ring.used).sum()
}
#[cfg(test)]
pub(crate) fn entry_count(&self) -> usize {
self.shards
.iter()
.map(|s| s.map.get().map_or(0, |m| m.len()))
.sum()
}
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn a_cache_entry_does_not_outgrow_its_charge() {
assert_eq!(
std::mem::size_of::<ClockEntry>(),
48,
"ClockEntry grew: either shrink it again or raise ENTRY_OVERHEAD \
and re-measure the hit rate, because admission changes with it"
);
}
use crate::engine::block::{BlockBuilder, RESTART_INTERVAL};
fn dummy_block(size: usize) -> Arc<Block> {
let mut builder = BlockBuilder::new(RESTART_INTERVAL);
let value = vec![0u8; size];
builder.add(b"k", &value);
Arc::new(Block::decode(builder.finish()).expect("decode"))
}
#[test]
fn single_insert_then_get() {
let cache = BlockCache::new(1024 * 1024);
let blk = dummy_block(256);
cache.insert(1, 0, blk.clone());
assert!(cache.get(1, 0).is_some());
assert!(cache.usage() >= 256);
}
#[test]
fn eviction_bounds_total_usage() {
let cache = BlockCache::with_config(4 * 1024, 6, false);
for i in 0..32u64 {
cache.insert(1, i * 100, dummy_block(1024));
}
let usage = cache.usage();
assert!(
usage <= cache.capacity(),
"usage {usage} exceeded capacity {}",
cache.capacity()
);
assert!(cache.get(1, 0).is_none());
}
#[test]
fn strict_capacity_rejects_oversized_entry() {
let cache = BlockCache::with_config(64 * 1024, 0, true);
let big = dummy_block(128 * 1024);
cache.insert(1, 0, big);
assert!(
cache.get(1, 0).is_none(),
"strict cache must reject oversized entries"
);
assert_eq!(cache.usage(), 0);
}
#[test]
fn non_strict_cache_admits_an_entry_bigger_than_one_shard() {
let cache = BlockCache::with_config(512 * 1024, 3, false);
assert_eq!(cache.num_shards(), 8);
cache.insert(1, 0, dummy_block(128 * 1024));
assert!(
cache.get(1, 0).is_some(),
"non-strict cache should admit an entry larger than one shard"
);
assert!(cache.usage() <= cache.capacity());
}
#[test]
fn non_strict_cache_refuses_an_entry_bigger_than_the_whole_budget() {
let cache = BlockCache::with_config(64 * 1024, 0, false);
cache.insert(1, 0, dummy_block(128 * 1024));
assert!(
cache.get(1, 0).is_none(),
"a block larger than the entire budget must not be cached"
);
assert_eq!(cache.usage(), 0);
}
#[test]
fn oversized_admissions_stay_inside_the_budget_at_every_shard_count() {
let budget = 256 * 64 * 1024;
let mut usages = Vec::new();
for bits in [0u32, 4, 8] {
let cache = BlockCache::with_config(budget, bits, false);
for file_id in 0..4096u64 {
cache.insert(file_id, 0, dummy_block(256 * 1024));
}
assert!(
cache.usage() <= cache.capacity(),
"shard_bits {bits}: usage {} over capacity {}",
cache.usage(),
cache.capacity()
);
usages.push(cache.usage());
}
assert_eq!(
usages[0], usages[2],
"resident bytes still track the shard count"
);
}
#[test]
fn sharding_distributes_inserts_across_shards() {
let cache = BlockCache::with_config(64 * 1024 * 1024, 6, false);
for i in 0..1024u64 {
cache.insert(i, i * 4096, dummy_block(1024));
}
let populated = cache.populated_shards();
assert_eq!(cache.num_shards(), 64);
assert!(
populated > 16,
"expected inserts to fan out across shards, populated = {populated}"
);
}
#[test]
fn evict_file_removes_only_that_files_blocks() {
let cache = BlockCache::with_config(64 * 1024 * 1024, 6, false);
for off in 0..16u64 {
cache.insert(1, off * 4096, dummy_block(1024));
cache.insert(2, off * 4096, dummy_block(1024));
}
cache.evict_file(1);
for off in 0..16u64 {
assert!(cache.get(1, off * 4096).is_none());
assert!(cache.get(2, off * 4096).is_some());
}
}
#[test]
fn clear_zeroes_usage() {
let cache = BlockCache::with_config(64 * 1024 * 1024, 6, false);
for i in 0..64u64 {
cache.insert(1, i * 4096, dummy_block(1024));
}
assert!(cache.usage() > 0);
cache.clear();
assert_eq!(cache.usage(), 0);
assert!(cache.get(1, 0).is_none());
}
#[test]
fn repeated_insert_at_same_key_does_not_double_count() {
let cache = BlockCache::with_config(64 * 1024, 0, false);
cache.insert(1, 0, dummy_block(1024));
let first_usage = cache.usage();
cache.insert(1, 0, dummy_block(1024));
cache.insert(1, 0, dummy_block(1024));
let final_usage = cache.usage();
assert_eq!(first_usage, final_usage);
}
#[test]
fn miss_on_absent_key_returns_none() {
let cache = BlockCache::with_config(64 * 1024, 0, false);
assert!(cache.get(99, 999).is_none());
}
#[test]
fn capacity_reflects_rounded_budget() {
let cache = BlockCache::with_config(100_000, 6, false);
assert!(cache.capacity() <= 100_000);
assert!(cache.capacity() > 0);
}
#[test]
fn resident_bytes_track_the_byte_budget_not_the_shard_count() {
let budget = 8 * 1024 * 1024;
let mut usages = Vec::new();
for bits in [0u32, 2, 4, 6] {
let cache = BlockCache::with_config(budget, bits, false);
assert_eq!(cache.usage(), 0, "a fresh cache holds nothing");
for i in 0..8192u64 {
cache.insert(1, i * 4096, dummy_block(4096));
}
assert!(
cache.usage() <= cache.capacity(),
"shard_bits {bits}: usage {} over capacity {}",
cache.usage(),
cache.capacity()
);
usages.push(cache.usage());
}
let spread =
usages.iter().max().copied().unwrap_or(0) - usages.iter().min().copied().unwrap_or(0);
assert!(
spread <= budget / 16,
"resident bytes moved with shard_bits: {usages:?}"
);
}
#[test]
fn per_entry_overhead_is_charged_against_the_budget() {
let cache = BlockCache::with_config(1024 * 1024, 0, false);
for i in 0..100_000u64 {
cache.insert(1, i * 64, dummy_block(0));
}
assert!(cache.usage() <= cache.capacity());
assert!(
cache.entry_count() <= cache.capacity() / ENTRY_OVERHEAD,
"held {} entries against a {}-byte budget",
cache.entry_count(),
cache.capacity()
);
}
#[test]
fn a_working_set_that_fits_the_budget_is_kept_whole() {
let cache = BlockCache::with_config(8 * 1024 * 1024, 0, false);
let mut offered = 0usize;
for i in 0..3500u64 {
let blk = dummy_block(1024);
offered += entry_charge(&CacheEntry::Data(Arc::clone(&blk)));
cache.insert(1, i * 4096, blk);
}
assert!(
offered <= cache.capacity(),
"test setup: the working set must fit the byte budget"
);
assert_eq!(
cache.entry_count(),
3500,
"the cache evicted entries that fit inside its byte budget"
);
assert_eq!(cache.usage(), offered);
}
#[test]
fn zero_budget_disables_the_cache() {
let cache = BlockCache::with_config(0, 6, false);
assert_eq!(cache.num_shards(), 0);
assert_eq!(cache.capacity(), 0);
cache.insert(1, 0, dummy_block(256));
assert!(cache.get(1, 0).is_none());
assert_eq!(cache.usage(), 0);
cache.evict_file(1);
cache.clear();
assert_eq!(cache.usage(), 0);
}
#[test]
fn zero_budget_strict_cache_is_also_disabled() {
let cache = BlockCache::with_config(0, 0, true);
cache.insert(1, 0, dummy_block(256));
assert!(cache.get(1, 0).is_none());
assert_eq!(cache.usage(), 0);
}
#[test]
fn tiny_budget_still_admits_a_block_that_fits() {
let cache = BlockCache::with_config(4096, 6, false);
cache.insert(1, 0, dummy_block(128));
assert!(cache.get(1, 0).is_some());
assert!(cache.usage() <= cache.capacity());
}
#[test]
fn byte_accounting_is_exact() {
let cache = BlockCache::with_config(64 * 1024 * 1024, 0, false);
let mut expected = 0usize;
for i in 0..64u64 {
let blk = dummy_block(512);
expected += entry_charge(&CacheEntry::Data(Arc::clone(&blk)));
cache.insert(1, i * 4096, blk);
}
assert_eq!(cache.usage(), expected);
}
#[test]
fn usage_does_not_drift_when_clear_races_insert() {
use std::sync::atomic::AtomicBool;
for _ in 0..50 {
let cache = Arc::new(BlockCache::with_config(64 * 1024 * 1024, 6, false));
let stop = Arc::new(AtomicBool::new(false));
let writer = {
let cache = Arc::clone(&cache);
let stop = Arc::clone(&stop);
std::thread::spawn(move || {
let mut i = 0u64;
while !stop.load(Ordering::Relaxed) {
cache.insert(i % 97, i * 4096, dummy_block(256));
i += 1;
}
})
};
for _ in 0..300 {
cache.clear();
}
stop.store(true, Ordering::Relaxed);
writer.join().expect("writer");
assert_eq!(
cache.usage(),
cache.true_usage(),
"usage() drifted away from the real byte total"
);
}
}
#[test]
fn concurrent_inserts_respect_the_budget() {
let cache = Arc::new(BlockCache::with_config(1024 * 1024, 2, false));
let mut handles = Vec::new();
for t in 0..8u64 {
let cache = Arc::clone(&cache);
handles.push(std::thread::spawn(move || {
for i in 0..4000u64 {
cache.insert(t, i * 64, dummy_block(64));
let _ = cache.get(t, (i / 2) * 64);
}
}));
}
for h in handles {
h.join().expect("worker");
}
assert!(
cache.true_usage() <= cache.capacity(),
"usage {} over capacity {}",
cache.true_usage(),
cache.capacity()
);
}
#[test]
fn evict_file_does_not_touch_other_files() {
let cache = BlockCache::with_config(64 * 1024 * 1024, 6, false);
cache.insert(7, 0, dummy_block(1024));
cache.insert(8, 0, dummy_block(1024));
let before = cache.usage();
cache.evict_file(99); assert_eq!(cache.usage(), before);
assert!(cache.get(7, 0).is_some());
assert!(cache.get(8, 0).is_some());
}
}
#[cfg(test)]
#[path = "block_cache_adversarial.rs"]
mod adversarial;