use std::{collections::HashMap, sync::Arc};
use reifydb_core::{common::CommitVersion, encoded::key::EncodedKey};
use reifydb_runtime::sync::mutex::Mutex;
use reifydb_value::util::cowvec::CowVec;
use crate::tier::VersionedGetResult;
#[derive(Clone)]
struct CacheEntry {
version: CommitVersion,
value: Option<CowVec<u8>>,
seq: u64,
}
struct Inner {
entries: HashMap<EncodedKey, CacheEntry>,
next_seq: u64,
capacity: usize,
}
#[derive(Clone)]
pub struct MultiReadBufferTier {
inner: Arc<Mutex<Inner>>,
}
impl MultiReadBufferTier {
pub fn new(capacity: usize) -> Self {
Self {
inner: Arc::new(Mutex::new(Inner {
entries: HashMap::new(),
next_seq: 0,
capacity: capacity.max(1),
})),
}
}
pub fn set_capacity(&self, capacity: usize) {
let mut inner = self.inner.lock();
inner.capacity = capacity.max(1);
inner.entries.clear();
}
pub fn get(&self, key: &EncodedKey, version: CommitVersion) -> VersionedGetResult {
let mut inner = self.inner.lock();
let Some(entry) = inner.entries.get(key) else {
return VersionedGetResult::NotFound;
};
if entry.version > version {
return VersionedGetResult::NotFound;
}
let stored_version = entry.version;
let value = entry.value.clone();
let seq = inner.next_seq;
inner.next_seq += 1;
if let Some(e) = inner.entries.get_mut(key) {
e.seq = seq;
}
match value {
Some(v) => VersionedGetResult::Value {
value: v,
version: stored_version,
},
None => VersionedGetResult::Tombstone,
}
}
pub fn insert(&self, key: EncodedKey, version: CommitVersion, value: Option<CowVec<u8>>) {
let mut inner = self.inner.lock();
match inner.entries.get(&key) {
Some(existing) if existing.version > version => return,
_ => {}
}
let seq = inner.next_seq;
inner.next_seq += 1;
inner.entries.insert(
key,
CacheEntry {
version,
value,
seq,
},
);
while inner.entries.len() > inner.capacity {
let Some(oldest) = inner.entries.iter().min_by_key(|(_, e)| e.seq).map(|(k, _)| k.clone())
else {
break;
};
inner.entries.remove(&oldest);
}
}
pub fn invalidate(&self, key: &EncodedKey) {
self.inner.lock().entries.remove(key);
}
pub fn clear(&self) {
self.inner.lock().entries.clear();
}
#[cfg(test)]
pub fn len(&self) -> usize {
self.inner.lock().entries.len()
}
}
#[cfg(test)]
mod tests {
use super::*;
fn key(s: &str) -> EncodedKey {
EncodedKey::new(s.as_bytes().to_vec())
}
fn val(s: &str) -> CowVec<u8> {
CowVec::new(s.as_bytes().to_vec())
}
#[test]
fn insert_then_get_returns_value_when_version_high_enough() {
let read = MultiReadBufferTier::new(8);
read.insert(key("k"), CommitVersion(5), Some(val("v5")));
match read.get(&key("k"), CommitVersion(5)) {
VersionedGetResult::Value {
value: v,
version: ver,
} => {
assert_eq!(v.as_ref(), b"v5");
assert_eq!(ver, CommitVersion(5));
}
VersionedGetResult::Tombstone => panic!("expected value, got tombstone"),
VersionedGetResult::NotFound => panic!("expected hit at exactly stored version"),
}
assert!(matches!(read.get(&key("k"), CommitVersion(9)), VersionedGetResult::Value { .. }));
}
#[test]
fn get_below_stored_version_misses_so_caller_reads_through() {
let read = MultiReadBufferTier::new(8);
read.insert(key("k"), CommitVersion(5), Some(val("v5")));
assert!(
matches!(read.get(&key("k"), CommitVersion(4)), VersionedGetResult::NotFound),
"must miss below the stored version"
);
}
#[test]
fn tombstone_is_cached_and_served() {
let read = MultiReadBufferTier::new(8);
read.insert(key("k"), CommitVersion(3), None);
assert!(matches!(read.get(&key("k"), CommitVersion(3)), VersionedGetResult::Tombstone));
}
#[test]
fn invalidate_removes_the_key() {
let read = MultiReadBufferTier::new(8);
read.insert(key("k"), CommitVersion(1), Some(val("v1")));
read.invalidate(&key("k"));
assert!(
matches!(read.get(&key("k"), CommitVersion(1)), VersionedGetResult::NotFound),
"invalidated key must miss"
);
}
#[test]
fn newer_insert_overwrites_but_older_insert_is_ignored() {
let read = MultiReadBufferTier::new(8);
read.insert(key("k"), CommitVersion(5), Some(val("v5")));
read.insert(key("k"), CommitVersion(2), Some(val("v2")));
match read.get(&key("k"), CommitVersion(5)) {
VersionedGetResult::Value {
value: v,
..
} => assert_eq!(v.as_ref(), b"v5", "older insert must not overwrite"),
VersionedGetResult::Tombstone => panic!("unexpected tombstone"),
VersionedGetResult::NotFound => panic!("unexpected miss"),
}
read.insert(key("k"), CommitVersion(7), Some(val("v7")));
match read.get(&key("k"), CommitVersion(7)) {
VersionedGetResult::Value {
value: v,
version: ver,
} => {
assert_eq!(v.as_ref(), b"v7");
assert_eq!(ver, CommitVersion(7));
}
VersionedGetResult::Tombstone => panic!("unexpected tombstone"),
VersionedGetResult::NotFound => panic!("unexpected miss"),
}
}
#[test]
fn eviction_bounds_size_and_never_changes_correctness() {
let read = MultiReadBufferTier::new(2);
read.insert(key("a"), CommitVersion(1), Some(val("a")));
read.insert(key("b"), CommitVersion(1), Some(val("b")));
read.insert(key("c"), CommitVersion(1), Some(val("c")));
assert!(read.len() <= 2, "read buffer must stay within capacity");
for k in ["a", "b", "c"] {
if let VersionedGetResult::Value {
value: v,
..
} = read.get(&key(k), CommitVersion(1))
{
assert_eq!(v.as_ref(), k.as_bytes());
}
}
}
#[test]
fn clone_shares_backing_storage() {
let a = MultiReadBufferTier::new(4);
let b = a.clone();
a.insert(key("k"), CommitVersion(1), Some(val("v")));
assert!(
matches!(b.get(&key("k"), CommitVersion(1)), VersionedGetResult::Value { .. }),
"clone observes writes from the original"
);
}
}