use std::collections::HashMap;
use std::sync::Arc;
#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash, PartialOrd, Ord)]
pub struct BlockKey {
pub sstable_id: u64,
pub block_offset: u64,
}
impl BlockKey {
pub fn new(sstable_id: u64, block_offset: u64) -> Self {
Self {
sstable_id,
block_offset,
}
}
}
pub type Block = Arc<[u8]>;
pub trait BlockCache: Send + Sync {
fn get(&self, key: &BlockKey) -> Option<Block>;
fn put(&self, key: BlockKey, block: Block);
fn len(&self) -> usize;
fn is_empty(&self) -> bool {
self.len() == 0
}
fn clear(&self);
}
pub struct LruBlockCache {
inner: std::sync::Mutex<LruInner>,
}
struct LruInner {
capacity: usize,
map: HashMap<BlockKey, usize>,
nodes: Vec<Node>,
head: Option<usize>,
tail: Option<usize>,
free: Vec<usize>,
hits: u64,
misses: u64,
}
struct Node {
key: BlockKey,
block: Block,
prev: Option<usize>,
next: Option<usize>,
}
impl LruBlockCache {
pub fn new(capacity: usize) -> Self {
let cap = capacity.max(1);
Self {
inner: std::sync::Mutex::new(LruInner {
capacity: cap,
map: HashMap::with_capacity(cap),
nodes: Vec::with_capacity(cap),
head: None,
tail: None,
free: Vec::new(),
hits: 0,
misses: 0,
}),
}
}
pub fn capacity(&self) -> usize {
self.inner.lock().unwrap().capacity
}
pub fn hits(&self) -> u64 {
self.inner.lock().unwrap().hits
}
pub fn misses(&self) -> u64 {
self.inner.lock().unwrap().misses
}
}
impl BlockCache for LruBlockCache {
fn get(&self, key: &BlockKey) -> Option<Block> {
let mut g = self.inner.lock().unwrap();
match g.map.get(key).copied() {
Some(idx) => {
let block = g.nodes[idx].block.clone();
g.move_to_front(idx);
g.hits += 1;
Some(block)
}
None => {
g.misses += 1;
None
}
}
}
fn put(&self, key: BlockKey, block: Block) {
let mut g = self.inner.lock().unwrap();
if let Some(&idx) = g.map.get(&key) {
g.nodes[idx].block = block;
g.move_to_front(idx);
return;
}
if g.map.len() >= g.capacity {
if let Some(tail_idx) = g.tail {
let tail_key = g.nodes[tail_idx].key;
g.detach(tail_idx);
g.map.remove(&tail_key);
g.free.push(tail_idx);
}
}
let idx = if let Some(slot) = g.free.pop() {
g.nodes[slot] = Node {
key,
block,
prev: None,
next: None,
};
slot
} else {
g.nodes.push(Node {
key,
block,
prev: None,
next: None,
});
g.nodes.len() - 1
};
g.map.insert(key, idx);
g.push_front(idx);
}
fn len(&self) -> usize {
self.inner.lock().unwrap().map.len()
}
fn clear(&self) {
let mut g = self.inner.lock().unwrap();
g.map.clear();
g.nodes.clear();
g.free.clear();
g.head = None;
g.tail = None;
}
}
impl LruInner {
fn push_front(&mut self, idx: usize) {
self.nodes[idx].prev = None;
self.nodes[idx].next = self.head;
if let Some(h) = self.head {
self.nodes[h].prev = Some(idx);
}
self.head = Some(idx);
if self.tail.is_none() {
self.tail = Some(idx);
}
}
fn detach(&mut self, idx: usize) {
let prev = self.nodes[idx].prev;
let next = self.nodes[idx].next;
if let Some(p) = prev {
self.nodes[p].next = next;
} else {
self.head = next;
}
if let Some(n) = next {
self.nodes[n].prev = prev;
} else {
self.tail = prev;
}
self.nodes[idx].prev = None;
self.nodes[idx].next = None;
}
fn move_to_front(&mut self, idx: usize) {
if self.head == Some(idx) {
return;
}
self.detach(idx);
self.push_front(idx);
}
}
#[cfg(test)]
mod tests {
use super::*;
fn block(bytes: &[u8]) -> Block {
Arc::from(bytes.to_vec().into_boxed_slice())
}
#[test]
fn miss_then_hit() {
let c = LruBlockCache::new(4);
let k = BlockKey::new(1, 0);
assert!(c.get(&k).is_none());
c.put(k, block(b"payload"));
let got = c.get(&k).unwrap();
assert_eq!(&*got, b"payload");
assert_eq!(c.hits(), 1);
assert_eq!(c.misses(), 1);
}
#[test]
fn lru_evicts_coldest() {
let c = LruBlockCache::new(2);
let a = BlockKey::new(1, 0);
let b = BlockKey::new(2, 0);
let z = BlockKey::new(3, 0);
c.put(a, block(b"A"));
c.put(b, block(b"B"));
c.get(&a);
c.put(z, block(b"Z"));
assert!(c.get(&a).is_some(), "recently-used 'a' stays");
assert!(c.get(&b).is_none(), "stale 'b' evicted");
assert!(c.get(&z).is_some(), "newest 'z' kept");
}
#[test]
fn put_of_existing_key_refreshes_value() {
let c = LruBlockCache::new(2);
let k = BlockKey::new(1, 0);
c.put(k, block(b"old"));
c.put(k, block(b"new"));
assert_eq!(&*c.get(&k).unwrap(), b"new");
assert_eq!(c.len(), 1);
}
#[test]
fn clear_drops_everything() {
let c = LruBlockCache::new(4);
c.put(BlockKey::new(1, 0), block(b"x"));
c.put(BlockKey::new(2, 0), block(b"y"));
assert_eq!(c.len(), 2);
c.clear();
assert_eq!(c.len(), 0);
assert!(c.is_empty());
}
#[test]
fn capacity_floor_is_one() {
let c = LruBlockCache::new(0);
assert_eq!(c.capacity(), 1);
c.put(BlockKey::new(1, 0), block(b"a"));
c.put(BlockKey::new(2, 0), block(b"b"));
assert_eq!(c.len(), 1, "cap-1 cache holds the latest only");
assert!(c.get(&BlockKey::new(2, 0)).is_some());
assert!(c.get(&BlockKey::new(1, 0)).is_none());
}
#[test]
fn block_arc_clone_is_shared() {
let c = LruBlockCache::new(2);
let k = BlockKey::new(1, 100);
c.put(k, block(b"shared"));
let h1 = c.get(&k).unwrap();
let h2 = c.get(&k).unwrap();
assert!(
Arc::ptr_eq(&h1, &h2),
"cache hits hand out Arc clones, not copies"
);
}
#[test]
fn hits_and_misses_are_counted() {
let c = LruBlockCache::new(4);
let k = BlockKey::new(1, 0);
c.put(k, block(b"v"));
c.get(&k);
c.get(&k);
c.get(&BlockKey::new(99, 0));
assert_eq!(c.hits(), 2);
assert_eq!(c.misses(), 1);
}
#[test]
fn distinct_keys_share_cache_when_under_capacity() {
let c = LruBlockCache::new(4);
for i in 0..4 {
c.put(BlockKey::new(i, 0), block(&[i as u8]));
}
for i in 0..4 {
assert!(c.get(&BlockKey::new(i, 0)).is_some());
}
assert_eq!(c.len(), 4);
}
}