use std::collections::HashMap;
use std::sync::Mutex;
use std::time::{Duration, Instant};
struct CacheEntry<T> {
value: T,
expires_at: Instant,
}
pub struct QueryCache<T> {
max_size: usize,
ttl: Duration,
store: Mutex<HashMap<String, CacheEntry<T>>>,
}
impl<T: Clone> QueryCache<T> {
pub fn new(max_size: usize, ttl: Duration) -> Self {
Self {
max_size,
ttl,
store: Mutex::new(HashMap::new()),
}
}
pub fn get(&self, key: &str) -> Option<T> {
let mut store = self.store.lock().unwrap();
let entry = store.get(key)?;
if Instant::now() > entry.expires_at {
store.remove(key);
return None;
}
Some(entry.value.clone())
}
pub fn set(&self, key: String, value: T) {
let mut store = self.store.lock().unwrap();
if store.len() >= self.max_size
&& !store.contains_key(&key)
&& let Some(oldest_key) = store.keys().next().cloned()
{
store.remove(&oldest_key);
}
store.insert(
key,
CacheEntry {
value,
expires_at: Instant::now() + self.ttl,
},
);
}
pub fn clear(&self) {
self.store.lock().unwrap().clear();
}
}
impl<T: Clone> Default for QueryCache<T> {
fn default() -> Self {
Self::new(500, Duration::from_secs(60))
}
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn returns_a_previously_set_value() {
let cache: QueryCache<String> = QueryCache::new(10, Duration::from_secs(60));
cache.set("key".to_string(), "value".to_string());
assert_eq!(cache.get("key"), Some("value".to_string()));
}
#[test]
fn expires_entries_past_their_ttl() {
let cache: QueryCache<String> = QueryCache::new(10, Duration::from_millis(1));
cache.set("key".to_string(), "value".to_string());
std::thread::sleep(Duration::from_millis(10));
assert_eq!(cache.get("key"), None);
}
#[test]
fn evicts_when_full() {
let cache: QueryCache<String> = QueryCache::new(1, Duration::from_secs(60));
cache.set("a".to_string(), "1".to_string());
cache.set("b".to_string(), "2".to_string());
let remaining = [cache.get("a"), cache.get("b")]
.into_iter()
.flatten()
.count();
assert_eq!(remaining, 1);
}
}