use std::borrow::Borrow;
use std::cmp::Ordering;
use std::collections::{BTreeSet, HashMap};
use std::hash::{Hash, Hasher};
use std::ops::Bound::{Excluded, Included};
use std::sync::Arc;
use std::time::{Duration, Instant};
#[derive(Eq)]
struct CacheArc<T>(Arc<T>);
impl<T> CacheArc<T> {
fn new(key: T) -> Self {
CacheArc(Arc::new(key))
}
}
impl<T> Clone for CacheArc<T> {
fn clone(&self) -> Self {
CacheArc(self.0.clone())
}
}
impl<T: PartialEq> PartialEq for CacheArc<T> {
fn eq(&self, other: &Self) -> bool {
self.0.eq(&other.0)
}
}
impl<T: PartialOrd> PartialOrd for CacheArc<T> {
fn partial_cmp(&self, other: &Self) -> Option<Ordering> {
self.0.partial_cmp(&other.0)
}
}
impl<T: Ord> Ord for CacheArc<T> {
fn cmp(&self, other: &Self) -> Ordering {
self.0.cmp(&other.0)
}
}
impl<T: Hash> Hash for CacheArc<T> {
fn hash<H: Hasher>(&self, state: &mut H) {
self.0.hash(state);
}
}
impl<T> Borrow<T> for CacheArc<T> {
fn borrow(&self) -> &T {
&self.0
}
}
impl Borrow<str> for CacheArc<String> {
fn borrow(&self) -> &str {
self.0.as_str()
}
}
impl<T> Borrow<[T]> for CacheArc<Vec<T>> {
fn borrow(&self) -> &[T] {
self.0.as_slice()
}
}
#[derive(Debug)]
pub enum Error {
TimeBounds,
}
#[derive(Hash, Eq, PartialEq, Ord, PartialOrd)]
struct Stamped<K> {
expiry: Instant,
key: Option<CacheArc<K>>,
}
impl<K> Clone for Stamped<K> {
fn clone(&self) -> Self {
Self {
expiry: self.expiry,
key: self.key.clone(),
}
}
}
impl<K> Stamped<K> {
fn bound(expiry: Instant) -> Stamped<K> {
Stamped { expiry, key: None }
}
}
struct Entry<K, V> {
expiry: Instant,
key: CacheArc<K>,
value: V,
}
impl<K, V> Entry<K, V> {
fn as_stamped(&self) -> Stamped<K> {
Stamped {
expiry: self.expiry,
key: Some(self.key.clone()),
}
}
fn is_expired(&self) -> bool {
self.expiry < Instant::now()
}
}
macro_rules! impl_get {
($_self:expr, $key:expr) => {{
$_self.map.get($key).and_then(|entry| {
if entry.is_expired() {
None
} else {
Some(&entry.value)
}
})
}};
}
pub struct ExpiringSizedCache<K, V> {
min_instant: Instant,
map: HashMap<CacheArc<K>, Entry<K, V>>,
keys: BTreeSet<Stamped<K>>,
pub ttl_millis: u64,
pub size_limit: Option<usize>,
}
impl<K: Hash + Eq + Ord, V> ExpiringSizedCache<K, V> {
pub fn new(ttl_millis: u64) -> Self {
Self {
min_instant: Instant::now(),
map: HashMap::new(),
keys: BTreeSet::new(),
ttl_millis,
size_limit: None,
}
}
pub fn with_capacity(ttl_millis: u64, size: usize) -> Self {
let mut new = Self::new(ttl_millis);
new.map.reserve(size);
new
}
pub fn size_limit(&mut self, size: usize) -> Option<usize> {
let prev = self.size_limit;
self.size_limit = Some(size);
prev
}
pub fn reserve(&mut self, more: usize) {
self.map.reserve(more);
}
pub fn ttl_millis(&mut self, ttl_millis: u64) -> u64 {
let prev = self.ttl_millis;
self.ttl_millis = ttl_millis;
prev
}
pub fn evict(&mut self) -> usize {
let cutoff = Instant::now();
let min = Stamped::bound(self.min_instant);
let max = Stamped::bound(cutoff);
let min = Included(&min);
let max = Excluded(&max);
let remove = self.keys.range((min, max)).count();
let mut count = 0;
while count < remove {
match self.keys.pop_first() {
None => break,
Some(stamped) => {
self.map.remove(
&stamped
.key
.expect("evicting: only artificial bounds are none"),
);
count += 1;
}
}
}
count
}
pub fn retain_latest(&mut self, count: usize, evict: bool) -> usize {
let retain_drop_count = self.len().saturating_sub(count);
let remove = if evict {
let cutoff = Instant::now();
let min = Stamped::bound(self.min_instant);
let max = Stamped::bound(cutoff);
let min = Included(&min);
let max = Excluded(&max);
let to_evict_count = self.keys.range((min, max)).count();
retain_drop_count.max(to_evict_count)
} else {
retain_drop_count
};
let mut count = 0;
while count < remove {
match self.keys.pop_first() {
None => break,
Some(stamped) => {
self.map.remove(
&stamped
.key
.expect("retaining: only artificial bounds are none"),
);
count += 1;
}
}
}
count
}
pub fn remove(&mut self, key: &K) -> Option<V> {
match self.map.remove(key) {
None => None,
Some(removed) => {
self.keys.remove(&removed.as_stamped());
if removed.is_expired() {
None
} else {
Some(removed.value)
}
}
}
}
pub fn insert(&mut self, key: K, value: V) -> Result<Option<V>, Error> {
self.insert_ttl_evict(key, value, None, false)
}
pub fn insert_ttl(&mut self, key: K, value: V, ttl_millis: u64) -> Result<Option<V>, Error> {
self.insert_ttl_evict(key, value, Some(ttl_millis), false)
}
pub fn insert_evict(&mut self, key: K, value: V, evict: bool) -> Result<Option<V>, Error> {
self.insert_ttl_evict(key, value, None, evict)
}
pub fn insert_ttl_evict(
&mut self,
key: K,
value: V,
ttl_millis: Option<u64>,
evict: bool,
) -> Result<Option<V>, Error> {
if let Some(size_limit) = self.size_limit {
if self.len() > size_limit - 1 {
self.retain_latest(size_limit - 1, evict);
}
} else if evict {
self.evict();
}
let key = CacheArc::new(key);
let expiry = Instant::now()
.checked_add(Duration::from_millis(ttl_millis.unwrap_or(self.ttl_millis)))
.ok_or(Error::TimeBounds)?;
let new_stamped = Stamped {
expiry,
key: Some(key.clone()),
};
self.keys.insert(new_stamped.clone());
let old = self.map.insert(key.clone(), Entry { expiry, key, value });
if let Some(old) = &old {
let old_stamped = old.as_stamped();
if old_stamped != new_stamped {
self.keys.remove(&old_stamped);
}
}
Ok(old.and_then(|entry| {
if entry.is_expired() {
None
} else {
Some(entry.value)
}
}))
}
pub fn clear(&mut self) {
self.map.clear();
self.keys.clear();
}
pub fn len(&self) -> usize {
self.map.len()
}
pub fn is_empty(&self) -> bool {
self.len() == 0
}
pub fn get(&self, key: &K) -> Option<&V> {
impl_get!(self, key)
}
}
impl<V> ExpiringSizedCache<String, V> {
pub fn get_borrowed(&self, key: &str) -> Option<&V> {
impl_get!(self, key)
}
}
impl<T: Hash + Eq + PartialEq, V> ExpiringSizedCache<Vec<T>, V> {
pub fn get_borrowed(&self, key: &[T]) -> Option<&V> {
impl_get!(self, key)
}
}
#[cfg(test)]
mod test {
use crate::stores::ExpiringSizedCache;
use std::time::Duration;
#[test]
fn borrow_keys() {
let mut cache = ExpiringSizedCache::with_capacity(100, 100);
cache.insert(String::from("a"), "a").unwrap();
assert_eq!(cache.get_borrowed("a").unwrap(), &"a");
let mut cache = ExpiringSizedCache::with_capacity(100, 100);
cache.insert(vec![0], "a").unwrap();
assert_eq!(cache.get_borrowed(&[0]).unwrap(), &"a");
}
#[test]
fn kitchen_sink() {
let mut cache = ExpiringSizedCache::with_capacity(100, 100);
assert_eq!(0, cache.evict());
assert_eq!(0, cache.retain_latest(100, true));
assert!(cache.get(&"a".into()).is_none());
cache.insert("a".to_string(), "A".to_string()).unwrap();
assert_eq!(cache.get(&"a".into()), Some("A".to_string()).as_ref());
assert_eq!(cache.len(), 1);
std::thread::sleep(Duration::from_millis(200));
assert_eq!(1, cache.evict());
assert!(cache.get(&"a".into()).is_none());
assert_eq!(cache.len(), 0);
cache.insert("a".to_string(), "A".to_string()).unwrap();
assert_eq!(cache.get(&"a".into()), Some("A".to_string()).as_ref());
assert_eq!(cache.len(), 1);
std::thread::sleep(Duration::from_millis(200));
assert_eq!(0, cache.retain_latest(1, false));
assert_eq!(cache.get(&"a".into()), None);
assert_eq!(cache.len(), 1);
assert_eq!(1, cache.retain_latest(1, true));
assert!(cache.get(&"a".into()).is_none());
assert_eq!(cache.len(), 0);
cache.insert("a".to_string(), "a".to_string()).unwrap();
cache.insert("b".to_string(), "b".to_string()).unwrap();
cache.insert("c".to_string(), "c".to_string()).unwrap();
cache.insert("d".to_string(), "d".to_string()).unwrap();
cache.insert("e".to_string(), "e".to_string()).unwrap();
assert_eq!(3, cache.retain_latest(2, false));
assert_eq!(2, cache.len());
assert_eq!(cache.get(&"a".into()), None);
assert_eq!(cache.get(&"b".into()), None);
assert_eq!(cache.get(&"c".into()), None);
assert_eq!(cache.get(&"d".into()), Some("d".to_string()).as_ref());
assert_eq!(cache.get(&"e".into()), Some("e".to_string()).as_ref());
cache.insert("a".to_string(), "a".to_string()).unwrap();
cache.insert("a".to_string(), "a".to_string()).unwrap();
cache.insert("b".to_string(), "b".to_string()).unwrap();
cache.insert("b".to_string(), "b".to_string()).unwrap();
assert_eq!(4, cache.len());
assert_eq!(2, cache.retain_latest(2, false));
assert_eq!(cache.get(&"d".into()), None);
assert_eq!(cache.get(&"e".into()), None);
assert_eq!(cache.get(&"a".into()), Some("a".to_string()).as_ref());
assert_eq!(cache.get(&"b".into()), Some("b".to_string()).as_ref());
assert_eq!(2, cache.len());
std::thread::sleep(Duration::from_millis(200));
assert_eq!(cache.remove(&"a".into()), None);
assert_eq!(1, cache.len());
cache.insert("a".to_string(), "a".to_string()).unwrap();
assert_eq!(cache.remove(&"a".into()), Some("a".to_string()));
assert_eq!(1, cache.len());
assert_eq!(1, cache.evict());
assert_eq!(0, cache.len());
cache
.insert_ttl("a".to_string(), "a".to_string(), 300)
.unwrap();
std::thread::sleep(Duration::from_millis(200));
assert_eq!(cache.get(&"a".into()), Some("a".to_string()).as_ref());
assert_eq!(1, cache.len());
std::thread::sleep(Duration::from_millis(200));
cache
.insert_ttl_evict("b".to_string(), "b".to_string(), Some(300), true)
.unwrap();
assert_eq!(1, cache.len());
assert_eq!(cache.get_borrowed("a"), None);
}
#[test]
fn size_limit() {
let mut cache = ExpiringSizedCache::with_capacity(100, 100);
cache.size_limit(2);
assert_eq!(0, cache.evict());
assert_eq!(0, cache.retain_latest(100, true));
assert!(cache.get(&"a".into()).is_none());
cache.insert("a".to_string(), "A".to_string()).unwrap();
assert_eq!(cache.get(&"a".into()), Some("A".to_string()).as_ref());
assert_eq!(cache.len(), 1);
cache.insert("b".to_string(), "B".to_string()).unwrap();
assert_eq!(cache.get(&"b".into()), Some("B".to_string()).as_ref());
assert_eq!(cache.len(), 2);
cache.insert("c".to_string(), "C".to_string()).unwrap();
assert_eq!(cache.len(), 2);
assert_eq!(cache.get(&"b".into()), Some("B".to_string()).as_ref());
assert_eq!(cache.get(&"c".into()), Some("C".to_string()).as_ref());
assert_eq!(cache.get(&"a".into()), None);
}
}