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