use std::any::Any;
use std::collections::HashMap;
use std::future::Future;
use std::hash::{Hash, Hasher};
use std::pin::Pin;
use std::sync::{Arc, Mutex};
use futures::FutureExt;
use futures::future::Shared;
use super::error::RequestCacheError;
type MemoShared<V> = Shared<Pin<Box<dyn Future<Output = Result<V, RequestCacheError>> + Send>>>;
#[derive(Default, Clone)]
pub struct RequestCache {
slots: Arc<Mutex<HashMap<u64, Box<dyn Any + Send + 'static>>>>,
}
impl RequestCache {
#[must_use]
pub fn new() -> Self {
Self::default()
}
pub async fn get_or_compute<K, V, F, Fut>(
&self,
name: &'static str,
key: &K,
compute: F,
) -> Result<V, RequestCacheError>
where
K: Hash + 'static,
V: Clone + Send + Sync + 'static,
F: FnOnce() -> Fut,
Fut: Future<Output = Result<V, RequestCacheError>> + Send + 'static,
{
let slot_id = slot_id::<K>(name, key);
let shared = self.slot_or_insert(slot_id, compute);
shared.await
}
fn slot_or_insert<F, Fut, V>(&self, slot_id: u64, compute: F) -> MemoShared<V>
where
F: FnOnce() -> Fut,
Fut: Future<Output = Result<V, RequestCacheError>> + Send + 'static,
V: Clone + Send + Sync + 'static,
{
let mut slots = match self.slots.lock() {
Ok(guard) => guard,
Err(poisoned) => poisoned.into_inner(),
};
if let Some(existing) = slots.get(&slot_id) {
let any = existing.downcast_ref::<MemoShared<V>>();
if let Some(shared) = any {
return shared.clone();
}
}
let future: Pin<Box<dyn Future<Output = Result<V, RequestCacheError>> + Send>> =
Box::pin(compute());
let shared = future.shared();
slots.insert(slot_id, Box::new(shared.clone()));
shared
}
#[must_use]
pub fn len(&self) -> usize {
let slots = match self.slots.lock() {
Ok(guard) => guard,
Err(poisoned) => poisoned.into_inner(),
};
slots.len()
}
#[must_use]
pub fn is_empty(&self) -> bool {
self.len() == 0
}
}
pub fn from_state<S>(state: &S) -> RequestCache
where
S: RequestCacheFactory,
{
state.request_cache()
}
pub trait RequestCacheFactory {
fn request_cache(&self) -> RequestCache;
}
impl<S> axum::extract::FromRequestParts<S> for RequestCache
where
S: Send + Sync,
{
type Rejection = std::convert::Infallible;
async fn from_request_parts(
parts: &mut axum::http::request::Parts,
_state: &S,
) -> Result<Self, Self::Rejection> {
if let Some(existing) = parts.extensions.get::<RequestCache>() {
return Ok(existing.clone());
}
let cache = RequestCache::new();
parts.extensions.insert(cache.clone());
Ok(cache)
}
}
fn slot_id<K: Hash + 'static>(name: &'static str, key: &K) -> u64 {
let mut hasher = std::collections::hash_map::DefaultHasher::new();
name.hash(&mut hasher);
key.hash(&mut hasher);
hasher.finish()
}
#[cfg(test)]
mod tests {
use super::*;
use crate::request_cache::MAX_KEY_BYTES;
use std::sync::atomic::{AtomicUsize, Ordering};
use std::time::Duration;
#[tokio::test]
async fn same_key_runs_once_for_concurrent_callers() {
let cache = RequestCache::new();
let calls = Arc::new(AtomicUsize::new(0));
let calls_clone = calls.clone();
let first = cache.get_or_compute("load_profile", &42u64, || {
let calls = calls_clone.clone();
async move {
calls.fetch_add(1, Ordering::SeqCst);
tokio::time::sleep(Duration::from_millis(50)).await;
Ok::<u64, RequestCacheError>(99)
}
});
let second = cache.get_or_compute("load_profile", &42u64, || {
let calls = calls.clone();
async move {
calls.fetch_add(1, Ordering::SeqCst);
Ok::<u64, RequestCacheError>(99)
}
});
let (a, b) = tokio::join!(first, second);
assert_eq!(a.unwrap(), 99);
assert_eq!(b.unwrap(), 99);
assert_eq!(
calls.load(Ordering::SeqCst),
1,
"resolver ran more than once"
);
}
#[tokio::test]
async fn different_keys_run_independently() {
let cache = RequestCache::new();
let calls = Arc::new(AtomicUsize::new(0));
let c1 = calls.clone();
let c2 = calls.clone();
let a = cache.get_or_compute("load_profile", &1u64, || {
let c = c1.clone();
async move {
c.fetch_add(1, Ordering::SeqCst);
tokio::time::sleep(Duration::from_millis(20)).await;
Ok::<u64, RequestCacheError>(1)
}
});
let b = cache.get_or_compute("load_profile", &2u64, || {
let c = c2.clone();
async move {
c.fetch_add(1, Ordering::SeqCst);
Ok(2)
}
});
let (ra, rb) = tokio::join!(a, b);
assert_eq!(ra.unwrap(), 1);
assert_eq!(rb.unwrap(), 2);
assert_eq!(
calls.load(Ordering::SeqCst),
2,
"expected two independent runs"
);
}
#[tokio::test]
async fn resolver_failure_is_cached_and_isolates_to_the_key() {
let cache = RequestCache::new();
let e1 = cache.get_or_compute::<_, u64, _, _>("load", &1u64, || async {
Err(RequestCacheError::Resolver(Arc::from("boom")))
});
let e1b = cache.get_or_compute::<_, u64, _, _>("load", &1u64, || async {
Err(RequestCacheError::Resolver(Arc::from("should-not-run")))
});
let ok2 = cache.get_or_compute::<_, u64, _, _>("load", &2u64, || async { Ok(2) });
let (r1, r1b, r2) = tokio::join!(e1, e1b, ok2);
assert!(matches!(r1, Err(RequestCacheError::Resolver(_))));
assert!(matches!(r1b, Err(RequestCacheError::Resolver(_))));
assert_eq!(r2.unwrap(), 2);
}
#[tokio::test]
async fn cancellation_drops_in_flight_work() {
let cache = RequestCache::new();
let calls = Arc::new(AtomicUsize::new(0));
let spawned_cache = cache.clone();
let c = calls.clone();
let handle = tokio::spawn(async move {
spawned_cache
.get_or_compute::<_, u64, _, _>("slow", &1u64, move || {
let c = c.clone();
async move {
c.fetch_add(1, Ordering::SeqCst);
tokio::time::sleep(Duration::from_millis(200)).await;
Ok(1)
}
})
.await
});
handle.abort();
let _ = handle.await;
let c2 = calls.clone();
let result = cache
.get_or_compute::<_, u64, _, _>("slow", &1u64, || {
let c = c2.clone();
async move {
c.fetch_add(1, Ordering::SeqCst);
Ok(1)
}
})
.await;
assert_eq!(result.unwrap(), 1);
assert!(
calls.load(Ordering::SeqCst) >= 1,
"expected at least one fresh run after cancellation"
);
}
#[tokio::test]
async fn empty_cache_is_empty() {
let cache = RequestCache::new();
assert!(cache.is_empty());
assert_eq!(cache.len(), 0);
let _ = cache
.get_or_compute::<_, u64, _, _>("x", &1u64, || async { Ok(1) })
.await;
assert_eq!(cache.len(), 1);
assert!(!cache.is_empty());
}
#[test]
fn request_cache_error_display_is_typed() {
let e = RequestCacheError::OversizedKey {
limit: MAX_KEY_BYTES,
actual: 100_000,
};
assert!(e.to_string().contains("too large"));
let e2 = RequestCacheError::from_display(&std::io::Error::other("db down"));
assert!(e2.to_string().contains("db down"));
}
#[test]
fn request_cache_error_is_clone_for_shared_failure() {
let e = RequestCacheError::Resolver(Arc::from("boom"));
let e2 = e.clone();
assert!(matches!(e2, RequestCacheError::Resolver(_)));
}
}