use std::any::Any; use std::collections::HashMap; use std::fmt; use std::future::Future; use std::hash::Hash; use std::sync::atomic::{AtomicBool, AtomicU64, Ordering}; use std::sync::Mutex as StdMutex; use std::time::Duration; use aether_cache::ExpiringMap; use serde_json::Value; use tokio::sync::futures::OwnedNotified; use tokio::sync::Notify; use crate::data::auth::GatewayAuthApiKeySnapshot; const AUTH_RUNTIME_CACHE_MAX_ENTRIES: usize = 16_384; #[derive(Debug)] struct CacheSingleflight { inflight: StdMutex>>, // Invalidations advance this generation. A load that started before an // invalidation must not publish its stale result after the clear. generation: AtomicU64, } struct CacheInflightState { completed: AtomicBool, error: StdMutex>>, notify: std::sync::Arc, } impl fmt::Debug for CacheInflightState { fn fmt(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result { formatter .debug_struct("CacheInflightState") .field("completed", &self.completed.load(Ordering::Acquire)) .field( "has_error", &self .error .lock() .map(|error| error.is_some()) .unwrap_or(true), ) .finish_non_exhaustive() } } impl CacheInflightState { fn new() -> Self { Self { completed: AtomicBool::new(false), error: StdMutex::new(None), notify: std::sync::Arc::new(Notify::new()), } } fn waiter(self: &std::sync::Arc) -> CacheInflightWaiter { CacheInflightWaiter { state: std::sync::Arc::clone(self), notified: std::sync::Arc::clone(&self.notify).notified_owned(), } } fn complete(&self) { self.completed.store(true, Ordering::Release); self.notify.notify_waiters(); } fn fail(&self, error: E) where E: Clone + Send + Sync + 'static, { if let Ok(mut current) = self.error.lock() { *current = Some(Box::new(error)); } self.complete(); } fn error(&self) -> Option where E: Clone + Send + Sync + 'static, { // Cache methods are generic over E, so retain the concrete error for // same-typed followers. A mismatched caller type safely falls back to // the existing retry path instead of receiving the wrong error type. self.error .lock() .ok() .and_then(|error| error.as_ref()?.downcast_ref::().cloned()) } } struct CacheInflightWaiter { state: std::sync::Arc, notified: OwnedNotified, } impl CacheInflightWaiter { async fn wait(self) -> Result<(), E> where E: Clone + Send + Sync + 'static, { let Self { state, notified } = self; if state.completed.load(Ordering::Acquire) { return state.error().map_or(Ok(()), Err); } tokio::pin!(notified); if notified.as_mut().enable() || state.completed.load(Ordering::Acquire) { return state.error().map_or(Ok(()), Err); } notified.await; state.error().map_or(Ok(()), Err) } } impl Default for CacheSingleflight { fn default() -> Self { Self { inflight: StdMutex::new(HashMap::new()), generation: AtomicU64::new(0), } } } enum CacheInflightRegistration<'a, K: Eq + Hash> { Leader(CacheInflightGuard<'a, K>), Follower(CacheInflightWaiter), Bypass, } struct CacheInflightGuard<'a, K: Eq + Hash> { singleflight: &'a CacheSingleflight, key: Option, state: std::sync::Arc, generation: u64, } struct CacheOwnedInflightGuard { singleflight: std::sync::Arc>, key: Option, state: std::sync::Arc, generation: u64, } impl CacheSingleflight where K: Eq + Hash, { fn generation(&self) -> u64 { self.generation.load(Ordering::Acquire) } fn generation_is_current(&self, generation: u64) -> bool { self.generation.load(Ordering::Acquire) == generation } fn finish(&self, key: &K, state: &std::sync::Arc) { let removed = self .inflight .lock() .map(|mut inflight| { inflight .get(key) .is_some_and(|current| std::sync::Arc::ptr_eq(current, state)) && inflight.remove(key).is_some() }) .unwrap_or(false); if removed { state.complete(); } } fn fail(&self, key: &K, state: &std::sync::Arc, error: E) where E: Clone + Send + Sync + 'static, { let removed_current = self .inflight .lock() .map(|mut inflight| { if inflight .get(key) .is_some_and(|current| std::sync::Arc::ptr_eq(current, state)) { inflight.remove(key); true } else { false } }) .unwrap_or(false); if removed_current { state.fail(error); } } fn clear(&self) { let states = self .inflight .lock() .map(|mut inflight| { self.generation.fetch_add(1, Ordering::AcqRel); inflight.drain().map(|(_, state)| state).collect::>() }) .unwrap_or_default(); for state in states { state.complete(); } } } impl CacheSingleflight where K: Clone + Eq + Hash, { fn register(&self, key: &K) -> CacheInflightRegistration<'_, K> { match self.inflight.lock() { Ok(mut inflight) => { if let Some(state) = inflight.get(key) { // Register while the map lock is held so leader completion // cannot race between lookup and waiter creation. CacheInflightRegistration::Follower(state.waiter()) } else { let state = std::sync::Arc::new(CacheInflightState::new()); inflight.insert(key.clone(), std::sync::Arc::clone(&state)); CacheInflightRegistration::Leader(CacheInflightGuard { singleflight: self, key: Some(key.clone()), state, generation: self.generation(), }) } } Err(_) => CacheInflightRegistration::Bypass, } } fn try_register_owned_leader( self: &std::sync::Arc, key: &K, ) -> Option> { let Ok(mut inflight) = self.inflight.lock() else { return None; }; if inflight.contains_key(key) { return None; } let state = std::sync::Arc::new(CacheInflightState::new()); inflight.insert(key.clone(), std::sync::Arc::clone(&state)); Some(CacheOwnedInflightGuard { singleflight: std::sync::Arc::clone(self), key: Some(key.clone()), state, generation: self.generation(), }) } } impl Drop for CacheInflightGuard<'_, K> where K: Eq + Hash, { fn drop(&mut self) { if let Some(key) = self.key.take() { self.singleflight.finish(&key, &self.state); } } } impl CacheInflightGuard<'_, K> where K: Eq + Hash, { fn generation_is_current(&self) -> bool { self.singleflight.generation_is_current(self.generation) } fn fail(&self, error: E) where E: Clone + Send + Sync + 'static, { if let Some(key) = self.key.as_ref() { self.singleflight.fail(key, &self.state, error); } } } impl Drop for CacheOwnedInflightGuard where K: Eq + Hash, { fn drop(&mut self) { if let Some(key) = self.key.take() { self.singleflight.finish(&key, &self.state); } } } impl CacheOwnedInflightGuard where K: Eq + Hash, { fn fail(&self, error: E) where E: Clone + Send + Sync + 'static, { if let Some(key) = self.key.as_ref() { self.singleflight.fail(key, &self.state, error); } } } #[derive(Debug, Clone, PartialEq, Eq, Hash)] pub(crate) struct AuthApiKeyIdentityCacheKey { user_id: String, api_key_id: String, } impl AuthApiKeyIdentityCacheKey { pub(crate) fn new(user_id: &str, api_key_id: &str) -> Self { Self { user_id: user_id.trim().to_string(), api_key_id: api_key_id.trim().to_string(), } } pub(crate) fn is_empty(&self) -> bool { self.user_id.is_empty() || self.api_key_id.is_empty() } } #[derive(Debug, Clone, PartialEq, Eq, Hash)] pub(crate) struct AuthApiKeyFeatureCacheKey { user_id: String, api_key_id: String, is_standalone: bool, } impl AuthApiKeyFeatureCacheKey { pub(crate) fn new(user_id: &str, api_key_id: &str, is_standalone: bool) -> Self { Self { user_id: user_id.trim().to_string(), api_key_id: api_key_id.trim().to_string(), is_standalone, } } pub(crate) fn is_empty(&self) -> bool { self.api_key_id.is_empty() || (!self.is_standalone && self.user_id.is_empty()) } } #[derive(Debug, Default)] pub(crate) struct AuthSnapshotCache { entries: ExpiringMap>, singleflight: CacheSingleflight, mutation: StdMutex<()>, } #[derive(Debug, Clone, PartialEq, Eq, Hash)] pub(crate) enum AuthSnapshotCacheKey { KeyHash(String), UserApiKeyIds(AuthApiKeyIdentityCacheKey), } impl AuthSnapshotCacheKey { pub(crate) fn key_hash(key_hash: &str) -> Self { Self::KeyHash(key_hash.trim().to_string()) } pub(crate) fn user_api_key_ids(user_id: &str, api_key_id: &str) -> Self { Self::UserApiKeyIds(AuthApiKeyIdentityCacheKey::new(user_id, api_key_id)) } pub(crate) fn is_empty(&self) -> bool { match self { Self::KeyHash(key_hash) => key_hash.is_empty(), Self::UserApiKeyIds(key) => key.is_empty(), } } } impl AuthSnapshotCache { pub(crate) fn generation(&self) -> u64 { self.singleflight.generation() } pub(crate) fn get( &self, key: &AuthSnapshotCacheKey, ttl: Duration, ) -> Option> { self.entries.get_fresh(key, ttl) } pub(crate) fn insert_if_generation( &self, key: AuthSnapshotCacheKey, value: Option, ttl: Duration, generation: u64, ) { let Ok(_mutation) = self.mutation.lock() else { return; }; if !self.singleflight.generation_is_current(generation) { return; } self.entries .insert(key, value, ttl, AUTH_RUNTIME_CACHE_MAX_ENTRIES); } pub(crate) async fn get_or_load( &self, key: AuthSnapshotCacheKey, ttl: Duration, mut load: F, ) -> Result, E> where E: Clone + Send + Sync + 'static, F: FnMut() -> Fut, Fut: Future, E>>, { if let Some(value) = self.get(&key, ttl) { return Ok(value); } loop { match self.singleflight.register(&key) { CacheInflightRegistration::Bypass => { let generation = self.singleflight.generation(); let value = load().await?; self.insert_if_generation(key, value.clone(), ttl, generation); return Ok(value); } CacheInflightRegistration::Follower(waiter) => { waiter.wait().await?; if let Some(value) = self.get(&key, ttl) { return Ok(value); } } CacheInflightRegistration::Leader(guard) => { let value = match load().await { Ok(value) => value, Err(error) => { guard.fail(error.clone()); return Err(error); } }; if guard.generation_is_current() { self.insert_if_generation(key, value.clone(), ttl, guard.generation); } return Ok(value); } } } } pub(crate) fn clear(&self) { let Ok(_mutation) = self.mutation.lock() else { return; }; self.entries.clear(); self.singleflight.clear(); } } #[derive(Debug)] pub(crate) struct JsonValueCache { entries: ExpiringMap>, singleflight: CacheSingleflight, mutation: StdMutex<()>, } impl Default for JsonValueCache { fn default() -> Self { Self { entries: ExpiringMap::default(), singleflight: CacheSingleflight::default(), mutation: StdMutex::new(()), } } } impl JsonValueCache where K: Clone + Eq + Hash, { pub(crate) fn get(&self, key: &K, ttl: Duration) -> Option> { self.entries.get_fresh(key, ttl) } pub(crate) fn insert(&self, key: K, value: Option, ttl: Duration) { let Ok(_mutation) = self.mutation.lock() else { return; }; self.entries .insert(key, value, ttl, AUTH_RUNTIME_CACHE_MAX_ENTRIES); self.singleflight.clear(); } fn insert_if_generation(&self, key: K, value: Option, ttl: Duration, generation: u64) { let Ok(_mutation) = self.mutation.lock() else { return; }; if !self.singleflight.generation_is_current(generation) { return; } self.entries .insert(key, value, ttl, AUTH_RUNTIME_CACHE_MAX_ENTRIES); } pub(crate) async fn get_or_load( &self, key: K, ttl: Duration, mut load: F, ) -> Result, E> where E: Clone + Send + Sync + 'static, F: FnMut() -> Fut, Fut: Future, E>>, { if let Some(value) = self.get(&key, ttl) { return Ok(value); } loop { match self.singleflight.register(&key) { CacheInflightRegistration::Bypass => { let generation = self.singleflight.generation(); let value = load().await?; self.insert_if_generation(key, value.clone(), ttl, generation); return Ok(value); } CacheInflightRegistration::Follower(waiter) => { waiter.wait().await?; if let Some(value) = self.get(&key, ttl) { return Ok(value); } } CacheInflightRegistration::Leader(guard) => { let value = match load().await { Ok(value) => value, Err(error) => { guard.fail(error.clone()); return Err(error); } }; self.insert_if_generation(key, value.clone(), ttl, guard.generation); return Ok(value); } } } } pub(crate) fn clear(&self) { let Ok(_mutation) = self.mutation.lock() else { return; }; self.entries.clear(); self.singleflight.clear(); } } #[derive(Debug)] pub(crate) struct ValueCache { entries: ExpiringMap>, singleflight: std::sync::Arc>, mutation: StdMutex<()>, } impl Default for ValueCache { fn default() -> Self { Self { entries: ExpiringMap::default(), singleflight: std::sync::Arc::new(CacheSingleflight::default()), mutation: StdMutex::new(()), } } } impl ValueCache where K: Clone + Eq + Hash + Send + Sync + 'static, V: Clone + Send + Sync + 'static, { pub(crate) fn get(&self, key: &K, ttl: Duration) -> Option> { self.entries.get_fresh(key, ttl) } pub(crate) fn insert(&self, key: K, value: Option, ttl: Duration) { let Ok(_mutation) = self.mutation.lock() else { return; }; self.entries .insert(key, value, ttl, AUTH_RUNTIME_CACHE_MAX_ENTRIES); self.singleflight.clear(); } fn insert_if_generation(&self, key: K, value: Option, ttl: Duration, generation: u64) { let Ok(_mutation) = self.mutation.lock() else { return; }; if !self.singleflight.generation_is_current(generation) { return; } self.entries .insert(key, value, ttl, AUTH_RUNTIME_CACHE_MAX_ENTRIES); } pub(crate) async fn get_or_load( &self, key: K, ttl: Duration, mut load: F, ) -> Result, E> where E: Clone + Send + Sync + 'static, F: FnMut() -> Fut, Fut: Future, E>>, { if let Some(value) = self.get(&key, ttl) { return Ok(value); } loop { match self.singleflight.register(&key) { CacheInflightRegistration::Bypass => { let generation = self.singleflight.generation(); let value = load().await?; self.insert_if_generation(key, value.clone(), ttl, generation); return Ok(value); } CacheInflightRegistration::Follower(waiter) => { waiter.wait().await?; if let Some(value) = self.get(&key, ttl) { return Ok(value); } } CacheInflightRegistration::Leader(guard) => { let value = match load().await { Ok(value) => value, Err(error) => { guard.fail(error.clone()); return Err(error); } }; self.insert_if_generation(key, value.clone(), ttl, guard.generation); return Ok(value); } } } } pub(crate) async fn get_or_load_once( &self, key: K, ttl: Duration, load: F, ) -> Result, E> where E: Clone + Send + Sync + 'static, F: FnOnce() -> Fut, Fut: Future, E>>, { self.get_or_load_once_with_observer(key, ttl, load, CacheLoadObserver::default()) .await } pub(crate) async fn get_or_load_once_with_observer( &self, key: K, ttl: Duration, load: F, observer: CacheLoadObserver, ) -> Result, E> where E: Clone + Send + Sync + 'static, F: FnOnce() -> Fut, Fut: Future, E>>, { if let Some(value) = self.get(&key, ttl) { observer.hit(); return Ok(value); } observer.miss(); let mut load = Some(load); loop { match self.singleflight.register(&key) { CacheInflightRegistration::Bypass => { let generation = self.singleflight.generation(); observer.load(); let value = load.take().expect("cache load closure should be available")().await?; self.insert_if_generation(key, value.clone(), ttl, generation); return Ok(value); } CacheInflightRegistration::Follower(waiter) => { observer.follower_wait(); waiter.wait().await?; if let Some(value) = self.get(&key, ttl) { observer.hit(); return Ok(value); } } CacheInflightRegistration::Leader(guard) => { observer.load(); let value = match load.take().expect("cache load closure should be available")() .await { Ok(value) => value, Err(error) => { guard.fail(error.clone()); return Err(error); } }; self.insert_if_generation(key, value.clone(), ttl, guard.generation); return Ok(value); } } } } pub(crate) async fn get_or_load_once_stale_while_refreshing( &self, key: K, ttl: Duration, stale_ttl: Duration, load: F, observer: CacheLoadObserver, ) -> Result, E> where E: Clone + Send + Sync + 'static, F: FnOnce() -> Fut, Fut: Future, E>>, { if let Some((value, age)) = self.entries.get_with_age(&key, stale_ttl) { if age <= ttl { observer.hit(); return Ok(value); } // Keep stale snapshots off the request critical path. The caller's // invalidation path clears entries when provider/catalog/routing // state changes, and the bounded stale TTL limits passive drift. observer.hit(); return Ok(value); } observer.miss(); let mut load = Some(load); loop { match self.singleflight.register(&key) { CacheInflightRegistration::Bypass => { let generation = self.singleflight.generation(); observer.load(); let value = load.take().expect("cache load closure should be available")().await?; self.insert_if_generation(key, value.clone(), stale_ttl, generation); return Ok(value); } CacheInflightRegistration::Follower(waiter) => { observer.follower_wait(); waiter.wait().await?; if let Some((value, _age)) = self.entries.get_with_age(&key, stale_ttl) { observer.hit(); return Ok(value); } } CacheInflightRegistration::Leader(guard) => { observer.load(); let value = match load.take().expect("cache load closure should be available")() .await { Ok(value) => value, Err(error) => { guard.fail(error.clone()); return Err(error); } }; self.insert_if_generation(key, value.clone(), stale_ttl, guard.generation); return Ok(value); } } } } pub(crate) async fn get_or_load_once_stale_while_revalidating< E, ColdLoad, ColdFuture, BackgroundRefresh, BackgroundFuture, >( self: &std::sync::Arc, key: K, ttl: Duration, stale_ttl: Duration, cold_load: ColdLoad, background_refresh: BackgroundRefresh, observer: CacheLoadObserver, ) -> Result, E> where K: Send + 'static, V: Send + 'static, E: Clone + Send + Sync + 'static, ColdLoad: FnOnce() -> ColdFuture, ColdFuture: Future, E>>, BackgroundRefresh: FnOnce() -> BackgroundFuture, BackgroundFuture: Future, E>> + Send + 'static, { let stale_ttl = stale_ttl.max(ttl); if let Some((value, age)) = self.entries.get_with_age(&key, stale_ttl) { observer.hit(); if age > ttl { if let Some(guard) = self.singleflight.try_register_owned_leader(&key) { observer.load(); let cache = std::sync::Arc::clone(self); let generation = guard.generation; // Build the owned refresh future only after this request // wins stale revalidation. Fresh hits never clone the // caller's plan/candidate inputs. let refresh = background_refresh(); tokio::spawn(async move { match refresh.await { Ok(refreshed) => { cache.insert_if_generation(key, refreshed, stale_ttl, generation); } Err(error) => guard.fail(error), } drop(guard); }); } } return Ok(value); } observer.miss(); let mut cold_load = Some(cold_load); loop { match self.singleflight.register(&key) { CacheInflightRegistration::Bypass => { let generation = self.singleflight.generation(); observer.load(); let cold_load = cold_load .take() .expect("cache cold-load closure should be available"); let value = cold_load().await?; self.insert_if_generation(key, value.clone(), stale_ttl, generation); return Ok(value); } CacheInflightRegistration::Follower(waiter) => { observer.follower_wait(); waiter.wait().await?; if let Some((value, _age)) = self.entries.get_with_age(&key, stale_ttl) { observer.hit(); return Ok(value); } } CacheInflightRegistration::Leader(guard) => { observer.load(); let cold_load = cold_load .take() .expect("cache cold-load closure should be available"); let value = match cold_load().await { Ok(value) => value, Err(error) => { guard.fail(error.clone()); return Err(error); } }; self.insert_if_generation(key, value.clone(), stale_ttl, guard.generation); return Ok(value); } } } } pub(crate) fn clear(&self) { let Ok(_mutation) = self.mutation.lock() else { return; }; self.entries.clear(); self.singleflight.clear(); } } #[derive(Clone, Copy, Default)] pub(crate) struct CacheLoadObserver { on_hit: Option, on_miss: Option, on_load: Option, on_follower_wait: Option, } impl CacheLoadObserver { pub(crate) fn new() -> Self { Self::default() } pub(crate) fn on_hit(mut self, callback: fn()) -> Self { self.on_hit = Some(callback); self } pub(crate) fn on_miss(mut self, callback: fn()) -> Self { self.on_miss = Some(callback); self } pub(crate) fn on_load(mut self, callback: fn()) -> Self { self.on_load = Some(callback); self } pub(crate) fn on_follower_wait(mut self, callback: fn()) -> Self { self.on_follower_wait = Some(callback); self } fn hit(self) { if let Some(callback) = self.on_hit { callback(); } } fn miss(self) { if let Some(callback) = self.on_miss { callback(); } } fn load(self) { if let Some(callback) = self.on_load { callback(); } } fn follower_wait(self) { if let Some(callback) = self.on_follower_wait { callback(); } } } #[cfg(test)] mod tests { use super::{CacheInflightRegistration, CacheLoadObserver, CacheSingleflight, ValueCache}; use std::sync::atomic::{AtomicUsize, Ordering}; use std::sync::Arc; use std::time::{Duration, Instant}; fn record_max(max_seen: &AtomicUsize, value: usize) { let mut current = max_seen.load(Ordering::Acquire); while value > current { match max_seen.compare_exchange(current, value, Ordering::AcqRel, Ordering::Acquire) { Ok(_) => break, Err(next) => current = next, } } } #[tokio::test] async fn cache_singleflight_notifies_only_followers_for_completed_key() { let singleflight = CacheSingleflight::::default(); let leader_a = match singleflight.register(&"key-a".to_string()) { CacheInflightRegistration::Leader(guard) => guard, _ => panic!("key-a should register a leader"), }; let leader_b = match singleflight.register(&"key-b".to_string()) { CacheInflightRegistration::Leader(guard) => guard, _ => panic!("key-b should register a leader"), }; let follower_a = match singleflight.register(&"key-a".to_string()) { CacheInflightRegistration::Follower(waiter) => waiter, _ => panic!("second key-a registration should follow"), }; let follower_b = match singleflight.register(&"key-b".to_string()) { CacheInflightRegistration::Follower(waiter) => waiter, _ => panic!("second key-b registration should follow"), }; drop(leader_b); tokio::time::timeout(Duration::from_millis(100), follower_b.wait::<()>()) .await .expect("key-b follower should wake when key-b completes") .expect("successful flight should not publish an error"); assert!( tokio::time::timeout(Duration::from_millis(20), follower_a.wait::<()>()) .await .is_err(), "key-a follower must not wake when unrelated key-b completes" ); drop(leader_a); assert!(singleflight.inflight.lock().unwrap().is_empty()); } #[tokio::test] async fn cache_singleflight_leader_drop_before_follower_first_poll_does_not_lose_wakeup() { let singleflight = CacheSingleflight::::default(); let key = "before-first-poll".to_string(); let leader = match singleflight.register(&key) { CacheInflightRegistration::Leader(guard) => guard, _ => panic!("first registration should lead"), }; let follower = match singleflight.register(&key) { CacheInflightRegistration::Follower(waiter) => waiter, _ => panic!("second registration should follow"), }; // Drop before the waiter future is ever polled, as happens when a // loader completes between registration and task scheduling. drop(leader); tokio::time::timeout(Duration::from_millis(100), follower.wait::<()>()) .await .expect("completion before first poll must still release the follower") .expect("successful flight should not publish an error"); } #[tokio::test] async fn cache_singleflight_leader_drop_broadcasts_to_all_followers() { const FOLLOWER_COUNT: usize = 2_048; let singleflight = CacheSingleflight::::default(); let key = "broadcast".to_string(); let leader = match singleflight.register(&key) { CacheInflightRegistration::Leader(guard) => guard, _ => panic!("first registration should lead"), }; let mut followers = Vec::with_capacity(FOLLOWER_COUNT); for _ in 0..FOLLOWER_COUNT { let waiter = match singleflight.register(&key) { CacheInflightRegistration::Follower(waiter) => waiter, _ => panic!("concurrent registration should follow"), }; followers.push(tokio::spawn(waiter.wait::<()>())); } tokio::task::yield_now().await; drop(leader); tokio::time::timeout(Duration::from_secs(2), async { for follower in followers { follower .await .expect("follower task should join") .expect("successful flight should not publish an error"); } }) .await .expect("one completion should broadcast to every follower"); assert!(singleflight.inflight.lock().unwrap().is_empty()); } #[tokio::test] async fn cache_singleflight_clear_wakes_old_followers_without_removing_replacement() { let singleflight = CacheSingleflight::::default(); let key = "clear-replacement".to_string(); let old_generation = singleflight.generation(); let old_leader = match singleflight.register(&key) { CacheInflightRegistration::Leader(guard) => guard, _ => panic!("first registration should lead"), }; let old_follower = match singleflight.register(&key) { CacheInflightRegistration::Follower(waiter) => waiter, _ => panic!("second registration should follow"), }; singleflight.clear(); assert_ne!(singleflight.generation(), old_generation); let replacement_leader = match singleflight.register(&key) { CacheInflightRegistration::Leader(guard) => guard, _ => panic!("clear should allow a replacement leader"), }; let replacement_follower = match singleflight.register(&key) { CacheInflightRegistration::Follower(waiter) => waiter, _ => panic!("registration behind replacement should follow"), }; drop(old_leader); assert_eq!(singleflight.inflight.lock().unwrap().len(), 1); tokio::time::timeout(Duration::from_millis(100), old_follower.wait::<()>()) .await .expect("clear should wake followers of the invalidated flight") .expect("cache clear should allow an immediate retry"); drop(replacement_leader); tokio::time::timeout( Duration::from_millis(100), replacement_follower.wait::<()>(), ) .await .expect("old guard drop must not remove or strand the replacement flight") .expect("successful replacement should not publish an error"); assert!(singleflight.inflight.lock().unwrap().is_empty()); } #[tokio::test] async fn value_cache_coalesces_concurrent_loads_for_same_key() { let cache = Arc::new(ValueCache::::default()); let calls = Arc::new(AtomicUsize::new(0)); let mut tasks = Vec::new(); for _ in 0..16 { let cache = Arc::clone(&cache); let calls = Arc::clone(&calls); tasks.push(tokio::spawn(async move { cache .get_or_load::<(), _, _>( "same-key".to_string(), Duration::from_secs(60), || async { calls.fetch_add(1, Ordering::AcqRel); tokio::time::sleep(Duration::from_millis(25)).await; Ok(Some(42)) }, ) .await .unwrap() })); } for task in tasks { assert_eq!(task.await.unwrap(), Some(42)); } assert_eq!(calls.load(Ordering::Acquire), 1); } #[tokio::test] async fn value_cache_shares_leader_failure_without_follower_reload() { let cache = Arc::new(ValueCache::::default()); let calls = Arc::new(AtomicUsize::new(0)); let leader_started = Arc::new(tokio::sync::Notify::new()); let leader_release = Arc::new(tokio::sync::Notify::new()); let leader_cache = Arc::clone(&cache); let leader_calls = Arc::clone(&calls); let started = Arc::clone(&leader_started); let release = Arc::clone(&leader_release); let leader = tokio::spawn(async move { leader_cache .get_or_load::( "failed-key".to_string(), Duration::from_secs(60), || { let leader_calls = Arc::clone(&leader_calls); let started = Arc::clone(&started); let release = Arc::clone(&release); async move { leader_calls.fetch_add(1, Ordering::AcqRel); started.notify_one(); release.notified().await; Err("forced cache load failure".to_string()) } }, ) .await }); tokio::time::timeout(Duration::from_secs(1), leader_started.notified()) .await .expect("leader load should start"); let follower_cache = Arc::clone(&cache); let follower_calls = Arc::clone(&calls); let follower = tokio::spawn(async move { follower_cache .get_or_load::( "failed-key".to_string(), Duration::from_secs(60), || { let follower_calls = Arc::clone(&follower_calls); async move { follower_calls.fetch_add(1, Ordering::AcqRel); Err("follower loader must not run".to_string()) } }, ) .await }); tokio::time::timeout(Duration::from_secs(1), async { loop { let follower_registered = cache .singleflight .inflight .lock() .ok() .and_then(|inflight| inflight.values().next().cloned()) .is_some_and(|state| Arc::strong_count(&state) >= 4); if follower_registered { break; } tokio::task::yield_now().await; } }) .await .expect("follower should register behind the failed leader"); leader_release.notify_one(); assert_eq!( leader .await .expect("leader task should join") .expect_err("leader should return the injected failure"), "forced cache load failure" ); assert_eq!( follower .await .expect("follower task should join") .expect_err("follower should receive the leader failure"), "forced cache load failure" ); assert_eq!(calls.load(Ordering::Acquire), 1); } #[tokio::test] async fn value_cache_loads_different_keys_without_global_blocking() { let cache = Arc::new(ValueCache::::default()); let calls = Arc::new(AtomicUsize::new(0)); let active = Arc::new(AtomicUsize::new(0)); let max_active = Arc::new(AtomicUsize::new(0)); let mut tasks = Vec::new(); for (key, value) in [("key-a", 1_u64), ("key-b", 2_u64)] { let cache = Arc::clone(&cache); let calls = Arc::clone(&calls); let active = Arc::clone(&active); let max_active = Arc::clone(&max_active); tasks.push(tokio::spawn(async move { cache .get_or_load::<(), _, _>(key.to_string(), Duration::from_secs(60), || { let calls = Arc::clone(&calls); let active = Arc::clone(&active); let max_active = Arc::clone(&max_active); async move { calls.fetch_add(1, Ordering::AcqRel); let current = active.fetch_add(1, Ordering::AcqRel) + 1; record_max(&max_active, current); tokio::time::sleep(Duration::from_millis(50)).await; active.fetch_sub(1, Ordering::AcqRel); Ok(Some(value)) } }) .await .unwrap() })); } let mut results = Vec::new(); for task in tasks { results.push(task.await.unwrap()); } results.sort_unstable(); assert_eq!(results, vec![Some(1), Some(2)]); assert_eq!(calls.load(Ordering::Acquire), 2); assert_eq!(max_active.load(Ordering::Acquire), 2); } #[tokio::test] async fn value_cache_legacy_stale_lookup_returns_without_invoking_loader() { let cache = Arc::new(ValueCache::::default()); let key = "hot-key".to_string(); let calls = Arc::new(AtomicUsize::new(0)); cache.insert(key.clone(), Some(1), Duration::from_millis(10)); tokio::time::sleep(Duration::from_millis(20)).await; let first_cache = Arc::clone(&cache); let first_key = key.clone(); let first_calls = Arc::clone(&calls); let first = tokio::spawn(async move { first_cache .get_or_load_once_stale_while_refreshing::<(), _, _>( first_key, Duration::from_millis(10), Duration::from_secs(1), || async move { first_calls.fetch_add(1, Ordering::AcqRel); tokio::time::sleep(Duration::from_millis(100)).await; Ok(Some(2)) }, CacheLoadObserver::default(), ) .await }); let follower_cache = Arc::clone(&cache); let follower_calls = Arc::clone(&calls); let follower = tokio::spawn(async move { follower_cache .get_or_load_once_stale_while_refreshing::<(), _, _>( key, Duration::from_millis(10), Duration::from_secs(1), || async move { follower_calls.fetch_add(1, Ordering::AcqRel); Ok(Some(3)) }, CacheLoadObserver::default(), ) .await }); assert_eq!(first.await.unwrap().unwrap(), Some(1)); assert_eq!(follower.await.unwrap().unwrap(), Some(1)); assert_eq!(calls.load(Ordering::Acquire), 0); } #[tokio::test] async fn value_cache_stale_revalidation_returns_immediately_and_refreshes_in_background() { let cache = Arc::new(ValueCache::::default()); let key = "revalidate-key".to_string(); let ttl = Duration::from_millis(5); let stale_ttl = Duration::from_secs(1); let cold_calls = Arc::new(AtomicUsize::new(0)); let refresh_started = Arc::new(tokio::sync::Notify::new()); let refresh_release = Arc::new(tokio::sync::Notify::new()); cache.insert(key.clone(), Some(1), stale_ttl); tokio::time::sleep(Duration::from_millis(15)).await; let cold_calls_for_load = Arc::clone(&cold_calls); let refresh_started_for_load = Arc::clone(&refresh_started); let refresh_release_for_load = Arc::clone(&refresh_release); let result = tokio::time::timeout( Duration::from_millis(100), cache.get_or_load_once_stale_while_revalidating::<(), _, _, _, _>( key.clone(), ttl, stale_ttl, || async move { cold_calls_for_load.fetch_add(1, Ordering::AcqRel); Ok(Some(99)) }, move || async move { refresh_started_for_load.notify_one(); refresh_release_for_load.notified().await; Ok(Some(2)) }, CacheLoadObserver::default(), ), ) .await .expect("stale request should not wait for background refresh") .expect("stale lookup should succeed"); assert_eq!(result, Some(1)); assert_eq!(cold_calls.load(Ordering::Acquire), 0); tokio::time::timeout(Duration::from_secs(1), refresh_started.notified()) .await .expect("background refresh should start"); refresh_release.notify_one(); tokio::time::timeout(Duration::from_secs(1), async { loop { if cache.get(&key, stale_ttl) == Some(Some(2)) { break; } tokio::task::yield_now().await; } }) .await .expect("background refresh should publish the new value"); } #[tokio::test] async fn value_cache_fresh_hit_does_not_build_background_refresh_future() { let cache = Arc::new(ValueCache::::default()); let key = "fresh-lazy-refresh-key".to_string(); let refresh_factory_calls = AtomicUsize::new(0); cache.insert(key.clone(), Some(1), Duration::from_secs(60)); let result = cache .get_or_load_once_stale_while_revalidating::<(), _, _, _, _>( key, Duration::from_secs(30), Duration::from_secs(60), || async { panic!("cold loader must not run for a fresh hit") }, || { refresh_factory_calls.fetch_add(1, Ordering::AcqRel); async { Ok(Some(2)) } }, CacheLoadObserver::default(), ) .await .expect("fresh lookup should succeed"); assert_eq!(result, Some(1)); assert_eq!(refresh_factory_calls.load(Ordering::Acquire), 0); } #[tokio::test] async fn value_cache_stale_revalidation_starts_only_one_refresh_per_key() { let cache = Arc::new(ValueCache::::default()); let key = "single-refresh-key".to_string(); let ttl = Duration::from_millis(5); let stale_ttl = Duration::from_secs(1); let refresh_calls = Arc::new(AtomicUsize::new(0)); let refresh_started = Arc::new(tokio::sync::Notify::new()); let refresh_release = Arc::new(tokio::sync::Notify::new()); cache.insert(key.clone(), Some(1), stale_ttl); tokio::time::sleep(Duration::from_millis(15)).await; let first_calls = Arc::clone(&refresh_calls); let first_started = Arc::clone(&refresh_started); let first_release = Arc::clone(&refresh_release); assert_eq!( cache .get_or_load_once_stale_while_revalidating::<(), _, _, _, _>( key.clone(), ttl, stale_ttl, || async { panic!("cold loader must not run for a stale hit") }, move || async move { first_calls.fetch_add(1, Ordering::AcqRel); first_started.notify_one(); first_release.notified().await; Ok(Some(2)) }, CacheLoadObserver::default(), ) .await .unwrap(), Some(1) ); tokio::time::timeout(Duration::from_secs(1), refresh_started.notified()) .await .expect("first refresh should start"); for _ in 0..32 { let calls = Arc::clone(&refresh_calls); assert_eq!( cache .get_or_load_once_stale_while_revalidating::<(), _, _, _, _>( key.clone(), ttl, stale_ttl, || async { panic!("cold loader must not run for a stale hit") }, move || async move { calls.fetch_add(1, Ordering::AcqRel); Ok(Some(3)) }, CacheLoadObserver::default(), ) .await .unwrap(), Some(1) ); } assert_eq!(refresh_calls.load(Ordering::Acquire), 1); refresh_release.notify_one(); tokio::time::timeout(Duration::from_secs(1), async { loop { if cache.get(&key, stale_ttl) == Some(Some(2)) { break; } tokio::task::yield_now().await; } }) .await .expect("the single refresh should publish"); } #[tokio::test] async fn value_cache_clear_rejects_old_stale_revalidation_result() { let cache = Arc::new(ValueCache::::default()); let key = "revalidate-generation-key".to_string(); let ttl = Duration::from_millis(5); let stale_ttl = Duration::from_secs(1); let refresh_started = Arc::new(tokio::sync::Notify::new()); let refresh_release = Arc::new(tokio::sync::Notify::new()); cache.insert(key.clone(), Some(1), stale_ttl); tokio::time::sleep(Duration::from_millis(15)).await; let started = Arc::clone(&refresh_started); let release = Arc::clone(&refresh_release); assert_eq!( cache .get_or_load_once_stale_while_revalidating::<(), _, _, _, _>( key.clone(), ttl, stale_ttl, || async { panic!("cold loader must not run for a stale hit") }, move || async move { started.notify_one(); release.notified().await; Ok(Some(2)) }, CacheLoadObserver::default(), ) .await .unwrap(), Some(1) ); tokio::time::timeout(Duration::from_secs(1), refresh_started.notified()) .await .expect("refresh should start before clear"); cache.clear(); cache.insert(key.clone(), Some(3), stale_ttl); refresh_release.notify_one(); tokio::time::timeout(Duration::from_secs(1), async { loop { if Arc::strong_count(&cache.singleflight) == 1 { break; } tokio::task::yield_now().await; } }) .await .expect("background refresh should finish"); assert_eq!(cache.get(&key, stale_ttl), Some(Some(3))); } #[tokio::test] async fn value_cache_cold_stale_followers_do_not_reload_after_fresh_ttl() { let cache = Arc::new(ValueCache::::default()); let key = "cold-hot-key".to_string(); let calls = Arc::new(AtomicUsize::new(0)); let leader_cache = Arc::clone(&cache); let leader_key = key.clone(); let leader_calls = Arc::clone(&calls); let leader = tokio::spawn(async move { leader_cache .get_or_load_once_stale_while_refreshing::<(), _, _>( leader_key, Duration::from_millis(10), Duration::from_secs(1), || async move { leader_calls.fetch_add(1, Ordering::AcqRel); Ok(Some(1)) }, CacheLoadObserver::default(), ) .await }); assert_eq!(leader.await.unwrap().unwrap(), Some(1)); tokio::time::sleep(Duration::from_millis(25)).await; let follower_started = Instant::now(); let follower_cache = Arc::clone(&cache); let follower_calls = Arc::clone(&calls); let follower = tokio::spawn(async move { follower_cache .get_or_load_once_stale_while_refreshing::<(), _, _>( key, Duration::from_millis(10), Duration::from_secs(1), || async move { follower_calls.fetch_add(1, Ordering::AcqRel); Ok(Some(2)) }, CacheLoadObserver::default(), ) .await }); assert_eq!(follower.await.unwrap().unwrap(), Some(1)); assert_eq!(calls.load(Ordering::Acquire), 1); assert!( follower_started.elapsed() < Duration::from_millis(50), "follower should reuse cold-loaded stale value without reloading" ); } #[tokio::test] async fn value_cache_clear_releases_same_key_followers() { let cache = Arc::new(ValueCache::::default()); let leader_cache = Arc::clone(&cache); let leader = tokio::spawn(async move { leader_cache .get_or_load::<(), _, _>( "stuck-key".to_string(), Duration::from_secs(60), || async { std::future::pending::<()>().await; Ok(Some(1)) }, ) .await }); tokio::time::sleep(Duration::from_millis(10)).await; let follower_cache = Arc::clone(&cache); let follower = tokio::spawn(async move { follower_cache .get_or_load::<(), _, _>( "stuck-key".to_string(), Duration::from_secs(60), || async { Ok(Some(2)) }, ) .await }); tokio::time::sleep(Duration::from_millis(10)).await; cache.clear(); let value = tokio::time::timeout(Duration::from_millis(200), follower) .await .expect("clear should wake followers") .unwrap() .unwrap(); assert_eq!(value, Some(2)); leader.abort(); let _ = leader.await; } #[tokio::test] async fn value_cache_direct_insert_rejects_older_inflight_publication() { let cache = Arc::new(ValueCache::::default()); let started = Arc::new(tokio::sync::Notify::new()); let release = Arc::new(tokio::sync::Notify::new()); let follower_loads = Arc::new(AtomicUsize::new(0)); let leader_cache = Arc::clone(&cache); let leader_started = Arc::clone(&started); let leader_release = Arc::clone(&release); let leader = tokio::spawn(async move { leader_cache .get_or_load::<(), _, _>("insert-key".to_string(), Duration::from_secs(60), || { let started = Arc::clone(&leader_started); let release = Arc::clone(&leader_release); async move { started.notify_one(); release.notified().await; Ok(Some(1)) } }) .await }); tokio::time::timeout(Duration::from_secs(1), started.notified()) .await .expect("leader should start loading"); let follower_cache = Arc::clone(&cache); let follower_loads_for_task = Arc::clone(&follower_loads); let follower = tokio::spawn(async move { follower_cache .get_or_load::<(), _, _>("insert-key".to_string(), Duration::from_secs(60), || { let follower_loads = Arc::clone(&follower_loads_for_task); async move { follower_loads.fetch_add(1, Ordering::AcqRel); Ok(Some(3)) } }) .await }); tokio::time::timeout(Duration::from_secs(1), async { loop { let follower_registered = cache .singleflight .inflight .lock() .ok() .and_then(|inflight| inflight.values().next().cloned()) .is_some_and(|state| Arc::strong_count(&state) >= 4); if follower_registered { break; } tokio::task::yield_now().await; } }) .await .expect("follower should register behind the old leader"); cache.insert("insert-key".to_string(), Some(2), Duration::from_secs(60)); assert_eq!( follower .await .expect("follower should join") .expect("follower should reuse the direct write"), Some(2) ); assert_eq!(follower_loads.load(Ordering::Acquire), 0); release.notify_one(); assert_eq!(leader.await.expect("leader should join").unwrap(), Some(1)); assert_eq!( cache.get(&"insert-key".to_string(), Duration::from_secs(60)), Some(Some(2)) ); } #[tokio::test] async fn value_cache_clear_prevents_old_leader_from_reinserting_stale_value() { let cache = Arc::new(ValueCache::::default()); let started = Arc::new(tokio::sync::Notify::new()); let release = Arc::new(tokio::sync::Notify::new()); let leader_cache = Arc::clone(&cache); let leader_started = Arc::clone(&started); let leader_release = Arc::clone(&release); let leader = tokio::spawn(async move { leader_cache .get_or_load::<(), _, _>("epoch-key".to_string(), Duration::from_secs(60), || { let started = Arc::clone(&leader_started); let release = Arc::clone(&leader_release); async move { started.notify_one(); release.notified().await; Ok(Some(1)) } }) .await }); tokio::time::timeout(Duration::from_secs(1), started.notified()) .await .expect("leader should start loading"); cache.clear(); let fresh = cache .get_or_load::<(), _, _>("epoch-key".to_string(), Duration::from_secs(60), || async { Ok(Some(2)) }) .await .expect("fresh load should succeed"); assert_eq!(fresh, Some(2)); release.notify_one(); assert_eq!(leader.await.expect("leader should join").unwrap(), Some(1)); assert_eq!( cache.get(&"epoch-key".to_string(), Duration::from_secs(60)), Some(Some(2)) ); } }