RustProxy/src/cache/collapse.rs

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<bool>,
}

pub enum FillClaim {
    Disabled,
    Leader(FillLeaderGuard),
    Waited,
    Timeout,
}

pub struct FillLeaderGuard {
    cache: std::sync::Weak<Cache>,
    key: Key,
    id: u64,
    done: tokio::sync::watch::Sender<bool>,
}

struct CollapseWaiterGuard {
    cache: std::sync::Weak<Cache>,
}

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<Self>, 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::<Vec<_>>()
        };
        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::<Vec<_>>();
            keys.into_iter()
                .filter_map(|key| flights.remove(&key))
                .collect::<Vec<_>>()
        };
        for flight in removed {
            flight.done.send_replace(true);
        }
    }
}

pub struct ProactiveScheduleGuard {
    cache: std::sync::Weak<Cache>,
    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<Cache>,
}
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<Self>,
        ticket: &Ticket,
        settings: &Settings,
    ) -> Option<ProactiveScheduleGuard> {
        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<Self>,
        settings: &Settings,
    ) -> Option<ProactiveInflightGuard> {
        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),
                });
            }
        }
    }
}