Skip to main content

tenferro_runtime/runtime/
cache.rs

1#![allow(dead_code)]
2// P4-C0 deliberately lands the complete generic cache before P4-C1 wires a
3// semantic `PreparedProgram` key into `RuntimeState`. The cache is exercised by
4// runtime::tests::cache under all-targets until the production owner is added.
5
6use std::any::TypeId;
7use std::cell::RefCell;
8use std::collections::{HashMap, VecDeque};
9use std::fmt::Debug;
10use std::mem::size_of;
11use std::num::NonZeroUsize;
12use std::sync::{Arc, Condvar, Mutex, MutexGuard};
13
14use lru::LruCache;
15
16use super::cache_owner::{FrozenCacheOwner, FrozenCacheOwnerKind};
17use super::{
18    CacheInFlightBehavior, CacheOwnerFailure, CacheStats, PreparationKeySummary, PrepareError,
19    RuntimeCacheError, RuntimeStateError, SpecializationRequirements,
20};
21
22const PREPARED_CACHE_LOCK: &str = "prepared-cache.state";
23const DEFAULT_MAX_PREPARED_ENTRIES: usize = 128;
24const DEFAULT_MAX_RETAINED_BYTES: usize = 64 * 1024 * 1024;
25const DEFAULT_MAX_IN_FLIGHT_ENTRIES: usize = 16;
26const DEFAULT_MAX_QUEUED_DISTINCT_KEYS: usize = 64;
27
28thread_local! {
29    static PREPARATION_STACK: RefCell<Vec<StackFrame>> = const { RefCell::new(Vec::new()) };
30}
31
32/// Bounded prepared-plan cache limits.
33///
34/// # Examples
35///
36/// ```
37/// use tenferro_runtime::PreparedPlanCacheLimits;
38///
39/// assert_eq!(PreparedPlanCacheLimits::default().max_entries.get(), 128);
40/// ```
41#[derive(Clone, Copy, Debug, Eq, PartialEq)]
42pub struct PreparedPlanCacheLimits {
43    /// Maximum retained terminal entries.
44    pub max_entries: NonZeroUsize,
45    /// Maximum retained logical bytes.
46    pub max_retained_bytes: NonZeroUsize,
47    /// Maximum distinct keys that may actively prepare at once.
48    pub max_in_flight_entries: NonZeroUsize,
49    /// Maximum distinct queued keys waiting for an in-flight slot.
50    pub max_queued_distinct_keys: NonZeroUsize,
51}
52
53impl Default for PreparedPlanCacheLimits {
54    fn default() -> Self {
55        Self {
56            max_entries: NonZeroUsize::new(DEFAULT_MAX_PREPARED_ENTRIES)
57                .expect("default entry limit is nonzero"),
58            max_retained_bytes: NonZeroUsize::new(DEFAULT_MAX_RETAINED_BYTES)
59                .expect("default byte limit is nonzero"),
60            max_in_flight_entries: NonZeroUsize::new(DEFAULT_MAX_IN_FLIGHT_ENTRIES)
61                .expect("default in-flight limit is nonzero"),
62            max_queued_distinct_keys: NonZeroUsize::new(DEFAULT_MAX_QUEUED_DISTINCT_KEYS)
63                .expect("default queue limit is nonzero"),
64        }
65    }
66}
67
68/// Prepared-plan cache statistics.
69///
70/// # Examples
71///
72/// ```
73/// use tenferro_runtime::PreparedPlanCacheStats;
74///
75/// assert_eq!(PreparedPlanCacheStats::default().entries, 0);
76/// ```
77#[derive(Clone, Copy, Debug, Default, Eq, PartialEq)]
78pub struct PreparedPlanCacheStats {
79    /// Retained terminal entries.
80    pub entries: usize,
81    /// Logical bytes retained by entries, active attempts, and queued records.
82    pub retained_bytes: usize,
83    /// Ready or redirect cache hits.
84    pub hits: u64,
85    /// Distinct producer starts.
86    pub misses: u64,
87    /// Waits on an existing active or queued key.
88    pub waits: u64,
89    /// Hits on retained deterministic failures.
90    pub negative_hits: u64,
91    /// Producer callbacks invoked.
92    pub preparations: u64,
93    /// Retained entries evicted by entry or byte limits.
94    pub evictions: u64,
95    /// Redirect terminals returned.
96    pub redirects: u64,
97    /// Active distinct producer count.
98    pub in_flight: usize,
99    /// Highest observed active producer count.
100    pub peak_in_flight: usize,
101    /// Distinct queued keys.
102    pub queued_distinct_keys: usize,
103    /// Capacity refusals.
104    pub capacity_refusals: u64,
105}
106
107/// Aggregate runtime cache statistics.
108///
109/// # Examples
110///
111/// ```
112/// use tenferro_runtime::RuntimeCacheStats;
113///
114/// assert_eq!(RuntimeCacheStats::default().prepared_plans.entries, 0);
115/// ```
116#[derive(Clone, Copy, Debug, Default, Eq, PartialEq)]
117pub struct RuntimeCacheStats {
118    /// Runtime-owned prepared-plan cache statistics.
119    pub prepared_plans: PreparedPlanCacheStats,
120    /// Extension module cache-owner statistics.
121    pub extensions: CacheStats,
122}
123
124pub(crate) trait PreparedCacheKey: Clone + Debug + Send + Sync + 'static {
125    type Shared: Debug + Send + Sync + 'static;
126
127    fn compact_digest(&self) -> u128;
128    fn exact_eq(&self, other: &Self) -> bool;
129    fn retained_bytes(&self) -> Option<usize>;
130    fn summary(&self) -> PreparationKeySummary;
131    fn shared_retention(&self) -> Option<SharedRetention<Self::Shared>>;
132}
133
134pub(crate) trait PreparedValue: Debug + Send + Sync + 'static {
135    fn retained_bytes(&self) -> Option<usize>;
136}
137
138#[derive(Debug)]
139pub(crate) struct SharedRetention<S: Debug + Send + Sync + 'static> {
140    pub(crate) value: Arc<S>,
141    pub(crate) retained_bytes: Option<usize>,
142}
143
144impl<S: Debug + Send + Sync + 'static> Clone for SharedRetention<S> {
145    fn clone(&self) -> Self {
146        Self {
147            value: Arc::clone(&self.value),
148            retained_bytes: self.retained_bytes,
149        }
150    }
151}
152
153#[derive(Debug)]
154pub(crate) enum CacheProduced<K: PreparedCacheKey, V: PreparedValue> {
155    Ready {
156        value: Arc<V>,
157        shared: Option<SharedRetention<K::Shared>>,
158    },
159    Redirect {
160        requirements: SpecializationRequirements,
161        shared: Option<SharedRetention<K::Shared>>,
162    },
163    FailedDeterministic {
164        error: Arc<PrepareError>,
165        shared: Option<SharedRetention<K::Shared>>,
166    },
167    FailedTransient(Arc<PrepareError>),
168}
169
170#[derive(Debug)]
171pub(crate) enum CacheLookup<K: PreparedCacheKey, V: PreparedValue> {
172    Ready(Arc<V>),
173    Redirect {
174        requirements: SpecializationRequirements,
175        shared: Option<Arc<K::Shared>>,
176    },
177    FailedDeterministic(Arc<PrepareError>),
178    FailedTransient(Arc<PrepareError>),
179}
180
181#[derive(Clone, Copy, Debug, Eq, Hash, Ord, PartialEq, PartialOrd)]
182struct AttemptId(u128);
183
184#[derive(Clone, Copy, Debug, Eq, Hash, Ord, PartialEq, PartialOrd)]
185struct EntryId(u128);
186
187#[derive(Clone, Copy, Debug, Eq, Hash, Ord, PartialEq, PartialOrd)]
188struct QueueId(u128);
189
190#[derive(Clone, Copy, Debug, Eq, Hash, Ord, PartialEq, PartialOrd)]
191struct SharedRetentionId(NonZeroUsize);
192
193impl SharedRetentionId {
194    fn of<S: Debug + Send + Sync + 'static>(value: &Arc<S>) -> Self {
195        Self(NonZeroUsize::new(Arc::as_ptr(value) as usize).expect("Arc data pointer is non-null"))
196    }
197}
198
199#[derive(Debug)]
200struct CacheGenerationToken;
201
202#[derive(Debug)]
203struct CacheRevisionToken;
204
205#[derive(Debug, thiserror::Error)]
206enum CacheInternalError {
207    #[error("prepared-cache retained-byte accounting overflow")]
208    AccountingOverflow,
209    #[error("prepared-cache shared retention pointer collision")]
210    SharedPointerCollision,
211}
212
213#[derive(Debug)]
214enum CacheTerminal<K: PreparedCacheKey, V: PreparedValue> {
215    Ready(Arc<V>),
216    Redirect {
217        requirements: SpecializationRequirements,
218        shared: Option<Arc<K::Shared>>,
219    },
220    FailedDeterministic(Arc<PrepareError>),
221    FailedTransient(Arc<PrepareError>),
222}
223
224impl<K: PreparedCacheKey, V: PreparedValue> Clone for CacheTerminal<K, V> {
225    fn clone(&self) -> Self {
226        match self {
227            Self::Ready(value) => Self::Ready(Arc::clone(value)),
228            Self::Redirect {
229                requirements,
230                shared,
231            } => Self::Redirect {
232                requirements: requirements.clone(),
233                shared: shared.as_ref().map(Arc::clone),
234            },
235            Self::FailedDeterministic(error) => Self::FailedDeterministic(Arc::clone(error)),
236            Self::FailedTransient(error) => Self::FailedTransient(Arc::clone(error)),
237        }
238    }
239}
240
241#[derive(Clone, Copy, Debug)]
242enum RetainedTerminalKind {
243    Ready,
244    Redirect,
245    FailedDeterministic,
246    FailedTransient,
247}
248
249impl RetainedTerminalKind {
250    fn from_terminal<K: PreparedCacheKey, V: PreparedValue>(
251        terminal: &CacheTerminal<K, V>,
252    ) -> Self {
253        match terminal {
254            CacheTerminal::Ready(_) => Self::Ready,
255            CacheTerminal::Redirect { .. } => Self::Redirect,
256            CacheTerminal::FailedDeterministic(_) => Self::FailedDeterministic,
257            CacheTerminal::FailedTransient(_) => Self::FailedTransient,
258        }
259    }
260}
261
262#[derive(Debug)]
263enum PreparedEntryState<K: PreparedCacheKey, V: PreparedValue> {
264    Preparing {
265        attempt_id: AttemptId,
266        waiting_readers: usize,
267    },
268    Retained(CacheTerminal<K, V>),
269    Ephemeral {
270        terminal: CacheTerminal<K, V>,
271        remaining_readers: usize,
272    },
273}
274
275#[derive(Debug)]
276struct EntryRecord<K: PreparedCacheKey, V: PreparedValue> {
277    id: EntryId,
278    digest: u128,
279    key: Arc<K>,
280    state: PreparedEntryState<K, V>,
281    generation: Arc<CacheGenerationToken>,
282    retained_bytes: usize,
283    shared_id: Option<SharedRetentionId>,
284}
285
286#[derive(Debug)]
287struct ActiveAttempt<K: PreparedCacheKey> {
288    id: AttemptId,
289    entry_id: Option<EntryId>,
290    digest: u128,
291    key: Arc<K>,
292    generation: Arc<CacheGenerationToken>,
293    key_bytes: usize,
294    signature_bytes: usize,
295    attempt_metadata_bytes: usize,
296    visible_entry_bytes: usize,
297}
298
299impl<K: PreparedCacheKey> ActiveAttempt<K> {
300    fn retained_bytes_without_visible_entry(&self) -> usize {
301        self.key_bytes
302            .saturating_add(self.signature_bytes)
303            .saturating_add(self.attempt_metadata_bytes)
304    }
305}
306
307#[derive(Debug)]
308struct QueueRecord<K: PreparedCacheKey> {
309    id: QueueId,
310    digest: u128,
311    key: Arc<K>,
312    generation: Arc<CacheGenerationToken>,
313    key_bytes: usize,
314    signature_bytes: usize,
315    metadata_bytes: usize,
316    tickets: usize,
317}
318
319impl<K: PreparedCacheKey> QueueRecord<K> {
320    fn retained_bytes(&self) -> usize {
321        self.key_bytes
322            .saturating_add(self.signature_bytes)
323            .saturating_add(self.metadata_bytes)
324    }
325}
326
327#[derive(Debug)]
328struct SharedCharge<S: Debug + Send + Sync + 'static> {
329    value: Arc<S>,
330    retained_bytes: usize,
331    references: usize,
332}
333
334#[derive(Debug)]
335struct CacheState<K: PreparedCacheKey, V: PreparedValue> {
336    entry_buckets: HashMap<u128, Vec<EntryId>>,
337    entries: HashMap<EntryId, EntryRecord<K, V>>,
338    lru: LruCache<EntryId, ()>,
339    queue_buckets: HashMap<u128, Vec<QueueId>>,
340    queue_order: VecDeque<QueueId>,
341    queue_records: HashMap<QueueId, QueueRecord<K>>,
342    attempts: HashMap<AttemptId, ActiveAttempt<K>>,
343    shared: HashMap<SharedRetentionId, SharedCharge<K::Shared>>,
344    generation: Arc<CacheGenerationToken>,
345    revision: Arc<CacheRevisionToken>,
346    next_attempt: u128,
347    next_entry: u128,
348    next_queue: u128,
349    limits: PreparedPlanCacheLimits,
350    retained_bytes: usize,
351    in_flight: usize,
352    stats: PreparedPlanCacheStats,
353}
354
355impl<K: PreparedCacheKey, V: PreparedValue> CacheState<K, V> {
356    fn new(limits: PreparedPlanCacheLimits) -> Self {
357        Self {
358            entry_buckets: HashMap::new(),
359            entries: HashMap::new(),
360            lru: LruCache::unbounded(),
361            queue_buckets: HashMap::new(),
362            queue_order: VecDeque::new(),
363            queue_records: HashMap::new(),
364            attempts: HashMap::new(),
365            shared: HashMap::new(),
366            generation: Arc::new(CacheGenerationToken),
367            revision: Arc::new(CacheRevisionToken),
368            next_attempt: 0,
369            next_entry: 0,
370            next_queue: 0,
371            limits,
372            retained_bytes: 0,
373            in_flight: 0,
374            stats: PreparedPlanCacheStats::default(),
375        }
376    }
377
378    fn snapshot_stats(&self) -> PreparedPlanCacheStats {
379        let mut stats = self.stats;
380        stats.entries = self.lru.len();
381        stats.retained_bytes = self.retained_bytes;
382        stats.in_flight = self.in_flight;
383        stats.queued_distinct_keys = self.queue_records.len();
384        stats
385    }
386
387    fn bump_revision(&mut self) {
388        self.revision = Arc::new(CacheRevisionToken);
389    }
390}
391
392pub(crate) struct PreparedPlanCache<K: PreparedCacheKey, V: PreparedValue> {
393    state: Mutex<CacheState<K, V>>,
394    changed: Condvar,
395}
396
397impl<K: PreparedCacheKey, V: PreparedValue> PreparedPlanCache<K, V> {
398    pub(crate) fn new(limits: PreparedPlanCacheLimits) -> Self {
399        Self {
400            state: Mutex::new(CacheState::new(limits)),
401            changed: Condvar::new(),
402        }
403    }
404
405    pub(crate) fn get_or_prepare(
406        &self,
407        key: K,
408        behavior: CacheInFlightBehavior,
409        signature_bytes: usize,
410        producer: impl FnOnce() -> CacheProduced<K, V>,
411    ) -> Result<CacheLookup<K, V>, Arc<PrepareError>> {
412        self.get_or_prepare_inner(key, behavior, signature_bytes, producer, || {}, || {})
413    }
414
415    #[cfg(test)]
416    pub(crate) fn get_or_prepare_with_entry_wait_hooks_for_test(
417        &self,
418        key: K,
419        behavior: CacheInFlightBehavior,
420        signature_bytes: usize,
421        producer: impl FnOnce() -> CacheProduced<K, V>,
422        before_entry_wait: impl FnMut(),
423        before_condvar_wait: impl FnMut(),
424    ) -> Result<CacheLookup<K, V>, Arc<PrepareError>> {
425        self.get_or_prepare_inner(
426            key,
427            behavior,
428            signature_bytes,
429            producer,
430            before_entry_wait,
431            before_condvar_wait,
432        )
433    }
434
435    fn get_or_prepare_inner(
436        &self,
437        key: K,
438        behavior: CacheInFlightBehavior,
439        signature_bytes: usize,
440        producer: impl FnOnce() -> CacheProduced<K, V>,
441        mut before_entry_wait: impl FnMut(),
442        mut before_condvar_wait: impl FnMut(),
443    ) -> Result<CacheLookup<K, V>, Arc<PrepareError>> {
444        let summary = key.summary();
445        check_preparation_stack::<K>(summary)?;
446        let digest = key.compact_digest();
447        let mut producer = Some(producer);
448        loop {
449            match self.probe(&key, digest)? {
450                ProbeResult::Entry(entry_id) => match self.handle_entry(entry_id)? {
451                    EntryAction::Return(lookup) => return Ok(lookup),
452                    EntryAction::Wait(waiter) => {
453                        before_entry_wait();
454                        match waiter.wait_once(&mut before_condvar_wait)? {
455                            WaitOutcome::RetryProbe => continue,
456                            WaitOutcome::Return(lookup) => return Ok(lookup),
457                        }
458                    }
459                    EntryAction::ReturnRetry => continue,
460                },
461                ProbeResult::Queue(queue_id) => {
462                    let ticket = self.add_queue_ticket(queue_id)?;
463                    match ticket.wait_for_turn()? {
464                        QueueOutcome::RetryProbe => continue,
465                        QueueOutcome::Produce(guard) => {
466                            return self.run_producer(
467                                summary,
468                                guard,
469                                producer
470                                    .take()
471                                    .expect("producer is consumed by only one cache miss"),
472                            );
473                        }
474                    }
475                }
476                ProbeResult::NoMatch {
477                    revision,
478                    generation,
479                } => {
480                    match self.reserve_miss(
481                        key.clone(),
482                        digest,
483                        behavior,
484                        signature_bytes,
485                        revision,
486                        generation,
487                    )? {
488                        MissReservation::Produce(guard) => {
489                            return self.run_producer(
490                                summary,
491                                guard,
492                                producer
493                                    .take()
494                                    .expect("producer is consumed by only one cache miss"),
495                            );
496                        }
497                        MissReservation::Wait(ticket) => match ticket.wait_for_turn()? {
498                            QueueOutcome::RetryProbe => continue,
499                            QueueOutcome::Produce(guard) => {
500                                return self.run_producer(
501                                    summary,
502                                    guard,
503                                    producer
504                                        .take()
505                                        .expect("producer is consumed by only one cache miss"),
506                                );
507                            }
508                        },
509                    }
510                }
511            }
512        }
513    }
514
515    pub(crate) fn limits(&self) -> Result<PreparedPlanCacheLimits, RuntimeStateError> {
516        self.state
517            .lock()
518            .map(|state| state.limits)
519            .map_err(|_| poisoned_state_error())
520    }
521
522    pub(crate) fn set_limits(
523        &self,
524        limits: PreparedPlanCacheLimits,
525    ) -> Result<(), RuntimeStateError> {
526        let mut state = self.state.lock().map_err(|_| poisoned_state_error())?;
527        state.limits = limits;
528        evict_to_limits(&mut state);
529        state.bump_revision();
530        drop(state);
531        self.changed.notify_all();
532        Ok(())
533    }
534
535    pub(crate) fn stats(&self) -> Result<PreparedPlanCacheStats, RuntimeStateError> {
536        self.state
537            .lock()
538            .map(|state| state.snapshot_stats())
539            .map_err(|_| poisoned_state_error())
540    }
541
542    pub(crate) fn clear(&self) -> Result<(), RuntimeStateError> {
543        let mut state = self.state.lock().map_err(|_| poisoned_state_error())?;
544        clear_state(&mut state);
545        drop(state);
546        self.changed.notify_all();
547        Ok(())
548    }
549
550    #[cfg(test)]
551    pub(crate) fn ready_charge_breakdown_for_test(
552        &self,
553        key: &K,
554        value: &V,
555        has_shared: bool,
556    ) -> Option<RetainedChargeBreakdown> {
557        let key_payload = key.retained_bytes()?;
558        let value_payload = value.retained_bytes()?;
559        let entry_record = size_of::<EntryRecord<K, V>>();
560        let compact_bucket_record = size_of::<EntryId>();
561        let lru_record = size_of::<EntryId>().checked_add(size_of::<()>())?;
562        let shared_record = if has_shared {
563            size_of::<SharedRetentionId>().checked_add(size_of::<SharedCharge<K::Shared>>())?
564        } else {
565            0
566        };
567        let total = checked_sum([
568            key_payload,
569            value_payload,
570            entry_record,
571            compact_bucket_record,
572            lru_record,
573            shared_record,
574        ])?;
575        Some(RetainedChargeBreakdown {
576            key_payload,
577            value_payload,
578            entry_record,
579            compact_bucket_record,
580            lru_record,
581            shared_record,
582            total,
583        })
584    }
585
586    fn run_producer(
587        &self,
588        summary: PreparationKeySummary,
589        guard: ProducerGuard<'_, K, V>,
590        producer: impl FnOnce() -> CacheProduced<K, V>,
591    ) -> Result<CacheLookup<K, V>, Arc<PrepareError>> {
592        let stack_guard = StackGuard::push::<K>(summary);
593        let produced = producer();
594        drop(stack_guard);
595        guard.publish(produced)
596    }
597
598    fn probe(&self, key: &K, digest: u128) -> Result<ProbeResult, Arc<PrepareError>> {
599        loop {
600            let (revision, generation, entry_ids, queue_ids) = {
601                let state = self.lock_state()?;
602                (
603                    Arc::clone(&state.revision),
604                    Arc::clone(&state.generation),
605                    state
606                        .entry_buckets
607                        .get(&digest)
608                        .cloned()
609                        .unwrap_or_default(),
610                    state
611                        .queue_buckets
612                        .get(&digest)
613                        .cloned()
614                        .unwrap_or_default(),
615                )
616            };
617
618            for entry_id in entry_ids {
619                let Some(candidate_key) =
620                    self.entry_key_if_generation_current(entry_id, digest, &revision, &generation)?
621                else {
622                    continue;
623                };
624                if key.exact_eq(candidate_key.as_ref()) {
625                    let state = self.lock_state()?;
626                    if !Arc::ptr_eq(&revision, &state.revision)
627                        || !Arc::ptr_eq(&generation, &state.generation)
628                    {
629                        continue;
630                    }
631                    if state
632                        .entries
633                        .get(&entry_id)
634                        .is_some_and(|entry| entry.digest == digest)
635                    {
636                        return Ok(ProbeResult::Entry(entry_id));
637                    }
638                    continue;
639                }
640            }
641
642            for queue_id in queue_ids {
643                let Some(candidate_key) =
644                    self.queue_key_if_generation_current(queue_id, digest, &revision, &generation)?
645                else {
646                    continue;
647                };
648                if key.exact_eq(candidate_key.as_ref()) {
649                    let state = self.lock_state()?;
650                    if !Arc::ptr_eq(&revision, &state.revision)
651                        || !Arc::ptr_eq(&generation, &state.generation)
652                    {
653                        continue;
654                    }
655                    if state
656                        .queue_records
657                        .get(&queue_id)
658                        .is_some_and(|record| record.digest == digest)
659                    {
660                        return Ok(ProbeResult::Queue(queue_id));
661                    }
662                    continue;
663                }
664            }
665
666            let state = self.lock_state()?;
667            if !Arc::ptr_eq(&revision, &state.revision)
668                || !Arc::ptr_eq(&generation, &state.generation)
669            {
670                continue;
671            }
672            return Ok(ProbeResult::NoMatch {
673                revision: Arc::clone(&state.revision),
674                generation: Arc::clone(&state.generation),
675            });
676        }
677    }
678
679    fn entry_key_if_generation_current(
680        &self,
681        entry_id: EntryId,
682        digest: u128,
683        revision: &Arc<CacheRevisionToken>,
684        generation: &Arc<CacheGenerationToken>,
685    ) -> Result<Option<Arc<K>>, Arc<PrepareError>> {
686        let state = self.lock_state()?;
687        if !Arc::ptr_eq(revision, &state.revision) || !Arc::ptr_eq(generation, &state.generation) {
688            return Ok(None);
689        }
690        Ok(state
691            .entries
692            .get(&entry_id)
693            .filter(|entry| entry.digest == digest)
694            .map(|entry| Arc::clone(&entry.key)))
695    }
696
697    fn queue_key_if_generation_current(
698        &self,
699        queue_id: QueueId,
700        digest: u128,
701        revision: &Arc<CacheRevisionToken>,
702        generation: &Arc<CacheGenerationToken>,
703    ) -> Result<Option<Arc<K>>, Arc<PrepareError>> {
704        let state = self.lock_state()?;
705        if !Arc::ptr_eq(revision, &state.revision) || !Arc::ptr_eq(generation, &state.generation) {
706            return Ok(None);
707        }
708        Ok(state
709            .queue_records
710            .get(&queue_id)
711            .filter(|record| record.digest == digest)
712            .map(|record| Arc::clone(&record.key)))
713    }
714
715    fn handle_entry(&self, entry_id: EntryId) -> Result<EntryAction<'_, K, V>, Arc<PrepareError>> {
716        let mut state = self.lock_state()?;
717        if let Some((lookup, terminal_kind)) = state.entries.get(&entry_id).and_then(|entry| {
718            if let PreparedEntryState::Retained(terminal) = &entry.state {
719                let lookup = lookup_from_terminal(terminal, entry.shared_id, &state.shared);
720                Some((lookup, RetainedTerminalKind::from_terminal(terminal)))
721            } else {
722                None
723            }
724        }) {
725            match terminal_kind {
726                RetainedTerminalKind::FailedDeterministic => {
727                    state.stats.negative_hits = state.stats.negative_hits.saturating_add(1);
728                }
729                RetainedTerminalKind::Redirect => {
730                    state.stats.hits = state.stats.hits.saturating_add(1);
731                    state.stats.redirects = state.stats.redirects.saturating_add(1);
732                }
733                RetainedTerminalKind::Ready => {
734                    state.stats.hits = state.stats.hits.saturating_add(1);
735                }
736                RetainedTerminalKind::FailedTransient => {}
737            }
738            let _ = state.lru.get(&entry_id);
739            return Ok(EntryAction::Return(lookup));
740        }
741
742        let Some(entry) = state.entries.get_mut(&entry_id) else {
743            return Ok(EntryAction::ReturnRetry);
744        };
745        match &mut entry.state {
746            PreparedEntryState::Retained(_) => Ok(EntryAction::ReturnRetry),
747            PreparedEntryState::Preparing {
748                attempt_id,
749                waiting_readers,
750            } => {
751                let attempt_id = *attempt_id;
752                *waiting_readers = waiting_readers
753                    .checked_add(1)
754                    .ok_or_else(accounting_prepare_error)?;
755                let _ = entry;
756                state.stats.waits = state.stats.waits.saturating_add(1);
757                Ok(EntryAction::Wait(EntryWaitGuard {
758                    cache: self,
759                    entry_id,
760                    attempt_id,
761                    armed: true,
762                }))
763            }
764            PreparedEntryState::Ephemeral { .. } => Ok(EntryAction::ReturnRetry),
765        }
766    }
767
768    fn add_queue_ticket(
769        &self,
770        queue_id: QueueId,
771    ) -> Result<QueueTicketGuard<'_, K, V>, Arc<PrepareError>> {
772        let mut state = self.lock_state()?;
773        let Some(record) = state.queue_records.get_mut(&queue_id) else {
774            return Ok(QueueTicketGuard {
775                cache: self,
776                queue_id,
777                armed: false,
778            });
779        };
780        record.tickets = record
781            .tickets
782            .checked_add(1)
783            .ok_or_else(accounting_prepare_error)?;
784        state.stats.waits = state.stats.waits.saturating_add(1);
785        Ok(QueueTicketGuard {
786            cache: self,
787            queue_id,
788            armed: true,
789        })
790    }
791
792    fn reserve_miss(
793        &self,
794        key: K,
795        digest: u128,
796        behavior: CacheInFlightBehavior,
797        signature_bytes: usize,
798        revision: Arc<CacheRevisionToken>,
799        generation: Arc<CacheGenerationToken>,
800    ) -> Result<MissReservation<'_, K, V>, Arc<PrepareError>> {
801        let mut state = self.lock_state()?;
802        if !Arc::ptr_eq(&revision, &state.revision) || !Arc::ptr_eq(&generation, &state.generation)
803        {
804            return Ok(MissReservation::Wait(QueueTicketGuard {
805                cache: self,
806                queue_id: QueueId(0),
807                armed: false,
808            }));
809        }
810        if state.in_flight < state.limits.max_in_flight_entries.get()
811            && state.queue_order.is_empty()
812        {
813            let guard = start_attempt(
814                &mut state,
815                self,
816                Arc::new(key),
817                digest,
818                signature_bytes,
819                None,
820            )?;
821            return Ok(MissReservation::Produce(guard));
822        }
823        if behavior == CacheInFlightBehavior::Refuse {
824            state.stats.capacity_refusals = state.stats.capacity_refusals.saturating_add(1);
825            return Err(capacity_error(&state));
826        }
827        if state.queue_records.len() >= state.limits.max_queued_distinct_keys.get() {
828            state.stats.capacity_refusals = state.stats.capacity_refusals.saturating_add(1);
829            return Err(capacity_error(&state));
830        }
831        let queue_id = allocate_queue_id(&mut state);
832        let key = Arc::new(key);
833        let key_bytes = key.retained_bytes().ok_or_else(accounting_prepare_error)?;
834        let metadata_bytes = checked_sum([
835            size_of::<QueueRecord<K>>(),
836            size_of::<QueueId>(),
837            size_of::<QueueId>(),
838        ])
839        .ok_or_else(accounting_prepare_error)?;
840        let retained_bytes = checked_sum([key_bytes, signature_bytes, metadata_bytes])
841            .ok_or_else(accounting_prepare_error)?;
842        let record = QueueRecord {
843            id: queue_id,
844            digest,
845            key,
846            generation: Arc::clone(&state.generation),
847            key_bytes,
848            signature_bytes,
849            metadata_bytes,
850            tickets: 1,
851        };
852        state
853            .queue_buckets
854            .entry(digest)
855            .or_default()
856            .push(queue_id);
857        state.queue_order.push_back(queue_id);
858        state.queue_records.insert(queue_id, record);
859        state.retained_bytes = state
860            .retained_bytes
861            .checked_add(retained_bytes)
862            .ok_or_else(accounting_prepare_error)?;
863        state.stats.waits = state.stats.waits.saturating_add(1);
864        state.bump_revision();
865        self.changed.notify_all();
866        Ok(MissReservation::Wait(QueueTicketGuard {
867            cache: self,
868            queue_id,
869            armed: true,
870        }))
871    }
872
873    fn lock_state(&self) -> Result<MutexGuard<'_, CacheState<K, V>>, Arc<PrepareError>> {
874        self.state.lock().map_err(|_| {
875            Arc::new(PrepareError::CacheState {
876                source: poisoned_state_error(),
877            })
878        })
879    }
880
881    fn recover_state(&self) -> MutexGuard<'_, CacheState<K, V>> {
882        match self.state.lock() {
883            Ok(state) => state,
884            Err(error) => {
885                // INVARIANT: Guard Drop cannot return Result, so it must finish guard-owned
886                // accounting and notify waiters; the poison bit remains set, making later
887                // Result-returning cache APIs report `prepared-cache.state`.
888                error.into_inner()
889            }
890        }
891    }
892
893    #[cfg(test)]
894    pub(crate) fn poison_state_for_test(&self) {
895        let _ = std::panic::catch_unwind(std::panic::AssertUnwindSafe(|| {
896            let _state = match self.state.lock() {
897                Ok(state) => state,
898                Err(_) => panic!("cache state should not already be poisoned in test setup"),
899            };
900            panic!("poison prepared cache state for test");
901        }));
902    }
903}
904
905#[derive(Debug)]
906enum ProbeResult {
907    Entry(EntryId),
908    Queue(QueueId),
909    NoMatch {
910        revision: Arc<CacheRevisionToken>,
911        generation: Arc<CacheGenerationToken>,
912    },
913}
914
915enum EntryAction<'a, K: PreparedCacheKey, V: PreparedValue> {
916    Return(CacheLookup<K, V>),
917    Wait(EntryWaitGuard<'a, K, V>),
918    ReturnRetry,
919}
920
921enum MissReservation<'a, K: PreparedCacheKey, V: PreparedValue> {
922    Produce(ProducerGuard<'a, K, V>),
923    Wait(QueueTicketGuard<'a, K, V>),
924}
925
926enum WaitOutcome<K: PreparedCacheKey, V: PreparedValue> {
927    RetryProbe,
928    Return(CacheLookup<K, V>),
929}
930
931enum QueueOutcome<'a, K: PreparedCacheKey, V: PreparedValue> {
932    RetryProbe,
933    Produce(ProducerGuard<'a, K, V>),
934}
935
936struct ProducerGuard<'a, K: PreparedCacheKey, V: PreparedValue> {
937    cache: &'a PreparedPlanCache<K, V>,
938    attempt_id: AttemptId,
939    entry_id: EntryId,
940    published: bool,
941}
942
943impl<K: PreparedCacheKey, V: PreparedValue> ProducerGuard<'_, K, V> {
944    fn publish(
945        mut self,
946        produced: CacheProduced<K, V>,
947    ) -> Result<CacheLookup<K, V>, Arc<PrepareError>> {
948        let mut state = self.cache.lock_state()?;
949        let Some(attempt) = state.attempts.remove(&self.attempt_id) else {
950            self.published = true;
951            return Ok(lookup_from_produced(produced));
952        };
953        debug_assert_eq!(attempt.id, self.attempt_id);
954        if let Some(active_entry_id) = attempt.entry_id {
955            debug_assert_eq!(active_entry_id, self.entry_id);
956            if let Some(entry) = state.entries.get(&active_entry_id) {
957                debug_assert_eq!(attempt.digest, entry.digest);
958            }
959        }
960        finish_attempt_accounting(&mut state, &attempt);
961        let Some(entry_id) = attempt.entry_id else {
962            state.bump_revision();
963            self.published = true;
964            drop(state);
965            self.cache.changed.notify_all();
966            return Ok(lookup_from_produced(produced));
967        };
968        let Some(entry) = state.entries.get(&entry_id) else {
969            state.bump_revision();
970            self.published = true;
971            drop(state);
972            self.cache.changed.notify_all();
973            return Ok(lookup_from_produced(produced));
974        };
975        debug_assert_eq!(entry.id, entry_id);
976        debug_assert!(Arc::ptr_eq(&entry.generation, &attempt.generation));
977        debug_assert_eq!(entry.retained_bytes, attempt.visible_entry_bytes);
978        let PreparedEntryState::Preparing {
979            attempt_id,
980            waiting_readers,
981        } = &entry.state
982        else {
983            state.bump_revision();
984            self.published = true;
985            drop(state);
986            self.cache.changed.notify_all();
987            return Ok(lookup_from_produced(produced));
988        };
989        if *attempt_id != self.attempt_id {
990            state.bump_revision();
991            self.published = true;
992            drop(state);
993            self.cache.changed.notify_all();
994            return Ok(lookup_from_produced(produced));
995        }
996        let waiting_readers = *waiting_readers;
997
998        if !Arc::ptr_eq(&attempt.generation, &state.generation) {
999            remove_entry(&mut state, entry_id, false);
1000            state.bump_revision();
1001            self.published = true;
1002            drop(state);
1003            self.cache.changed.notify_all();
1004            return Ok(lookup_from_produced(produced));
1005        }
1006
1007        let publication = build_publication(attempt.key.as_ref(), produced);
1008        let producer_lookup = publication.lookup();
1009        match publication.retention {
1010            PublicationRetention::Transient | PublicationRetention::Uncacheable => {
1011                publish_ephemeral_or_remove(
1012                    &mut state,
1013                    entry_id,
1014                    publication.terminal,
1015                    waiting_readers,
1016                );
1017            }
1018            PublicationRetention::Retain {
1019                entry_bytes,
1020                shared,
1021            } => {
1022                let shared_new_bytes = shared
1023                    .as_ref()
1024                    .map(|shared| shared_new_charge(&state, shared))
1025                    .transpose()?
1026                    .unwrap_or(0);
1027                let total_new_entry = entry_bytes
1028                    .checked_add(shared_new_bytes)
1029                    .ok_or_else(accounting_prepare_error)?;
1030                if total_new_entry > state.limits.max_retained_bytes.get() {
1031                    publish_ephemeral_or_remove(
1032                        &mut state,
1033                        entry_id,
1034                        publication.terminal,
1035                        waiting_readers,
1036                    );
1037                } else {
1038                    let shared_id = if let Some(shared) = shared {
1039                        Some(retain_shared(&mut state, shared)?)
1040                    } else {
1041                        None
1042                    };
1043                    retain_entry(
1044                        &mut state,
1045                        entry_id,
1046                        publication.terminal,
1047                        entry_bytes,
1048                        shared_id,
1049                    )?;
1050                    evict_to_limits(&mut state);
1051                }
1052            }
1053        }
1054        state.bump_revision();
1055        self.published = true;
1056        drop(state);
1057        self.cache.changed.notify_all();
1058        Ok(producer_lookup)
1059    }
1060}
1061
1062impl<K: PreparedCacheKey, V: PreparedValue> Drop for ProducerGuard<'_, K, V> {
1063    fn drop(&mut self) {
1064        if self.published {
1065            return;
1066        }
1067        let mut state = self.cache.recover_state();
1068        if let Some(attempt) = state.attempts.remove(&self.attempt_id) {
1069            finish_attempt_accounting(&mut state, &attempt);
1070            if let Some(entry_id) = attempt.entry_id {
1071                remove_entry(&mut state, entry_id, false);
1072            }
1073            state.bump_revision();
1074        } else {
1075            remove_entry(&mut state, self.entry_id, false);
1076            state.bump_revision();
1077        }
1078        drop(state);
1079        self.cache.changed.notify_all();
1080    }
1081}
1082
1083struct EntryWaitGuard<'a, K: PreparedCacheKey, V: PreparedValue> {
1084    cache: &'a PreparedPlanCache<K, V>,
1085    entry_id: EntryId,
1086    attempt_id: AttemptId,
1087    armed: bool,
1088}
1089
1090impl<K: PreparedCacheKey, V: PreparedValue> EntryWaitGuard<'_, K, V> {
1091    fn wait_once(
1092        mut self,
1093        before_condvar_wait: &mut impl FnMut(),
1094    ) -> Result<WaitOutcome<K, V>, Arc<PrepareError>> {
1095        let mut state = self.cache.lock_state()?;
1096        while self.same_attempt_is_preparing(&state) {
1097            before_condvar_wait();
1098            state = match self.cache.changed.wait(state) {
1099                Ok(state) => state,
1100                Err(error) => {
1101                    drop(error.into_inner());
1102                    return Err(Arc::new(PrepareError::CacheState {
1103                        source: poisoned_state_error(),
1104                    }));
1105                }
1106            };
1107        }
1108        let outcome = self.complete_with_locked_state(&mut state);
1109        drop(state);
1110        if matches!(outcome, WaitOutcome::RetryProbe) {
1111            self.cache.changed.notify_all();
1112        }
1113        Ok(outcome)
1114    }
1115
1116    fn same_attempt_is_preparing(&self, state: &CacheState<K, V>) -> bool {
1117        matches!(
1118            state.entries.get(&self.entry_id),
1119            Some(EntryRecord {
1120                state: PreparedEntryState::Preparing { attempt_id, .. },
1121                ..
1122            }) if *attempt_id == self.attempt_id
1123        )
1124    }
1125
1126    fn complete_with_locked_state(&mut self, state: &mut CacheState<K, V>) -> WaitOutcome<K, V> {
1127        if !self.armed {
1128            return WaitOutcome::RetryProbe;
1129        }
1130        let Some(entry) = state.entries.get_mut(&self.entry_id) else {
1131            self.armed = false;
1132            return WaitOutcome::RetryProbe;
1133        };
1134        match &mut entry.state {
1135            PreparedEntryState::Preparing {
1136                attempt_id,
1137                waiting_readers,
1138            } if *attempt_id == self.attempt_id => {
1139                *waiting_readers = waiting_readers.saturating_sub(1);
1140                self.armed = false;
1141                WaitOutcome::RetryProbe
1142            }
1143            PreparedEntryState::Retained(_) => {
1144                self.armed = false;
1145                WaitOutcome::RetryProbe
1146            }
1147            PreparedEntryState::Ephemeral {
1148                terminal,
1149                remaining_readers,
1150            } => {
1151                let lookup = lookup_from_terminal(terminal, entry.shared_id, &state.shared);
1152                *remaining_readers = remaining_readers.saturating_sub(1);
1153                let remove = *remaining_readers == 0;
1154                self.armed = false;
1155                if remove {
1156                    remove_entry(state, self.entry_id, false);
1157                    state.bump_revision();
1158                }
1159                WaitOutcome::Return(lookup)
1160            }
1161            _ => {
1162                self.armed = false;
1163                WaitOutcome::RetryProbe
1164            }
1165        }
1166    }
1167}
1168
1169impl<K: PreparedCacheKey, V: PreparedValue> Drop for EntryWaitGuard<'_, K, V> {
1170    fn drop(&mut self) {
1171        if !self.armed {
1172            return;
1173        }
1174        let mut state = self.cache.recover_state();
1175        match state.entries.get_mut(&self.entry_id) {
1176            Some(EntryRecord {
1177                state:
1178                    PreparedEntryState::Preparing {
1179                        attempt_id,
1180                        waiting_readers,
1181                    },
1182                ..
1183            }) if *attempt_id == self.attempt_id => {
1184                *waiting_readers = waiting_readers.saturating_sub(1);
1185            }
1186            Some(EntryRecord {
1187                state:
1188                    PreparedEntryState::Ephemeral {
1189                        remaining_readers, ..
1190                    },
1191                ..
1192            }) => {
1193                *remaining_readers = remaining_readers.saturating_sub(1);
1194                if *remaining_readers == 0 {
1195                    remove_entry(&mut state, self.entry_id, false);
1196                    state.bump_revision();
1197                }
1198            }
1199            _ => {}
1200        }
1201        drop(state);
1202        self.cache.changed.notify_all();
1203    }
1204}
1205
1206struct QueueTicketGuard<'a, K: PreparedCacheKey, V: PreparedValue> {
1207    cache: &'a PreparedPlanCache<K, V>,
1208    queue_id: QueueId,
1209    armed: bool,
1210}
1211
1212impl<'a, K: PreparedCacheKey, V: PreparedValue> QueueTicketGuard<'a, K, V> {
1213    fn wait_for_turn(mut self) -> Result<QueueOutcome<'a, K, V>, Arc<PrepareError>> {
1214        if !self.armed {
1215            return Ok(QueueOutcome::RetryProbe);
1216        }
1217        let mut state = self.cache.lock_state()?;
1218        loop {
1219            let Some(record) = state.queue_records.get(&self.queue_id) else {
1220                self.armed = false;
1221                return Ok(QueueOutcome::RetryProbe);
1222            };
1223            debug_assert_eq!(record.id, self.queue_id);
1224            debug_assert!(Arc::ptr_eq(&record.generation, &state.generation));
1225            let at_front = state.queue_order.front().copied() == Some(self.queue_id);
1226            if at_front && state.in_flight < state.limits.max_in_flight_entries.get() {
1227                let key = Arc::clone(&record.key);
1228                let digest = record.digest;
1229                let signature_bytes = record.signature_bytes;
1230                let guard = start_attempt(
1231                    &mut state,
1232                    self.cache,
1233                    key,
1234                    digest,
1235                    signature_bytes,
1236                    Some(self.queue_id),
1237                )?;
1238                self.armed = false;
1239                return Ok(QueueOutcome::Produce(guard));
1240            }
1241            state = match self.cache.changed.wait(state) {
1242                Ok(state) => state,
1243                Err(error) => {
1244                    drop(error.into_inner());
1245                    return Err(Arc::new(PrepareError::CacheState {
1246                        source: poisoned_state_error(),
1247                    }));
1248                }
1249            };
1250        }
1251    }
1252}
1253
1254impl<K: PreparedCacheKey, V: PreparedValue> Drop for QueueTicketGuard<'_, K, V> {
1255    fn drop(&mut self) {
1256        if !self.armed {
1257            return;
1258        }
1259        let mut state = self.cache.recover_state();
1260        if let Some(record) = state.queue_records.get_mut(&self.queue_id) {
1261            record.tickets = record.tickets.saturating_sub(1);
1262            if record.tickets == 0 {
1263                remove_queue_record(&mut state, self.queue_id);
1264                state.bump_revision();
1265            }
1266        }
1267        drop(state);
1268        self.cache.changed.notify_all();
1269    }
1270}
1271
1272#[derive(Clone, Copy, Debug)]
1273struct StackFrame {
1274    type_id: TypeId,
1275    summary: PreparationKeySummary,
1276}
1277
1278struct StackGuard;
1279
1280impl StackGuard {
1281    fn push<K: PreparedCacheKey>(summary: PreparationKeySummary) -> Self {
1282        PREPARATION_STACK.with(|stack| {
1283            stack.borrow_mut().push(StackFrame {
1284                type_id: TypeId::of::<K>(),
1285                summary,
1286            });
1287        });
1288        Self
1289    }
1290}
1291
1292impl Drop for StackGuard {
1293    fn drop(&mut self) {
1294        PREPARATION_STACK.with(|stack| {
1295            let popped = stack.borrow_mut().pop();
1296            debug_assert!(popped.is_some());
1297        });
1298    }
1299}
1300
1301pub(crate) struct RuntimeCacheSet<K: PreparedCacheKey, V: PreparedValue> {
1302    prepared: PreparedPlanCache<K, V>,
1303}
1304
1305impl<K: PreparedCacheKey, V: PreparedValue> RuntimeCacheSet<K, V> {
1306    pub(crate) fn new(limits: PreparedPlanCacheLimits) -> Self {
1307        Self {
1308            prepared: PreparedPlanCache::new(limits),
1309        }
1310    }
1311
1312    pub(crate) fn cache_stats(
1313        &self,
1314        owners: &[FrozenCacheOwner],
1315    ) -> Result<RuntimeCacheStats, RuntimeCacheError> {
1316        let mut runtime_error = None;
1317        let prepared = match self.prepared.stats() {
1318            Ok(stats) => stats,
1319            Err(error) => {
1320                runtime_error = Some(error);
1321                PreparedPlanCacheStats::default()
1322            }
1323        };
1324        let mut extensions = CacheStats::default();
1325        let mut failures = Vec::new();
1326        for owner in owners {
1327            match owner.owner.cache_stats() {
1328                Ok(stats) => match owner.kind {
1329                    FrozenCacheOwnerKind::Extension => {
1330                        saturating_add_cache_stats(&mut extensions, stats);
1331                    }
1332                },
1333                Err(source) => failures.push(CacheOwnerFailure {
1334                    owner: owner.id.clone(),
1335                    source,
1336                }),
1337            }
1338        }
1339        if runtime_error.is_none() && failures.is_empty() {
1340            Ok(RuntimeCacheStats {
1341                prepared_plans: prepared,
1342                extensions,
1343            })
1344        } else {
1345            Err(RuntimeCacheError::Aggregate {
1346                runtime: runtime_error,
1347                owners: failures.into_boxed_slice(),
1348            })
1349        }
1350    }
1351
1352    pub(crate) fn clear_caches(
1353        &self,
1354        owners: &[FrozenCacheOwner],
1355    ) -> Result<(), RuntimeCacheError> {
1356        let runtime_error = self.prepared.clear().err();
1357        let mut failures = Vec::new();
1358        for owner in owners {
1359            if let Err(source) = owner.owner.clear_caches() {
1360                failures.push(CacheOwnerFailure {
1361                    owner: owner.id.clone(),
1362                    source,
1363                });
1364            }
1365        }
1366        if runtime_error.is_none() && failures.is_empty() {
1367            Ok(())
1368        } else {
1369            Err(RuntimeCacheError::Aggregate {
1370                runtime: runtime_error,
1371                owners: failures.into_boxed_slice(),
1372            })
1373        }
1374    }
1375
1376    #[allow(dead_code)]
1377    pub(crate) fn prepared(&self) -> &PreparedPlanCache<K, V> {
1378        &self.prepared
1379    }
1380}
1381
1382#[cfg(test)]
1383pub(crate) struct RetainedChargeBreakdown {
1384    pub key_payload: usize,
1385    pub value_payload: usize,
1386    pub entry_record: usize,
1387    pub compact_bucket_record: usize,
1388    pub lru_record: usize,
1389    pub shared_record: usize,
1390    pub total: usize,
1391}
1392
1393#[derive(Debug)]
1394struct Publication<K: PreparedCacheKey, V: PreparedValue> {
1395    terminal: CacheTerminal<K, V>,
1396    lookup: CacheLookup<K, V>,
1397    retention: PublicationRetention<K>,
1398}
1399
1400impl<K: PreparedCacheKey, V: PreparedValue> Publication<K, V> {
1401    fn lookup(&self) -> CacheLookup<K, V> {
1402        match &self.lookup {
1403            CacheLookup::Ready(value) => CacheLookup::Ready(Arc::clone(value)),
1404            CacheLookup::Redirect {
1405                requirements,
1406                shared,
1407            } => CacheLookup::Redirect {
1408                requirements: requirements.clone(),
1409                shared: shared.as_ref().map(Arc::clone),
1410            },
1411            CacheLookup::FailedDeterministic(error) => {
1412                CacheLookup::FailedDeterministic(Arc::clone(error))
1413            }
1414            CacheLookup::FailedTransient(error) => CacheLookup::FailedTransient(Arc::clone(error)),
1415        }
1416    }
1417}
1418
1419#[derive(Debug)]
1420enum PublicationRetention<K: PreparedCacheKey> {
1421    Retain {
1422        entry_bytes: usize,
1423        shared: Option<SharedRetention<K::Shared>>,
1424    },
1425    Transient,
1426    Uncacheable,
1427}
1428
1429fn build_publication<K: PreparedCacheKey, V: PreparedValue>(
1430    key: &K,
1431    produced: CacheProduced<K, V>,
1432) -> Publication<K, V> {
1433    match produced {
1434        CacheProduced::Ready { value, shared } => {
1435            let effective_shared = shared.or_else(|| key.shared_retention());
1436            let terminal = CacheTerminal::Ready(Arc::clone(&value));
1437            let lookup = CacheLookup::Ready(value);
1438            let retention = ready_entry_bytes::<K, V>(
1439                key,
1440                match &terminal {
1441                    CacheTerminal::Ready(value) => value,
1442                    _ => unreachable!(),
1443                },
1444            )
1445            .map(|entry_bytes| PublicationRetention::Retain {
1446                entry_bytes,
1447                shared: effective_shared,
1448            })
1449            .unwrap_or(PublicationRetention::Uncacheable);
1450            Publication {
1451                terminal,
1452                lookup,
1453                retention,
1454            }
1455        }
1456        CacheProduced::Redirect {
1457            requirements,
1458            shared,
1459        } => {
1460            let effective_shared = shared.or_else(|| key.shared_retention());
1461            let shared_arc = effective_shared
1462                .as_ref()
1463                .map(|shared| Arc::clone(&shared.value));
1464            let terminal = CacheTerminal::Redirect {
1465                requirements: requirements.clone(),
1466                shared: shared_arc.clone(),
1467            };
1468            let lookup = CacheLookup::Redirect {
1469                requirements: requirements.clone(),
1470                shared: shared_arc,
1471            };
1472            let retention = redirect_entry_bytes::<K, V>(key, &requirements)
1473                .map(|entry_bytes| PublicationRetention::Retain {
1474                    entry_bytes,
1475                    shared: effective_shared,
1476                })
1477                .unwrap_or(PublicationRetention::Uncacheable);
1478            Publication {
1479                terminal,
1480                lookup,
1481                retention,
1482            }
1483        }
1484        CacheProduced::FailedDeterministic { error, shared } => {
1485            let effective_shared = shared.or_else(|| key.shared_retention());
1486            let terminal = CacheTerminal::FailedDeterministic(Arc::clone(&error));
1487            let lookup = CacheLookup::FailedDeterministic(error);
1488            let retention = deterministic_error_entry_bytes::<K, V>(key)
1489                .map(|entry_bytes| PublicationRetention::Retain {
1490                    entry_bytes,
1491                    shared: effective_shared,
1492                })
1493                .unwrap_or(PublicationRetention::Uncacheable);
1494            Publication {
1495                terminal,
1496                lookup,
1497                retention,
1498            }
1499        }
1500        CacheProduced::FailedTransient(error) => Publication {
1501            terminal: CacheTerminal::FailedTransient(Arc::clone(&error)),
1502            lookup: CacheLookup::FailedTransient(error),
1503            retention: PublicationRetention::Transient,
1504        },
1505    }
1506}
1507
1508fn start_attempt<'a, K: PreparedCacheKey, V: PreparedValue>(
1509    state: &mut CacheState<K, V>,
1510    cache: &'a PreparedPlanCache<K, V>,
1511    key: Arc<K>,
1512    digest: u128,
1513    signature_bytes: usize,
1514    promoted_queue: Option<QueueId>,
1515) -> Result<ProducerGuard<'a, K, V>, Arc<PrepareError>> {
1516    if let Some(queue_id) = promoted_queue {
1517        remove_queue_record(state, queue_id);
1518    }
1519    let key_bytes = key.retained_bytes().ok_or_else(accounting_prepare_error)?;
1520    let attempt_metadata_bytes = size_of::<ActiveAttempt<K>>();
1521    let visible_entry_bytes = checked_sum([size_of::<EntryRecord<K, V>>(), size_of::<EntryId>()])
1522        .ok_or_else(accounting_prepare_error)?;
1523    let active_bytes = checked_sum([
1524        key_bytes,
1525        signature_bytes,
1526        attempt_metadata_bytes,
1527        visible_entry_bytes,
1528    ])
1529    .ok_or_else(accounting_prepare_error)?;
1530    let attempt_id = allocate_attempt_id(state);
1531    let entry_id = allocate_entry_id(state);
1532    let attempt = ActiveAttempt {
1533        id: attempt_id,
1534        entry_id: Some(entry_id),
1535        digest,
1536        key: Arc::clone(&key),
1537        generation: Arc::clone(&state.generation),
1538        key_bytes,
1539        signature_bytes,
1540        attempt_metadata_bytes,
1541        visible_entry_bytes,
1542    };
1543    let entry = EntryRecord {
1544        id: entry_id,
1545        digest,
1546        key,
1547        state: PreparedEntryState::Preparing {
1548            attempt_id,
1549            waiting_readers: 0,
1550        },
1551        generation: Arc::clone(&state.generation),
1552        retained_bytes: visible_entry_bytes,
1553        shared_id: None,
1554    };
1555    state.attempts.insert(attempt_id, attempt);
1556    state.entries.insert(entry_id, entry);
1557    state
1558        .entry_buckets
1559        .entry(digest)
1560        .or_default()
1561        .push(entry_id);
1562    state.retained_bytes = state
1563        .retained_bytes
1564        .checked_add(active_bytes)
1565        .ok_or_else(accounting_prepare_error)?;
1566    state.in_flight = state
1567        .in_flight
1568        .checked_add(1)
1569        .ok_or_else(accounting_prepare_error)?;
1570    state.stats.preparations = state.stats.preparations.saturating_add(1);
1571    state.stats.misses = state.stats.misses.saturating_add(1);
1572    state.stats.peak_in_flight = state.stats.peak_in_flight.max(state.in_flight);
1573    state.bump_revision();
1574    cache.changed.notify_all();
1575    Ok(ProducerGuard {
1576        cache,
1577        attempt_id,
1578        entry_id,
1579        published: false,
1580    })
1581}
1582
1583fn retain_entry<K: PreparedCacheKey, V: PreparedValue>(
1584    state: &mut CacheState<K, V>,
1585    entry_id: EntryId,
1586    terminal: CacheTerminal<K, V>,
1587    entry_bytes: usize,
1588    shared_id: Option<SharedRetentionId>,
1589) -> Result<(), Arc<PrepareError>> {
1590    let Some(entry) = state.entries.get_mut(&entry_id) else {
1591        return Ok(());
1592    };
1593    let old_bytes = entry.retained_bytes;
1594    entry.state = PreparedEntryState::Retained(terminal);
1595    entry.retained_bytes = entry_bytes;
1596    entry.shared_id = shared_id;
1597    state.retained_bytes = state.retained_bytes.saturating_sub(old_bytes);
1598    state.retained_bytes = state
1599        .retained_bytes
1600        .checked_add(entry_bytes)
1601        .ok_or_else(accounting_prepare_error)?;
1602    state.lru.put(entry_id, ());
1603    Ok(())
1604}
1605
1606fn publish_ephemeral_or_remove<K: PreparedCacheKey, V: PreparedValue>(
1607    state: &mut CacheState<K, V>,
1608    entry_id: EntryId,
1609    terminal: CacheTerminal<K, V>,
1610    waiting_readers: usize,
1611) {
1612    if waiting_readers == 0 {
1613        remove_entry(state, entry_id, false);
1614        return;
1615    }
1616    let Some(digest) = state.entries.get(&entry_id).map(|entry| entry.digest) else {
1617        return;
1618    };
1619    remove_entry_index(state, digest, entry_id);
1620    if let Some(entry) = state.entries.get_mut(&entry_id) {
1621        state.retained_bytes = state.retained_bytes.saturating_sub(entry.retained_bytes);
1622        let ephemeral_bytes = size_of::<EntryRecord<K, V>>();
1623        entry.retained_bytes = ephemeral_bytes;
1624        entry.shared_id = None;
1625        entry.state = PreparedEntryState::Ephemeral {
1626            terminal,
1627            remaining_readers: waiting_readers,
1628        };
1629        state.retained_bytes = state.retained_bytes.saturating_add(ephemeral_bytes);
1630    }
1631}
1632
1633fn evict_to_limits<K: PreparedCacheKey, V: PreparedValue>(state: &mut CacheState<K, V>) {
1634    while state.lru.len() > state.limits.max_entries.get()
1635        || state.retained_bytes > state.limits.max_retained_bytes.get()
1636    {
1637        let Some((entry_id, ())) = state.lru.pop_lru() else {
1638            break;
1639        };
1640        remove_entry(state, entry_id, true);
1641        state.stats.evictions = state.stats.evictions.saturating_add(1);
1642    }
1643}
1644
1645fn remove_entry<K: PreparedCacheKey, V: PreparedValue>(
1646    state: &mut CacheState<K, V>,
1647    entry_id: EntryId,
1648    lru_already_removed: bool,
1649) {
1650    let Some(entry) = state.entries.remove(&entry_id) else {
1651        return;
1652    };
1653    remove_entry_index(state, entry.digest, entry_id);
1654    if !lru_already_removed {
1655        let _ = state.lru.pop(&entry_id);
1656    }
1657    if let Some(shared_id) = entry.shared_id {
1658        release_shared(state, shared_id);
1659    }
1660    state.retained_bytes = state.retained_bytes.saturating_sub(entry.retained_bytes);
1661}
1662
1663fn remove_entry_index<K: PreparedCacheKey, V: PreparedValue>(
1664    state: &mut CacheState<K, V>,
1665    digest: u128,
1666    entry_id: EntryId,
1667) {
1668    if let Some(bucket) = state.entry_buckets.get_mut(&digest) {
1669        bucket.retain(|id| *id != entry_id);
1670        if bucket.is_empty() {
1671            state.entry_buckets.remove(&digest);
1672        }
1673    }
1674}
1675
1676fn remove_queue_record<K: PreparedCacheKey, V: PreparedValue>(
1677    state: &mut CacheState<K, V>,
1678    queue_id: QueueId,
1679) {
1680    let Some(record) = state.queue_records.remove(&queue_id) else {
1681        return;
1682    };
1683    if let Some(bucket) = state.queue_buckets.get_mut(&record.digest) {
1684        bucket.retain(|id| *id != queue_id);
1685        if bucket.is_empty() {
1686            state.queue_buckets.remove(&record.digest);
1687        }
1688    }
1689    state.queue_order.retain(|id| *id != queue_id);
1690    state.retained_bytes = state.retained_bytes.saturating_sub(record.retained_bytes());
1691}
1692
1693fn clear_state<K: PreparedCacheKey, V: PreparedValue>(state: &mut CacheState<K, V>) {
1694    let mut retained_bytes = 0usize;
1695    for attempt in state.attempts.values_mut() {
1696        attempt.entry_id = None;
1697        retained_bytes =
1698            retained_bytes.saturating_add(attempt.retained_bytes_without_visible_entry());
1699    }
1700    state.entries.clear();
1701    state.entry_buckets.clear();
1702    state.lru.clear();
1703    state.queue_records.clear();
1704    state.queue_buckets.clear();
1705    state.queue_order.clear();
1706    state.shared.clear();
1707    state.retained_bytes = retained_bytes;
1708    state.generation = Arc::new(CacheGenerationToken);
1709    state.bump_revision();
1710}
1711
1712fn finish_attempt_accounting<K: PreparedCacheKey, V: PreparedValue>(
1713    state: &mut CacheState<K, V>,
1714    attempt: &ActiveAttempt<K>,
1715) {
1716    state.retained_bytes = state
1717        .retained_bytes
1718        .saturating_sub(attempt.retained_bytes_without_visible_entry());
1719    state.in_flight = state.in_flight.saturating_sub(1);
1720}
1721
1722fn retain_shared<K: PreparedCacheKey, V: PreparedValue>(
1723    state: &mut CacheState<K, V>,
1724    shared: SharedRetention<K::Shared>,
1725) -> Result<SharedRetentionId, Arc<PrepareError>> {
1726    let id = SharedRetentionId::of(&shared.value);
1727    if let Some(charge) = state.shared.get_mut(&id) {
1728        if !Arc::ptr_eq(&charge.value, &shared.value) {
1729            return Err(Arc::new(PrepareError::Engine {
1730                source: Arc::new(CacheInternalError::SharedPointerCollision),
1731            }));
1732        }
1733        charge.references = charge
1734            .references
1735            .checked_add(1)
1736            .ok_or_else(accounting_prepare_error)?;
1737        return Ok(id);
1738    }
1739    let retained_bytes = shared_new_charge(state, &shared)?;
1740    state.retained_bytes = state
1741        .retained_bytes
1742        .checked_add(retained_bytes)
1743        .ok_or_else(accounting_prepare_error)?;
1744    state.shared.insert(
1745        id,
1746        SharedCharge {
1747            value: shared.value,
1748            retained_bytes,
1749            references: 1,
1750        },
1751    );
1752    Ok(id)
1753}
1754
1755fn shared_new_charge<K: PreparedCacheKey, V: PreparedValue>(
1756    state: &CacheState<K, V>,
1757    shared: &SharedRetention<K::Shared>,
1758) -> Result<usize, Arc<PrepareError>> {
1759    let id = SharedRetentionId::of(&shared.value);
1760    if let Some(charge) = state.shared.get(&id) {
1761        if Arc::ptr_eq(&charge.value, &shared.value) {
1762            return Ok(0);
1763        }
1764        return Err(Arc::new(PrepareError::Engine {
1765            source: Arc::new(CacheInternalError::SharedPointerCollision),
1766        }));
1767    }
1768    checked_sum([
1769        shared.retained_bytes.ok_or_else(accounting_prepare_error)?,
1770        size_of::<SharedRetentionId>(),
1771        size_of::<SharedCharge<K::Shared>>(),
1772    ])
1773    .ok_or_else(accounting_prepare_error)
1774}
1775
1776fn release_shared<K: PreparedCacheKey, V: PreparedValue>(
1777    state: &mut CacheState<K, V>,
1778    shared_id: SharedRetentionId,
1779) {
1780    let Some(charge) = state.shared.get_mut(&shared_id) else {
1781        return;
1782    };
1783    charge.references = charge.references.saturating_sub(1);
1784    if charge.references == 0 {
1785        let charge = state
1786            .shared
1787            .remove(&shared_id)
1788            .expect("shared charge exists");
1789        state.retained_bytes = state.retained_bytes.saturating_sub(charge.retained_bytes);
1790    }
1791}
1792
1793fn lookup_from_terminal<K: PreparedCacheKey, V: PreparedValue>(
1794    terminal: &CacheTerminal<K, V>,
1795    shared_id: Option<SharedRetentionId>,
1796    shared_table: &HashMap<SharedRetentionId, SharedCharge<K::Shared>>,
1797) -> CacheLookup<K, V> {
1798    match terminal {
1799        CacheTerminal::Ready(value) => CacheLookup::Ready(Arc::clone(value)),
1800        CacheTerminal::Redirect {
1801            requirements,
1802            shared,
1803        } => CacheLookup::Redirect {
1804            requirements: requirements.clone(),
1805            shared: shared.as_ref().map(Arc::clone).or_else(|| {
1806                shared_id.and_then(|id| shared_table.get(&id).map(|s| Arc::clone(&s.value)))
1807            }),
1808        },
1809        CacheTerminal::FailedDeterministic(error) => {
1810            CacheLookup::FailedDeterministic(Arc::clone(error))
1811        }
1812        CacheTerminal::FailedTransient(error) => CacheLookup::FailedTransient(Arc::clone(error)),
1813    }
1814}
1815
1816fn lookup_from_produced<K: PreparedCacheKey, V: PreparedValue>(
1817    produced: CacheProduced<K, V>,
1818) -> CacheLookup<K, V> {
1819    match produced {
1820        CacheProduced::Ready { value, .. } => CacheLookup::Ready(value),
1821        CacheProduced::Redirect {
1822            requirements,
1823            shared,
1824        } => CacheLookup::Redirect {
1825            requirements,
1826            shared: shared.map(|shared| shared.value),
1827        },
1828        CacheProduced::FailedDeterministic { error, .. } => CacheLookup::FailedDeterministic(error),
1829        CacheProduced::FailedTransient(error) => CacheLookup::FailedTransient(error),
1830    }
1831}
1832
1833fn ready_entry_bytes<K: PreparedCacheKey, V: PreparedValue>(key: &K, value: &V) -> Option<usize> {
1834    checked_sum([
1835        key.retained_bytes()?,
1836        value.retained_bytes()?,
1837        size_of::<EntryRecord<K, V>>(),
1838        size_of::<EntryId>(),
1839        size_of::<EntryId>().checked_add(size_of::<()>())?,
1840    ])
1841}
1842
1843fn redirect_entry_bytes<K: PreparedCacheKey, V: PreparedValue>(
1844    key: &K,
1845    requirements: &SpecializationRequirements,
1846) -> Option<usize> {
1847    checked_sum([
1848        key.retained_bytes()?,
1849        specialization_requirements_retained_bytes(requirements)?,
1850        size_of::<EntryRecord<K, V>>(),
1851        size_of::<EntryId>(),
1852        size_of::<EntryId>().checked_add(size_of::<()>())?,
1853    ])
1854}
1855
1856fn deterministic_error_entry_bytes<K: PreparedCacheKey, V: PreparedValue>(
1857    key: &K,
1858) -> Option<usize> {
1859    checked_sum([
1860        key.retained_bytes()?,
1861        size_of::<PrepareError>(),
1862        size_of::<EntryRecord<K, V>>(),
1863        size_of::<EntryId>(),
1864        size_of::<EntryId>().checked_add(size_of::<()>())?,
1865    ])
1866}
1867
1868fn specialization_requirements_retained_bytes(
1869    requirements: &SpecializationRequirements,
1870) -> Option<usize> {
1871    checked_sum([
1872        size_of::<SpecializationRequirements>(),
1873        requirements
1874            .inputs()
1875            .len()
1876            .checked_mul(size_of::<super::InputSpecializationRequirements>())?,
1877    ])
1878}
1879
1880fn allocate_attempt_id<K: PreparedCacheKey, V: PreparedValue>(
1881    state: &mut CacheState<K, V>,
1882) -> AttemptId {
1883    loop {
1884        state.next_attempt = next_nonzero_id(state.next_attempt);
1885        let id = AttemptId(state.next_attempt);
1886        if !state.attempts.contains_key(&id) {
1887            return id;
1888        }
1889    }
1890}
1891
1892fn allocate_entry_id<K: PreparedCacheKey, V: PreparedValue>(
1893    state: &mut CacheState<K, V>,
1894) -> EntryId {
1895    loop {
1896        state.next_entry = next_nonzero_id(state.next_entry);
1897        let id = EntryId(state.next_entry);
1898        if !state.entries.contains_key(&id) {
1899            return id;
1900        }
1901    }
1902}
1903
1904fn allocate_queue_id<K: PreparedCacheKey, V: PreparedValue>(
1905    state: &mut CacheState<K, V>,
1906) -> QueueId {
1907    loop {
1908        state.next_queue = next_nonzero_id(state.next_queue);
1909        let id = QueueId(state.next_queue);
1910        if !state.queue_records.contains_key(&id) {
1911            return id;
1912        }
1913    }
1914}
1915
1916fn next_nonzero_id(current: u128) -> u128 {
1917    let next = current.wrapping_add(1);
1918    if next == 0 {
1919        1
1920    } else {
1921        next
1922    }
1923}
1924
1925fn checked_sum(values: impl IntoIterator<Item = usize>) -> Option<usize> {
1926    values
1927        .into_iter()
1928        .try_fold(0usize, |sum, value| sum.checked_add(value))
1929}
1930
1931fn check_preparation_stack<K: PreparedCacheKey>(
1932    requested: PreparationKeySummary,
1933) -> Result<(), Arc<PrepareError>> {
1934    PREPARATION_STACK.with(|stack| {
1935        let stack = stack.borrow();
1936        let Some(parent) = stack.last().copied() else {
1937            return Ok(());
1938        };
1939        if parent.type_id == TypeId::of::<K>() && parent.summary == requested {
1940            Err(Arc::new(PrepareError::PreparationCycle { key: requested }))
1941        } else {
1942            Err(Arc::new(PrepareError::NestedPreparationUnsupported {
1943                parent: parent.summary,
1944                requested,
1945            }))
1946        }
1947    })
1948}
1949
1950fn poisoned_state_error() -> RuntimeStateError {
1951    RuntimeStateError::Poisoned {
1952        lock: PREPARED_CACHE_LOCK,
1953    }
1954}
1955
1956fn accounting_prepare_error() -> Arc<PrepareError> {
1957    Arc::new(PrepareError::Engine {
1958        source: Arc::new(CacheInternalError::AccountingOverflow),
1959    })
1960}
1961
1962fn capacity_error<K: PreparedCacheKey, V: PreparedValue>(
1963    state: &CacheState<K, V>,
1964) -> Arc<PrepareError> {
1965    Arc::new(PrepareError::CacheInFlightCapacityExceeded {
1966        in_flight: state.in_flight,
1967        queued_distinct_keys: state.queue_records.len(),
1968    })
1969}
1970
1971fn saturating_add_cache_stats(total: &mut CacheStats, value: CacheStats) {
1972    total.entries = total.entries.saturating_add(value.entries);
1973    total.retained_bytes = total.retained_bytes.saturating_add(value.retained_bytes);
1974    total.hits = total.hits.saturating_add(value.hits);
1975    total.misses = total.misses.saturating_add(value.misses);
1976    total.evictions = total.evictions.saturating_add(value.evictions);
1977    total.clears = total.clears.saturating_add(value.clears);
1978}