Skip to main content

acyclic_fs/
cache.rs

1//! Bounded disposable immutable-object caching and single-flight reads.
2//!
3//! This layer owns no authority. Every entry is keyed by a canonical typed
4//! object identity and may be discarded without changing filesystem results.
5
6use crate::async_storage::{
7    AsyncObjectStore, DecodedCacheAdmission, DecodedCacheKey, DecodedCacheValue,
8};
9use crate::cancellation::CancellationToken;
10use crate::performance::{WorkBudget, WorkCounters};
11use crate::storage::{
12    ObjectFailure, ObjectId, ObjectRead, ObjectReadRequest, ObjectReadRetention, ObjectReceipt,
13    ObjectResult, ObjectStoreError, ObjectWrite,
14};
15use bytes::Bytes;
16use std::collections::{BTreeMap, BTreeSet};
17use std::future::{Future, poll_fn};
18use std::mem::size_of;
19use std::pin::pin;
20use std::sync::{Arc, Mutex, MutexGuard};
21use std::task::{Context, Poll, Waker};
22use thiserror::Error;
23
24/// Hard process-local bounds for one disposable immutable-object cache.
25#[derive(Clone, Copy, Debug, Eq, PartialEq)]
26pub struct ObjectCacheOptions {
27    /// Maximum resident immutable objects.
28    pub maximum_entries: u32,
29    /// Maximum resident canonical object bytes.
30    pub maximum_bytes: u64,
31    /// Maximum distinct in-flight object identities.
32    pub maximum_in_flight: u32,
33    /// Maximum followers retained behind one in-flight read.
34    pub maximum_waiters_per_object: u32,
35}
36
37impl Default for ObjectCacheOptions {
38    fn default() -> Self {
39        Self {
40            maximum_entries: 4_096,
41            maximum_bytes: 256 * 1024 * 1024,
42            maximum_in_flight: 1_024,
43            maximum_waiters_per_object: 1_024,
44        }
45    }
46}
47
48/// Invalid cache configuration.
49#[derive(Clone, Copy, Debug, Eq, Error, PartialEq)]
50pub enum ObjectCacheConfigError {
51    /// Every capacity must admit forward progress.
52    #[error("every immutable-object cache bound must be positive")]
53    ZeroBound,
54}
55
56impl ObjectCacheOptions {
57    fn validate(self) -> Result<Self, ObjectCacheConfigError> {
58        if self.maximum_entries == 0
59            || self.maximum_bytes == 0
60            || self.maximum_in_flight == 0
61            || self.maximum_waiters_per_object == 0
62        {
63            return Err(ObjectCacheConfigError::ZeroBound);
64        }
65        Ok(self)
66    }
67}
68
69/// Exact process-local accelerator observations.
70#[derive(Clone, Copy, Debug, Default, Eq, PartialEq)]
71pub struct ObjectCacheStats {
72    /// Complete object reads served from resident immutable bytes.
73    pub hits: u64,
74    /// Authenticated decoded pages reused without canonical decode.
75    pub decoded_hits: u64,
76    /// Object reads requiring a backend request or in-flight wait.
77    pub misses: u64,
78    /// Reads that reused one concurrent backend result.
79    pub coalesced_reads: u64,
80    /// Entries removed by deterministic least-recently-used eviction.
81    pub evictions: u64,
82    /// Current resident entries.
83    pub resident_entries: u64,
84    /// Current canonical plus decoded logical bytes.
85    pub resident_bytes: u64,
86    /// Current resident canonical immutable objects.
87    pub resident_canonical_objects: u64,
88    /// Current resident canonical immutable-object bytes.
89    pub resident_canonical_bytes: u64,
90    /// Current resident authenticated decoded pages.
91    pub resident_decoded_pages: u64,
92    /// Current resident decoded-page logical bytes.
93    pub resident_decoded_bytes: u64,
94    /// Current distinct in-flight identities.
95    pub in_flight: u64,
96}
97
98struct CacheEntry {
99    bytes: Bytes,
100    touched: u64,
101}
102
103struct DecodedCacheEntry {
104    value: DecodedCacheValue,
105    touched: u64,
106}
107
108#[derive(Clone, Copy, Debug, Eq, Ord, PartialEq, PartialOrd)]
109enum RecencyKey {
110    Canonical(ObjectId),
111    Decoded(DecodedCacheKey),
112}
113
114#[derive(Default)]
115struct MutableStats {
116    hits: u64,
117    decoded_hits: u64,
118    misses: u64,
119    coalesced_reads: u64,
120    evictions: u64,
121}
122
123struct CacheState {
124    entries: BTreeMap<ObjectId, CacheEntry>,
125    decoded: BTreeMap<DecodedCacheKey, DecodedCacheEntry>,
126    recency: BTreeSet<(u64, RecencyKey)>,
127    flights: BTreeMap<ObjectId, Arc<Flight>>,
128    resident_bytes: u64,
129    resident_canonical_bytes: u64,
130    resident_decoded_bytes: u64,
131    clock: u64,
132    stats: MutableStats,
133}
134
135impl Default for CacheState {
136    fn default() -> Self {
137        Self {
138            entries: BTreeMap::new(),
139            decoded: BTreeMap::new(),
140            recency: BTreeSet::new(),
141            flights: BTreeMap::new(),
142            resident_bytes: 0,
143            resident_canonical_bytes: 0,
144            resident_decoded_bytes: 0,
145            clock: 1,
146            stats: MutableStats::default(),
147        }
148    }
149}
150
151struct Shared {
152    options: ObjectCacheOptions,
153    state: Mutex<CacheState>,
154}
155
156enum FlightState {
157    Running(Vec<Waker>),
158    Complete(Option<Bytes>),
159}
160
161struct Flight {
162    state: Mutex<FlightState>,
163    maximum_waiters: u32,
164}
165
166impl Flight {
167    fn new(maximum_waiters: u32) -> Self {
168        Self {
169            state: Mutex::new(FlightState::Running(Vec::new())),
170            maximum_waiters,
171        }
172    }
173
174    fn finish(&self, bytes: Option<Bytes>) {
175        let waiters = match self.state.lock() {
176            Ok(mut state) => match std::mem::replace(&mut *state, FlightState::Complete(bytes)) {
177                FlightState::Running(waiters) => waiters,
178                FlightState::Complete(_) => return,
179            },
180            Err(poisoned) => {
181                let mut state = poisoned.into_inner();
182                match std::mem::replace(&mut *state, FlightState::Complete(None)) {
183                    FlightState::Running(waiters) => waiters,
184                    FlightState::Complete(_) => return,
185                }
186            }
187        };
188        for waiter in waiters {
189            waiter.wake();
190        }
191    }
192
193    fn poll(&self, context: &Context<'_>) -> Poll<Result<Option<Bytes>, ObjectStoreError>> {
194        let mut state = self.state.lock().map_err(|_| ObjectStoreError::Corrupt)?;
195        match &mut *state {
196            FlightState::Complete(bytes) => Poll::Ready(Ok(bytes.clone())),
197            FlightState::Running(waiters) => {
198                if waiters
199                    .iter()
200                    .any(|waiter| waiter.will_wake(context.waker()))
201                {
202                    return Poll::Pending;
203                }
204                if waiters.len() >= usize::try_from(self.maximum_waiters).unwrap_or(usize::MAX) {
205                    return Poll::Ready(Err(ObjectStoreError::Rejected(
206                        "immutable-object single-flight waiter bound exhausted".to_owned(),
207                    )));
208                }
209                waiters.push(context.waker().clone());
210                Poll::Pending
211            }
212        }
213    }
214}
215
216struct ActiveFlightGuard<'a, S> {
217    store: &'a CachedObjectStore<S>,
218    object_id: ObjectId,
219    flight: Arc<Flight>,
220    finished: bool,
221}
222
223impl<'a, S> ActiveFlightGuard<'a, S> {
224    fn new(store: &'a CachedObjectStore<S>, object_id: ObjectId, flight: Arc<Flight>) -> Self {
225        Self {
226            store,
227            object_id,
228            flight,
229            finished: false,
230        }
231    }
232
233    fn finish(mut self, bytes: Option<Bytes>) -> Result<(), ObjectStoreError> {
234        self.flight.finish(bytes);
235        self.store.remove_flight(self.object_id, &self.flight)?;
236        self.finished = true;
237        Ok(())
238    }
239}
240
241impl<S> Drop for ActiveFlightGuard<'_, S> {
242    fn drop(&mut self) {
243        if !self.finished {
244            self.flight.finish(None);
245            let _ = self.store.remove_flight(self.object_id, &self.flight);
246        }
247    }
248}
249
250enum ReadPlan {
251    Hit(Bytes),
252    Leader(Arc<Flight>),
253    Follower(Arc<Flight>),
254}
255
256enum BatchSlot {
257    Hit(Bytes),
258    LocalLeader(usize),
259    ExternalFollower(Arc<Flight>),
260}
261
262struct BatchLeader {
263    object_id: ObjectId,
264    maximum_bytes: u64,
265    flight: Arc<Flight>,
266}
267
268struct PreparedBatch {
269    slots: Vec<BatchSlot>,
270    leaders: Vec<BatchLeader>,
271    work: WorkCounters,
272    planning_bytes: u64,
273    result_bytes: u64,
274    item_count: u64,
275}
276
277/// Async immutable-object accelerator over any conforming backend.
278pub struct CachedObjectStore<S> {
279    inner: S,
280    shared: Arc<Shared>,
281}
282
283impl<S: Clone> Clone for CachedObjectStore<S> {
284    fn clone(&self) -> Self {
285        Self {
286            inner: self.inner.clone(),
287            shared: Arc::clone(&self.shared),
288        }
289    }
290}
291
292impl<S> CachedObjectStore<S> {
293    /// Wraps one backend with an independently bounded disposable cache.
294    ///
295    /// # Errors
296    ///
297    /// Rejects a zero bound.
298    pub fn new(inner: S, options: ObjectCacheOptions) -> Result<Self, ObjectCacheConfigError> {
299        Ok(Self {
300            inner,
301            shared: Arc::new(Shared {
302                options: options.validate()?,
303                state: Mutex::new(CacheState::default()),
304            }),
305        })
306    }
307
308    /// Returns the wrapped correctness backend.
309    #[must_use]
310    pub const fn inner(&self) -> &S {
311        &self.inner
312    }
313
314    /// Returns exact current and cumulative accelerator observations.
315    ///
316    /// # Errors
317    ///
318    /// Fails closed if accelerator state was poisoned.
319    pub fn stats(&self) -> Result<ObjectCacheStats, ObjectStoreError> {
320        let state = self.lock_state()?;
321        Ok(ObjectCacheStats {
322            hits: state.stats.hits,
323            decoded_hits: state.stats.decoded_hits,
324            misses: state.stats.misses,
325            coalesced_reads: state.stats.coalesced_reads,
326            evictions: state.stats.evictions,
327            resident_entries: u64::try_from(
328                state.entries.len().saturating_add(state.decoded.len()),
329            )
330            .unwrap_or(u64::MAX),
331            resident_bytes: state.resident_bytes,
332            resident_canonical_objects: u64::try_from(state.entries.len()).unwrap_or(u64::MAX),
333            resident_canonical_bytes: state.resident_canonical_bytes,
334            resident_decoded_pages: u64::try_from(state.decoded.len()).unwrap_or(u64::MAX),
335            resident_decoded_bytes: state.resident_decoded_bytes,
336            in_flight: u64::try_from(state.flights.len()).unwrap_or(u64::MAX),
337        })
338    }
339
340    /// Discards every resident byte. In-flight reads remain valid and may
341    /// repopulate after this call.
342    ///
343    /// # Errors
344    ///
345    /// Fails closed if accelerator state was poisoned.
346    pub fn clear(&self) -> Result<(), ObjectStoreError> {
347        let mut state = self.lock_state()?;
348        state.entries.clear();
349        state.decoded.clear();
350        state.recency.clear();
351        state.resident_bytes = 0;
352        state.resident_canonical_bytes = 0;
353        state.resident_decoded_bytes = 0;
354        Ok(())
355    }
356
357    fn lock_state(&self) -> Result<MutexGuard<'_, CacheState>, ObjectStoreError> {
358        self.shared
359            .state
360            .lock()
361            .map_err(|_| ObjectStoreError::Corrupt)
362    }
363
364    fn read_plan(&self, object_id: ObjectId) -> Result<ReadPlan, ObjectStoreError> {
365        let mut state = self.lock_state()?;
366        let tick = next_tick(&mut state)?;
367        if let Some(entry) = state.entries.get_mut(&object_id) {
368            let prior = entry.touched;
369            entry.touched = tick;
370            let bytes = entry.bytes.clone();
371            state
372                .recency
373                .remove(&(prior, RecencyKey::Canonical(object_id)));
374            state
375                .recency
376                .insert((tick, RecencyKey::Canonical(object_id)));
377            state.stats.hits = increment(state.stats.hits)?;
378            return Ok(ReadPlan::Hit(bytes));
379        }
380        state.stats.misses = increment(state.stats.misses)?;
381        if let Some(flight) = state.flights.get(&object_id) {
382            return Ok(ReadPlan::Follower(Arc::clone(flight)));
383        }
384        if state.flights.len()
385            >= usize::try_from(self.shared.options.maximum_in_flight).unwrap_or(usize::MAX)
386        {
387            return Err(ObjectStoreError::Rejected(
388                "immutable-object single-flight identity bound exhausted".to_owned(),
389            ));
390        }
391        let flight = Arc::new(Flight::new(self.shared.options.maximum_waiters_per_object));
392        state.flights.insert(object_id, Arc::clone(&flight));
393        Ok(ReadPlan::Leader(flight))
394    }
395
396    fn cached_contains(&self, object_id: ObjectId) -> Result<bool, ObjectStoreError> {
397        let mut state = self.lock_state()?;
398        let tick = next_tick(&mut state)?;
399        let Some(entry) = state.entries.get_mut(&object_id) else {
400            state.stats.misses = increment(state.stats.misses)?;
401            return Ok(false);
402        };
403        let prior = entry.touched;
404        entry.touched = tick;
405        state
406            .recency
407            .remove(&(prior, RecencyKey::Canonical(object_id)));
408        state
409            .recency
410            .insert((tick, RecencyKey::Canonical(object_id)));
411        state.stats.hits = increment(state.stats.hits)?;
412        Ok(true)
413    }
414
415    fn insert(&self, object_id: ObjectId, bytes: Bytes) -> Result<bool, ObjectStoreError> {
416        let length = u64::try_from(bytes.len()).unwrap_or(u64::MAX);
417        if length > self.shared.options.maximum_bytes {
418            return Ok(false);
419        }
420        let mut state = self.lock_state()?;
421        if let Some(existing) = state.entries.get(&object_id) {
422            if existing.bytes != bytes {
423                return Err(ObjectStoreError::Corrupt);
424            }
425            return Ok(true);
426        }
427        while state.entries.len().saturating_add(state.decoded.len())
428            >= usize::try_from(self.shared.options.maximum_entries).unwrap_or(usize::MAX)
429            || state
430                .resident_bytes
431                .checked_add(length)
432                .is_none_or(|total| total > self.shared.options.maximum_bytes)
433        {
434            evict_oldest(&mut state)?;
435        }
436        let tick = next_tick(&mut state)?;
437        state.resident_bytes = state
438            .resident_bytes
439            .checked_add(length)
440            .ok_or(ObjectStoreError::Work(crate::WorkError::Overflow))?;
441        state.resident_canonical_bytes = state
442            .resident_canonical_bytes
443            .checked_add(length)
444            .ok_or(ObjectStoreError::Work(crate::WorkError::Overflow))?;
445        state.entries.insert(
446            object_id,
447            CacheEntry {
448                bytes,
449                touched: tick,
450            },
451        );
452        state
453            .recency
454            .insert((tick, RecencyKey::Canonical(object_id)));
455        Ok(true)
456    }
457
458    fn decoded_get(
459        &self,
460        key: DecodedCacheKey,
461    ) -> Result<Option<DecodedCacheValue>, ObjectStoreError> {
462        let mut state = self.lock_state()?;
463        let tick = next_tick(&mut state)?;
464        let Some(entry) = state.decoded.get_mut(&key) else {
465            return Ok(None);
466        };
467        let prior = entry.touched;
468        entry.touched = tick;
469        let value = entry.value.clone();
470        state.recency.remove(&(prior, RecencyKey::Decoded(key)));
471        state.recency.insert((tick, RecencyKey::Decoded(key)));
472        state.stats.hits = increment(state.stats.hits)?;
473        state.stats.decoded_hits = increment(state.stats.decoded_hits)?;
474        Ok(Some(value))
475    }
476
477    fn decoded_admit(
478        &self,
479        key: DecodedCacheKey,
480        value: DecodedCacheValue,
481    ) -> Result<DecodedCacheAdmission, ObjectStoreError> {
482        if value.logical_bytes > self.shared.options.maximum_bytes {
483            return Ok(DecodedCacheAdmission::Uncached(value));
484        }
485        let mut state = self.lock_state()?;
486        if let Some(existing) = state.decoded.get(&key) {
487            return Ok(DecodedCacheAdmission::Shared(existing.value.clone()));
488        }
489        while state.entries.len().saturating_add(state.decoded.len())
490            >= usize::try_from(self.shared.options.maximum_entries).unwrap_or(usize::MAX)
491            || state
492                .resident_bytes
493                .checked_add(value.logical_bytes)
494                .is_none_or(|total| total > self.shared.options.maximum_bytes)
495        {
496            evict_oldest(&mut state)?;
497        }
498        let tick = next_tick(&mut state)?;
499        state.resident_bytes = state
500            .resident_bytes
501            .checked_add(value.logical_bytes)
502            .ok_or(ObjectStoreError::Work(crate::WorkError::Overflow))?;
503        state.resident_decoded_bytes = state
504            .resident_decoded_bytes
505            .checked_add(value.logical_bytes)
506            .ok_or(ObjectStoreError::Work(crate::WorkError::Overflow))?;
507        state.decoded.insert(
508            key,
509            DecodedCacheEntry {
510                value: value.clone(),
511                touched: tick,
512            },
513        );
514        state.recency.insert((tick, RecencyKey::Decoded(key)));
515        Ok(DecodedCacheAdmission::Shared(value))
516    }
517
518    fn remove_flight(
519        &self,
520        object_id: ObjectId,
521        flight: &Arc<Flight>,
522    ) -> Result<(), ObjectStoreError> {
523        let mut state = self.lock_state()?;
524        if state
525            .flights
526            .get(&object_id)
527            .is_some_and(|active| Arc::ptr_eq(active, flight))
528        {
529            state.flights.remove(&object_id);
530        }
531        Ok(())
532    }
533
534    fn record_coalesced(&self) -> Result<(), ObjectStoreError> {
535        let mut state = self.lock_state()?;
536        state.stats.coalesced_reads = increment(state.stats.coalesced_reads)?;
537        Ok(())
538    }
539}
540
541fn increment(value: u64) -> Result<u64, ObjectStoreError> {
542    value
543        .checked_add(1)
544        .ok_or(ObjectStoreError::Work(crate::WorkError::Overflow))
545}
546
547fn next_tick(state: &mut CacheState) -> Result<u64, ObjectStoreError> {
548    let tick = state.clock;
549    state.clock = increment(state.clock)?;
550    Ok(tick)
551}
552
553fn evict_oldest(state: &mut CacheState) -> Result<(), ObjectStoreError> {
554    let &(touched, key) = state.recency.first().ok_or(ObjectStoreError::Corrupt)?;
555    state.recency.remove(&(touched, key));
556    let removed_bytes = match key {
557        RecencyKey::Canonical(object_id) => {
558            let bytes = state
559                .entries
560                .remove(&object_id)
561                .map(|entry| u64::try_from(entry.bytes.len()).unwrap_or(u64::MAX))
562                .ok_or(ObjectStoreError::Corrupt)?;
563            state.resident_canonical_bytes = state
564                .resident_canonical_bytes
565                .checked_sub(bytes)
566                .ok_or(ObjectStoreError::Corrupt)?;
567            bytes
568        }
569        RecencyKey::Decoded(decoded) => {
570            let bytes = state
571                .decoded
572                .remove(&decoded)
573                .map(|entry| entry.value.logical_bytes)
574                .ok_or(ObjectStoreError::Corrupt)?;
575            state.resident_decoded_bytes = state
576                .resident_decoded_bytes
577                .checked_sub(bytes)
578                .ok_or(ObjectStoreError::Corrupt)?;
579            bytes
580        }
581    };
582    state.resident_bytes = state
583        .resident_bytes
584        .checked_sub(removed_bytes)
585        .ok_or(ObjectStoreError::Corrupt)?;
586    state.stats.evictions = increment(state.stats.evictions)?;
587    Ok(())
588}
589
590fn cache_probe_work() -> WorkCounters {
591    WorkCounters {
592        object_probes: 1,
593        items_examined: 1,
594        ..WorkCounters::default()
595    }
596}
597
598fn combine_failure(prior: WorkCounters, failure: ObjectFailure) -> ObjectFailure {
599    match prior.checked_add(*failure.work) {
600        Ok(work) => ObjectFailure::new(failure.error, work),
601        Err(error) => ObjectFailure::new(error.into(), prior),
602    }
603}
604
605fn shared_read(bytes: Bytes, work: WorkCounters) -> ObjectReceipt<ObjectRead> {
606    ObjectReceipt {
607        value: ObjectRead {
608            bytes,
609            retention: ObjectReadRetention::Shared,
610        },
611        work,
612    }
613}
614
615async fn wait_for_flight(
616    flight: &Flight,
617    cancellation: &CancellationToken,
618) -> Result<Option<Bytes>, ObjectStoreError> {
619    let mut cancelled = pin!(cancellation.cancelled());
620    poll_fn(|context| {
621        if cancelled.as_mut().poll(context).is_ready() {
622            return Poll::Ready(Err(ObjectStoreError::Cancelled));
623        }
624        flight.poll(context)
625    })
626    .await
627}
628
629impl<S: AsyncObjectStore> CachedObjectStore<S> {
630    fn validate_cached_read(
631        _object_id: ObjectId,
632        bytes: Bytes,
633        maximum_bytes: u64,
634        work: WorkCounters,
635    ) -> ObjectResult<ObjectRead> {
636        let observed = u64::try_from(bytes.len()).unwrap_or(u64::MAX);
637        if observed > maximum_bytes {
638            return Err(ObjectFailure::new(
639                ObjectStoreError::TooLarge {
640                    observed,
641                    maximum: maximum_bytes,
642                },
643                work,
644            ));
645        }
646        Ok(shared_read(bytes, work))
647    }
648
649    async fn load_leader(
650        &self,
651        object_id: ObjectId,
652        maximum_bytes: u64,
653        budget: WorkBudget,
654        cancellation: &CancellationToken,
655        prior: WorkCounters,
656        flight: Arc<Flight>,
657    ) -> ObjectResult<ObjectRead> {
658        let guard = ActiveFlightGuard::new(self, object_id, Arc::clone(&flight));
659        let receipt = match self
660            .inner
661            .read(
662                object_id,
663                maximum_bytes,
664                prior
665                    .remaining(budget)
666                    .map_err(|error| ObjectFailure::new(error.into(), prior))?,
667                cancellation,
668            )
669            .await
670        {
671            Ok(receipt) => receipt,
672            Err(failure) => {
673                guard
674                    .finish(None)
675                    .map_err(|error| ObjectFailure::new(error, prior))?;
676                return Err(combine_failure(prior, failure));
677            }
678        };
679        let work = prior
680            .checked_add(receipt.work)
681            .map_err(|error| ObjectFailure::new(error.into(), prior))?;
682        let bytes = receipt.value.bytes.clone();
683        if self
684            .insert(object_id, bytes.clone())
685            .map_err(|error| ObjectFailure::new(error, work))?
686        {
687            guard
688                .finish(Some(bytes.clone()))
689                .map_err(|error| ObjectFailure::new(error, work))?;
690            return Ok(shared_read(bytes, work));
691        }
692        guard
693            .finish(None)
694            .map_err(|error| ObjectFailure::new(error, work))?;
695        Ok(ObjectReceipt {
696            value: receipt.value,
697            work,
698        })
699    }
700
701    fn prepare_batch(
702        &self,
703        requests: &[ObjectReadRequest],
704        budget: WorkBudget,
705        cancellation: &CancellationToken,
706    ) -> Result<PreparedBatch, ObjectFailure> {
707        cancellation
708            .check()
709            .map_err(|_| ObjectFailure::before_work(ObjectStoreError::Cancelled))?;
710        if requests.is_empty() {
711            return Err(ObjectFailure::before_work(ObjectStoreError::Rejected(
712                "object read batch is empty".to_owned(),
713            )));
714        }
715        let item_count = u64::try_from(requests.len()).unwrap_or(u64::MAX);
716        let result_bytes = item_count
717            .checked_mul(u64::try_from(size_of::<ObjectRead>()).unwrap_or(u64::MAX))
718            .ok_or_else(|| ObjectFailure::before_work(crate::WorkError::Overflow.into()))?;
719        let planning_bytes = item_count
720            .checked_mul(
721                u64::try_from(
722                    size_of::<BatchSlot>()
723                        + size_of::<BatchLeader>()
724                        + size_of::<ObjectReadRequest>(),
725                )
726                .unwrap_or(u64::MAX),
727            )
728            .ok_or_else(|| ObjectFailure::before_work(crate::WorkError::Overflow.into()))?;
729        let mut work = WorkCounters {
730            object_probes: item_count,
731            items_examined: item_count,
732            items_returned: item_count,
733            allocation_operations: 4,
734            peak_allocation_bytes: planning_bytes.saturating_add(result_bytes),
735            ..WorkCounters::default()
736        };
737        work.verify(budget)
738            .map_err(|error| ObjectFailure::before_work(error.into()))?;
739        work.items_returned = 0;
740        let mut slots = Vec::new();
741        let mut leaders = Vec::new();
742        let mut local_leaders = BTreeMap::<ObjectId, usize>::new();
743        slots.try_reserve_exact(requests.len()).map_err(|_| {
744            ObjectFailure::before_work(ObjectStoreError::Rejected(
745                "object cache batch plan allocation failed".to_owned(),
746            ))
747        })?;
748        leaders.try_reserve_exact(requests.len()).map_err(|_| {
749            ObjectFailure::before_work(ObjectStoreError::Rejected(
750                "object cache leader allocation failed".to_owned(),
751            ))
752        })?;
753        for request in requests {
754            match self.prepare_batch_slot(request, &mut leaders, &mut local_leaders, work) {
755                Ok(slot) => slots.push(slot),
756                Err(failure) => {
757                    self.abandon_unstarted_leaders(&leaders);
758                    return Err(failure);
759                }
760            }
761        }
762        Ok(PreparedBatch {
763            slots,
764            leaders,
765            work,
766            planning_bytes,
767            result_bytes,
768            item_count,
769        })
770    }
771
772    fn prepare_batch_slot(
773        &self,
774        request: &ObjectReadRequest,
775        leaders: &mut Vec<BatchLeader>,
776        local_leaders: &mut BTreeMap<ObjectId, usize>,
777        work: WorkCounters,
778    ) -> Result<BatchSlot, ObjectFailure> {
779        match self
780            .read_plan(request.object_id)
781            .map_err(|error| ObjectFailure::new(error, work))?
782        {
783            ReadPlan::Hit(bytes) => Ok(BatchSlot::Hit(bytes)),
784            ReadPlan::Leader(flight) => {
785                let index = leaders.len();
786                leaders.push(BatchLeader {
787                    object_id: request.object_id,
788                    maximum_bytes: request.maximum_bytes,
789                    flight: Arc::clone(&flight),
790                });
791                local_leaders.insert(request.object_id, index);
792                Ok(BatchSlot::LocalLeader(index))
793            }
794            ReadPlan::Follower(flight) => {
795                if let Some(&index) = local_leaders.get(&request.object_id)
796                    && leaders
797                        .get(index)
798                        .is_some_and(|leader| Arc::ptr_eq(&leader.flight, &flight))
799                {
800                    if let Some(leader) = leaders.get_mut(index) {
801                        leader.maximum_bytes = leader.maximum_bytes.max(request.maximum_bytes);
802                    }
803                    Ok(BatchSlot::LocalLeader(index))
804                } else {
805                    Ok(BatchSlot::ExternalFollower(flight))
806                }
807            }
808        }
809    }
810
811    async fn load_batch_leaders(
812        &self,
813        prepared: &PreparedBatch,
814        budget: WorkBudget,
815        cancellation: &CancellationToken,
816    ) -> ObjectResult<Vec<ObjectRead>> {
817        if prepared.leaders.is_empty() {
818            return Ok(ObjectReceipt {
819                value: Vec::new(),
820                work: prepared.work,
821            });
822        }
823        let mut requests = Vec::new();
824        requests
825            .try_reserve_exact(prepared.leaders.len())
826            .map_err(|_| {
827                ObjectFailure::new(
828                    ObjectStoreError::Rejected(
829                        "object cache backend batch allocation failed".to_owned(),
830                    ),
831                    prepared.work,
832                )
833            })?;
834        requests.extend(prepared.leaders.iter().map(|leader| ObjectReadRequest {
835            object_id: leader.object_id,
836            maximum_bytes: leader.maximum_bytes,
837        }));
838        let mut guards = prepared
839            .leaders
840            .iter()
841            .map(|leader| {
842                Some(ActiveFlightGuard::new(
843                    self,
844                    leader.object_id,
845                    Arc::clone(&leader.flight),
846                ))
847            })
848            .collect::<Vec<_>>();
849        let receipt = match self
850            .inner
851            .read_many(
852                &requests,
853                prepared
854                    .work
855                    .remaining(budget)
856                    .map_err(|error| ObjectFailure::new(error.into(), prepared.work))?,
857                cancellation,
858            )
859            .await
860        {
861            Ok(receipt) => receipt,
862            Err(failure) => {
863                Self::abort_batch_leaders(&mut guards);
864                return Err(combine_failure(prepared.work, failure));
865            }
866        };
867        let mut work = prepared
868            .work
869            .checked_add(receipt.work)
870            .map_err(|error| ObjectFailure::new(error.into(), prepared.work))?;
871        if receipt.value.len() != prepared.leaders.len() {
872            Self::abort_batch_leaders(&mut guards);
873            return Err(ObjectFailure::new(ObjectStoreError::Corrupt, work));
874        }
875        let mut values = receipt.value;
876        for ((leader, value), guard) in prepared
877            .leaders
878            .iter()
879            .zip(values.iter_mut())
880            .zip(guards.iter_mut())
881        {
882            let bytes = value.bytes.clone();
883            let retained = self
884                .insert(leader.object_id, bytes.clone())
885                .map_err(|error| ObjectFailure::new(error, work))?;
886            if retained {
887                value.retention = ObjectReadRetention::Shared;
888            }
889            if let Some(guard) = guard.take() {
890                guard
891                    .finish(retained.then_some(bytes))
892                    .map_err(|error| ObjectFailure::new(error, work))?;
893            }
894        }
895        work.peak_allocation_bytes = work.peak_allocation_bytes.max(
896            prepared
897                .planning_bytes
898                .saturating_add(prepared.result_bytes),
899        );
900        Ok(ObjectReceipt {
901            value: values,
902            work,
903        })
904    }
905
906    fn abort_batch_leaders(guards: &mut [Option<ActiveFlightGuard<'_, S>>]) {
907        for guard in guards {
908            if let Some(guard) = guard.take() {
909                let _ = guard.finish(None);
910            }
911        }
912    }
913
914    fn abandon_unstarted_leaders(&self, leaders: &[BatchLeader]) {
915        for leader in leaders {
916            leader.flight.finish(None);
917            let _ = self.remove_flight(leader.object_id, &leader.flight);
918        }
919    }
920
921    async fn resolve_prepared_batch(
922        &self,
923        requests: &[ObjectReadRequest],
924        prepared: PreparedBatch,
925        leaders: ObjectReceipt<Vec<ObjectRead>>,
926        budget: WorkBudget,
927        cancellation: &CancellationToken,
928    ) -> ObjectResult<Vec<ObjectRead>> {
929        let mut work = leaders.work;
930        let mut values = Vec::new();
931        values.try_reserve_exact(requests.len()).map_err(|_| {
932            ObjectFailure::new(
933                ObjectStoreError::Rejected("object cache result allocation failed".to_owned()),
934                work,
935            )
936        })?;
937        for (request, slot) in requests.iter().zip(prepared.slots) {
938            cancellation
939                .check()
940                .map_err(|_| ObjectFailure::new(ObjectStoreError::Cancelled, work))?;
941            let value = self
942                .resolve_batch_slot(
943                    request,
944                    slot,
945                    &leaders.value,
946                    &mut work,
947                    budget,
948                    cancellation,
949                )
950                .await?;
951            values.push(value);
952        }
953        work.items_returned = prepared.item_count;
954        work.verify(budget)
955            .map_err(|error| ObjectFailure::new(error.into(), work))?;
956        Ok(ObjectReceipt {
957            value: values,
958            work,
959        })
960    }
961
962    async fn resolve_batch_slot(
963        &self,
964        request: &ObjectReadRequest,
965        slot: BatchSlot,
966        leader_values: &[ObjectRead],
967        work: &mut WorkCounters,
968        budget: WorkBudget,
969        cancellation: &CancellationToken,
970    ) -> Result<ObjectRead, ObjectFailure> {
971        let bytes = match slot {
972            BatchSlot::Hit(bytes) => bytes,
973            BatchSlot::LocalLeader(index) => leader_values
974                .get(index)
975                .ok_or_else(|| ObjectFailure::new(ObjectStoreError::Corrupt, *work))?
976                .bytes
977                .clone(),
978            BatchSlot::ExternalFollower(flight) => {
979                let result = wait_for_flight(&flight, cancellation)
980                    .await
981                    .map_err(|error| ObjectFailure::new(error, *work))?;
982                self.record_coalesced()
983                    .map_err(|error| ObjectFailure::new(error, *work))?;
984                if let Some(bytes) = result {
985                    bytes
986                } else {
987                    self.remove_flight(request.object_id, &flight)
988                        .map_err(|error| ObjectFailure::new(error, *work))?;
989                    let receipt = self
990                        .read(
991                            request.object_id,
992                            request.maximum_bytes,
993                            work.remaining(budget)
994                                .map_err(|error| ObjectFailure::new(error.into(), *work))?,
995                            cancellation,
996                        )
997                        .await
998                        .map_err(|failure| combine_failure(*work, failure))?;
999                    *work = work
1000                        .checked_add(receipt.work)
1001                        .map_err(|error| ObjectFailure::new(error.into(), *work))?;
1002                    return Ok(receipt.value);
1003                }
1004            }
1005        };
1006        Ok(Self::validate_cached_read(
1007            request.object_id,
1008            bytes,
1009            request.maximum_bytes,
1010            WorkCounters::default(),
1011        )?
1012        .value)
1013    }
1014}
1015
1016impl<S: AsyncObjectStore> AsyncObjectStore for CachedObjectStore<S> {
1017    fn decoded_cache_get(
1018        &self,
1019        key: DecodedCacheKey,
1020    ) -> Result<Option<DecodedCacheValue>, ObjectStoreError> {
1021        self.decoded_get(key)
1022    }
1023
1024    fn decoded_cache_admit(
1025        &self,
1026        key: DecodedCacheKey,
1027        value: DecodedCacheValue,
1028    ) -> Result<DecodedCacheAdmission, ObjectStoreError> {
1029        self.decoded_admit(key, value)
1030    }
1031
1032    async fn put(
1033        &self,
1034        object_id: ObjectId,
1035        bytes: Bytes,
1036        budget: WorkBudget,
1037        cancellation: &CancellationToken,
1038    ) -> ObjectResult<()> {
1039        cancellation
1040            .check()
1041            .map_err(|_| ObjectFailure::before_work(ObjectStoreError::Cancelled))?;
1042        let retained = bytes.clone();
1043        let receipt = self
1044            .inner
1045            .put(object_id, bytes, budget, cancellation)
1046            .await?;
1047        self.insert(object_id, retained)
1048            .map_err(|error| ObjectFailure::new(error, receipt.work))?;
1049        Ok(receipt)
1050    }
1051
1052    async fn put_many(
1053        &self,
1054        writes: &[ObjectWrite],
1055        budget: WorkBudget,
1056        cancellation: &CancellationToken,
1057    ) -> ObjectResult<()> {
1058        cancellation
1059            .check()
1060            .map_err(|_| ObjectFailure::before_work(ObjectStoreError::Cancelled))?;
1061        let receipt = self.inner.put_many(writes, budget, cancellation).await?;
1062        for write in writes {
1063            self.insert(write.object_id, write.bytes.clone())
1064                .map_err(|error| ObjectFailure::new(error, receipt.work))?;
1065        }
1066        Ok(receipt)
1067    }
1068
1069    async fn read(
1070        &self,
1071        object_id: ObjectId,
1072        maximum_bytes: u64,
1073        budget: WorkBudget,
1074        cancellation: &CancellationToken,
1075    ) -> ObjectResult<ObjectRead> {
1076        cancellation
1077            .check()
1078            .map_err(|_| ObjectFailure::before_work(ObjectStoreError::Cancelled))?;
1079        let mut prior = cache_probe_work();
1080        prior
1081            .verify(budget)
1082            .map_err(|error| ObjectFailure::before_work(error.into()))?;
1083        loop {
1084            cancellation
1085                .check()
1086                .map_err(|_| ObjectFailure::new(ObjectStoreError::Cancelled, prior))?;
1087            match self
1088                .read_plan(object_id)
1089                .map_err(|error| ObjectFailure::new(error, prior))?
1090            {
1091                ReadPlan::Hit(bytes) => {
1092                    return Self::validate_cached_read(object_id, bytes, maximum_bytes, prior);
1093                }
1094                ReadPlan::Leader(flight) => {
1095                    return self
1096                        .load_leader(
1097                            object_id,
1098                            maximum_bytes,
1099                            budget,
1100                            cancellation,
1101                            prior,
1102                            flight,
1103                        )
1104                        .await;
1105                }
1106                ReadPlan::Follower(flight) => {
1107                    let result = wait_for_flight(&flight, cancellation)
1108                        .await
1109                        .map_err(|error| ObjectFailure::new(error, prior))?;
1110                    self.record_coalesced()
1111                        .map_err(|error| ObjectFailure::new(error, prior))?;
1112                    if let Some(bytes) = result {
1113                        return Self::validate_cached_read(object_id, bytes, maximum_bytes, prior);
1114                    }
1115                    self.remove_flight(object_id, &flight)
1116                        .map_err(|error| ObjectFailure::new(error, prior))?;
1117                    prior = prior
1118                        .checked_add(cache_probe_work())
1119                        .map_err(|error| ObjectFailure::new(error.into(), prior))?;
1120                    prior
1121                        .verify(budget)
1122                        .map_err(|error| ObjectFailure::new(error.into(), prior))?;
1123                }
1124            }
1125        }
1126    }
1127
1128    async fn read_many(
1129        &self,
1130        requests: &[ObjectReadRequest],
1131        budget: WorkBudget,
1132        cancellation: &CancellationToken,
1133    ) -> ObjectResult<Vec<ObjectRead>> {
1134        let prepared = self.prepare_batch(requests, budget, cancellation)?;
1135        let leaders = self
1136            .load_batch_leaders(&prepared, budget, cancellation)
1137            .await?;
1138        self.resolve_prepared_batch(requests, prepared, leaders, budget, cancellation)
1139            .await
1140    }
1141
1142    async fn contains(
1143        &self,
1144        object_id: ObjectId,
1145        budget: WorkBudget,
1146        cancellation: &CancellationToken,
1147    ) -> ObjectResult<bool> {
1148        cancellation
1149            .check()
1150            .map_err(|_| ObjectFailure::before_work(ObjectStoreError::Cancelled))?;
1151        let prior = cache_probe_work();
1152        prior
1153            .verify(budget)
1154            .map_err(|error| ObjectFailure::before_work(error.into()))?;
1155        if self
1156            .cached_contains(object_id)
1157            .map_err(|error| ObjectFailure::new(error, prior))?
1158        {
1159            return Ok(ObjectReceipt {
1160                value: true,
1161                work: prior,
1162            });
1163        }
1164        let receipt = self
1165            .inner
1166            .contains(
1167                object_id,
1168                prior
1169                    .remaining(budget)
1170                    .map_err(|error| ObjectFailure::new(error.into(), prior))?,
1171                cancellation,
1172            )
1173            .await
1174            .map_err(|failure| combine_failure(prior, failure))?;
1175        Ok(ObjectReceipt {
1176            value: receipt.value,
1177            work: prior
1178                .checked_add(receipt.work)
1179                .map_err(|error| ObjectFailure::new(error.into(), prior))?,
1180        })
1181    }
1182}
1183
1184#[cfg(all(test, feature = "memory"))]
1185#[path = "tests/cache.rs"]
1186mod tests;