use std::collections::VecDeque;
use crate::Cache;
pub struct BoundedCache<C> {
inner: C,
max_bytes: usize,
current_bytes: usize,
order: VecDeque<Vec<u8>>,
}
impl<C> BoundedCache<C>
where
C: Cache<Vec<u8>, Vec<u8>>,
{
pub fn new(inner: C, max_bytes: usize) -> Self {
BoundedCache {
inner,
max_bytes,
current_bytes: 0,
order: VecDeque::new(),
}
}
#[must_use]
pub fn current_bytes(&self) -> usize {
self.current_bytes
}
#[must_use]
pub fn max_bytes(&self) -> usize {
self.max_bytes
}
fn evict_until_fits(&mut self, needed: usize) {
while self.current_bytes + needed > self.max_bytes {
let oldest = match self.order.pop_front() {
Some(k) => k,
None => break,
};
if let Some(old_val) = self.inner.remove(&oldest) {
let freed = oldest.len() + old_val.len();
self.current_bytes = self.current_bytes.saturating_sub(freed);
}
}
}
}
impl<C> Cache<Vec<u8>, Vec<u8>> for BoundedCache<C>
where
C: Cache<Vec<u8>, Vec<u8>>,
{
fn get(&mut self, key: &Vec<u8>) -> Option<&Vec<u8>> {
self.inner.get(key)
}
fn put(&mut self, key: Vec<u8>, value: Vec<u8>) -> Option<Vec<u8>> {
let entry_size = key.len() + value.len();
self.evict_until_fits(entry_size);
if let Some(existing) = self.inner.peek(&key) {
let old_size = key.len() + existing.len();
self.current_bytes = self.current_bytes.saturating_sub(old_size);
self.order.retain(|k| k != &key);
}
self.current_bytes += entry_size;
self.order.push_back(key.clone());
self.inner.put(key, value)
}
fn put_with_ttl(
&mut self,
key: Vec<u8>,
value: Vec<u8>,
ttl: std::time::Duration,
) -> Option<Vec<u8>> {
let entry_size = key.len() + value.len();
self.evict_until_fits(entry_size);
if let Some(existing) = self.inner.peek(&key) {
let old_size = key.len() + existing.len();
self.current_bytes = self.current_bytes.saturating_sub(old_size);
self.order.retain(|k| k != &key);
}
self.current_bytes += entry_size;
self.order.push_back(key.clone());
self.inner.put_with_ttl(key, value, ttl)
}
fn len(&self) -> usize {
self.inner.len()
}
fn cap(&self) -> usize {
self.inner.cap()
}
fn remove(&mut self, key: &Vec<u8>) -> Option<Vec<u8>> {
if let Some(val) = self.inner.remove(key) {
let freed = key.len() + val.len();
self.current_bytes = self.current_bytes.saturating_sub(freed);
self.order.retain(|k| k != key);
Some(val)
} else {
None
}
}
fn clear(&mut self) {
self.inner.clear();
self.current_bytes = 0;
self.order.clear();
}
fn peek(&self, key: &Vec<u8>) -> Option<&Vec<u8>> {
self.inner.peek(key)
}
fn contains_key(&self, key: &Vec<u8>) -> bool {
self.inner.contains_key(key)
}
fn resize(&mut self, new_cap: usize) {
self.inner.resize(new_cap);
}
}
#[cfg(test)]
mod tests {
use super::*;
use crate::LruCache;
fn make_bounded(max_bytes: usize, cap: usize) -> BoundedCache<LruCache<Vec<u8>, Vec<u8>>> {
BoundedCache::new(LruCache::new(cap), max_bytes)
}
#[test]
fn bounded_under_budget() {
let mut cache = make_bounded(100, 10);
cache.put(b"key1".to_vec(), b"val1".to_vec());
cache.put(b"key2".to_vec(), b"val2".to_vec());
assert!(cache.current_bytes() <= 100);
assert_eq!(cache.get(&b"key1".to_vec()), Some(&b"val1".to_vec()));
}
#[test]
fn bounded_evicts_when_over_budget() {
let mut cache = make_bounded(16, 100);
cache.put(b"key1".to_vec(), b"val1".to_vec()); cache.put(b"key2".to_vec(), b"val2".to_vec()); assert_eq!(cache.current_bytes(), 16);
cache.put(b"key3".to_vec(), b"val3".to_vec());
assert!(cache.current_bytes() <= 16);
assert!(cache.get(&b"key1".to_vec()).is_none());
}
#[test]
fn bounded_remove_updates_bytes() {
let mut cache = make_bounded(100, 10);
cache.put(b"hello".to_vec(), b"world".to_vec());
let before = cache.current_bytes();
cache.remove(&b"hello".to_vec());
assert_eq!(cache.current_bytes(), before - 10); }
#[test]
fn bounded_clear_resets_bytes() {
let mut cache = make_bounded(100, 10);
cache.put(b"a".to_vec(), b"b".to_vec());
cache.clear();
assert_eq!(cache.current_bytes(), 0);
}
#[test]
fn bounded_update_existing_key() {
let mut cache = make_bounded(100, 10);
cache.put(b"key".to_vec(), b"short".to_vec()); cache.put(b"key".to_vec(), b"longer_value".to_vec()); assert!(cache.current_bytes() <= 100);
assert_eq!(cache.get(&b"key".to_vec()), Some(&b"longer_value".to_vec()));
}
}