use std::{
collections::{HashMap, VecDeque},
future::Future,
hash::Hash,
sync::Arc,
time::{Duration, Instant},
};
use thiserror::Error;
use tokio::sync::Mutex;
#[derive(Clone, Copy, Debug, Eq, PartialEq)]
pub struct CacheConfig {
pub hard_ttl: Duration,
pub max_bytes: usize,
pub max_flights: usize,
}
#[derive(Clone, Copy, Debug, Eq, PartialEq)]
pub enum CacheOutcomeKind {
Hit,
Miss,
}
#[derive(Clone, Debug, Eq, PartialEq)]
pub struct CacheOutcome<V> {
pub value: V,
pub kind: CacheOutcomeKind,
}
#[derive(Clone, Copy, Debug, Eq, PartialEq)]
pub struct CacheSnapshot {
pub entries: usize,
pub bytes: usize,
pub flights: usize,
}
#[derive(Debug, Error)]
pub enum CacheError<E> {
#[error("cache flight limit reached")]
FlightLimit,
#[error(transparent)]
Loader(E),
}
struct Entry<V> {
value: V,
expires_at: Instant,
bytes: usize,
generation: u64,
}
struct State<K, V> {
entries: HashMap<K, Entry<V>>,
order: VecDeque<K>,
generations: HashMap<K, u64>,
bytes: usize,
}
pub struct GenerationCache<K, V> {
state: Arc<Mutex<State<K, V>>>,
flights: Arc<Mutex<HashMap<K, Arc<Mutex<()>>>>>,
config: CacheConfig,
}
struct CancelledFlightCleanup<K: Eq + Hash + Send + 'static, V: Send + 'static> {
key: Option<K>,
flight: Option<Arc<Mutex<()>>>,
flights: Arc<Mutex<HashMap<K, Arc<Mutex<()>>>>>,
state: Arc<Mutex<State<K, V>>>,
}
impl<K: Eq + Hash + Send + 'static, V: Send + 'static> CancelledFlightCleanup<K, V> {
fn disarm(&mut self) {
self.key.take();
self.flight.take();
}
}
impl<K, V> Drop for CancelledFlightCleanup<K, V>
where
K: Eq + Hash + Send + 'static,
V: Send + 'static,
{
fn drop(&mut self) {
let (Some(key), Some(flight)) = (self.key.take(), self.flight.take()) else {
return;
};
let flights = Arc::clone(&self.flights);
let state = Arc::clone(&self.state);
let Ok(runtime) = tokio::runtime::Handle::try_current() else {
return;
};
runtime.spawn(async move {
let mut flights = flights.lock().await;
if !flights
.get(&key)
.is_some_and(|current| Arc::ptr_eq(current, &flight) && Arc::strong_count(&flight) <= 2)
{
return;
}
flights.remove(&key);
drop(flights);
let mut state = state.lock().await;
if !state.entries.contains_key(&key) {
state.generations.remove(&key);
}
});
}
}
impl<K, V> GenerationCache<K, V>
where
K: Clone + Eq + Hash + Send + 'static,
V: Clone + Send + 'static,
{
pub fn new(config: CacheConfig) -> Self {
Self {
state: Arc::new(Mutex::new(State {
entries: HashMap::new(),
order: VecDeque::new(),
generations: HashMap::new(),
bytes: 0,
})),
flights: Arc::new(Mutex::new(HashMap::new())),
config,
}
}
pub async fn get_or_load<E, F, Fut>(&self, key: K, bytes: usize, load: F) -> Result<CacheOutcome<V>, CacheError<E>>
where
F: FnOnce() -> Fut,
Fut: Future<Output = Result<V, E>>,
{
if let Some(value) = self.get(&key).await {
return Ok(CacheOutcome {
value,
kind: CacheOutcomeKind::Hit,
});
}
let flight = {
let mut flights = self.flights.lock().await;
if let Some(flight) = flights.get(&key) {
Arc::clone(flight)
} else {
if flights.len() >= self.config.max_flights {
return Err(CacheError::FlightLimit);
}
let flight = Arc::new(Mutex::new(()));
flights.insert(key.clone(), Arc::clone(&flight));
flight
}
};
let mut cancelled_cleanup = CancelledFlightCleanup {
key: Some(key.clone()),
flight: Some(Arc::clone(&flight)),
flights: Arc::clone(&self.flights),
state: Arc::clone(&self.state),
};
let _guard = flight.lock().await;
if let Some(value) = self.get(&key).await {
cancelled_cleanup.disarm();
self.release_flight(&key, &flight).await;
return Ok(CacheOutcome {
value,
kind: CacheOutcomeKind::Miss,
});
}
let generation = self.generation(&key).await;
let result = load().await;
if let Ok(value) = &result {
self.insert(key.clone(), value.clone(), bytes, generation).await;
}
cancelled_cleanup.disarm();
self.release_flight(&key, &flight).await;
result
.map(|value| CacheOutcome {
value,
kind: CacheOutcomeKind::Miss,
})
.map_err(CacheError::Loader)
}
pub async fn invalidate(&self, key: &K) {
let in_flight = self.flights.lock().await.contains_key(key);
let mut state = self.state.lock().await;
if let Some(entry) = state.entries.remove(key) {
state.bytes -= entry.bytes;
}
state.order.retain(|candidate| candidate != key);
if in_flight {
*state.generations.entry(key.clone()).or_default() += 1;
} else {
state.generations.remove(key);
}
}
pub async fn snapshot(&self) -> CacheSnapshot {
let state = self.state.lock().await;
let entries = state.entries.len();
let bytes = state.bytes;
drop(state);
CacheSnapshot {
entries,
bytes,
flights: self.flights.lock().await.len(),
}
}
async fn release_flight(&self, key: &K, flight: &Arc<Mutex<()>>) {
let mut flights = self.flights.lock().await;
if flights
.get(key)
.is_some_and(|current| Arc::ptr_eq(current, flight) && Arc::strong_count(flight) <= 3)
{
flights.remove(key);
let mut state = self.state.lock().await;
if !state.entries.contains_key(key) {
state.generations.remove(key);
}
}
}
async fn get(&self, key: &K) -> Option<V> {
let in_flight = self.flights.lock().await.contains_key(key);
let mut state = self.state.lock().await;
let current_generation = state.generations.get(key).copied().unwrap_or_default();
let valid = state
.entries
.get(key)
.is_some_and(|entry| entry.expires_at > Instant::now() && entry.generation == current_generation);
if !valid {
if let Some(entry) = state.entries.remove(key) {
state.bytes -= entry.bytes;
}
state.order.retain(|candidate| candidate != key);
if !in_flight {
state.generations.remove(key);
}
return None;
}
let value = state.entries.get(key)?.value.clone();
state.order.retain(|candidate| candidate != key);
state.order.push_back(key.clone());
Some(value)
}
async fn generation(&self, key: &K) -> u64 {
self
.state
.lock()
.await
.generations
.get(key)
.copied()
.unwrap_or_default()
}
async fn insert(&self, key: K, value: V, bytes: usize, generation: u64) {
let mut state = self.state.lock().await;
if state.generations.get(&key).copied().unwrap_or_default() != generation {
return;
}
if let Some(previous) = state.entries.remove(&key) {
state.bytes -= previous.bytes;
}
state.order.retain(|candidate| candidate != &key);
state.bytes += bytes;
state.order.push_back(key.clone());
state.entries.insert(
key,
Entry {
value,
expires_at: Instant::now() + self.config.hard_ttl,
bytes,
generation,
},
);
while state.bytes > self.config.max_bytes {
let Some(oldest) = state.order.pop_front() else {
break;
};
if let Some(entry) = state.entries.remove(&oldest) {
state.bytes -= entry.bytes;
state.generations.remove(&oldest);
}
}
}
}
#[cfg(test)]
#[path = "../tests/cache/generation/tests.rs"]
mod tests;