const UNCACHEABLE_MEMO_TTL: Duration = Duration::from_secs(30); const UNCACHEABLE_MEMO_CAPACITY: usize = 8192; struct FillFlight { id: u64, scope: String, scheme: String, authority: String, target: String, done: tokio::sync::watch::Sender, } pub enum FillClaim { Disabled, Leader(FillLeaderGuard), Waited, Timeout, } pub struct FillLeaderGuard { cache: std::sync::Weak, key: Key, id: u64, done: tokio::sync::watch::Sender, } struct CollapseWaiterGuard { cache: std::sync::Weak, } impl Drop for CollapseWaiterGuard { fn drop(&mut self) { if let Some(cache) = self.cache.upgrade() { cache.collapse_waiters.fetch_sub(1, Ordering::Relaxed); } } } impl Drop for FillLeaderGuard { fn drop(&mut self) { if let Some(cache) = self.cache.upgrade() { let mut flights = cache.fill_flights.lock().unwrap_or_else(|e| e.into_inner()); let current = flights .get(&self.key) .is_some_and(|flight| flight.id == self.id); if current { flights.remove(&self.key); } } self.done.send_replace(true); } } impl Cache { pub async fn claim_fill(self: &Arc, ticket: &Ticket, settings: &Settings) -> FillClaim { if !settings.collapse_enabled || self.recently_uncacheable(&ticket.key) { return FillClaim::Disabled; } let existing = { let mut flights = self.fill_flights.lock().unwrap_or_else(|e| e.into_inner()); if let Some(flight) = flights.get(&ticket.key) { Some(flight.clone()) } else { let id = self .fill_flight_sequence .fetch_add(1, Ordering::Relaxed) .wrapping_add(1); let (done, _) = tokio::sync::watch::channel(false); flights.insert( ticket.key, Arc::new(FillFlight { id, scope: ticket.scope.clone(), scheme: ticket.scheme.clone(), authority: ticket.authority.clone(), target: ticket.target.clone(), done: done.clone(), }), ); return FillClaim::Leader(FillLeaderGuard { cache: Arc::downgrade(self), key: ticket.key, id, done, }); } }; let Some(flight) = existing else { unreachable!("existing fill resolved above") }; let mut done = flight.done.subscribe(); self.collapse_waiters.fetch_add(1, Ordering::Relaxed); let _waiter = CollapseWaiterGuard { cache: Arc::downgrade(self), }; let waited = if *done.borrow() { Ok(()) } else { tokio::time::timeout( Duration::from_millis(settings.collapse_wait_timeout_ms), done.changed(), ) .await .map(|_| ()) }; match waited { Ok(()) => { self.counters .collapsed_fills .fetch_add(1, Ordering::Relaxed); FillClaim::Waited } Err(_) => { self.counters .collapse_timeouts .fetch_add(1, Ordering::Relaxed); FillClaim::Timeout } } } /// Marks a key as uncacheable so requests skip coalescing. pub fn remember_uncacheable(&self, ticket: &Ticket) { let now = std::time::Instant::now(); let mut memo = self .uncacheable_until .lock() .unwrap_or_else(|e| e.into_inner()); if memo.len() >= UNCACHEABLE_MEMO_CAPACITY { memo.retain(|_, until| *until > now); if memo.len() >= UNCACHEABLE_MEMO_CAPACITY { return; } } memo.insert(ticket.key, now + UNCACHEABLE_MEMO_TTL); } fn recently_uncacheable(&self, key: &Key) -> bool { let mut memo = self .uncacheable_until .lock() .unwrap_or_else(|e| e.into_inner()); match memo.get(key) { Some(until) if *until > std::time::Instant::now() => true, Some(_) => { memo.remove(key); false } None => false, } } fn forget_uncacheable(&self) { self.uncacheable_until .lock() .unwrap_or_else(|e| e.into_inner()) .clear(); } fn cancel_all_fill_flights(&self) { self.forget_uncacheable(); let flights = { let mut flights = self.fill_flights.lock().unwrap_or_else(|e| e.into_inner()); flights .drain() .map(|(_, flight)| flight) .collect::>() }; for flight in flights { flight.done.send_replace(true); } } fn cancel_scope_fill_flights(&self, scope: &str) { self.forget_uncacheable(); let removed = { let mut flights = self.fill_flights.lock().unwrap_or_else(|e| e.into_inner()); let keys = flights .iter() .filter(|(_, flight)| flight.scope == scope) .map(|(key, _)| *key) .collect::>(); keys.into_iter() .filter_map(|key| flights.remove(&key)) .collect::>() }; for flight in removed { flight.done.send_replace(true); } } } pub struct ProactiveScheduleGuard { cache: std::sync::Weak, key: Key, } impl Drop for ProactiveScheduleGuard { fn drop(&mut self) { if let Some(cache) = self.cache.upgrade() { cache .proactive_scheduled_keys .lock() .unwrap_or_else(|e| e.into_inner()) .remove(&self.key); } } } pub struct ProactiveInflightGuard { cache: std::sync::Weak, } impl Drop for ProactiveInflightGuard { fn drop(&mut self) { if let Some(cache) = self.cache.upgrade() { cache.proactive_inflight.fetch_sub(1, Ordering::Relaxed); } } } impl Cache { pub fn begin_proactive_schedule( self: &Arc, ticket: &Ticket, settings: &Settings, ) -> Option { if !settings.proactive_refresh_enabled || settings.proactive_refresh_percent == 0 { return None; } let sample = u16::from_be_bytes([ticket.key[0], ticket.key[1]]) % 100; if sample >= settings.proactive_refresh_percent as u16 { self.counters .proactive_skipped .fetch_add(1, Ordering::Relaxed); return None; } let mut scheduled = self .proactive_scheduled_keys .lock() .unwrap_or_else(|e| e.into_inner()); let pending_limit = settings .proactive_refresh_max_inflight .saturating_mul(64) .clamp(64, 256); if scheduled.len() >= pending_limit && !scheduled.contains(&ticket.key) { self.counters .proactive_skipped .fetch_add(1, Ordering::Relaxed); return None; } if !scheduled.insert(ticket.key) { return None; } self.counters .proactive_scheduled .fetch_add(1, Ordering::Relaxed); Some(ProactiveScheduleGuard { cache: Arc::downgrade(self), key: ticket.key, }) } pub fn try_begin_proactive_refresh( self: &Arc, settings: &Settings, ) -> Option { let limit = settings.proactive_refresh_max_inflight as u64; loop { let current = self.proactive_inflight.load(Ordering::Relaxed); if current >= limit { self.counters .proactive_skipped .fetch_add(1, Ordering::Relaxed); return None; } if self .proactive_inflight .compare_exchange_weak(current, current + 1, Ordering::AcqRel, Ordering::Relaxed) .is_ok() { self.counters .proactive_started .fetch_add(1, Ordering::Relaxed); return Some(ProactiveInflightGuard { cache: Arc::downgrade(self), }); } } } }