1use 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#[derive(Clone, Copy, Debug, Eq, PartialEq)]
26pub struct ObjectCacheOptions {
27 pub maximum_entries: u32,
29 pub maximum_bytes: u64,
31 pub maximum_in_flight: u32,
33 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#[derive(Clone, Copy, Debug, Eq, Error, PartialEq)]
50pub enum ObjectCacheConfigError {
51 #[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#[derive(Clone, Copy, Debug, Default, Eq, PartialEq)]
71pub struct ObjectCacheStats {
72 pub hits: u64,
74 pub decoded_hits: u64,
76 pub misses: u64,
78 pub coalesced_reads: u64,
80 pub evictions: u64,
82 pub resident_entries: u64,
84 pub resident_bytes: u64,
86 pub resident_canonical_objects: u64,
88 pub resident_canonical_bytes: u64,
90 pub resident_decoded_pages: u64,
92 pub resident_decoded_bytes: u64,
94 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
277pub 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 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 #[must_use]
310 pub const fn inner(&self) -> &S {
311 &self.inner
312 }
313
314 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 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;