1use super::{
27 CHECKSUM_SIZE, get_page_with_checksum_from_blob, physical_page, validate_physical_page,
28};
29use crate::{BLOB_PAGE_SIZE, Blob, Error, IoBuf, IoBufMut, ReadOptions};
30use ahash::{AHashMap, AHashSet};
31use commonware_utils::{
32 Widen, cache,
33 futures::{Shared, WeakShared},
34 sync::{Mutex, RwLock, RwLockWriteGuard},
35};
36use futures::{FutureExt, Stream, TryStreamExt as _, future::BoxFuture, stream::FuturesUnordered};
37use std::{
38 collections::{VecDeque, hash_map::Entry},
39 mem,
40 num::{NonZeroU16, NonZeroUsize},
41 pin::{Pin, pin},
42 sync::{
43 Arc, Weak,
44 atomic::{AtomicU64, Ordering},
45 },
46 task::{Context, Poll},
47};
48use tracing::{debug, error, trace};
49
50type PageFetch = Shared<BoxFuture<'static, Result<IoBuf, Error>>>;
54
55struct PageFetchEntry {
63 fetch: PageFetch,
65 waiters: usize,
67}
68
69struct PageFetchGuard<'a> {
73 cache: &'a RwLock<Cache>,
74 key: (u64, u64),
75 fetch: PageFetch,
76 armed: bool,
77}
78
79impl<'a> PageFetchGuard<'a> {
80 const fn new(cache: &'a RwLock<Cache>, key: (u64, u64), fetch: PageFetch) -> Self {
81 Self {
82 cache,
83 key,
84 fetch,
85 armed: true,
86 }
87 }
88
89 const fn disarm(&mut self) {
90 self.armed = false;
91 }
92
93 async fn wait(mut self) -> Result<IoBuf, Error> {
97 let result = (&mut self.fetch).await;
98 self.disarm();
99 result
100 }
101}
102
103impl Drop for PageFetchGuard<'_> {
104 fn drop(&mut self) {
105 if !self.armed {
106 return;
107 }
108
109 let mut cache = self.cache.write();
116 let Entry::Occupied(mut current) = cache.page_fetches.entry(self.key) else {
117 return;
118 };
119 if !current.get().fetch.ptr_eq(&self.fetch) {
120 return;
121 }
122 if current.get().waiters == 1 {
123 current.remove();
124 } else {
125 current.get_mut().waiters -= 1;
126 }
127 }
128}
129
130const MISSES_PER_LOCK: usize = 64;
135
136type BatchFuture = BoxFuture<'static, Result<(), Error>>;
138
139type BatchDriver = Shared<BatchFuture>;
141
142type Delivery = (usize, Result<IoBuf, Error>);
145
146const STARTER: usize = 0;
148
149type Joiner = (usize, usize);
152
153struct BatchState {
157 joiners: Box<[Option<Vec<Joiner>>]>,
160 inboxes: Vec<Option<VecDeque<Delivery>>>,
164 error: Option<Error>,
166}
167
168struct Batch {
170 blob_id: u64,
171 pages: Box<[u64]>,
173 state: Mutex<BatchState>,
174}
175
176impl Batch {
177 fn new(blob_id: u64, pages: Vec<u64>) -> Self {
178 let state = BatchState {
179 joiners: vec![Some(Vec::new()); pages.len()].into(),
180 inboxes: vec![Some(VecDeque::new())],
181 error: None,
182 };
183 Self {
184 blob_id,
185 pages: pages.into(),
186 state: Mutex::new(state),
187 }
188 }
189
190 async fn read<B: Blob>(
197 self: Arc<Self>,
198 cache: Arc<RwLock<Cache>>,
199 blob: Arc<B>,
200 physical: Vec<(u64, usize)>,
201 page_size: NonZeroU16,
202 _retire: RetireGuard,
203 ) -> Result<(), Error> {
204 let result = async {
205 let mut stream = pin!(blob.read_many(&physical, ReadOptions::DONT_CACHE));
208 while let Some((index, bufs)) = stream.try_next().await? {
209 let Some(&page_num) = self.pages.get(index) else {
211 return Err(Error::ReadFailed);
212 };
213 let page = validate_physical_page(bufs.coalesce())
214 .and_then(|(page, _)| cacheable_page(page, page_num, page_size))
215 .inspect_err(|err| error!(page_num, ?err, "Batched page failed validation"));
216 self.publish(&cache, index, page)?;
217 }
218 if self.state.lock().joiners.iter().any(Option::is_some) {
219 return Err(Error::ReadFailed);
220 }
221 Ok(())
222 }
223 .await;
224 if let Err(err) = &result {
225 self.fail(&cache, err);
226 }
227 result
228 }
229
230 fn publish(
233 &self,
234 cache: &RwLock<Cache>,
235 index: usize,
236 result: Result<IoBuf, Error>,
237 ) -> Result<(), Error> {
238 let mut cache = cache.write();
239 let mut state = self.state.lock();
240 let BatchState {
241 joiners, inboxes, ..
242 } = &mut *state;
243 let Some(joiners) = joiners[index].take() else {
244 return Err(Error::ReadFailed);
245 };
246 if let Ok(page) = &result {
247 cache.cache(self.blob_id, page.as_ref(), self.pages[index]);
248 }
249 drop(cache);
250 for (tag, reader) in joiners {
251 if let Some(inbox) = &mut inboxes[reader] {
252 inbox.push_back((tag, result.clone()));
253 }
254 }
255 if let Some(inbox) = &mut inboxes[STARTER] {
256 inbox.push_back((index, result));
257 }
258 Ok(())
259 }
260
261 fn fail(&self, cache: &RwLock<Cache>, err: &Error) {
264 error!(blob_id = self.blob_id, ?err, "Batched page read failed");
265 let _cache = cache.write();
266 let mut state = self.state.lock();
267 state.error = Some(err.clone());
268 state.joiners.fill(None);
269 }
270}
271
272type Retired = Arc<Mutex<Vec<(u64, Weak<Batch>)>>>;
274
275struct RetireGuard {
278 retired: Retired,
279 blob_id: u64,
280 batch: Weak<Batch>,
281}
282
283impl Drop for RetireGuard {
284 fn drop(&mut self) {
285 let batch = std::mem::take(&mut self.batch);
286 self.retired.lock().push((self.blob_id, batch));
287 }
288}
289
290type Registration = (Weak<Batch>, WeakShared<BatchFuture>);
292
293#[derive(Default)]
299struct BatchRegistry {
300 by_blob: AHashMap<u64, Vec<Registration>>,
301 retired: Retired,
302}
303
304impl BatchRegistry {
305 fn find_unpublished(
308 &self,
309 blob_id: u64,
310 page_num: u64,
311 ) -> Option<(Arc<Batch>, usize, BatchDriver)> {
312 self.by_blob
313 .get(&blob_id)?
314 .iter()
315 .find_map(|(batch, driver)| {
316 let batch = batch.upgrade()?;
317 let index = batch.pages.binary_search(&page_num).ok();
318 let index = index.filter(|&index| batch.state.lock().joiners[index].is_some())?;
319 Some((batch, index, driver.upgrade()?))
321 })
322 }
323
324 fn retire_on_drop(&self, batch: &Arc<Batch>) -> RetireGuard {
326 RetireGuard {
327 retired: self.retired.clone(),
328 blob_id: batch.blob_id,
329 batch: Arc::downgrade(batch),
330 }
331 }
332
333 fn register(&mut self, batch: &Arc<Batch>, driver: &BatchDriver) {
335 self.by_blob
336 .entry(batch.blob_id)
337 .or_default()
338 .push((Arc::downgrade(batch), driver.downgrade()));
339 }
340
341 fn unregister_retired(&mut self) {
343 let mut retired = std::mem::take(&mut *self.retired.lock());
344 if retired.is_empty() {
345 return;
346 }
347
348 let batches: AHashSet<*const Batch> =
351 retired.iter().map(|(_, batch)| batch.as_ptr()).collect();
352 retired.sort_unstable_by_key(|(blob_id, _)| *blob_id);
353 retired.dedup_by_key(|(blob_id, _)| *blob_id);
354 for (blob_id, _) in retired {
355 if let Entry::Occupied(mut registered) = self.by_blob.entry(blob_id) {
356 registered
357 .get_mut()
358 .retain(|(batch, _)| !batches.contains(&batch.as_ptr()));
359 if registered.get().is_empty() {
360 registered.remove();
361 }
362 }
363 }
364 }
365}
366
367struct Subscription {
370 batch: Arc<Batch>,
371 reader: usize,
373 taken: VecDeque<Delivery>,
375 remaining: usize,
377 driver: Option<BatchDriver>,
379}
380
381impl Subscription {
382 fn starting(batch: Arc<Batch>, driver: BatchDriver) -> Self {
384 Self {
385 remaining: batch.pages.len(),
386 driver: Some(driver),
387 batch,
388 reader: STARTER,
389 taken: VecDeque::new(),
390 }
391 }
392
393 fn joining(batch: Arc<Batch>, driver: BatchDriver) -> Self {
395 let reader = {
396 let mut state = batch.state.lock();
397 state.inboxes.push(Some(VecDeque::new()));
398 state.inboxes.len() - 1
399 };
400 Self {
401 driver: Some(driver),
402 batch,
403 reader,
404 taken: VecDeque::new(),
405 remaining: 0,
406 }
407 }
408
409 fn select(&mut self, index: usize, tag: usize) {
412 let mut state = self.batch.state.lock();
413 let BatchState {
414 joiners, inboxes, ..
415 } = &mut *state;
416 let joiners = joiners[index]
417 .as_mut()
418 .expect("only unpublished pages are selected");
419
420 if joiners.len() == joiners.capacity() {
423 joiners.retain(|&(_, reader)| inboxes[reader].is_some());
424 joiners.reserve(joiners.len());
425 }
426 joiners.push((tag, self.reader));
427 self.remaining += 1;
428 }
429
430 fn take(&mut self) -> Option<Delivery> {
432 if self.taken.is_empty() {
433 let mut state = self.batch.state.lock();
434 if let Some(inbox) = &mut state.inboxes[self.reader] {
435 mem::swap(inbox, &mut self.taken);
436 }
437 }
438 self.taken.pop_front()
439 }
440}
441
442impl Drop for Subscription {
443 fn drop(&mut self) {
444 let inbox = self.batch.state.lock().inboxes[self.reader].take();
446 drop(inbox);
447 }
448}
449
450impl Stream for Subscription {
451 type Item = Result<(usize, IoBuf), Error>;
452
453 fn poll_next(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Option<Self::Item>> {
454 let this = &mut *self;
455 if this.remaining == 0 {
456 return Poll::Ready(None);
457 }
458
459 let delivery = match this.take() {
465 Some(delivery) => delivery,
466 None => {
467 let driver = this.driver.as_mut();
468 if driver.is_some_and(|driver| driver.poll_unpin(cx).is_ready()) {
469 this.driver = None;
470 }
471 match this.take() {
472 Some(delivery) => delivery,
473 None if this.driver.is_some() => return Poll::Pending,
474 None => {
475 this.remaining = 0;
477 let error = this.batch.state.lock().error.clone();
478 return Poll::Ready(Some(Err(error.unwrap_or(Error::ReadFailed))));
479 }
480 }
481 }
482 };
483 let (tag, result) = delivery;
484 this.remaining -= 1;
485 if this.remaining == 0 {
486 this.driver = None;
487 }
488 Poll::Ready(Some(result.map(|page| (tag, page))))
489 }
490}
491
492enum PageSource {
495 Fetch(PageFetch),
497 Batch(Subscription),
499}
500
501impl PageSource {
502 async fn wait(self, cache: &RwLock<Cache>, key: (u64, u64)) -> Result<IoBuf, Error> {
504 match self {
505 Self::Fetch(fetch) => PageFetchGuard::new(cache, key, fetch).wait().await,
506 Self::Batch(mut subscription) => {
507 let delivery = subscription.try_next().await?;
508 let (_, page) = delivery.expect("a subscription yields the page it selected");
509 Ok(page)
510 }
511 }
512 }
513}
514
515type Pieces<'a, 'b> = &'a mut [(u64, usize, &'b mut [u8])];
518
519struct BatchRead<'a, 'b> {
522 subscription: Subscription,
523 pieces: Vec<Pieces<'a, 'b>>,
525 completion: Option<BatchDriver>,
528}
529
530impl<'a, 'b> BatchRead<'a, 'b> {
531 fn starting(batch: Arc<Batch>, driver: BatchDriver, pieces: Vec<Pieces<'a, 'b>>) -> Self {
533 Self {
534 completion: Some(driver.clone()),
535 subscription: Subscription::starting(batch, driver),
536 pieces,
537 }
538 }
539
540 fn joining(batch: Arc<Batch>, driver: BatchDriver) -> Self {
542 Self {
543 subscription: Subscription::joining(batch, driver),
544 pieces: Vec::new(),
545 completion: None,
546 }
547 }
548
549 fn select(&mut self, index: usize, pieces: Pieces<'a, 'b>) {
552 self.subscription.select(index, self.pieces.len());
553 self.pieces.push(pieces);
554 }
555
556 async fn complete(mut self) -> Result<(), Error> {
559 while let Some((tag, page)) = self.subscription.try_next().await? {
560 copy_page(page.as_ref(), self.pieces[tag]);
561 }
562 match self.completion {
563 Some(driver) => driver.await,
564 None => Ok(()),
565 }
566 }
567}
568
569fn copy_page(page: &[u8], pieces: Pieces<'_, '_>) {
571 for (_, offset, piece) in pieces {
572 piece.copy_from_slice(&page[*offset..*offset + piece.len()]);
573 }
574}
575
576struct Cache {
593 cache: cache::Cache<(u64, u64), IoBufMut>,
595
596 hints: Vec<usize>,
602
603 page_size: NonZeroU16,
606
607 page_fetches: AHashMap<(u64, u64), PageFetchEntry>,
610
611 batches: BatchRegistry,
613}
614
615#[derive(Clone)]
619pub struct CacheRef {
620 page_size: NonZeroU16,
631
632 next_id: Arc<AtomicU64>,
634
635 cache: Arc<RwLock<Cache>>,
637}
638
639impl CacheRef {
640 pub fn new(page_size: NonZeroU16, capacity: NonZeroUsize) -> Self {
651 let page_size_u64: u64 = page_size.widen();
652 let physical_page_size = page_size_u64 + CHECKSUM_SIZE;
653 if !physical_page_size.is_multiple_of(u64::from(BLOB_PAGE_SIZE))
654 && !u64::from(BLOB_PAGE_SIZE).is_multiple_of(physical_page_size)
655 {
656 debug!(
657 page_size = page_size.get(),
658 physical_page_size, "physical pages do not align with blob pages"
659 );
660 }
661
662 Self {
663 page_size,
664 next_id: Arc::new(AtomicU64::new(0)),
665 cache: Arc::new(RwLock::new(Cache::new(page_size, capacity))),
666 }
667 }
668
669 #[inline]
672 pub const fn page_size(&self) -> NonZeroU16 {
673 self.page_size
674 }
675
676 pub fn next_id(&self) -> u64 {
678 self.next_id.fetch_add(1, Ordering::Relaxed)
679 }
680
681 pub fn offset_to_page(&self, offset: u64) -> (u64, u64) {
684 Cache::offset_to_page(self.page_size, offset)
685 }
686
687 pub(super) fn read_cached(
690 &self,
691 blob_id: u64,
692 mut buf: &mut [u8],
693 mut logical_offset: u64,
694 ) -> usize {
695 let original_len = buf.len();
696 let page_cache = self.cache.read();
697 while !buf.is_empty() {
698 let count = page_cache.read_at(blob_id, buf, logical_offset);
699 if count == 0 {
700 break;
702 }
703 logical_offset += count as u64;
704 buf = &mut buf[count..];
705 }
706 original_len - buf.len()
707 }
708
709 pub(super) fn read_cached_many(&self, blob_id: u64, ranges: &mut Vec<(&mut [u8], u64)>) {
716 let page_cache = self.cache.read();
717 let page_size = page_cache.page_size;
718
719 let mut srcs: Vec<Option<&[u8]>> = Vec::with_capacity(ranges.len());
723 for (buf, offset) in ranges.iter() {
724 let (page_num, offset_in_page, remaining) = Cache::locate(page_size, *offset);
725 let seg = std::cmp::min(buf.len(), remaining);
726 srcs.push(
727 page_cache
728 .get_page(blob_id, page_num)
729 .map(|page| &page.as_ref()[offset_in_page..offset_in_page + seg]),
730 );
731 }
732
733 let mut next = 0;
735 ranges.retain_mut(|(buf, offset)| {
736 let src = srcs[next];
737 next += 1;
738 if buf.is_empty() {
739 return false;
740 }
741 let Some(src) = src else {
742 return true;
743 };
744 buf[..src.len()].copy_from_slice(src);
745 let mut done = src.len();
746 while done < buf.len() {
747 let count = page_cache.read_at(blob_id, &mut buf[done..], *offset + done as u64);
748 if count == 0 {
749 *offset += done as u64;
750 *buf = std::mem::take(buf).split_at_mut(done).1;
751 return true;
752 }
753 done += count;
754 }
755 false
756 });
757 }
758
759 pub(super) async fn read_after_miss<B: Blob>(
763 &self,
764 blob: &Arc<B>,
765 blob_id: u64,
766 mut buf: &mut [u8],
767 mut offset: u64,
768 ) -> Result<(), Error> {
769 while !buf.is_empty() {
772 let count = self
774 .read_after_page_fault(blob, blob_id, buf, offset)
775 .await?;
776 offset += count as u64;
777 buf = &mut buf[count..];
778 if buf.is_empty() {
779 break;
780 }
781
782 let cached = self.read_cached(blob_id, buf, offset);
784 offset += cached as u64;
785 buf = &mut buf[cached..];
786 }
787
788 Ok(())
789 }
790
791 pub(super) async fn read_after_misses<B: Blob>(
803 &self,
804 blob: &Arc<B>,
805 blob_id: u64,
806 ranges: Vec<(&mut [u8], u64)>,
807 ) -> Result<(), Error> {
808 let mut waiting = Vec::new();
811 {
812 let cache = self.cache.read();
813 for (mut buf, mut offset) in ranges {
814 while !buf.is_empty() {
815 let (page_num, offset_in_page, remaining) =
816 Cache::locate(self.page_size, offset);
817 let len = remaining.min(buf.len());
818 let (piece, rest) = std::mem::take(&mut buf).split_at_mut(len);
819 if cache.read_at(blob_id, piece, offset) == 0 {
820 waiting.push((page_num, offset_in_page, piece));
821 }
822 offset += Widen::widen(len);
823 buf = rest;
824 }
825 }
826 }
827
828 let page_size: u64 = self.page_size.widen();
831 let mut misses = Vec::new();
832 for pieces in waiting.chunk_by_mut(|(a, _, _), (b, _, _)| a == b) {
833 misses.push((physical_page(pieces[0].0, page_size)?, pieces));
834 }
835
836 let mut fetches = Vec::new();
840 let mut reads: Vec<BatchRead<'_, '_>> = Vec::new();
841 {
842 let mut unread = Vec::new();
843 let mut cache = self.cache.write();
844 cache.batches.unregister_retired();
845 for (i, (physical, pieces)) in misses.into_iter().enumerate() {
846 if i > 0 && i % MISSES_PER_LOCK == 0 {
847 RwLockWriteGuard::bump(&mut cache);
848 }
849 let page_num = pieces[0].0;
850 let key = (blob_id, page_num);
851 if let Some(page) = cache.get_page(blob_id, page_num) {
852 copy_page(page.as_ref(), pieces);
853 } else if let Some(entry) = cache.page_fetches.get_mut(&key) {
854 entry.waiters += 1;
855 let fetch = PageFetchGuard::new(&self.cache, key, entry.fetch.clone());
856 fetches.push((fetch, pieces));
857 } else if let Some((batch, index, driver)) =
858 cache.batches.find_unpublished(blob_id, page_num)
859 {
860 let joined = reads
861 .iter()
862 .rposition(|read| Arc::ptr_eq(&read.subscription.batch, &batch));
863 let joined = joined.unwrap_or_else(|| {
864 reads.push(BatchRead::joining(batch, driver));
865 reads.len() - 1
866 });
867 reads[joined].select(index, pieces);
868 } else {
869 unread.push((physical, pieces));
870 }
871 }
872 if !unread.is_empty() {
873 let (physical, pieces): (Vec<_>, Vec<_>) = unread.into_iter().unzip();
874 let pages = pieces.iter().map(|pieces| pieces[0].0).collect();
875 let batch = Arc::new(Batch::new(blob_id, pages));
876 let retire = cache.batches.retire_on_drop(&batch);
877 let read = batch.clone().read(
878 self.cache.clone(),
879 blob.clone(),
880 physical,
881 self.page_size,
882 retire,
883 );
884 let driver = Shared::new(read.boxed());
885 cache.batches.register(&batch, &driver);
886 reads.push(BatchRead::starting(batch, driver, pieces));
887 }
888 }
889
890 let fetches = fetches
893 .into_iter()
894 .map(|(fetch, pieces)| async move {
895 copy_page(fetch.wait().await?.as_ref(), pieces);
896 Ok(())
897 })
898 .collect::<FuturesUnordered<_>>()
899 .try_collect::<()>();
900 let reads = reads
901 .into_iter()
902 .map(BatchRead::complete)
903 .collect::<FuturesUnordered<_>>()
904 .try_collect::<()>();
905 futures::try_join!(fetches, reads)?;
906 Ok(())
907 }
908
909 pub(super) async fn read_after_page_fault<B: Blob>(
913 &self,
914 blob: &Arc<B>,
915 blob_id: u64,
916 buf: &mut [u8],
917 offset: u64,
918 ) -> Result<usize, Error> {
919 assert!(!buf.is_empty());
920
921 let (page_num, offset_in_page, _) = Cache::locate(self.page_size, offset);
922 trace!(page_num, blob_id, "page fault");
923
924 let key = (blob_id, page_num);
928 let source = {
929 let mut cache = self.cache.write();
930
931 let count = cache.read_at(blob_id, buf, offset);
934 if count != 0 {
935 return Ok(count);
936 }
937
938 let Cache {
939 page_fetches,
940 batches,
941 ..
942 } = &mut *cache;
943 match page_fetches.entry(key) {
944 Entry::Occupied(o) => {
945 let entry = o.into_mut();
947 entry.waiters += 1;
948 PageSource::Fetch(entry.fetch.clone())
949 }
950 Entry::Vacant(v) => {
951 batches.unregister_retired();
952 if let Some((batch, index, driver)) =
953 batches.find_unpublished(blob_id, page_num)
954 {
955 let mut subscription = Subscription::joining(batch, driver);
957 subscription.select(index, 0);
958 PageSource::Batch(subscription)
959 } else {
960 let blob = blob.clone();
964 let cache = Arc::clone(&self.cache);
965 let page_size = self.page_size;
966 let future = async move {
967 let result =
968 fetch_cacheable_page(blob.as_ref(), page_num, page_size).await;
969 if let Err(err) = &result {
970 error!(page_num, ?err, "Page fetch failed");
971 }
972
973 let mut cache = cache.write();
980 if let Ok(page) = &result {
981 cache.cache(blob_id, page.as_ref(), page_num);
982 }
983 let _ = cache.page_fetches.remove(&key);
984 result
985 };
986
987 let fetch = Shared::new(future.boxed());
989 v.insert(PageFetchEntry {
990 fetch: fetch.clone(),
991 waiters: 1,
992 });
993 PageSource::Fetch(fetch)
994 }
995 }
996 }
997 };
998 let page_buf = source.wait(&self.cache, key).await?;
999
1000 let bytes_to_copy = std::cmp::min(buf.len(), page_buf.len() - offset_in_page);
1002 buf[..bytes_to_copy]
1003 .copy_from_slice(&page_buf.as_ref()[offset_in_page..offset_in_page + bytes_to_copy]);
1004
1005 Ok(bytes_to_copy)
1006 }
1007
1008 pub fn cache(&self, blob_id: u64, buf: &[u8], offset: u64) -> usize {
1015 let page_size: usize = self.page_size.widen();
1016 buf.len() - self.cache_pages(blob_id, buf.chunks_exact(page_size), offset)
1017 }
1018
1019 pub(super) fn cache_pages<'a>(
1021 &self,
1022 blob_id: u64,
1023 pages: impl Iterator<Item = &'a [u8]>,
1024 offset: u64,
1025 ) -> usize {
1026 let (mut page_num, offset_in_page) = self.offset_to_page(offset);
1027 assert_eq!(offset_in_page, 0);
1028 let mut cached = 0;
1029 let mut page_cache = self.cache.write();
1030 for page in pages {
1031 page_cache.cache(blob_id, page, page_num);
1032 cached += page.len();
1033 page_num = match page_num.checked_add(1) {
1034 Some(next) => next,
1035 None => break,
1036 };
1037 }
1038 cached
1039 }
1040
1041 #[cfg(any(test, feature = "test-utils"))]
1045 pub fn clear(&self) {
1046 self.cache.write().clear();
1047 }
1048}
1049
1050impl Cache {
1051 pub fn new(page_size: NonZeroU16, capacity: NonZeroUsize) -> Self {
1054 let slot_size: usize = page_size.widen();
1055 let mut cache = cache::Cache::new(capacity);
1056 cache.prefill(|| IoBufMut::zeroed(slot_size));
1057 let hints = capacity.get().saturating_mul(2).next_power_of_two();
1058 Self {
1059 cache,
1060 hints: vec![usize::MAX; hints],
1061 page_size,
1062 page_fetches: AHashMap::new(),
1063 batches: BatchRegistry::default(),
1064 }
1065 }
1066
1067 fn offset_to_page(page_size: NonZeroU16, offset: u64) -> (u64, u64) {
1070 let page_size: u64 = page_size.widen();
1071 (offset / page_size, offset % page_size)
1072 }
1073
1074 fn locate(page_size: NonZeroU16, offset: u64) -> (u64, usize, usize) {
1077 let (page_num, offset_in_page) = Self::offset_to_page(page_size, offset);
1078 let offset_in_page = offset_in_page as usize;
1079 let width: usize = page_size.widen();
1080 let remaining = width - offset_in_page;
1081 (page_num, offset_in_page, remaining)
1082 }
1083
1084 fn read_at(&self, blob_id: u64, buf: &mut [u8], logical_offset: u64) -> usize {
1090 let (page_num, offset_in_page, remaining) = Self::locate(self.page_size, logical_offset);
1091 let Some(page) = self.get_page(blob_id, page_num) else {
1092 return 0;
1093 };
1094 let page = page.as_ref();
1095
1096 let bytes_to_copy = std::cmp::min(buf.len(), remaining);
1097 buf[..bytes_to_copy].copy_from_slice(&page[offset_in_page..offset_in_page + bytes_to_copy]);
1098
1099 bytes_to_copy
1100 }
1101
1102 fn cache(&mut self, blob_id: u64, page: &[u8], page_num: u64) {
1104 let page_size: usize = self.page_size.widen();
1105 assert_eq!(page.len(), page_size);
1106 let (slot, buf) = self
1107 .cache
1108 .get_or_insert_mut((blob_id, page_num), || IoBufMut::zeroed(page_size));
1109 buf.as_mut().copy_from_slice(page);
1110 let hint = self.hint_index(blob_id, page_num);
1111 self.hints[hint] = slot;
1112 }
1113
1114 #[inline]
1123 const fn hint_index(&self, blob_id: u64, page_num: u64) -> usize {
1124 let salted = page_num.wrapping_add(blob_id.wrapping_mul(commonware_utils::GOLDEN_RATIO));
1125 (salted & (self.hints.len() as u64 - 1)) as usize
1126 }
1127
1128 #[inline]
1130 fn get_page(&self, blob_id: u64, page_num: u64) -> Option<&IoBufMut> {
1131 let key = (blob_id, page_num);
1132 let slot = self.hints[self.hint_index(blob_id, page_num)];
1133 if let Some(page) = self.cache.get_at(slot, &key) {
1134 return Some(page);
1135 }
1136 self.cache.get(&key)
1137 }
1138
1139 #[cfg(any(test, feature = "test-utils"))]
1141 fn clear(&mut self) {
1142 self.cache.retain(|_, _| false);
1143 self.page_fetches.clear();
1144 self.batches = BatchRegistry::default();
1145 }
1146}
1147
1148async fn fetch_cacheable_page(
1150 blob: &impl Blob,
1151 page_num: u64,
1152 page_size: NonZeroU16,
1153) -> Result<IoBuf, Error> {
1154 let width: u64 = page_size.widen();
1156 let (page, _) =
1157 get_page_with_checksum_from_blob(blob, page_num, width, ReadOptions::DONT_CACHE).await?;
1158 cacheable_page(page, page_num, page_size)
1159}
1160
1161fn cacheable_page(page: IoBuf, page_num: u64, page_size: NonZeroU16) -> Result<IoBuf, Error> {
1164 let len = page.len();
1167 let expected: usize = page_size.widen();
1168 if len != expected {
1169 error!(
1170 page_num,
1171 expected = page_size,
1172 actual = len,
1173 "attempted to fetch partial page from blob"
1174 );
1175 return Err(Error::InvalidChecksum);
1176 }
1177 Ok(page)
1178}
1179
1180#[cfg(test)]
1181mod tests {
1182 use super::{
1183 super::{
1184 Checksum,
1185 view::{Tail, View},
1186 },
1187 *,
1188 };
1189 use crate::{
1190 Clock as _, Handle, IoBufMut, IoBufs, IoBufsMut, Runner as _, Spawner as _, Storage as _,
1191 Supervisor as _, WriteOptions, buffer::paged::CHECKSUM_SIZE, deterministic,
1192 mocks::MigratingReadBlob,
1193 };
1194 use commonware_cryptography::Crc32;
1195 use commonware_macros::{select, test_traced};
1196 use commonware_utils::{
1197 NZU16, NZUsize,
1198 channel::{
1199 mpsc::{UnboundedReceiver, unbounded_channel},
1200 oneshot,
1201 },
1202 sync::Mutex,
1203 };
1204 use futures::{future::pending, poll};
1205 use rand::{RngExt as _, seq::SliceRandom as _};
1206 use rstest::rstest;
1207 use std::{
1208 future::poll_fn,
1209 num::NonZeroU16,
1210 sync::{
1211 Arc,
1212 atomic::{AtomicUsize, Ordering},
1213 },
1214 task::Waker,
1215 time::Duration,
1216 };
1217
1218 const PAGE_SIZE: NonZeroU16 = NZU16!(1024);
1220 const PAGE_SIZE_U64: u64 = PAGE_SIZE.get() as u64;
1221
1222 fn expected_cached_bytes(logical_offset: u64, len: usize) -> Vec<u8> {
1223 (0..len)
1224 .map(|i| {
1225 let page = (logical_offset + i as u64) / PAGE_SIZE_U64;
1226 page as u8 + 1
1227 })
1228 .collect()
1229 }
1230
1231 struct BlockingBlob {
1233 started: Mutex<Option<oneshot::Sender<()>>>,
1234 }
1235
1236 impl Blob for BlockingBlob {
1237 async fn read_at(
1238 &self,
1239 offset: u64,
1240 len: usize,
1241 options: ReadOptions,
1242 ) -> Result<IoBufsMut, Error> {
1243 self.read_at_buf(offset, len, IoBufMut::with_capacity(len), options)
1244 .await
1245 }
1246
1247 async fn read_at_buf(
1248 &self,
1249 _offset: u64,
1250 _len: usize,
1251 _bufs: impl Into<IoBufsMut> + Send,
1252 _options: ReadOptions,
1253 ) -> Result<IoBufsMut, Error> {
1254 let sender = self
1255 .started
1256 .lock()
1257 .take()
1258 .expect("blocking blob read started more than once");
1259 let _ = sender.send(());
1260 pending::<()>().await;
1261 unreachable!()
1262 }
1263
1264 async fn write_at(
1265 &self,
1266 _offset: u64,
1267 _bufs: impl Into<crate::IoBufs> + Send,
1268 _options: WriteOptions,
1269 ) -> Result<(), Error> {
1270 Ok(())
1271 }
1272
1273 async fn resize(&self, _len: u64) -> Result<(), Error> {
1274 Ok(())
1275 }
1276
1277 async fn sync(&self) -> Result<(), Error> {
1278 Ok(())
1279 }
1280
1281 async fn start_sync(&self) -> Handle<()> {
1282 Handle::ready(self.sync().await)
1283 }
1284 }
1285
1286 enum ControlledBlobResult {
1287 Success(Vec<u8>),
1288 Error,
1289 }
1290
1291 struct ControlledBlob {
1293 started: Mutex<Option<oneshot::Sender<()>>>,
1294 release: Mutex<Option<oneshot::Receiver<()>>>,
1295 reads: Arc<AtomicUsize>,
1296 result: ControlledBlobResult,
1297 }
1298
1299 impl Blob for ControlledBlob {
1300 async fn read_at(
1301 &self,
1302 offset: u64,
1303 len: usize,
1304 options: ReadOptions,
1305 ) -> Result<IoBufsMut, Error> {
1306 self.read_at_buf(offset, len, IoBufMut::with_capacity(len), options)
1307 .await
1308 }
1309
1310 async fn read_at_buf(
1311 &self,
1312 _offset: u64,
1313 _len: usize,
1314 _bufs: impl Into<IoBufsMut> + Send,
1315 _options: ReadOptions,
1316 ) -> Result<IoBufsMut, Error> {
1317 self.reads.fetch_add(1, Ordering::Relaxed);
1318 let release = self.release.lock().take();
1319 if let Some(release) = release {
1320 let sender = self
1321 .started
1322 .lock()
1323 .take()
1324 .expect("controlled blob start signal missing");
1325 let _ = sender.send(());
1326 release.await.expect("release signal dropped");
1327 }
1328
1329 match &self.result {
1330 ControlledBlobResult::Success(page) => Ok(IoBufsMut::from(page.clone())),
1331 ControlledBlobResult::Error => Err(Error::ReadFailed),
1332 }
1333 }
1334
1335 async fn write_at(
1336 &self,
1337 _offset: u64,
1338 _bufs: impl Into<crate::IoBufs> + Send,
1339 _options: WriteOptions,
1340 ) -> Result<(), Error> {
1341 Ok(())
1342 }
1343
1344 async fn resize(&self, _len: u64) -> Result<(), Error> {
1345 Ok(())
1346 }
1347
1348 async fn sync(&self) -> Result<(), Error> {
1349 Ok(())
1350 }
1351
1352 async fn start_sync(&self) -> Handle<()> {
1353 Handle::ready(self.sync().await)
1354 }
1355 }
1356
1357 #[test_traced]
1358 fn test_cache_basic() {
1359 let mut cache: Cache = Cache::new(PAGE_SIZE, NZUsize!(10));
1360
1361 let mut buf = vec![0; PAGE_SIZE.get() as usize];
1363 let bytes_read = cache.read_at(0, &mut buf, 0);
1364 assert_eq!(bytes_read, 0);
1365
1366 cache.cache(0, &[1; PAGE_SIZE.get() as usize], 0);
1367 let bytes_read = cache.read_at(0, &mut buf, 0);
1368 assert_eq!(bytes_read, PAGE_SIZE.get() as usize);
1369 assert_eq!(buf, [1; PAGE_SIZE.get() as usize]);
1370
1371 cache.cache(0, &[2; PAGE_SIZE.get() as usize], 0);
1373 let bytes_read = cache.read_at(0, &mut buf, 0);
1374 assert_eq!(bytes_read, PAGE_SIZE.get() as usize);
1375 assert_eq!(buf, [2; PAGE_SIZE.get() as usize]);
1376
1377 for i in 0u64..11 {
1379 cache.cache(0, &[i as u8; PAGE_SIZE.get() as usize], i);
1380 }
1381 let bytes_read = cache.read_at(0, &mut buf, 0);
1383 assert_eq!(bytes_read, 0);
1384 for i in 1u64..11 {
1386 let bytes_read = cache.read_at(0, &mut buf, i * PAGE_SIZE_U64);
1387 assert_eq!(bytes_read, PAGE_SIZE.get() as usize);
1388 assert_eq!(buf, [i as u8; PAGE_SIZE.get() as usize]);
1389 }
1390
1391 let mut buf = vec![0; PAGE_SIZE.get() as usize];
1394 let bytes_read = cache.read_at(0, &mut buf, PAGE_SIZE_U64 + 2);
1395 assert_eq!(bytes_read, PAGE_SIZE.get() as usize - 2);
1396 assert_eq!(
1397 &buf[..PAGE_SIZE.get() as usize - 2],
1398 [1; PAGE_SIZE.get() as usize - 2]
1399 );
1400 }
1401
1402 #[test_traced]
1403 fn test_clear_does_not_orphan_recached_page() {
1404 let mut cache: Cache = Cache::new(PAGE_SIZE, NZUsize!(2));
1406 let blob_id = 0u64;
1407 let page_size = PAGE_SIZE.get() as usize;
1408
1409 cache.cache(blob_id, &vec![0xAA; page_size], 0);
1411 cache.cache(blob_id, &vec![0xBB; page_size], 1);
1412 cache.clear();
1413
1414 cache.cache(blob_id, &vec![0xCC; page_size], 1);
1416 let mut buf = vec![0u8; page_size];
1417 assert_eq!(
1418 cache.read_at(blob_id, &mut buf, PAGE_SIZE_U64),
1419 page_size,
1420 "page 1 should be readable after re-cache"
1421 );
1422 assert_eq!(buf, vec![0xCC; page_size]);
1423
1424 cache.cache(blob_id, &vec![0xDD; page_size], 2);
1426
1427 let mut buf = vec![0u8; page_size];
1429 assert_eq!(
1430 cache.read_at(blob_id, &mut buf, PAGE_SIZE_U64),
1431 page_size,
1432 "live page 1 was orphaned by stale-slot eviction"
1433 );
1434 assert_eq!(buf, vec![0xCC; page_size]);
1435
1436 let mut buf = vec![0u8; page_size];
1438 assert_eq!(
1439 cache.read_at(blob_id, &mut buf, PAGE_SIZE_U64 * 2),
1440 page_size
1441 );
1442 assert_eq!(buf, vec![0xDD; page_size]);
1443 }
1444
1445 #[test_traced]
1446 fn test_cache_read_with_blob() {
1447 let executor = deterministic::Runner::default();
1449 executor.start(|context| async move {
1451 let physical_page_size = PAGE_SIZE_U64 + CHECKSUM_SIZE;
1453
1454 let (blob, size) = context
1456 .open("test", "blob".as_bytes())
1457 .await
1458 .expect("Failed to open blob");
1459 assert_eq!(size, 0);
1460 let blob = Arc::new(blob);
1461 for i in 0..11 {
1462 let logical_data = vec![i as u8; PAGE_SIZE.get() as usize];
1464 let crc = Crc32::checksum(&logical_data);
1465 let record = Checksum::new(PAGE_SIZE.get(), crc);
1466 let mut page_data = logical_data;
1467 page_data.extend_from_slice(&record.to_bytes());
1468 blob.write_at(i * physical_page_size, page_data, WriteOptions::default())
1469 .await
1470 .unwrap();
1471 }
1472
1473 let cache_ref = CacheRef::new(PAGE_SIZE, NZUsize!(10));
1475 assert_eq!(cache_ref.next_id(), 0);
1476 assert_eq!(cache_ref.next_id(), 1);
1477 for i in 0..11 {
1478 let mut buf = vec![0; PAGE_SIZE.get() as usize];
1480 cache_ref
1481 .read_after_miss(&blob, 0, &mut buf, i * PAGE_SIZE_U64)
1482 .await
1483 .unwrap();
1484 assert_eq!(buf, [i as u8; PAGE_SIZE.get() as usize]);
1485 }
1486
1487 for i in 1..11 {
1490 let mut buf = vec![0; PAGE_SIZE.get() as usize];
1491 cache_ref
1492 .read_after_miss(&blob, 0, &mut buf, i * PAGE_SIZE_U64)
1493 .await
1494 .unwrap();
1495 assert_eq!(buf, [i as u8; PAGE_SIZE.get() as usize]);
1496 }
1497
1498 blob.sync().await.unwrap();
1500 });
1501 }
1502
1503 #[rstest]
1504 #[case::deterministic(deterministic::Runner::default(), false)]
1505 #[case::tokio(crate::tokio::Runner::default(), false)]
1506 #[cfg_attr(
1507 target_os = "linux",
1508 case::iouring(crate::iouring::Runner::default(), true)
1509 )]
1510 fn test_page_fetch_migrates_between_dedicated_workers<R: crate::Runner>(
1511 #[case] runner: R,
1512 #[case] require_pending: bool,
1513 ) where
1514 R::Context: crate::Clock + crate::Spawner + crate::Storage,
1515 {
1516 runner.start(|context| async move {
1517 let logical_page = vec![7u8; PAGE_SIZE.get() as usize];
1519 let crc = Crc32::checksum(&logical_page);
1520 let mut physical_page = logical_page.clone();
1521 physical_page.extend_from_slice(&Checksum::new(PAGE_SIZE.get(), crc).to_bytes());
1522
1523 let partition = "cache_page_fetch_migration";
1524 let (blob, size) = context.open(partition, b"blob").await.unwrap();
1525 assert_eq!(size, 0);
1526 blob.write_at(0, physical_page, WriteOptions::default())
1527 .await
1528 .unwrap();
1529
1530 let blob = Arc::new(MigratingReadBlob::new(blob, require_pending));
1532 let cache_ref = CacheRef::new(PAGE_SIZE, NZUsize!(2));
1533 let blob_id = cache_ref.next_id();
1534
1535 let concurrent_reads = async {
1539 let (source_polled_tx, source_polled_rx) = oneshot::channel();
1540 let (source_release_tx, source_release_rx) = oneshot::channel();
1541 let first_cache = cache_ref.clone();
1542 let first_blob = blob.clone();
1543 let first = context.child("first").dedicated().spawn(move |_| async move {
1544 let mut buf = vec![0; PAGE_SIZE.get() as usize];
1545 let mut read = Box::pin(first_cache.read_after_miss(
1546 &first_blob,
1547 blob_id,
1548 &mut buf,
1549 0,
1550 ));
1551
1552 assert!(poll!(&mut read).is_pending());
1555 source_polled_tx
1556 .send(())
1557 .expect("source poll receiver dropped");
1558 source_release_rx
1559 .await
1560 .expect("source release signal dropped");
1561 read.await.unwrap();
1562 buf
1563 });
1564 source_polled_rx
1565 .await
1566 .expect("source cache poll never completed");
1567
1568 let second_cache = cache_ref.clone();
1571 let second_blob = blob.clone();
1572 let second = context
1573 .child("second")
1574 .dedicated()
1575 .spawn(move |_| async move {
1576 let mut buf = vec![0; PAGE_SIZE.get() as usize];
1577 second_cache
1578 .read_after_miss(&second_blob, blob_id, &mut buf, 0)
1579 .await
1580 .unwrap();
1581 buf
1582 });
1583
1584 let second_buf = second.await.unwrap();
1586 source_release_tx
1587 .send(())
1588 .expect("source worker exited before release");
1589 (first.await.unwrap(), second_buf)
1590 };
1591
1592 let (first_buf, second_buf) = select! {
1595 result = concurrent_reads => result,
1596 _ = context.sleep(Duration::from_secs(10)) => panic!("concurrent page fetch timed out"),
1597 };
1598
1599 assert_eq!(first_buf, logical_page);
1601 assert_eq!(second_buf, logical_page);
1602 assert_eq!(blob.reads(), 1);
1603 context.remove(partition, None).await.unwrap();
1604 });
1605 }
1606
1607 #[test_traced]
1608 fn test_cache_clear_forces_uncached_blob_read() {
1609 struct CountingBlob {
1610 reads: AtomicUsize,
1611 read_options: Mutex<Vec<ReadOptions>>,
1612 page: Vec<u8>,
1613 }
1614
1615 impl Blob for CountingBlob {
1616 async fn read_at(
1617 &self,
1618 offset: u64,
1619 len: usize,
1620 options: ReadOptions,
1621 ) -> Result<IoBufsMut, Error> {
1622 self.read_at_buf(offset, len, IoBufsMut::default(), options)
1623 .await
1624 }
1625
1626 async fn read_at_buf(
1627 &self,
1628 _offset: u64,
1629 _len: usize,
1630 _bufs: impl Into<IoBufsMut> + Send,
1631 options: ReadOptions,
1632 ) -> Result<IoBufsMut, Error> {
1633 self.reads.fetch_add(1, Ordering::Relaxed);
1634 self.read_options.lock().push(options);
1635 Ok(IoBufsMut::from(self.page.clone()))
1636 }
1637
1638 async fn write_at(
1639 &self,
1640 _offset: u64,
1641 _bufs: impl Into<IoBufs> + Send,
1642 _options: WriteOptions,
1643 ) -> Result<(), Error> {
1644 Ok(())
1645 }
1646
1647 async fn resize(&self, _len: u64) -> Result<(), Error> {
1648 Ok(())
1649 }
1650
1651 async fn sync(&self) -> Result<(), Error> {
1652 Ok(())
1653 }
1654
1655 async fn start_sync(&self) -> Handle<()> {
1656 Handle::ready(self.sync().await)
1657 }
1658 }
1659
1660 let executor = deterministic::Runner::default();
1661 executor.start(|_| async move {
1662 let page = vec![7u8; PAGE_SIZE.get() as usize];
1663 let crc = Crc32::checksum(&page);
1664 let record = Checksum::new(PAGE_SIZE.get(), crc);
1665 let mut physical_page = page.clone();
1666 physical_page.extend_from_slice(&record.to_bytes());
1667 let blob = Arc::new(CountingBlob {
1668 reads: AtomicUsize::new(0),
1669 read_options: Mutex::new(Vec::new()),
1670 page: physical_page,
1671 });
1672 let cache_ref = CacheRef::new(PAGE_SIZE, NZUsize!(2));
1673
1674 let mut buf = vec![0u8; page.len()];
1675 cache_ref
1676 .read_after_miss(&blob, 0, &mut buf, 0)
1677 .await
1678 .unwrap();
1679 assert_eq!(buf, page);
1680 assert_eq!(blob.reads.load(Ordering::Relaxed), 1);
1681
1682 let mut buf = vec![0u8; page.len()];
1683 cache_ref
1684 .read_after_miss(&blob, 0, &mut buf, 0)
1685 .await
1686 .unwrap();
1687 assert_eq!(buf, page);
1688 assert_eq!(blob.reads.load(Ordering::Relaxed), 1);
1689
1690 cache_ref.clear();
1691
1692 let mut buf = vec![0u8; page.len()];
1693 cache_ref
1694 .read_after_miss(&blob, 0, &mut buf, 0)
1695 .await
1696 .unwrap();
1697 assert_eq!(buf, page);
1698 assert_eq!(blob.reads.load(Ordering::Relaxed), 2);
1699 assert_eq!(
1700 *blob.read_options.lock(),
1701 vec![ReadOptions::DONT_CACHE, ReadOptions::DONT_CACHE]
1702 );
1703 });
1704 }
1705
1706 #[test_traced]
1707 fn test_cache_max_page() {
1708 let executor = deterministic::Runner::default();
1709 executor.start(|_| async move {
1710 let cache_ref = CacheRef::new(PAGE_SIZE, NZUsize!(2));
1711
1712 let aligned_max_offset = u64::MAX - (u64::MAX % PAGE_SIZE_U64);
1714
1715 let logical_data = vec![42u8; PAGE_SIZE.get() as usize];
1717
1718 let remaining = cache_ref.cache(0, logical_data.as_slice(), aligned_max_offset);
1720 assert_eq!(remaining, 0);
1721
1722 let mut buf = vec![0u8; PAGE_SIZE.get() as usize];
1724 let page_cache = cache_ref.cache.read();
1725 let bytes_read = page_cache.read_at(0, &mut buf, aligned_max_offset);
1726 assert_eq!(bytes_read, PAGE_SIZE.get() as usize);
1727 assert!(buf.iter().all(|b| *b == 42));
1728 });
1729 }
1730
1731 async fn checksummed_blob<S: crate::Storage>(storage: &S, pages: u64) -> Arc<S::Blob> {
1734 let physical_page_size = PAGE_SIZE_U64 + CHECKSUM_SIZE;
1735 let (blob, size) = storage.open("fill", b"blob").await.unwrap();
1736 assert_eq!(size, 0);
1737 for i in 0..pages {
1738 let logical = vec![i as u8; PAGE_SIZE.get() as usize];
1739 let crc = Crc32::checksum(&logical);
1740 let mut page = logical;
1741 page.extend_from_slice(&Checksum::new(PAGE_SIZE.get(), crc).to_bytes());
1742 blob.write_at(i * physical_page_size, page, WriteOptions::default())
1743 .await
1744 .unwrap();
1745 }
1746 Arc::new(blob)
1747 }
1748
1749 #[test_traced]
1750 fn test_read_after_misses_serves_ranges() {
1751 let executor = deterministic::Runner::default();
1752 executor.start(|context| async move {
1753 let (recording, recordings) = crate::mocks::RecordingContext::new(context);
1754 let blob = checksummed_blob(&recording, 8).await;
1755 let cache_ref = CacheRef::new(PAGE_SIZE, NZUsize!(16));
1756 let page = PAGE_SIZE.get() as usize;
1757
1758 let mut buf = vec![0u8; page];
1761 cache_ref
1762 .read_after_miss(&blob, 0, &mut buf, 2 * PAGE_SIZE_U64)
1763 .await
1764 .unwrap();
1765 recordings.clear();
1766 let ranges = [
1767 (5u64, 3usize),
1768 (2 * PAGE_SIZE_U64 + 1, 2),
1769 (3 * PAGE_SIZE_U64 + page as u64 - 1, 2),
1770 (5 * PAGE_SIZE_U64, page + 1),
1771 ];
1772 let mut bufs: Vec<Vec<u8>> = ranges.iter().map(|&(_, len)| vec![0; len]).collect();
1773 let pending: Vec<(&mut [u8], u64)> = bufs
1774 .iter_mut()
1775 .zip(ranges)
1776 .map(|(buf, (offset, _))| (buf.as_mut_slice(), offset))
1777 .collect();
1778 cache_ref
1779 .read_after_misses(&blob, 0, pending)
1780 .await
1781 .unwrap();
1782 assert_eq!(recordings.snapshot().reads.len(), 5);
1783 for ((offset, len), buf) in ranges.into_iter().zip(&bufs) {
1784 let expected: Vec<u8> = (offset..offset + len as u64)
1785 .map(|o| (o / PAGE_SIZE_U64) as u8)
1786 .collect();
1787 assert_eq!(buf, &expected);
1788 }
1789
1790 recordings.clear();
1792 for (offset, len) in ranges {
1793 let mut buf = vec![0u8; len];
1794 cache_ref
1795 .read_after_miss(&blob, 0, &mut buf, offset)
1796 .await
1797 .unwrap();
1798 }
1799 assert!(recordings.snapshot().reads.is_empty());
1800
1801 let mut inside = [0u8; 1];
1803 let mut straddling = [0u8; 2];
1804 cache_ref
1805 .read_after_misses(
1806 &blob,
1807 0,
1808 vec![
1809 (inside.as_mut_slice(), 7),
1810 (
1811 straddling.as_mut_slice(),
1812 3 * PAGE_SIZE_U64 + page as u64 - 1,
1813 ),
1814 ],
1815 )
1816 .await
1817 .unwrap();
1818 assert_eq!((inside, straddling), ([0], [3, 4]));
1819 assert!(recordings.snapshot().reads.is_empty());
1820 });
1821 }
1822
1823 #[test_traced]
1824 fn test_read_after_misses_rejects_corruption() {
1825 let executor = deterministic::Runner::default();
1826 executor.start(|context| async move {
1827 let blob = checksummed_blob(&context, 3).await;
1828 let physical_page_size = PAGE_SIZE_U64 + CHECKSUM_SIZE;
1829 blob.write_at(physical_page_size + 1, vec![0xFF], WriteOptions::default())
1830 .await
1831 .unwrap();
1832 let cache_ref = CacheRef::new(PAGE_SIZE, NZUsize!(4));
1833
1834 let mut bufs = [[0u8; 1]; 3];
1837 let pending: Vec<(&mut [u8], u64)> = bufs
1838 .iter_mut()
1839 .zip([0, PAGE_SIZE_U64, 2 * PAGE_SIZE_U64])
1840 .map(|(buf, offset)| (buf.as_mut_slice(), offset))
1841 .collect();
1842 assert!(matches!(
1843 cache_ref.read_after_misses(&blob, 0, pending).await,
1844 Err(Error::InvalidChecksum)
1845 ));
1846 let mut buf = [0u8; 1];
1847 assert_eq!(cache_ref.read_cached(0, &mut buf, PAGE_SIZE_U64), 0);
1848
1849 let mut bufs = [[0u8; 1]; 2];
1851 let pending: Vec<(&mut [u8], u64)> = bufs
1852 .iter_mut()
1853 .zip([0, 2 * PAGE_SIZE_U64])
1854 .map(|(buf, offset)| (buf.as_mut_slice(), offset))
1855 .collect();
1856 cache_ref
1857 .read_after_misses(&blob, 0, pending)
1858 .await
1859 .unwrap();
1860 assert_eq!(bufs, [[0], [2]]);
1861 });
1862 }
1863
1864 #[test_traced]
1865 fn test_read_many_into_reads_each_missing_page_once() {
1866 let executor = deterministic::Runner::default();
1867 executor.start(|context| async move {
1868 let (recording, recordings) = crate::mocks::RecordingContext::new(context);
1869 let blob = checksummed_blob(&recording, 40).await;
1870
1871 let cache_ref = CacheRef::new(PAGE_SIZE, NZUsize!(20));
1873 for page in 20..40 {
1874 let mut buf = [0u8; 1];
1875 cache_ref
1876 .read_after_miss(&blob, 0, &mut buf, page * PAGE_SIZE_U64)
1877 .await
1878 .unwrap();
1879 }
1880 recordings.clear();
1881
1882 let size = 40 * PAGE_SIZE_U64;
1885 let view = View {
1886 blob: &blob,
1887 cache_ref: &cache_ref,
1888 id: 0,
1889 size,
1890 tail_offset: size,
1891 tail: Tail::Sealed(&[]),
1892 };
1893 let offsets: Vec<u64> = (0..10).map(|page| page * PAGE_SIZE_U64).collect();
1894 let mut buf = vec![0u8; offsets.len()];
1895 let served = view
1896 .read_many_into(&mut buf, &offsets, NZUsize!(1))
1897 .await
1898 .unwrap();
1899 assert_eq!(served, 0);
1900 assert_eq!(buf, (0u8..10).collect::<Vec<_>>());
1901 assert_eq!(recordings.snapshot().reads.len(), offsets.len());
1902 });
1903 }
1904
1905 #[test_traced]
1906 fn test_read_many_into_copies_resident_pages_first() {
1907 let executor = deterministic::Runner::default();
1908 executor.start(|context| async move {
1909 let (recording, recordings) = crate::mocks::RecordingContext::new(context);
1910 let blob = checksummed_blob(&recording, 6).await;
1911
1912 let cache_ref = CacheRef::new(PAGE_SIZE, NZUsize!(2));
1915 let mut buf = [0u8; 1];
1916 cache_ref
1917 .read_after_miss(&blob, 0, &mut buf, PAGE_SIZE_U64)
1918 .await
1919 .unwrap();
1920 recordings.clear();
1921
1922 let size = 6 * PAGE_SIZE_U64;
1923 let view = View {
1924 blob: &blob,
1925 cache_ref: &cache_ref,
1926 id: 0,
1927 size,
1928 tail_offset: size,
1929 tail: Tail::Sealed(&[]),
1930 };
1931 let mut buf = vec![0u8; size as usize];
1932 let served = view
1933 .read_many_into(&mut buf, &[0], NZUsize!(size as usize))
1934 .await
1935 .unwrap();
1936 assert_eq!(served, 0);
1937 let expected: Vec<u8> = (0..size).map(|o| (o / PAGE_SIZE_U64) as u8).collect();
1938 assert_eq!(buf, expected);
1939 assert_eq!(recordings.snapshot().reads.len(), 5);
1940 });
1941 }
1942
1943 fn page_bytes(value: u8, len: u16) -> Vec<u8> {
1946 let mut page = vec![value; usize::from(len)];
1947 let crc = Crc32::checksum(&page);
1948 page.resize(PAGE_SIZE.get() as usize, 0);
1949 page.extend_from_slice(&Checksum::new(len, crc).to_bytes());
1950 page
1951 }
1952
1953 fn full_page(page: usize) -> Vec<u8> {
1955 page_bytes(page as u8, PAGE_SIZE.get())
1956 }
1957
1958 fn corrupt_page(page: usize) -> Vec<u8> {
1960 let mut bytes = full_page(page);
1961 bytes[0] ^= 0xFF;
1962 bytes
1963 }
1964
1965 struct GatedBlob {
1967 gates: Mutex<Vec<Option<oneshot::Receiver<()>>>>,
1968 reads: Mutex<Vec<usize>>,
1969 corrupt: Mutex<Option<usize>>,
1971 polls: AtomicUsize,
1973 }
1974
1975 impl GatedBlob {
1976 fn new(pages: usize) -> (Arc<Self>, Vec<Option<oneshot::Sender<()>>>) {
1978 let (opens, gates) = (0..pages)
1979 .map(|_| {
1980 let (open, gate) = oneshot::channel();
1981 (Some(open), Some(gate))
1982 })
1983 .unzip();
1984 let blob = Self {
1985 gates: Mutex::new(gates),
1986 reads: Mutex::new(Vec::new()),
1987 corrupt: Mutex::new(None),
1988 polls: AtomicUsize::new(0),
1989 };
1990 (Arc::new(blob), opens)
1991 }
1992
1993 fn reads_of(&self, page: usize) -> usize {
1995 self.reads
1996 .lock()
1997 .iter()
1998 .filter(|&&read| read == page)
1999 .count()
2000 }
2001 }
2002
2003 impl Blob for GatedBlob {
2004 async fn read_at(
2005 &self,
2006 offset: u64,
2007 len: usize,
2008 options: ReadOptions,
2009 ) -> Result<IoBufsMut, Error> {
2010 self.read_at_buf(offset, len, IoBufsMut::default(), options)
2011 .await
2012 }
2013
2014 async fn read_at_buf(
2015 &self,
2016 offset: u64,
2017 _len: usize,
2018 _bufs: impl Into<IoBufsMut> + Send,
2019 _options: ReadOptions,
2020 ) -> Result<IoBufsMut, Error> {
2021 let page = (offset / (PAGE_SIZE_U64 + CHECKSUM_SIZE)) as usize;
2022 self.reads.lock().push(page);
2023 let corrupted = self.corrupt.lock().take_if(|corrupt| *corrupt == page);
2024 let bytes = if corrupted.is_some() {
2025 corrupt_page(page)
2026 } else {
2027 full_page(page)
2028 };
2029 let gate = self.gates.lock()[page].take();
2030 if let Some(gate) = gate {
2031 let _ = gate.await;
2032 }
2033 Ok(IoBufsMut::from(bytes))
2034 }
2035
2036 fn read_many(
2037 &self,
2038 ranges: &[(u64, usize)],
2039 options: ReadOptions,
2040 ) -> impl futures::Stream<Item = Result<(usize, IoBufsMut), Error>> + Send {
2041 let mut reads = Box::pin(crate::read_each(self, ranges, options));
2042 futures::stream::poll_fn(move |cx| {
2043 self.polls.fetch_add(1, Ordering::Relaxed);
2044 reads.as_mut().poll_next(cx)
2045 })
2046 }
2047
2048 async fn write_at(
2049 &self,
2050 _offset: u64,
2051 _bufs: impl Into<IoBufs> + Send,
2052 _options: WriteOptions,
2053 ) -> Result<(), Error> {
2054 Ok(())
2055 }
2056
2057 async fn resize(&self, _len: u64) -> Result<(), Error> {
2058 Ok(())
2059 }
2060
2061 async fn sync(&self) -> Result<(), Error> {
2062 Ok(())
2063 }
2064
2065 async fn start_sync(&self) -> Handle<()> {
2066 Handle::ready(self.sync().await)
2067 }
2068 }
2069
2070 type Hook = Box<dyn FnMut(&mut Context<'_>) -> Poll<()> + Send>;
2072
2073 enum Step {
2075 Yield(usize, Vec<u8>),
2077 Fail(Error),
2079 Wait(oneshot::Receiver<()>),
2081 Call(Hook),
2083 }
2084
2085 impl Step {
2086 fn page(index: usize) -> Self {
2088 Self::Yield(index, full_page(index))
2089 }
2090 }
2091
2092 struct ScriptedBlob(Mutex<Vec<Step>>);
2095
2096 impl ScriptedBlob {
2097 fn new(script: Vec<Step>) -> Arc<Self> {
2098 Arc::new(Self(Mutex::new(script)))
2099 }
2100 }
2101
2102 impl Blob for ScriptedBlob {
2103 async fn read_at(
2104 &self,
2105 _offset: u64,
2106 _len: usize,
2107 _options: ReadOptions,
2108 ) -> Result<IoBufsMut, Error> {
2109 unreachable!("only batched reads are scripted")
2110 }
2111
2112 async fn read_at_buf(
2113 &self,
2114 _offset: u64,
2115 _len: usize,
2116 _bufs: impl Into<IoBufsMut> + Send,
2117 _options: ReadOptions,
2118 ) -> Result<IoBufsMut, Error> {
2119 unreachable!("only batched reads are scripted")
2120 }
2121
2122 fn read_many(
2123 &self,
2124 _ranges: &[(u64, usize)],
2125 _options: ReadOptions,
2126 ) -> impl futures::Stream<Item = Result<(usize, IoBufsMut), Error>> + Send {
2127 let steps = std::mem::take(&mut *self.0.lock());
2128 futures::stream::unfold(steps.into_iter(), |mut steps| async move {
2129 loop {
2130 match steps.next()? {
2131 Step::Yield(index, bytes) => {
2132 return Some((Ok((index, IoBufsMut::from(bytes))), steps));
2133 }
2134 Step::Fail(err) => return Some((Err(err), steps)),
2135 Step::Wait(gate) => {
2136 let _ = gate.await;
2137 }
2138 Step::Call(mut call) => poll_fn(|cx| call(cx)).await,
2139 }
2140 }
2141 })
2142 }
2143
2144 async fn write_at(
2145 &self,
2146 _offset: u64,
2147 _bufs: impl Into<IoBufs> + Send,
2148 _options: WriteOptions,
2149 ) -> Result<(), Error> {
2150 Ok(())
2151 }
2152
2153 async fn resize(&self, _len: u64) -> Result<(), Error> {
2154 Ok(())
2155 }
2156
2157 async fn sync(&self) -> Result<(), Error> {
2158 Ok(())
2159 }
2160
2161 async fn start_sync(&self) -> Handle<()> {
2162 Handle::ready(self.sync().await)
2163 }
2164 }
2165
2166 fn open(opens: &mut [Option<oneshot::Sender<()>>]) {
2168 for open in opens.iter_mut().filter_map(Option::take) {
2169 let _ = open.send(());
2170 }
2171 }
2172
2173 async fn read_first_bytes<B: Blob>(
2176 blob: &Arc<B>,
2177 cache_ref: &CacheRef,
2178 pages: &[u64],
2179 ) -> Result<Vec<u8>, Error> {
2180 let size = pages.last().map_or(0, |last| (last + 1) * PAGE_SIZE_U64);
2181 let view = View {
2182 blob,
2183 cache_ref,
2184 id: 0,
2185 size,
2186 tail_offset: size,
2187 tail: Tail::Sealed(&[]),
2188 };
2189 let offsets: Vec<u64> = pages.iter().map(|page| page * PAGE_SIZE_U64).collect();
2190 let mut buf = vec![0u8; pages.len()];
2191 view.read_many_into(&mut buf, &offsets, NZUsize!(1)).await?;
2192 Ok(buf)
2193 }
2194
2195 async fn read_first_byte<B: Blob>(
2197 blob: &Arc<B>,
2198 cache_ref: &CacheRef,
2199 page: u64,
2200 ) -> Result<u8, Error> {
2201 let mut buf = [0u8; 1];
2202 cache_ref
2203 .read_after_miss(blob, 0, &mut buf, page * PAGE_SIZE_U64)
2204 .await?;
2205 Ok(buf[0])
2206 }
2207
2208 fn assert_idle(cache_ref: &CacheRef) {
2210 let cache = cache_ref.cache.read();
2211 assert!(cache.page_fetches.is_empty());
2212 let mut batches = cache.batches.by_blob.values().flatten();
2213 assert!(batches.all(|(batch, _)| batch.strong_count() == 0));
2214 }
2215
2216 fn joiners(cache_ref: &CacheRef) -> Vec<Option<(usize, usize)>> {
2219 let cache = cache_ref.cache.read();
2220 let batches = cache.batches.by_blob.values().flatten();
2221 let batches = batches.filter_map(|(batch, _)| batch.upgrade());
2222 batches
2223 .flat_map(|batch| {
2224 let state = batch.state.lock();
2225 let pages = state.joiners.iter().map(|joiners| {
2226 joiners.as_ref().map(|joiners| {
2227 let live = joiners
2228 .iter()
2229 .filter(|&&(_, reader)| state.inboxes[reader].is_some());
2230 (joiners.len(), live.count())
2231 })
2232 });
2233 pages.collect::<Vec<_>>()
2234 })
2235 .collect()
2236 }
2237
2238 fn selected(cache_ref: &CacheRef) -> usize {
2240 let pages = joiners(cache_ref);
2241 pages.iter().flatten().map(|(_, live)| live).sum()
2242 }
2243
2244 #[derive(Default)]
2246 struct Wakes(AtomicUsize);
2247
2248 impl std::task::Wake for Wakes {
2249 fn wake(self: Arc<Self>) {
2250 self.0.fetch_add(1, Ordering::SeqCst);
2251 }
2252 }
2253
2254 #[test_traced]
2255 fn test_read_many_into_joins_fetch_in_flight() {
2256 let executor = deterministic::Runner::default();
2257 executor.start(|_| async move {
2258 const PAGES: u64 = 10;
2259 let (blob, mut opens) = GatedBlob::new(PAGES as usize);
2260 let cache_ref = CacheRef::new(PAGE_SIZE, NZUsize!(4));
2261
2262 let mut first = pin!(read_first_byte(&blob, &cache_ref, 0));
2264 assert!(poll!(first.as_mut()).is_pending());
2265
2266 let pages: Vec<u64> = (0..PAGES).collect();
2269 let mut second = pin!(read_first_bytes(&blob, &cache_ref, &pages));
2270 assert!(poll!(second.as_mut()).is_pending());
2271 opens[0].take().unwrap().send(()).unwrap();
2272 assert_eq!(first.await.unwrap(), 0);
2273 open(&mut opens);
2274 assert_eq!(second.await.unwrap(), (0..PAGES as u8).collect::<Vec<_>>());
2275 assert_eq!(blob.reads_of(0), 1);
2276 });
2277 }
2278
2279 #[test_traced]
2280 fn test_overlapping_reads_read_each_page_once() {
2281 let executor = deterministic::Runner::default();
2282 executor.start(|_| async move {
2283 const PAGES: u64 = 7;
2284 let (blob, mut opens) = GatedBlob::new(PAGES as usize);
2285 let cache_ref = CacheRef::new(PAGE_SIZE, NZUsize!(16));
2286
2287 let mut first = pin!(read_first_bytes(&blob, &cache_ref, &[0, 1, 2, 3]));
2290 assert!(poll!(first.as_mut()).is_pending());
2291 let mut second = pin!(read_first_bytes(&blob, &cache_ref, &[2, 3, 4, 5]));
2292 assert!(poll!(second.as_mut()).is_pending());
2293
2294 let mut third = pin!(read_first_bytes(&blob, &cache_ref, &[1, 5, 6]));
2297 assert!(poll!(third.as_mut()).is_pending());
2298 let mut single = pin!(read_first_byte(&blob, &cache_ref, 3));
2299 assert!(poll!(single.as_mut()).is_pending());
2300 assert_eq!(selected(&cache_ref), 2 + (1 + 1) + 1);
2301
2302 open(&mut opens);
2303 assert_eq!(first.await.unwrap(), [0, 1, 2, 3]);
2304 assert_eq!(second.await.unwrap(), [2, 3, 4, 5]);
2305 assert_eq!(third.await.unwrap(), [1, 5, 6]);
2306 assert_eq!(single.await.unwrap(), 3);
2307 for page in 0..PAGES as usize {
2308 assert_eq!(blob.reads_of(page), 1);
2309 }
2310 assert_idle(&cache_ref);
2311 });
2312 }
2313
2314 #[test_traced]
2315 fn test_joining_many_pages_selects_all_before_driving() {
2316 let executor = deterministic::Runner::default();
2317 executor.start(|_| async move {
2318 const PAGES: u64 = 3 * MISSES_PER_LOCK as u64;
2319 let (blob, mut opens) = GatedBlob::new(PAGES as usize);
2320 let cache_ref = CacheRef::new(PAGE_SIZE, NZUsize!(2 * PAGES as usize));
2321 let pages: Vec<u64> = (0..PAGES).collect();
2322 let expected: Vec<u8> = pages.iter().map(|&page| page as u8).collect();
2323
2324 let mut first = pin!(read_first_bytes(&blob, &cache_ref, &pages));
2326 assert!(poll!(first.as_mut()).is_pending());
2327
2328 let mut second = pin!(read_first_bytes(&blob, &cache_ref, &pages));
2331 assert!(poll!(second.as_mut()).is_pending());
2332 assert_eq!(selected(&cache_ref), PAGES as usize);
2333
2334 open(&mut opens);
2335 assert_eq!(second.await.unwrap(), expected);
2336 assert_eq!(first.await.unwrap(), expected);
2337 assert_eq!(blob.reads.lock().len(), PAGES as usize);
2338 assert_idle(&cache_ref);
2339 });
2340 }
2341
2342 #[test_traced]
2343 fn test_read_joins_unpolled_batch_under_eviction() {
2344 let executor = deterministic::Runner::default();
2345 executor.start(|_| async move {
2346 let (blob, mut opens) = GatedBlob::new(40);
2348 open(&mut opens[20..]);
2349 let cache_ref = CacheRef::new(PAGE_SIZE, NZUsize!(20));
2350 for page in 20..40 {
2351 read_first_byte(&blob, &cache_ref, page).await.unwrap();
2352 }
2353 blob.reads.lock().clear();
2354
2355 let pages: Vec<u64> = (0..10).collect();
2357 let mut first = pin!(read_first_bytes(&blob, &cache_ref, &pages));
2358 assert!(poll!(first.as_mut()).is_pending());
2359 open(&mut opens);
2360
2361 let expected: Vec<u8> = (0..10).collect();
2364 assert_eq!(
2365 read_first_bytes(&blob, &cache_ref, &pages).await.unwrap(),
2366 expected
2367 );
2368 assert_eq!(first.await.unwrap(), expected);
2369 assert_eq!(blob.reads.lock().len(), pages.len());
2370 assert_idle(&cache_ref);
2371 });
2372 }
2373
2374 #[test_traced]
2375 fn test_read_joining_batch_outlives_its_starter() {
2376 let executor = deterministic::Runner::default();
2377 executor.start(|_| async move {
2378 let (blob, mut opens) = GatedBlob::new(2);
2379 let cache_ref = CacheRef::new(PAGE_SIZE, NZUsize!(16));
2380
2381 let mut first = Box::pin(read_first_bytes(&blob, &cache_ref, &[0, 1]));
2383 assert!(poll!(first.as_mut()).is_pending());
2384 let mut second = pin!(read_first_byte(&blob, &cache_ref, 1));
2385 assert!(poll!(second.as_mut()).is_pending());
2386
2387 drop(first);
2390 open(&mut opens);
2391 assert_eq!(second.await.unwrap(), 1);
2392 assert_eq!(blob.reads.lock().len(), 2);
2393 assert_idle(&cache_ref);
2394 });
2395 }
2396
2397 #[test_traced]
2398 fn test_joined_read_finishes_before_its_batch() {
2399 let executor = deterministic::Runner::default();
2400 executor.start(|_| async move {
2401 let (blob, mut opens) = GatedBlob::new(3);
2402 let cache_ref = CacheRef::new(PAGE_SIZE, NZUsize!(16));
2403 let mut first = pin!(read_first_bytes(&blob, &cache_ref, &[0, 1, 2]));
2404 assert!(poll!(first.as_mut()).is_pending());
2405
2406 let mut second = pin!(read_first_byte(&blob, &cache_ref, 1));
2409 assert!(poll!(second.as_mut()).is_pending());
2410 opens[1].take().unwrap().send(()).unwrap();
2411 assert_eq!(second.await.unwrap(), 1);
2412 assert_eq!(joiners(&cache_ref), [Some((0, 0)), None, Some((0, 0))]);
2413
2414 open(&mut opens);
2415 assert_eq!(first.await.unwrap(), [0, 1, 2]);
2416 assert_idle(&cache_ref);
2417 });
2418 }
2419
2420 #[test_traced]
2421 fn test_batch_read_polled_once_per_completion() {
2422 const PAGES: usize = 32;
2423 const JOINERS: usize = 4;
2424
2425 let polls = |joiners: usize| {
2428 deterministic::Runner::default().start(|context| async move {
2429 let (blob, mut opens) = GatedBlob::new(PAGES);
2430 let cache_ref = CacheRef::new(PAGE_SIZE, NZUsize!(2 * PAGES));
2431 let pages: Vec<u64> = (0..PAGES as u64).collect();
2432 let reads: Vec<_> = (0..=joiners)
2433 .map(|_| {
2434 let (blob, cache_ref, pages) =
2435 (blob.clone(), cache_ref.clone(), pages.clone());
2436 context.child("reader").spawn(move |_| async move {
2437 read_first_bytes(&blob, &cache_ref, &pages).await.unwrap()
2438 })
2439 })
2440 .collect();
2441 context.sleep(Duration::from_millis(10)).await;
2442 let before = blob.polls.load(Ordering::Relaxed);
2443 for open in &mut opens {
2444 let _ = open.take().unwrap().send(());
2445 context.sleep(Duration::from_millis(1)).await;
2446 }
2447 let expected: Vec<u8> = (0..PAGES as u8).collect();
2448 for read in futures::future::join_all(reads).await {
2449 assert_eq!(read.unwrap(), expected);
2450 }
2451 for page in 0..PAGES {
2452 assert_eq!(blob.reads_of(page), 1);
2453 }
2454 blob.polls.load(Ordering::Relaxed) - before
2455 })
2456 };
2457
2458 let alone = polls(0);
2461 let joined = polls(JOINERS);
2462 assert!(
2463 joined <= alone + PAGES,
2464 "{joined} stream polls with {JOINERS} joiners, {alone} alone"
2465 );
2466 }
2467
2468 #[test]
2469 fn test_reader_that_stops_driving_rewakes_waiting_reader() {
2470 type Read = Pin<Box<dyn Future<Output = Result<Vec<u8>, Error>> + Send>>;
2471 let starter: Arc<Mutex<Option<Read>>> = Arc::default();
2472 let starter_wakes = Arc::new(Wakes::default());
2473 let starter_waker = Waker::from(starter_wakes.clone());
2474 let woken_before = Arc::new(AtomicUsize::new(0));
2475 let stalled: Arc<Mutex<Option<Waker>>> = Arc::default();
2476
2477 let mut first_call = true;
2481 let stall = {
2482 let stalled = stalled.clone();
2483 move |cx: &mut Context<'_>| {
2484 if !std::mem::take(&mut first_call) {
2485 return Poll::Ready(());
2486 }
2487 *stalled.lock() = Some(cx.waker().clone());
2488 Poll::Pending
2489 }
2490 };
2491 let mut first_call = true;
2492 let race = {
2493 let (starter, waker) = (starter.clone(), starter_waker.clone());
2494 let (wakes, woken_before) = (starter_wakes.clone(), woken_before.clone());
2495 move |cx: &mut Context<'_>| {
2496 if !std::mem::take(&mut first_call) {
2497 return Poll::Ready(());
2498 }
2499 cx.waker().wake_by_ref();
2500 let mut starter = starter.lock();
2501 let read = starter.as_mut().unwrap().as_mut();
2502 assert!(read.poll(&mut Context::from_waker(&waker)).is_pending());
2503 woken_before.store(wakes.0.load(Ordering::SeqCst), Ordering::SeqCst);
2504 Poll::Pending
2505 }
2506 };
2507 let blob = ScriptedBlob::new(vec![
2508 Step::Call(Box::new(stall)),
2509 Step::page(0),
2510 Step::Call(Box::new(race)),
2511 Step::page(1),
2512 ]);
2513 let cache_ref = CacheRef::new(PAGE_SIZE, NZUsize!(16));
2514
2515 *starter.lock() = Some(Box::pin({
2517 let (blob, cache_ref) = (blob.clone(), cache_ref.clone());
2518 async move { read_first_bytes(&blob, &cache_ref, &[0, 1]).await }
2519 }));
2520 let poll_starter = || {
2521 let mut starter = starter.lock();
2522 let read = starter.as_mut().unwrap().as_mut();
2523 read.poll(&mut Context::from_waker(&starter_waker))
2524 };
2525 assert!(poll_starter().is_pending());
2526
2527 stalled.lock().take().unwrap().wake();
2530 let mut joiner = pin!(read_first_byte(&blob, &cache_ref, 0));
2531 let joined = joiner
2532 .as_mut()
2533 .poll(&mut Context::from_waker(Waker::noop()));
2534 assert!(matches!(joined, Poll::Ready(Ok(0))));
2535
2536 assert!(starter_wakes.0.load(Ordering::SeqCst) > woken_before.load(Ordering::SeqCst));
2539 assert!(matches!(poll_starter(), Poll::Ready(Ok(ref bytes)) if bytes == &[0, 1]));
2540 }
2541
2542 const SCHEDULE_DEADLINE: Duration = Duration::from_secs(30);
2546
2547 type Completion = Result<(usize, IoBufsMut), Error>;
2549
2550 struct CompletionBlob {
2553 completions: Mutex<Option<UnboundedReceiver<Completion>>>,
2554 paused: std::sync::mpsc::Sender<()>,
2556 resume: Mutex<std::sync::mpsc::Receiver<()>>,
2558 yielded: AtomicUsize,
2560 }
2561
2562 impl Blob for CompletionBlob {
2563 async fn read_at(
2564 &self,
2565 _offset: u64,
2566 _len: usize,
2567 _options: ReadOptions,
2568 ) -> Result<IoBufsMut, Error> {
2569 unreachable!("pages are read in batches")
2570 }
2571
2572 async fn read_at_buf(
2573 &self,
2574 _offset: u64,
2575 _len: usize,
2576 _bufs: impl Into<IoBufsMut> + Send,
2577 _options: ReadOptions,
2578 ) -> Result<IoBufsMut, Error> {
2579 unreachable!("pages are read in batches")
2580 }
2581
2582 fn read_many(
2583 &self,
2584 ranges: &[(u64, usize)],
2585 _options: ReadOptions,
2586 ) -> impl futures::Stream<Item = Completion> + Send {
2587 assert_eq!(ranges.len(), 2);
2588 let mut completions = self.completions.lock().take().unwrap();
2589 let mut paused = false;
2590 futures::stream::poll_fn(move |cx| {
2591 if self.yielded.load(Ordering::SeqCst) == 2 {
2592 return Poll::Ready(None);
2593 }
2594 let result = completions.poll_recv(cx);
2595 if result.is_ready() {
2596 self.yielded.fetch_add(1, Ordering::SeqCst);
2597 }
2598
2599 if result.is_pending() && self.yielded.load(Ordering::SeqCst) == 1 && !paused {
2602 paused = true;
2603 self.paused.send(()).unwrap();
2604 self.resume.lock().recv_timeout(SCHEDULE_DEADLINE).unwrap();
2605 }
2606 result
2607 })
2608 }
2609
2610 async fn write_at(
2611 &self,
2612 _offset: u64,
2613 _bufs: impl Into<IoBufs> + Send,
2614 _options: WriteOptions,
2615 ) -> Result<(), Error> {
2616 Ok(())
2617 }
2618
2619 async fn resize(&self, _len: u64) -> Result<(), Error> {
2620 Ok(())
2621 }
2622
2623 async fn sync(&self) -> Result<(), Error> {
2624 Ok(())
2625 }
2626
2627 async fn start_sync(&self) -> Handle<()> {
2628 Handle::ready(self.sync().await)
2629 }
2630 }
2631
2632 #[test]
2638 fn test_reader_that_stops_driving_rewakes_waiting_reader_on_tokio() {
2639 let runner = crate::tokio::Runner::new(crate::tokio::Config::new().with_worker_threads(4));
2640 runner.start(|context| async move {
2641 for _ in 0..20 {
2642 let (completions, receiver) = unbounded_channel();
2643 let (paused_tx, paused_rx) = std::sync::mpsc::channel();
2644 let (resume_stream_tx, resume_stream_rx) = std::sync::mpsc::channel();
2645 let blob = Arc::new(CompletionBlob {
2646 completions: Mutex::new(Some(receiver)),
2647 paused: paused_tx,
2648 resume: Mutex::new(resume_stream_rx),
2649 yielded: AtomicUsize::new(0),
2650 });
2651 let cache_ref = CacheRef::new(PAGE_SIZE, NZUsize!(4));
2652
2653 let (first_polled_tx, first_polled_rx) = std::sync::mpsc::channel();
2657 let (second_poll_tx, second_poll_rx) = std::sync::mpsc::channel();
2658 let (resume_starter_tx, resume_starter_rx) = std::sync::mpsc::channel();
2659 let (second_polled_tx, second_polled_rx) = std::sync::mpsc::channel();
2660 let starter = context.child("starter").spawn({
2661 let (blob, cache_ref) = (blob.clone(), cache_ref.clone());
2662 move |_| async move {
2663 let mut read = pin!(read_first_bytes(&blob, &cache_ref, &[0, 1]));
2664 let mut polls = 0;
2665 poll_fn(move |cx| {
2666 polls += 1;
2667 if polls == 2 {
2668 second_poll_tx.send(()).unwrap();
2669 resume_starter_rx.recv_timeout(SCHEDULE_DEADLINE).unwrap();
2670 }
2671 let result = read.as_mut().poll(cx);
2672 if polls <= 2 {
2673 assert!(result.is_pending());
2674 let polled = if polls == 1 {
2675 &first_polled_tx
2676 } else {
2677 &second_polled_tx
2678 };
2679 polled.send(()).unwrap();
2680 }
2681 result
2682 })
2683 .await
2684 }
2685 });
2686
2687 let (start_joiner_tx, start_joiner_rx) = oneshot::channel();
2691 let schedule = std::thread::spawn(move || {
2692 first_polled_rx.recv_timeout(SCHEDULE_DEADLINE).unwrap();
2693 completions
2694 .send(Ok((0, IoBufsMut::from(full_page(0)))))
2695 .unwrap();
2696 second_poll_rx.recv_timeout(SCHEDULE_DEADLINE).unwrap();
2697 start_joiner_tx.send(()).unwrap();
2698 paused_rx.recv_timeout(SCHEDULE_DEADLINE).unwrap();
2699 completions
2700 .send(Ok((1, IoBufsMut::from(full_page(1)))))
2701 .unwrap();
2702 resume_starter_tx.send(()).unwrap();
2703 second_polled_rx.recv_timeout(SCHEDULE_DEADLINE).unwrap();
2704 resume_stream_tx.send(()).unwrap();
2705 });
2706 select! {
2707 result = start_joiner_rx => result.unwrap(),
2708 _ = context.sleep(SCHEDULE_DEADLINE) => panic!("starter was not polled again"),
2709 }
2710
2711 let joiner = context.child("joiner").spawn({
2713 let (blob, cache_ref) = (blob.clone(), cache_ref.clone());
2714 move |_| async move { read_first_byte(&blob, &cache_ref, 0).await }
2715 });
2716 let joined = select! {
2717 result = joiner => result.unwrap(),
2718 _ = context.sleep(SCHEDULE_DEADLINE) => panic!("joiner stalled"),
2719 };
2720 assert_eq!(joined.unwrap(), 0);
2721 schedule.join().unwrap();
2722
2723 let started = select! {
2725 result = starter => result.unwrap(),
2726 _ = context.sleep(SCHEDULE_DEADLINE) => {
2727 panic!("starter stalled after the joiner left")
2728 },
2729 };
2730 assert_eq!(started.unwrap(), [0, 1]);
2731 assert_eq!(blob.yielded.load(Ordering::SeqCst), 2);
2732 assert_idle(&cache_ref);
2733 }
2734 });
2735 }
2736
2737 #[rstest]
2738 #[case::ends(false)]
2739 #[case::fails(true)]
2740 fn test_starting_read_waits_for_its_batch_to_end(#[case] fail: bool) {
2741 let executor = deterministic::Runner::default();
2742 executor.start(|_| async move {
2743 let (start, started) = oneshot::channel();
2744 let (end, ended) = oneshot::channel();
2745 let mut script = vec![
2746 Step::Wait(started),
2747 Step::page(0),
2748 Step::page(1),
2749 Step::Wait(ended),
2750 ];
2751 if fail {
2752 script.push(Step::Fail(Error::ReadFailed));
2753 }
2754 let blob = ScriptedBlob::new(script);
2755 let cache_ref = CacheRef::new(PAGE_SIZE, NZUsize!(16));
2756
2757 let mut read = pin!(read_first_bytes(&blob, &cache_ref, &[0, 1]));
2759 assert!(poll!(read.as_mut()).is_pending());
2760 let mut joiner = pin!(read_first_byte(&blob, &cache_ref, 0));
2761 assert!(poll!(joiner.as_mut()).is_pending());
2762
2763 start.send(()).unwrap();
2766 assert_eq!(joiner.await.unwrap(), 0);
2767 assert!(poll!(read.as_mut()).is_pending());
2768 end.send(()).unwrap();
2769 let result = read.await;
2770 if fail {
2771 assert!(matches!(result, Err(Error::ReadFailed)));
2772 } else {
2773 assert_eq!(result.unwrap(), [0, 1]);
2774 }
2775 assert_idle(&cache_ref);
2776 });
2777 }
2778
2779 #[test_traced]
2780 fn test_failed_page_fails_its_readers() {
2781 let executor = deterministic::Runner::default();
2782 executor.start(|_| async move {
2783 let (blob, mut opens) = GatedBlob::new(2);
2784 *blob.corrupt.lock() = Some(1);
2785 let cache_ref = CacheRef::new(PAGE_SIZE, NZUsize!(16));
2786
2787 let mut first = pin!(read_first_bytes(&blob, &cache_ref, &[0, 1]));
2789 assert!(poll!(first.as_mut()).is_pending());
2790 let mut zero = pin!(read_first_byte(&blob, &cache_ref, 0));
2791 assert!(poll!(zero.as_mut()).is_pending());
2792 let mut one = pin!(read_first_byte(&blob, &cache_ref, 1));
2793 assert!(poll!(one.as_mut()).is_pending());
2794
2795 open(&mut opens[1..]);
2797 assert!(matches!(first.await, Err(Error::InvalidChecksum)));
2798 assert!(matches!(one.await, Err(Error::InvalidChecksum)));
2799
2800 assert_eq!(read_first_byte(&blob, &cache_ref, 1).await.unwrap(), 1);
2803 open(&mut opens);
2804 assert_eq!(zero.await.unwrap(), 0);
2805 assert_eq!(blob.reads_of(1), 2);
2806 assert_idle(&cache_ref);
2807 });
2808 }
2809
2810 #[test_traced]
2811 fn test_failed_read_fails_readers_still_waiting() {
2812 let executor = deterministic::Runner::default();
2813 executor.start(|_| async move {
2814 let (start, started) = oneshot::channel();
2815 let (fail, failed) = oneshot::channel();
2816 let blob = ScriptedBlob::new(vec![
2817 Step::Wait(started),
2818 Step::page(0),
2819 Step::Wait(failed),
2820 Step::Fail(Error::ReadFailed),
2821 ]);
2822 let cache_ref = CacheRef::new(PAGE_SIZE, NZUsize!(16));
2823
2824 let mut first = pin!(read_first_bytes(&blob, &cache_ref, &[0, 1, 2]));
2827 assert!(poll!(first.as_mut()).is_pending());
2828 let mut zero = pin!(read_first_byte(&blob, &cache_ref, 0));
2829 assert!(poll!(zero.as_mut()).is_pending());
2830 let mut two = pin!(read_first_byte(&blob, &cache_ref, 2));
2831 assert!(poll!(two.as_mut()).is_pending());
2832 let mut second = pin!(read_first_bytes(&blob, &cache_ref, &[0, 2]));
2833 assert!(poll!(second.as_mut()).is_pending());
2834
2835 start.send(()).unwrap();
2839 assert_eq!(zero.await.unwrap(), 0);
2840 fail.send(()).unwrap();
2841 assert!(matches!(first.await, Err(Error::ReadFailed)));
2842 assert!(
2843 cache_ref
2844 .cache
2845 .read()
2846 .batches
2847 .find_unpublished(0, 2)
2848 .is_none()
2849 );
2850 assert!(matches!(two.await, Err(Error::ReadFailed)));
2851 assert!(matches!(second.await, Err(Error::ReadFailed)));
2852 assert_idle(&cache_ref);
2853 });
2854 }
2855
2856 #[rstest]
2857 #[case::out_of_order(vec![Step::page(2), Step::page(0), Step::page(1)], None)]
2858 #[case::out_of_range(vec![Step::page(0), Step::page(3)], Some(Error::ReadFailed))]
2859 #[case::duplicate(
2860 vec![Step::page(0), Step::page(0), Step::page(1), Step::page(2)],
2861 Some(Error::ReadFailed)
2862 )]
2863 #[case::missing(vec![Step::page(0), Step::page(1)], Some(Error::ReadFailed))]
2864 #[case::corrupt(
2865 vec![Step::page(0), Step::Yield(1, corrupt_page(1)), Step::page(2)],
2866 Some(Error::InvalidChecksum)
2867 )]
2868 #[case::partial(
2869 vec![Step::page(0), Step::Yield(1, page_bytes(1, 100)), Step::page(2)],
2870 Some(Error::InvalidChecksum)
2871 )]
2872 #[case::error(
2873 vec![Step::page(0), Step::Fail(Error::BlobInsufficientLength)],
2874 Some(Error::BlobInsufficientLength)
2875 )]
2876 fn test_batch_checks_read_many_output(
2877 #[case] script: Vec<Step>,
2878 #[case] expected: Option<Error>,
2879 ) {
2880 let executor = deterministic::Runner::default();
2881 executor.start(|_| async move {
2882 let (start, started) = oneshot::channel();
2883 let blob =
2884 ScriptedBlob::new(std::iter::once(Step::Wait(started)).chain(script).collect());
2885 let cache_ref = CacheRef::new(PAGE_SIZE, NZUsize!(16));
2886
2887 let mut first = pin!(read_first_bytes(&blob, &cache_ref, &[0, 1, 2]));
2889 assert!(poll!(first.as_mut()).is_pending());
2890 let mut second = pin!(read_first_bytes(&blob, &cache_ref, &[0, 1, 2]));
2891 assert!(poll!(second.as_mut()).is_pending());
2892
2893 start.send(()).unwrap();
2895 let (first, second) = futures::join!(first, second);
2896 for result in [first, second] {
2897 match &expected {
2898 None => assert_eq!(result.unwrap(), [0, 1, 2]),
2899 Some(expected) => assert_eq!(
2900 std::mem::discriminant(&result.unwrap_err()),
2901 std::mem::discriminant(expected)
2902 ),
2903 }
2904 }
2905 assert_idle(&cache_ref);
2906 });
2907 }
2908
2909 #[rstest]
2913 #[case::thirty_fetches(30, false)]
2914 #[case::thirty_one_fetches(31, false)]
2915 #[case::thirty_batches(30, true)]
2916 #[case::thirty_one_batches(31, true)]
2917 fn test_read_fails_as_soon_as_any_source_fails(#[case] count: usize, #[case] batched: bool) {
2918 let executor = deterministic::Runner::default();
2919 executor.start(|_| async move {
2920 let (blob, mut opens) = GatedBlob::new(count);
2921 *blob.corrupt.lock() = Some(count - 1);
2922 let cache_ref = CacheRef::new(PAGE_SIZE, NZUsize!(4));
2923
2924 let (blob, cache_ref) = (&blob, &cache_ref);
2926 let mut sources: Vec<_> = (0..count as u64)
2927 .map(|page| {
2928 Box::pin(async move {
2929 if batched {
2930 read_first_bytes(blob, cache_ref, &[page]).await.map(|_| ())
2931 } else {
2932 read_first_byte(blob, cache_ref, page).await.map(|_| ())
2933 }
2934 })
2935 })
2936 .collect();
2937 for source in &mut sources {
2938 assert!(poll!(source.as_mut()).is_pending());
2939 }
2940
2941 let pages: Vec<u64> = (0..count as u64).collect();
2944 let mut read = pin!(read_first_bytes(blob, cache_ref, &pages));
2945 assert!(poll!(read.as_mut()).is_pending());
2946 opens[count - 1].take().unwrap().send(()).unwrap();
2947 assert!(matches!(
2948 poll!(read.as_mut()),
2949 Poll::Ready(Err(Error::InvalidChecksum))
2950 ));
2951 });
2952 }
2953
2954 #[test_traced]
2955 fn test_each_page_routes_to_its_joiners_only() {
2956 let executor = deterministic::Runner::default();
2957 executor.start(|_| async move {
2958 const PAGES: u64 = 4;
2959 let (blob, mut opens) = GatedBlob::new(PAGES as usize);
2960 let cache_ref = CacheRef::new(PAGE_SIZE, NZUsize!(16));
2961 let pages: Vec<u64> = (0..PAGES).collect();
2962 let mut first = pin!(read_first_bytes(&blob, &cache_ref, &pages));
2963 assert!(poll!(first.as_mut()).is_pending());
2964
2965 let mut singles: Vec<_> = pages
2967 .iter()
2968 .map(|&page| Box::pin(read_first_byte(&blob, &cache_ref, page)))
2969 .collect();
2970 for single in &mut singles {
2971 assert!(poll!(single.as_mut()).is_pending());
2972 }
2973 assert_eq!(joiners(&cache_ref), [Some((1, 1)); PAGES as usize]);
2974
2975 open(&mut opens);
2976 for (page, single) in singles.into_iter().enumerate() {
2977 assert_eq!(single.await.unwrap(), page as u8);
2978 }
2979 assert_eq!(first.await.unwrap(), [0, 1, 2, 3]);
2980 assert_eq!(blob.reads.lock().len(), PAGES as usize);
2981 assert_idle(&cache_ref);
2982 });
2983 }
2984
2985 #[test_traced]
2986 fn test_many_reads_join_one_page() {
2987 let executor = deterministic::Runner::default();
2988 executor.start(|_| async move {
2989 const READS: usize = 128;
2990 let (blob, mut opens) = GatedBlob::new(1);
2991 let cache_ref = CacheRef::new(PAGE_SIZE, NZUsize!(16));
2992 let mut first = pin!(read_first_bytes(&blob, &cache_ref, &[0]));
2993 assert!(poll!(first.as_mut()).is_pending());
2994
2995 let mut reads: Vec<_> = (0..READS)
2997 .map(|_| Box::pin(read_first_byte(&blob, &cache_ref, 0)))
2998 .collect();
2999 for read in &mut reads {
3000 assert!(poll!(read.as_mut()).is_pending());
3001 }
3002 assert_eq!(joiners(&cache_ref), [Some((READS, READS))]);
3003
3004 open(&mut opens);
3005 for read in reads {
3006 assert_eq!(read.await.unwrap(), 0);
3007 }
3008 assert_eq!(first.await.unwrap(), [0]);
3009 assert_eq!(blob.reads.lock().len(), 1);
3010 assert_idle(&cache_ref);
3011 });
3012 }
3013
3014 #[test_traced]
3015 fn test_subscription_dropped_unpolled_gets_no_pages() {
3016 let executor = deterministic::Runner::default();
3017 executor.start(|_| async move {
3018 let (blob, mut opens) = GatedBlob::new(2);
3019 let cache_ref = CacheRef::new(PAGE_SIZE, NZUsize!(16));
3020
3021 let mut first = pin!(read_first_bytes(&blob, &cache_ref, &[0, 1]));
3024 assert!(poll!(first.as_mut()).is_pending());
3025 let batch = {
3026 let cache = cache_ref.cache.write();
3027 let (batch, index, driver) = cache.batches.find_unpublished(0, 0).unwrap();
3028 let mut subscription = Subscription::joining(batch.clone(), driver);
3029 subscription.select(index, 0);
3030 let reader = subscription.reader;
3031 drop(subscription);
3032 assert!(batch.state.lock().inboxes[reader].is_none());
3033 batch
3034 };
3035
3036 open(&mut opens);
3039 assert_eq!(first.await.unwrap(), [0, 1]);
3040 let state = batch.state.lock();
3041 assert!(state.inboxes[1..].iter().all(Option::is_none));
3042 drop(state);
3043 drop(batch);
3044 assert_eq!(blob.reads_of(0), 1);
3045 assert_idle(&cache_ref);
3046 });
3047 }
3048
3049 #[test_traced]
3050 fn test_page_delivered_before_batch_fails_still_serves_its_reader() {
3051 let executor = deterministic::Runner::default();
3052 executor.start(|_| async move {
3053 let (start, started) = oneshot::channel();
3054 let blob = ScriptedBlob::new(vec![
3055 Step::Wait(started),
3056 Step::page(0),
3057 Step::Fail(Error::ReadFailed),
3058 ]);
3059 let cache_ref = CacheRef::new(PAGE_SIZE, NZUsize!(16));
3060
3061 let mut first = pin!(read_first_bytes(&blob, &cache_ref, &[0, 1]));
3064 assert!(poll!(first.as_mut()).is_pending());
3065 let mut zero = pin!(read_first_byte(&blob, &cache_ref, 0));
3066 assert!(poll!(zero.as_mut()).is_pending());
3067 let mut both = pin!(read_first_bytes(&blob, &cache_ref, &[0, 1]));
3068 assert!(poll!(both.as_mut()).is_pending());
3069
3070 start.send(()).unwrap();
3073 assert!(matches!(first.await, Err(Error::ReadFailed)));
3074
3075 assert_eq!(zero.await.unwrap(), 0);
3078 assert!(matches!(both.await, Err(Error::ReadFailed)));
3079 assert_idle(&cache_ref);
3080 });
3081 }
3082
3083 #[test_traced]
3084 fn test_cancelled_joins_do_not_disturb_other_reads() {
3085 let executor = deterministic::Runner::default();
3086 executor.start(|_| async move {
3087 let (blob, mut opens) = GatedBlob::new(3);
3088 let cache_ref = CacheRef::new(PAGE_SIZE, NZUsize!(16));
3089 let mut first = pin!(read_first_bytes(&blob, &cache_ref, &[0, 1, 2]));
3090 assert!(poll!(first.as_mut()).is_pending());
3091
3092 let mut kept = pin!(read_first_bytes(&blob, &cache_ref, &[1, 2]));
3094 assert!(poll!(kept.as_mut()).is_pending());
3095 assert_eq!(selected(&cache_ref), 2);
3096
3097 {
3099 let mut joined = pin!(read_first_bytes(&blob, &cache_ref, &[0, 1, 2]));
3100 assert!(poll!(joined.as_mut()).is_pending());
3101 let mut single = pin!(read_first_byte(&blob, &cache_ref, 1));
3102 assert!(poll!(single.as_mut()).is_pending());
3103 assert_eq!(selected(&cache_ref), 2 + 3 + 1);
3104 }
3105 assert_eq!(selected(&cache_ref), 2);
3106
3107 for _ in 0..64 {
3110 let mut cancelled = pin!(read_first_byte(&blob, &cache_ref, 1));
3111 assert!(poll!(cancelled.as_mut()).is_pending());
3112 }
3113 let (selections, subscribed) = joiners(&cache_ref)[1].unwrap();
3114 assert_eq!(subscribed, 1);
3115 assert!(selections < 8, "{selections} selections remain");
3116
3117 open(&mut opens);
3118 assert_eq!(kept.await.unwrap(), [1, 2]);
3119 assert_eq!(first.await.unwrap(), [0, 1, 2]);
3120 assert_eq!(blob.reads.lock().len(), 3);
3121 assert_idle(&cache_ref);
3122 });
3123 }
3124
3125 #[test_traced]
3126 fn test_dropped_batches_are_unregistered() {
3127 let executor = deterministic::Runner::default();
3128 executor.start(|_| async move {
3129 let (blob, mut opens) = GatedBlob::new(8);
3130 let cache_ref = CacheRef::new(PAGE_SIZE, NZUsize!(16));
3131
3132 let started = (0..4)
3134 .map(|id| (id, id))
3135 .chain((4..8).map(|page| (0, page)));
3136 for (id, page) in started {
3137 let mut buf = [0u8; 1];
3138 let ranges = vec![(buf.as_mut_slice(), page * PAGE_SIZE_U64)];
3139 let mut read = pin!(cache_ref.read_after_misses(&blob, id, ranges));
3140 assert!(poll!(read.as_mut()).is_pending());
3141 }
3142 assert_idle(&cache_ref);
3143 open(&mut opens);
3144
3145 assert_eq!(read_first_byte(&blob, &cache_ref, 0).await.unwrap(), 0);
3148 assert_eq!(blob.reads_of(0), 2);
3149 assert!(cache_ref.cache.read().batches.by_blob.is_empty());
3150 });
3151 }
3152
3153 #[test_traced]
3154 fn test_batch_whose_driver_was_dropped_is_not_joined() {
3155 let executor = deterministic::Runner::default();
3156 executor.start(|_| async move {
3157 let (blob, mut opens) = GatedBlob::new(1);
3158 open(&mut opens);
3159 let cache_ref = CacheRef::new(PAGE_SIZE, NZUsize!(16));
3160
3161 let batch = Arc::new(Batch::new(0, vec![0]));
3164 let driver: BatchDriver = Shared::new(pending::<Result<(), Error>>().boxed());
3165 cache_ref.cache.write().batches.register(&batch, &driver);
3166 drop(driver);
3167
3168 assert_eq!(read_first_byte(&blob, &cache_ref, 0).await.unwrap(), 0);
3170 assert_eq!(blob.reads_of(0), 1);
3171 });
3172 }
3173
3174 #[test_traced]
3175 fn test_overlapping_reads_with_cancellations() {
3176 for seed in 0..16 {
3177 let executor = deterministic::Runner::seeded(seed);
3178 executor.start(|mut context| async move {
3179 const PAGES: u64 = MISSES_PER_LOCK as u64 + 8;
3180 let (blob, mut opens) = GatedBlob::new(PAGES as usize);
3181 let cache_ref = CacheRef::new(PAGE_SIZE, NZUsize!(4));
3182
3183 let mut readers = Vec::new();
3186 for _ in 0..8 {
3187 let mut pages: Vec<u64> = if context.random_bool(0.25) {
3188 (0..PAGES).collect()
3189 } else {
3190 (0..PAGES).filter(|_| context.random_bool(0.3)).collect()
3191 };
3192 if pages.is_empty() {
3193 pages.push(context.random_range(0..PAGES));
3194 }
3195 let single = pages.len() == 1 && context.random_bool(0.5);
3196 let expected: Vec<u8> = pages.iter().map(|&page| page as u8).collect();
3197 let (blob, cache_ref) = (blob.clone(), cache_ref.clone());
3198 let reader = context.child("reader").spawn(move |_| async move {
3199 if single {
3200 read_first_byte(&blob, &cache_ref, pages[0])
3201 .await
3202 .map(|byte| vec![byte])
3203 } else {
3204 read_first_bytes(&blob, &cache_ref, &pages).await
3205 }
3206 });
3207 readers.push((reader, expected));
3208 let delay = context.random_range(0..3);
3209 context.sleep(Duration::from_millis(delay)).await;
3210 }
3211
3212 let mut order: Vec<usize> = (0..PAGES as usize).collect();
3214 order.shuffle(&mut context);
3215 for page in order {
3216 open(&mut opens[page..=page]);
3217 if context.random_bool(0.05) {
3218 readers[context.random_range(0..readers.len())].0.abort();
3219 }
3220 context.sleep(Duration::from_millis(1)).await;
3221 }
3222
3223 for (reader, expected) in readers {
3225 if let Ok(result) = reader.await {
3226 assert_eq!(result.unwrap(), expected);
3227 }
3228 }
3229 assert_idle(&cache_ref);
3230 });
3231 }
3232 }
3233
3234 #[test_traced]
3235 fn test_concurrent_reads_read_each_page_once() {
3236 let executor = deterministic::Runner::default();
3237 executor.start(|context| async move {
3238 const PAGES: u64 = 3 * MISSES_PER_LOCK as u64;
3241 let (recording, recordings) = crate::mocks::RecordingContext::new(context);
3242 let blob = checksummed_blob(&recording, PAGES).await;
3243 let context = recording.inner;
3244 recordings.clear();
3245 let cache_ref = CacheRef::new(PAGE_SIZE, NZUsize!(2 * PAGES as usize));
3246 let readers = (0..4).map(|_| {
3247 let (blob, cache_ref) = (blob.clone(), cache_ref.clone());
3248 context.child("reader").spawn(move |_| async move {
3249 let pages: Vec<u64> = (0..PAGES).collect();
3250 read_first_bytes(&blob, &cache_ref, &pages).await.unwrap()
3251 })
3252 });
3253 let expected: Vec<u8> = (0..PAGES).map(|page| page as u8).collect();
3254 for reader in futures::future::join_all(readers).await {
3255 assert_eq!(reader.unwrap(), expected);
3256 }
3257
3258 assert_eq!(recordings.snapshot().reads.len(), PAGES as usize);
3261 });
3262 }
3263
3264 #[rstest]
3265 #[case::deterministic(deterministic::Runner::default())]
3266 #[case::tokio(crate::tokio::Runner::default())]
3267 #[cfg_attr(target_os = "linux", case::iouring(crate::iouring::Runner::default()))]
3268 fn test_read_many_into_serves_misses<R: crate::Runner>(#[case] runner: R)
3269 where
3270 R::Context: crate::Storage,
3271 {
3272 runner.start(|context| async move {
3273 let blob = checksummed_blob(&context, 40).await;
3276 let cache_ref = CacheRef::new(PAGE_SIZE, NZUsize!(4));
3277 let size = 40 * PAGE_SIZE_U64;
3278 let view = View {
3279 blob: &blob,
3280 cache_ref: &cache_ref,
3281 id: 0,
3282 size,
3283 tail_offset: size,
3284 tail: Tail::Sealed(&[]),
3285 };
3286 let offsets: Vec<u64> = (1..40).map(|page| page * PAGE_SIZE_U64 - 2).collect();
3287 let mut buf = vec![0u8; offsets.len() * 4];
3288 view.read_many_into(&mut buf, &offsets, NZUsize!(4))
3289 .await
3290 .unwrap();
3291 let expected: Vec<u8> = offsets
3292 .iter()
3293 .flat_map(|&offset| (offset..offset + 4).map(|o| (o / PAGE_SIZE_U64) as u8))
3294 .collect();
3295 assert_eq!(buf, expected);
3296 });
3297 }
3298
3299 #[rstest]
3300 #[case::deterministic(deterministic::Runner::default())]
3301 #[case::tokio(crate::tokio::Runner::default())]
3302 #[cfg_attr(target_os = "linux", case::iouring(crate::iouring::Runner::default()))]
3303 fn test_concurrent_reads_with_eviction_and_cancellation<R: crate::Runner>(#[case] runner: R)
3304 where
3305 R::Context: crate::Clock + crate::Spawner + crate::Storage,
3306 {
3307 runner.start(|context| async move {
3308 const PAGES: u64 = 3 * MISSES_PER_LOCK as u64;
3311 let blob = checksummed_blob(&context, PAGES).await;
3312 let cache_ref = CacheRef::new(PAGE_SIZE, NZUsize!(4));
3313 let readers: Vec<_> = (0..8)
3314 .map(|reader| {
3315 let pages: Vec<u64> = (reader * 8..PAGES).collect();
3316 let (blob, cache_ref) = (blob.clone(), cache_ref.clone());
3317 let handle = context.child("reader").spawn(move |_| async move {
3318 read_first_bytes(&blob, &cache_ref, &pages).await
3319 });
3320 (reader, handle)
3321 })
3322 .collect();
3323 context.sleep(Duration::from_millis(1)).await;
3324 for (_, handle) in readers.iter().step_by(3) {
3325 handle.abort();
3326 }
3327
3328 for (reader, handle) in readers {
3330 select! {
3331 result = handle => {
3332 if reader % 3 != 0 {
3333 let expected: Vec<u8> =
3334 (reader * 8..PAGES).map(|page| page as u8).collect();
3335 assert_eq!(result.unwrap().unwrap(), expected);
3336 }
3337 },
3338 _ = context.sleep(Duration::from_secs(10)) => panic!("a read hung"),
3339 }
3340 }
3341 });
3342 }
3343
3344 #[test_traced]
3345 fn test_cache_at_high_offset() {
3346 let executor = deterministic::Runner::default();
3347 executor.start(|_| async move {
3348 const MIN_PAGE_SIZE: u64 = CHECKSUM_SIZE + 1;
3350 let cache_ref = CacheRef::new(NZU16!(MIN_PAGE_SIZE as u16), NZUsize!(2));
3351
3352 let data = vec![1u8; MIN_PAGE_SIZE as usize * 2];
3355
3356 let aligned_max_offset = u64::MAX - (u64::MAX % MIN_PAGE_SIZE);
3359 let high_offset = aligned_max_offset - (MIN_PAGE_SIZE * 2);
3360 let remaining = cache_ref.cache(0, &data, high_offset);
3361 assert_eq!(remaining, 0);
3363
3364 let mut buf = vec![0u8; MIN_PAGE_SIZE as usize];
3366 let page_cache = cache_ref.cache.read();
3367 assert_eq!(
3368 page_cache.read_at(0, &mut buf, high_offset),
3369 MIN_PAGE_SIZE as usize
3370 );
3371 assert!(buf.iter().all(|b| *b == 1));
3372
3373 assert_eq!(
3375 page_cache.read_at(0, &mut buf, high_offset + MIN_PAGE_SIZE),
3376 MIN_PAGE_SIZE as usize
3377 );
3378 assert!(buf.iter().all(|b| *b == 1));
3379 });
3380 }
3381
3382 #[test_traced]
3383 fn test_page_fetches_entry_removed_when_first_fetcher_cancelled() {
3384 let executor = deterministic::Runner::default();
3385 executor.start(|context| async move {
3386 let blob_id = 0;
3388 let cache_ref = CacheRef::new(PAGE_SIZE, NZUsize!(10));
3389 let (started_tx, started_rx) = oneshot::channel();
3390 let blob = Arc::new(BlockingBlob {
3391 started: Mutex::new(Some(started_tx)),
3392 });
3393 let mut read_buf = vec![0u8; PAGE_SIZE.get() as usize];
3394
3395 let cache_ref_for_task = cache_ref.clone();
3397 let blob_for_task = blob.clone();
3398 let handle = context.spawn(move |_| async move {
3399 let _ = cache_ref_for_task
3400 .read_after_miss(&blob_for_task, blob_id, &mut read_buf, 0)
3401 .await;
3402 });
3403
3404 started_rx.await.expect("blocking read never started");
3406 {
3407 let page_cache = cache_ref.cache.read();
3408 assert!(page_cache.page_fetches.contains_key(&(blob_id, 0)));
3409 }
3410
3411 handle.abort();
3413 assert!(matches!(handle.await, Err(Error::Closed)));
3414
3415 let page_cache = cache_ref.cache.read();
3417 assert!(
3418 !page_cache.page_fetches.contains_key(&(blob_id, 0)),
3419 "cancelled first fetcher should not leave stale page_fetches entry"
3420 );
3421 });
3422 }
3423
3424 #[test_traced]
3425 fn test_followers_keep_single_flight_after_first_fetcher_cancellation() {
3426 let executor = deterministic::Runner::default();
3427 executor.start(|context| async move {
3428 let blob_id = 0;
3429 let cache_ref = CacheRef::new(PAGE_SIZE, NZUsize!(10));
3430
3431 let logical_page = vec![7u8; PAGE_SIZE.get() as usize];
3433 let crc = Crc32::checksum(&logical_page);
3434 let mut physical_page = logical_page.clone();
3435 physical_page.extend_from_slice(&Checksum::new(PAGE_SIZE.get(), crc).to_bytes());
3436 let (started_tx, started_rx) = oneshot::channel();
3437 let (release_tx, release_rx) = oneshot::channel();
3438 let reads = Arc::new(AtomicUsize::new(0));
3439 let blob = Arc::new(ControlledBlob {
3440 started: Mutex::new(Some(started_tx)),
3441 release: Mutex::new(Some(release_rx)),
3442 reads: reads.clone(),
3443 result: ControlledBlobResult::Success(physical_page),
3444 });
3445
3446 let mut first_buf = vec![0u8; PAGE_SIZE.get() as usize];
3448 let cache_ref_for_first = cache_ref.clone();
3449 let blob_for_first = blob.clone();
3450 let first = context.child("first").spawn(move |_| async move {
3451 let _ = cache_ref_for_first
3452 .read_after_miss(&blob_for_first, blob_id, &mut first_buf, 0)
3453 .await;
3454 });
3455 started_rx.await.expect("first read never started");
3456
3457 let mut second_buf = vec![0u8; PAGE_SIZE.get() as usize];
3459 let cache_ref_for_second = cache_ref.clone();
3460 let blob_for_second = blob.clone();
3461 let second = context.child("second").spawn(move |_| async move {
3462 cache_ref_for_second
3463 .read_after_miss(&blob_for_second, blob_id, &mut second_buf, 0)
3464 .await
3465 .expect("second read failed");
3466 second_buf
3467 });
3468
3469 loop {
3471 let joined = {
3472 let page_cache = cache_ref.cache.read();
3473 page_cache
3474 .page_fetches
3475 .get(&(blob_id, 0))
3476 .map(|fetch| fetch.waiters == 2)
3477 .unwrap_or(false)
3478 };
3479 if joined {
3480 break;
3481 }
3482 context.sleep(Duration::from_millis(1)).await;
3483 }
3484
3485 first.abort();
3487 assert!(matches!(first.await, Err(Error::Closed)));
3488
3489 let mut third_buf = vec![0u8; PAGE_SIZE.get() as usize];
3492 let cache_ref_for_third = cache_ref.clone();
3493 let blob_for_third = blob.clone();
3494 let third = context.child("third").spawn(move |_| async move {
3495 cache_ref_for_third
3496 .read_after_miss(&blob_for_third, blob_id, &mut third_buf, 0)
3497 .await
3498 .expect("third read failed");
3499 third_buf
3500 });
3501
3502 loop {
3505 let third_entered = {
3506 let page_cache = cache_ref.cache.read();
3507 reads.load(Ordering::Relaxed) > 1
3508 || page_cache
3509 .page_fetches
3510 .get(&(blob_id, 0))
3511 .map(|fetch| fetch.waiters == 2)
3512 .unwrap_or(false)
3513 };
3514 if third_entered {
3515 break;
3516 }
3517 context.sleep(Duration::from_millis(1)).await;
3518 }
3519
3520 let _ = release_tx.send(());
3522 let second_buf = second.await.expect("second task failed");
3523 let third_buf = third.await.expect("third task failed");
3524 assert_eq!(second_buf, logical_page);
3525 assert_eq!(third_buf, logical_page);
3526
3527 assert_eq!(reads.load(Ordering::Relaxed), 1);
3529
3530 let mut cached = vec![0u8; PAGE_SIZE.get() as usize];
3532 assert_eq!(
3533 cache_ref.read_cached(blob_id, &mut cached, 0),
3534 PAGE_SIZE.get() as usize
3535 );
3536 assert_eq!(cached, logical_page);
3537
3538 let mut fourth_buf = vec![0u8; PAGE_SIZE.get() as usize];
3540 cache_ref
3541 .read_after_miss(&blob, blob_id, &mut fourth_buf, 0)
3542 .await
3543 .unwrap();
3544 assert_eq!(fourth_buf, logical_page);
3545 assert_eq!(reads.load(Ordering::Relaxed), 1);
3546
3547 let page_cache = cache_ref.cache.read();
3548 assert!(
3549 !page_cache.page_fetches.contains_key(&(blob_id, 0)),
3550 "completed fetch should leave no stale page_fetches entry"
3551 );
3552 });
3553 }
3554
3555 #[test_traced]
3556 fn test_page_fetch_error_removes_entry_for_all_waiters() {
3557 let executor = deterministic::Runner::default();
3558 executor.start(|context| async move {
3559 let blob_id = 0;
3560 let cache_ref = CacheRef::new(PAGE_SIZE, NZUsize!(10));
3561
3562 let (started_tx, started_rx) = oneshot::channel();
3564 let (release_tx, release_rx) = oneshot::channel();
3565 let reads = Arc::new(AtomicUsize::new(0));
3566 let blob = Arc::new(ControlledBlob {
3567 started: Mutex::new(Some(started_tx)),
3568 release: Mutex::new(Some(release_rx)),
3569 reads: reads.clone(),
3570 result: ControlledBlobResult::Error,
3571 });
3572
3573 let mut first_buf = vec![0u8; PAGE_SIZE.get() as usize];
3575 let cache_ref_for_first = cache_ref.clone();
3576 let blob_for_first = blob.clone();
3577 let first = context.child("first").spawn(move |_| async move {
3578 cache_ref_for_first
3579 .read_after_miss(&blob_for_first, blob_id, &mut first_buf, 0)
3580 .await
3581 });
3582 started_rx.await.expect("first erroring read never started");
3583
3584 let mut second_buf = vec![0u8; PAGE_SIZE.get() as usize];
3586 let cache_ref_for_second = cache_ref.clone();
3587 let blob_for_second = blob.clone();
3588 let second = context.child("second").spawn(move |_| async move {
3589 cache_ref_for_second
3590 .read_after_miss(&blob_for_second, blob_id, &mut second_buf, 0)
3591 .await
3592 });
3593
3594 loop {
3596 let joined = {
3597 let page_cache = cache_ref.cache.read();
3598 page_cache
3599 .page_fetches
3600 .get(&(blob_id, 0))
3601 .map(|fetch| fetch.waiters == 2)
3602 .unwrap_or(false)
3603 };
3604 if joined {
3605 break;
3606 }
3607 context.sleep(Duration::from_millis(1)).await;
3608 }
3609
3610 let _ = release_tx.send(());
3612
3613 assert!(matches!(first.await, Ok(Err(Error::ReadFailed))));
3614 assert!(matches!(second.await, Ok(Err(Error::ReadFailed))));
3615 assert_eq!(reads.load(Ordering::Relaxed), 1);
3617
3618 {
3620 let page_cache = cache_ref.cache.read();
3621 assert!(
3622 !page_cache.page_fetches.contains_key(&(blob_id, 0)),
3623 "erroring fetch should leave no stale page_fetches entry"
3624 );
3625 }
3626 let mut cached = vec![0u8; PAGE_SIZE.get() as usize];
3627 assert_eq!(cache_ref.read_cached(blob_id, &mut cached, 0), 0);
3628
3629 let mut third_buf = vec![0u8; PAGE_SIZE.get() as usize];
3631 assert!(matches!(
3632 cache_ref
3633 .read_after_miss(&blob, blob_id, &mut third_buf, 0)
3634 .await,
3635 Err(Error::ReadFailed)
3636 ));
3637 assert_eq!(reads.load(Ordering::Relaxed), 2);
3638 });
3639 }
3640
3641 #[test_traced]
3642 fn test_cancelled_stale_waiter_preserves_new_fetch() {
3643 let executor = deterministic::Runner::default();
3644 executor.start(|_| async move {
3645 let blob_id = 0;
3646 let cache_ref = CacheRef::new(PAGE_SIZE, NZUsize!(10));
3647 let (started_tx, started_rx) = oneshot::channel();
3648 let (release_tx, release_rx) = oneshot::channel();
3649 let reads = Arc::new(AtomicUsize::new(0));
3650 let blob = Arc::new(ControlledBlob {
3651 started: Mutex::new(Some(started_tx)),
3652 release: Mutex::new(Some(release_rx)),
3653 reads: reads.clone(),
3654 result: ControlledBlobResult::Error,
3655 });
3656
3657 let mut first_buf = vec![0; PAGE_SIZE.get() as usize];
3659 let mut first = Box::pin(cache_ref.read_after_miss(&blob, blob_id, &mut first_buf, 0));
3660 assert!(poll!(&mut first).is_pending());
3661 started_rx.await.unwrap();
3662 let mut stale_buf = vec![0; PAGE_SIZE.get() as usize];
3663 let mut stale = Box::pin(cache_ref.read_after_miss(&blob, blob_id, &mut stale_buf, 0));
3664 assert!(poll!(&mut stale).is_pending());
3665 assert_eq!(reads.load(Ordering::Relaxed), 1);
3666 release_tx.send(()).unwrap();
3667 assert!(matches!(first.await, Err(Error::ReadFailed)));
3668
3669 let (started_tx, started_rx) = oneshot::channel();
3671 let (release_tx, release_rx) = oneshot::channel();
3672 *blob.started.lock() = Some(started_tx);
3673 *blob.release.lock() = Some(release_rx);
3674 let mut retry_buf = vec![0; PAGE_SIZE.get() as usize];
3675 let mut retry = Box::pin(cache_ref.read_after_miss(&blob, blob_id, &mut retry_buf, 0));
3676 assert!(poll!(&mut retry).is_pending());
3677 started_rx.await.unwrap();
3678 assert_eq!(reads.load(Ordering::Relaxed), 2);
3679
3680 drop(stale);
3682 let mut joined_buf = vec![0; PAGE_SIZE.get() as usize];
3683 let mut joined =
3684 Box::pin(cache_ref.read_after_miss(&blob, blob_id, &mut joined_buf, 0));
3685 assert!(poll!(&mut joined).is_pending());
3686 assert_eq!(reads.load(Ordering::Relaxed), 2);
3687 release_tx.send(()).unwrap();
3688 assert!(matches!(retry.await, Err(Error::ReadFailed)));
3689 assert!(matches!(joined.await, Err(Error::ReadFailed)));
3690 assert_eq!(reads.load(Ordering::Relaxed), 2);
3691 });
3692 }
3693
3694 #[test_traced]
3695 fn test_read_cached_many_all_cached() {
3696 let cache_ref = CacheRef::new(PAGE_SIZE, NZUsize!(10));
3697 let blob_id = cache_ref.next_id();
3698 let page0 = vec![0xAA; PAGE_SIZE.get() as usize];
3699 let page1 = vec![0xBB; PAGE_SIZE.get() as usize];
3700
3701 {
3703 let mut cache = cache_ref.cache.write();
3704 cache.cache(blob_id, &page0, 0);
3705 cache.cache(blob_id, &page1, 1);
3706 }
3707
3708 let mut buf0 = vec![0u8; PAGE_SIZE_U64 as usize];
3709 let mut buf1 = vec![0u8; PAGE_SIZE_U64 as usize];
3710 let mut ranges: Vec<(&mut [u8], u64)> = vec![(&mut buf0, 0), (&mut buf1, PAGE_SIZE_U64)];
3711
3712 cache_ref.read_cached_many(blob_id, &mut ranges);
3713
3714 assert!(ranges.is_empty());
3716 drop(ranges);
3717
3718 assert!(buf0 == page0);
3720 assert!(buf1 == page1);
3721 }
3722
3723 #[test_traced]
3724 fn test_read_cached_many_resumes_after_evicted_prefix() {
3725 let executor = deterministic::Runner::default();
3726 executor.start(|_| async move {
3727 let cache_ref = CacheRef::new(PAGE_SIZE, NZUsize!(3));
3728 let blob_id = cache_ref.next_id();
3729 let page_size = PAGE_SIZE.get() as usize;
3730 cache_ref.cache(blob_id, &vec![0x11; page_size], 0);
3731 cache_ref.cache(blob_id, &vec![0x22; page_size], PAGE_SIZE_U64);
3732
3733 let mut buf = vec![0; page_size + 8];
3734 let mut ranges = vec![(buf.as_mut_slice(), PAGE_SIZE_U64 - 3)];
3735 cache_ref.read_cached_many(blob_id, &mut ranges);
3736 assert_eq!(ranges.len(), 1);
3737 assert_eq!(ranges[0].1, PAGE_SIZE_U64 * 2);
3738 assert_eq!(ranges[0].0.len(), 5);
3739
3740 cache_ref.clear();
3742 let mut physical_page = vec![0x33; page_size];
3743 let crc = Crc32::checksum(&physical_page);
3744 physical_page.extend_from_slice(&Checksum::new(PAGE_SIZE.get(), crc).to_bytes());
3745 let reads = Arc::new(AtomicUsize::new(0));
3746 let blob = Arc::new(ControlledBlob {
3747 started: Mutex::new(None),
3748 release: Mutex::new(None),
3749 reads: reads.clone(),
3750 result: ControlledBlobResult::Success(physical_page),
3751 });
3752 for (remaining, offset) in ranges {
3753 cache_ref
3754 .read_after_miss(&blob, blob_id, remaining, offset)
3755 .await
3756 .unwrap();
3757 }
3758 assert_eq!(reads.load(Ordering::Relaxed), 1);
3759 assert_eq!(&buf[..3], &[0x11; 3]);
3760 assert!(buf[3..3 + page_size].iter().all(|&byte| byte == 0x22));
3761 assert_eq!(&buf[3 + page_size..], &[0x33; 5]);
3762 });
3763 }
3764
3765 #[test_traced]
3766 fn test_observed_miss_filled_before_completion() {
3767 let executor = deterministic::Runner::default();
3768 executor.start(|_| async move {
3769 let cache_ref = CacheRef::new(PAGE_SIZE, NZUsize!(1));
3770 let blob_id = cache_ref.next_id();
3771 let mut buf = [0; 8];
3772 assert_eq!(cache_ref.read_cached(blob_id, &mut buf, 3), 0);
3773 let reads = Arc::new(AtomicUsize::new(0));
3774 let blob = Arc::new(ControlledBlob {
3775 started: Mutex::new(None),
3776 release: Mutex::new(None),
3777 reads: reads.clone(),
3778 result: ControlledBlobResult::Error,
3779 });
3780
3781 let completion = cache_ref.read_after_miss(&blob, blob_id, &mut buf, 3);
3783 cache_ref.cache(blob_id, &vec![0x44; PAGE_SIZE.get() as usize], 0);
3784 completion.await.unwrap();
3785 assert_eq!(reads.load(Ordering::Relaxed), 0);
3786 assert_eq!(buf, [0x44; 8]);
3787 });
3788 }
3789
3790 #[test_traced]
3791 fn test_read_cached_many_none_cached() {
3792 let cache_ref = CacheRef::new(PAGE_SIZE, NZUsize!(10));
3793 let blob_id = cache_ref.next_id();
3794
3795 let mut buf0 = vec![0u8; PAGE_SIZE_U64 as usize];
3796 let mut buf1 = vec![0u8; PAGE_SIZE_U64 as usize];
3797 let mut ranges: Vec<(&mut [u8], u64)> = vec![(&mut buf0, 0), (&mut buf1, PAGE_SIZE_U64)];
3798
3799 cache_ref.read_cached_many(blob_id, &mut ranges);
3801 assert_eq!(ranges.len(), 2);
3802 assert_eq!(ranges[0].1, 0);
3803 assert_eq!(ranges[1].1, PAGE_SIZE_U64);
3804 }
3805
3806 #[test_traced]
3807 fn test_read_cached_many_scattered_misses() {
3808 let cache_ref = CacheRef::new(PAGE_SIZE, NZUsize!(10));
3811 let blob_id = cache_ref.next_id();
3812
3813 let page0 = vec![0x11; PAGE_SIZE.get() as usize];
3814 let page2 = vec![0x33; PAGE_SIZE.get() as usize];
3815 {
3816 let mut cache = cache_ref.cache.write();
3817 cache.cache(blob_id, &page0, 0);
3818 cache.cache(blob_id, &page2, 2);
3820 }
3821
3822 let mut buf0 = vec![0u8; PAGE_SIZE_U64 as usize];
3823 let mut buf1 = vec![0u8; PAGE_SIZE_U64 as usize];
3824 let mut buf2 = vec![0u8; PAGE_SIZE_U64 as usize];
3825 let mut ranges: Vec<(&mut [u8], u64)> = vec![
3826 (&mut buf0, 0),
3827 (&mut buf1, PAGE_SIZE_U64),
3828 (&mut buf2, PAGE_SIZE_U64 * 2),
3829 ];
3830
3831 cache_ref.read_cached_many(blob_id, &mut ranges);
3832
3833 assert_eq!(ranges.len(), 1);
3836 assert_eq!(ranges[0].1, PAGE_SIZE_U64);
3837 drop(ranges);
3838
3839 assert!(buf0 == page0);
3841 assert!(buf2 == page2);
3842 assert!(buf1.iter().all(|b| *b == 0));
3844 }
3845
3846 #[test_traced]
3847 fn test_read_cached_many_stale_hint_after_eviction() {
3848 let cache_ref = CacheRef::new(PAGE_SIZE, NZUsize!(2));
3852 let blob_id = cache_ref.next_id();
3853 let page_size = PAGE_SIZE.get() as usize;
3854 {
3855 let mut cache = cache_ref.cache.write();
3856 for page in 0u64..3 {
3857 cache.cache(blob_id, &vec![page as u8 + 1; page_size], page);
3858 }
3859 }
3860
3861 let mut bufs: Vec<Vec<u8>> = (0..3).map(|_| vec![0u8; page_size]).collect();
3862 let mut iter = bufs.iter_mut();
3863 let mut ranges: Vec<(&mut [u8], u64)> = (0..3u64)
3864 .map(|page| (iter.next().unwrap().as_mut_slice(), page * PAGE_SIZE_U64))
3865 .collect();
3866 cache_ref.read_cached_many(blob_id, &mut ranges);
3867
3868 assert_eq!(ranges.len(), 1);
3870 assert_eq!(ranges[0].1, 0);
3871 drop(ranges);
3872 assert!(bufs[0].iter().all(|b| *b == 0));
3873 assert_eq!(bufs[1], vec![2u8; page_size]);
3874 assert_eq!(bufs[2], vec![3u8; page_size]);
3875 }
3876
3877 #[test_traced]
3878 fn test_read_cached_many_cross_blob_hint_collision() {
3879 let cache_ref = CacheRef::new(PAGE_SIZE, NZUsize!(4));
3884 let blob_a = cache_ref.next_id();
3885 let blob_b = cache_ref.next_id();
3886 let page_size = PAGE_SIZE.get() as usize;
3887 let page_a = 5u64;
3888 let page_b = {
3889 let mut cache = cache_ref.cache.write();
3890
3891 let mask = cache.hints.len() as u64 - 1;
3893 let page_b = page_a
3894 .wrapping_add(blob_a.wrapping_mul(commonware_utils::GOLDEN_RATIO))
3895 .wrapping_sub(blob_b.wrapping_mul(commonware_utils::GOLDEN_RATIO))
3896 & mask;
3897 assert_eq!(
3898 cache.hint_index(blob_a, page_a),
3899 cache.hint_index(blob_b, page_b)
3900 );
3901 cache.cache(blob_a, &vec![0xAA; page_size], page_a);
3902 cache.cache(blob_b, &vec![0xBB; page_size], page_b);
3903 page_b
3904 };
3905
3906 for (blob, page, byte) in [(blob_a, page_a, 0xAAu8), (blob_b, page_b, 0xBB)] {
3907 let mut buf = vec![0u8; page_size];
3908 let mut ranges: Vec<(&mut [u8], u64)> = vec![(&mut buf, page * PAGE_SIZE_U64)];
3909 cache_ref.read_cached_many(blob, &mut ranges);
3910 assert!(
3911 ranges.is_empty(),
3912 "blob {blob} page {page} should be cached"
3913 );
3914 drop(ranges);
3915 assert_eq!(buf, vec![byte; page_size]);
3916 }
3917 }
3918
3919 #[test_traced]
3920 fn test_read_cached_many_sparse_page_number_keeps_hints_fixed() {
3921 let cache_ref = CacheRef::new(PAGE_SIZE, NZUsize!(2));
3924 let blob_id = cache_ref.next_id();
3925 let page_size = PAGE_SIZE.get() as usize;
3926 let page_num = u64::MAX / PAGE_SIZE_U64 - 1;
3927 {
3928 let mut cache = cache_ref.cache.write();
3929 let hints = cache.hints.len();
3930 cache.cache(blob_id, &vec![0x5A; page_size], page_num);
3931 assert_eq!(cache.hints.len(), hints);
3932 }
3933
3934 let mut buf = vec![0u8; page_size];
3935 let mut ranges: Vec<(&mut [u8], u64)> = vec![(&mut buf, page_num * PAGE_SIZE_U64)];
3936 cache_ref.read_cached_many(blob_id, &mut ranges);
3937 assert!(ranges.is_empty());
3938 drop(ranges);
3939 assert_eq!(buf, vec![0x5A; page_size]);
3940 }
3941
3942 #[test_traced]
3943 fn test_read_cached_many_evicted_page_is_a_miss() {
3944 let cache_ref = CacheRef::new(PAGE_SIZE, NZUsize!(1));
3945 let blob_id = cache_ref.next_id();
3946 let page_size = PAGE_SIZE.get() as usize;
3947 let read_page = |page: u64| {
3948 let mut buf = vec![0u8; page_size];
3949 let mut ranges: Vec<(&mut [u8], u64)> = vec![(&mut buf, page * PAGE_SIZE_U64)];
3950 cache_ref.read_cached_many(blob_id, &mut ranges);
3951 let hit = ranges.is_empty();
3952 drop(ranges);
3953 hit.then_some(buf)
3954 };
3955
3956 cache_ref.cache(blob_id, &vec![1; page_size], 0);
3957 assert_eq!(read_page(0), Some(vec![1; page_size]));
3958
3959 cache_ref.cache(blob_id, &vec![2; page_size], PAGE_SIZE_U64);
3961 assert_eq!(read_page(0), None);
3962 assert_eq!(read_page(1), Some(vec![2; page_size]));
3963
3964 cache_ref.cache(blob_id, &vec![0xCC; page_size], 0);
3966 assert_eq!(read_page(0), Some(vec![0xCC; page_size]));
3967 assert_eq!(read_page(1), None);
3968 }
3969
3970 #[rstest]
3971 #[case::empty_read(vec![], 0, 0, 0)]
3972 #[case::single_cached_page(vec![0], 3, 5, 5)]
3973 #[case::cached_range_can_cross_pages(vec![0, 1], PAGE_SIZE_U64 - 2, 4, 4)]
3974 #[case::missing_first_page_reads_nothing(vec![1], 0, 4, 0)]
3975 #[case::missing_later_page_truncates_read(vec![0], PAGE_SIZE_U64 - 2, 4, 2)]
3976 fn test_read_cached(
3977 #[case] cached_pages: Vec<u64>,
3978 #[case] logical_offset: u64,
3979 #[case] len: usize,
3980 #[case] expected_count: usize,
3981 ) {
3982 let cache_ref = CacheRef::new(PAGE_SIZE, NZUsize!(10));
3983 let blob_id = cache_ref.next_id();
3984 let sentinel = 0xEE;
3985 let page_size = PAGE_SIZE.get() as usize;
3986
3987 {
3988 let mut cache = cache_ref.cache.write();
3989 for page in cached_pages {
3990 cache.cache(blob_id, &vec![page as u8 + 1; page_size], page);
3992 }
3993 }
3994
3995 let mut buf = vec![sentinel; len];
3996 let count = cache_ref.read_cached(blob_id, &mut buf, logical_offset);
3997 assert_eq!(count, expected_count);
3998
3999 assert_eq!(buf[..count], expected_cached_bytes(logical_offset, count));
4001 assert!(buf[count..].iter().all(|b| *b == sentinel));
4002 }
4003}