1#![allow(dead_code)]
2use 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#[derive(Clone, Copy, Debug, Eq, PartialEq)]
42pub struct PreparedPlanCacheLimits {
43 pub max_entries: NonZeroUsize,
45 pub max_retained_bytes: NonZeroUsize,
47 pub max_in_flight_entries: NonZeroUsize,
49 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#[derive(Clone, Copy, Debug, Default, Eq, PartialEq)]
78pub struct PreparedPlanCacheStats {
79 pub entries: usize,
81 pub retained_bytes: usize,
83 pub hits: u64,
85 pub misses: u64,
87 pub waits: u64,
89 pub negative_hits: u64,
91 pub preparations: u64,
93 pub evictions: u64,
95 pub redirects: u64,
97 pub in_flight: usize,
99 pub peak_in_flight: usize,
101 pub queued_distinct_keys: usize,
103 pub capacity_refusals: u64,
105}
106
107#[derive(Clone, Copy, Debug, Default, Eq, PartialEq)]
117pub struct RuntimeCacheStats {
118 pub prepared_plans: PreparedPlanCacheStats,
120 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 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}