use std::collections::{HashMap, VecDeque};
use std::sync::atomic::{AtomicBool, Ordering};
#[derive(Default)]
pub struct LRUCache<T> {
cache: HashMap<u64, (T, AtomicBool, u32)>,
lru_queue: VecDeque<(u64, u32)>,
next_seq: u32,
}
impl<T> LRUCache<T> {
pub(crate) fn new() -> Self {
Self {
cache: Default::default(),
lru_queue: Default::default(),
next_seq: 0,
}
}
pub(crate) fn len(&self) -> usize {
self.cache.len()
}
pub(crate) fn insert(&mut self, key: u64, value: T) -> Option<T> {
if let Some((existing, second_chance, _)) = self.cache.get_mut(&key) {
second_chance.store(false, Ordering::Release);
return Some(std::mem::replace(existing, value));
}
let seq = self.next_seq;
self.next_seq = self.next_seq.wrapping_add(1);
self.cache.insert(key, (value, AtomicBool::new(false), seq));
self.lru_queue.push_back((key, seq));
None
}
pub(crate) fn remove(&mut self, key: u64) -> Option<T> {
let (value, _, _) = self.cache.remove(&key)?;
if self.lru_queue.len() > 2 * self.cache.len() {
for _ in 0..2 {
if let Some((removed_key, seq)) = self.lru_queue.pop_front() {
if let Some(second_chance) = self.live_entry(removed_key, seq) {
second_chance.store(false, Ordering::Release);
self.lru_queue.push_back((removed_key, seq));
}
}
}
}
Some(value)
}
fn live_entry(&self, key: u64, seq: u32) -> Option<&AtomicBool> {
self.cache
.get(&key)
.filter(|(_, _, live_seq)| *live_seq == seq)
.map(|(_, second_chance, _)| second_chance)
}
pub(crate) fn get(&self, key: u64) -> Option<&T> {
if let Some((value, second_chance, _)) = self.cache.get(&key) {
second_chance.store(true, Ordering::Release);
Some(value)
} else {
None
}
}
pub(crate) fn get_mut(&mut self, key: u64) -> Option<&mut T> {
if let Some((value, second_chance, _)) = self.cache.get_mut(&key) {
second_chance.store(true, Ordering::Release);
Some(value)
} else {
None
}
}
pub(crate) fn iter(&self) -> impl ExactSizeIterator<Item = (&u64, &T)> {
self.cache.iter().map(|(k, (v, _, _))| (k, v))
}
pub(crate) fn iter_mut(&mut self) -> impl ExactSizeIterator<Item = (&u64, &mut T)> {
self.cache.iter_mut().map(|(k, (v, _, _))| (k, v))
}
pub(crate) fn pop_lowest_priority(&mut self) -> Option<(u64, T)> {
while let Some((key, seq)) = self.lru_queue.pop_front() {
if let Some(second_chance) = self.live_entry(key, seq) {
if second_chance
.compare_exchange(true, false, Ordering::AcqRel, Ordering::Acquire)
.is_ok()
{
self.lru_queue.push_back((key, seq));
} else {
let (value, _, _) = self.cache.remove(&key).unwrap();
return Some((key, value));
}
}
}
None
}
pub(crate) fn clear(&mut self) {
self.cache.clear();
self.lru_queue.clear();
}
}
#[cfg(test)]
mod test {
use super::LRUCache;
#[test]
fn queue_stays_bounded_without_eviction() {
const KEYS: u64 = 64;
let mut cache: LRUCache<u64> = LRUCache::new();
for k in 0..KEYS {
cache.insert(k, k);
}
for _ in 0..5_000 {
for k in 0..KEYS {
cache.remove(k);
cache.insert(k, k);
}
}
let keys = usize::try_from(KEYS).unwrap();
assert_eq!(cache.len(), keys);
assert!(cache.lru_queue.len() <= 3 * keys);
let mut evicted = vec![];
while let Some((k, _)) = cache.pop_lowest_priority() {
evicted.push(k);
}
evicted.sort_unstable();
assert_eq!(evicted, (0..KEYS).collect::<Vec<_>>());
assert_eq!(cache.len(), 0);
assert!(cache.lru_queue.is_empty());
}
#[test]
fn queue_stays_bounded_with_hot_key() {
const KEYS: u64 = 1024;
let mut cache: LRUCache<u64> = LRUCache::new();
for k in 0..KEYS {
cache.insert(k, k);
}
for _ in 0..100_000 {
cache.remove(KEYS - 1);
cache.insert(KEYS - 1, KEYS - 1);
}
let keys = usize::try_from(KEYS).unwrap();
assert_eq!(cache.len(), keys);
assert!(cache.lru_queue.len() <= 3 * keys);
}
#[test]
fn reinsert_moves_key_to_the_back() {
let mut cache: LRUCache<u64> = LRUCache::new();
for k in 1..=3 {
cache.insert(k, k);
}
cache.remove(1);
cache.insert(1, 1);
let mut evicted = vec![];
while let Some((k, _)) = cache.pop_lowest_priority() {
evicted.push(k);
}
assert_eq!(evicted, vec![2, 3, 1]);
}
#[test]
fn second_chance_survives_one_eviction_pass() {
let mut cache: LRUCache<u64> = LRUCache::new();
for k in 1..=3 {
cache.insert(k, k);
}
assert_eq!(cache.get(1), Some(&1));
assert_eq!(cache.pop_lowest_priority(), Some((2, 2)));
assert_eq!(cache.pop_lowest_priority(), Some((3, 3)));
assert_eq!(cache.pop_lowest_priority(), Some((1, 1)));
assert_eq!(cache.pop_lowest_priority(), None);
}
#[test]
fn replacing_a_value_keeps_its_queue_position() {
let mut cache: LRUCache<u64> = LRUCache::new();
cache.insert(1, 1);
cache.insert(2, 2);
assert_eq!(cache.insert(1, 10), Some(1));
assert_eq!(cache.lru_queue.len(), 2);
assert_eq!(cache.pop_lowest_priority(), Some((1, 10)));
assert_eq!(cache.pop_lowest_priority(), Some((2, 2)));
}
}