1use super::Checksum;
2use crate::{Blob, Error, Handle, IoBuf, IoBufsMut, ReadOptions};
3use bytes::{BufMut, Bytes, BytesMut, TryGetError};
4use commonware_codec::{Buf, FixedSize};
5use commonware_utils::Widen;
6use futures::FutureExt as _;
7use std::{collections::VecDeque, num::NonZeroU16, ops::Range, sync::Arc};
8use tracing::error;
9
10pub(super) struct BufferState {
15 buffer: Bytes,
17 num_pages: usize,
19 last_page_len: usize,
21}
22
23const MAX_READ_AHEAD: usize = 3;
29
30struct Inflight {
32 first_page: u64,
34 pages: usize,
36 read: Handle<Vec<IoBufsMut>>,
38}
39
40impl Inflight {
41 fn end(&self) -> u64 {
43 self.first_page + Widen::widen(self.pages)
44 }
45
46 async fn wait(self) -> Result<IoBufsMut, Error> {
48 self.read.await?.pop().ok_or(Error::ReadFailed)
49 }
50}
51
52pub(super) struct PageReader<B: Blob> {
66 blob: Arc<B>,
68 physical_page_size: usize,
70 page_size: usize,
72 physical_blob_size: u64,
74 logical_blob_size: u64,
76 partial_page: Option<Bytes>,
78 blob_page: u64,
80 batch_pages: usize,
82 inflight: VecDeque<Inflight>,
85 read_options: ReadOptions,
87}
88
89impl<B: Blob> PageReader<B> {
90 pub(super) fn new(
102 blob: Arc<B>,
103 mut physical_blob_size: u64,
104 logical_blob_size: u64,
105 partial_page: Option<IoBuf>,
106 buffer_pages: usize,
107 page_size: NonZeroU16,
108 read_options: ReadOptions,
109 ) -> Self {
110 let page_size = page_size.get() as usize;
111 let physical_page_size = page_size + Checksum::SIZE;
112 let physical_pages = physical_blob_size / physical_page_size as u64;
113 let logical_pages = if logical_blob_size == 0 {
114 0
115 } else {
116 ((logical_blob_size - 1) / page_size as u64) + 1
117 };
118 assert_eq!(physical_blob_size % physical_page_size as u64, 0);
119 assert_eq!(physical_pages, logical_pages);
120 if let Some(partial_page) = &partial_page {
121 assert!(
122 !partial_page.is_empty()
123 && Widen::widen(partial_page.len())
124 == logical_blob_size % Widen::widen(page_size),
125 "frozen tail must match the final partial page"
126 );
127 physical_blob_size -= Widen::widen(physical_page_size);
128 }
129
130 Self {
131 blob,
132 physical_page_size,
133 page_size,
134 physical_blob_size,
135 logical_blob_size,
136 partial_page: partial_page.map(Bytes::from),
137 blob_page: 0,
138 batch_pages: buffer_pages.div_ceil(MAX_READ_AHEAD + 1),
139 inflight: VecDeque::new(),
140 read_options,
141 }
142 }
143
144 fn seek_to_page(&mut self, page: u64) {
146 self.blob_page = page;
147 self.inflight.clear();
148 }
149
150 async fn top_up(&mut self, consumed_end: u64) -> Result<(), Error> {
153 let mut next = self.inflight.back().map_or(consumed_end, Inflight::end);
154 while self.inflight.len() < MAX_READ_AHEAD {
155 let Some(batch) = self.start_batch(next).await? else {
156 break;
157 };
158 next = batch.end();
159 self.inflight.push_back(batch);
160 }
161 Ok(())
162 }
163
164 async fn start_batch(&self, first_page: u64) -> Result<Option<Inflight>, Error> {
167 let physical_page_size: u64 = Widen::widen(self.physical_page_size);
168 let start = first_page
169 .checked_mul(physical_page_size)
170 .ok_or(Error::OffsetOverflow)?;
171 if start >= self.physical_blob_size {
172 return Ok(None);
173 }
174 let pages = ((self.physical_blob_size - start) / physical_page_size)
175 .min(Widen::widen(self.batch_pages)) as usize;
176 let len = pages * self.physical_page_size;
177
178 let read = self
182 .blob
183 .start_read_many(&[(start, len)], self.read_options)
184 .boxed()
185 .await;
186 Ok(Some(Inflight {
187 first_page,
188 pages,
189 read,
190 }))
191 }
192
193 pub(super) const fn blob_size(&self) -> u64 {
195 self.logical_blob_size
196 }
197
198 pub(super) const fn physical_page_size(&self) -> usize {
200 self.physical_page_size
201 }
202
203 pub(super) const fn page_size(&self) -> usize {
205 self.page_size
206 }
207
208 async fn read_batch(&mut self) -> Result<(usize, IoBufsMut), Error> {
216 if let Some(batch) = self
217 .inflight
218 .pop_front_if(|batch| batch.first_page == self.blob_page)
219 {
220 let pages = batch.pages;
221 self.top_up(batch.end()).await?;
222 if let Ok(bufs) = batch.wait().await {
223 return Ok((pages, bufs));
224 }
225 }
226
227 self.inflight.clear();
228 let batch = self
229 .start_batch(self.blob_page)
230 .await?
231 .expect("pages remain to be issued");
232 let pages = batch.pages;
233 self.top_up(batch.end()).await?;
234 Ok((pages, batch.wait().await?))
235 }
236
237 pub(super) async fn fill(&mut self) -> Result<Option<(BufferState, usize)>, Error> {
242 let start_offset = match self.blob_page.checked_mul(self.physical_page_size as u64) {
244 Some(o) => o,
245 None => return Err(Error::OffsetOverflow),
246 };
247 if start_offset == self.physical_blob_size
248 && let Some(partial_page) = &self.partial_page
249 {
250 self.blob_page += 1;
251 let len = partial_page.len();
252 return Ok(Some((
253 BufferState {
254 buffer: partial_page.clone(),
255 num_pages: 1,
256 last_page_len: len,
257 },
258 len,
259 )));
260 }
261 if start_offset >= self.physical_blob_size {
262 return Ok(None); }
264
265 let (pages_to_read, physical_buf) = self.read_batch().await?;
266 let physical_buf = Bytes::from(physical_buf.coalesce().freeze());
267
268 let mut total_logical = 0usize;
270 let mut last_len = 0usize;
271 for page_idx in 0..pages_to_read {
272 let page_start = page_idx * self.physical_page_size;
273 let page_slice =
274 &physical_buf.as_ref()[page_start..page_start + self.physical_page_size];
275 let Some(checksum) = Checksum::validate_page(page_slice) else {
276 error!(page = self.blob_page + page_idx as u64, "CRC mismatch");
277 return Err(Error::InvalidChecksum);
278 };
279 let len = checksum.len as usize;
280
281 if len < self.page_size {
283 error!(
284 page = self.blob_page + Widen::widen(page_idx),
285 expected = self.page_size,
286 actual = len,
287 "page validates fewer bytes than the replay reads"
288 );
289 return Err(Error::InvalidChecksum);
290 }
291 total_logical += self.page_size;
292 last_len = self.page_size;
293 }
294 self.blob_page += Widen::widen(pages_to_read);
295
296 let state = BufferState {
297 buffer: physical_buf,
298 num_pages: pages_to_read,
299 last_page_len: last_len,
300 };
301
302 Ok(Some((state, total_logical)))
303 }
304}
305
306struct ReplayBuf {
312 physical_page_size: usize,
314 page_size: usize,
316 buffers: VecDeque<BufferState>,
318 current_page: usize,
320 offset_in_page: usize,
322 remaining: usize,
324}
325
326impl ReplayBuf {
327 const fn new(physical_page_size: usize, page_size: usize) -> Self {
329 Self {
330 physical_page_size,
331 page_size,
332 buffers: VecDeque::new(),
333 current_page: 0,
334 offset_in_page: 0,
335 remaining: 0,
336 }
337 }
338
339 fn clear(&mut self) {
341 self.buffers.clear();
342 self.current_page = 0;
343 self.offset_in_page = 0;
344 self.remaining = 0;
345 }
346
347 fn push(&mut self, state: BufferState, logical_bytes: usize) {
349 let skip = if self.buffers.is_empty() {
352 self.offset_in_page
353 } else {
354 0
355 };
356 self.buffers.push_back(state);
357 self.remaining += logical_bytes.saturating_sub(skip);
358 }
359
360 const fn page_len(buf: &BufferState, page_idx: usize, page_size: usize) -> usize {
362 if page_idx + 1 == buf.num_pages {
363 buf.last_page_len
364 } else {
365 page_size
366 }
367 }
368
369 fn chunk_range(&self) -> Option<(&BufferState, Range<usize>)> {
372 let buf = self.buffers.front()?;
373 if self.current_page >= buf.num_pages {
374 return None;
375 }
376 let page_len = Self::page_len(buf, self.current_page, self.page_size);
377 let page_start = self.current_page * self.physical_page_size;
378 Some((buf, page_start + self.offset_in_page..page_start + page_len))
379 }
380}
381
382impl Buf for ReplayBuf {}
383
384impl bytes::Buf for ReplayBuf {
385 fn copy_to_bytes(&mut self, len: usize) -> Bytes {
386 assert!(len <= self.remaining, "copy_to_bytes out of bounds");
387 if len == 0 {
388 return Bytes::new();
389 }
390 if let Some((buf, range)) = self.chunk_range()
391 && len <= range.len()
392 {
393 let bytes = buf.buffer.slice(range.start..range.start + len);
394 self.advance(len);
395 return bytes;
396 }
397
398 let mut bytes = BytesMut::with_capacity(len);
400 bytes.put(self.take(len));
401 bytes.freeze()
402 }
403
404 fn remaining(&self) -> usize {
405 self.remaining
406 }
407
408 #[inline(always)]
409 fn try_copy_to_slice(&mut self, mut dst: &mut [u8]) -> Result<(), TryGetError> {
410 if dst.len() > self.remaining {
411 return Err(TryGetError {
412 requested: dst.len(),
413 available: self.remaining,
414 });
415 }
416
417 let chunk = self.chunk();
420 if dst.len() < chunk.len() {
421 dst.copy_from_slice(&chunk[..dst.len()]);
422 self.offset_in_page += dst.len();
423 self.remaining -= dst.len();
424 return Ok(());
425 }
426
427 while !dst.is_empty() {
428 let src = self.chunk();
429 let cnt = usize::min(src.len(), dst.len());
430 dst[..cnt].copy_from_slice(&src[..cnt]);
431 dst = &mut dst[cnt..];
432 self.advance(cnt);
433 }
434 Ok(())
435 }
436
437 fn chunk(&self) -> &[u8] {
438 self.chunk_range()
439 .map_or(&[], |(buf, range)| &buf.buffer.as_ref()[range])
440 }
441
442 fn advance(&mut self, mut cnt: usize) {
443 self.remaining = self.remaining.saturating_sub(cnt);
444
445 while cnt > 0 {
446 let Some(buf) = self.buffers.front() else {
447 break;
448 };
449
450 while cnt > 0 && self.current_page < buf.num_pages {
452 let page_len = Self::page_len(buf, self.current_page, self.page_size);
453 let available = page_len - self.offset_in_page;
454 if cnt < available {
455 self.offset_in_page += cnt;
456 return;
457 }
458 cnt -= available;
459 self.current_page += 1;
460 self.offset_in_page = 0;
461 }
462
463 if self.current_page >= buf.num_pages {
465 self.buffers.pop_front();
466 self.current_page = 0;
467 self.offset_in_page = 0;
468 }
469 }
470 }
471}
472
473pub struct Replay<B: Blob> {
482 reader: PageReader<B>,
484 buffer: ReplayBuf,
486 exhausted: bool,
488}
489
490impl<B: Blob> Replay<B> {
491 pub(super) const fn new(reader: PageReader<B>) -> Self {
493 let physical_page_size = reader.physical_page_size();
494 let page_size = reader.page_size();
495 Self {
496 reader,
497 buffer: ReplayBuf::new(physical_page_size, page_size),
498 exhausted: false,
499 }
500 }
501
502 pub const fn blob_size(&self) -> u64 {
504 self.reader.blob_size()
505 }
506
507 pub const fn is_exhausted(&self) -> bool {
512 self.exhausted
513 }
514
515 pub async fn ensure(&mut self, n: usize) -> Result<bool, Error> {
526 while self.buffer.remaining < n && !self.exhausted {
527 match self.reader.fill().await? {
528 Some((state, logical_bytes)) => {
529 self.buffer.push(state, logical_bytes);
530 }
531 None => {
532 self.exhausted = true;
533 }
534 }
535 }
536 Ok(self.buffer.remaining >= n)
537 }
538
539 pub fn seek_to(&mut self, offset: u64) -> Result<(), Error> {
542 if offset > self.reader.blob_size() {
543 return Err(Error::BlobInsufficientLength);
544 }
545
546 self.buffer.clear();
547 self.exhausted = false;
548
549 let page_size = self.reader.page_size as u64;
550 self.reader.seek_to_page(offset / page_size);
551 self.buffer.current_page = 0;
552 self.buffer.offset_in_page = (offset % page_size) as usize;
553
554 Ok(())
555 }
556}
557
558impl<B: Blob> Buf for Replay<B> {}
559
560impl<B: Blob> bytes::Buf for Replay<B> {
561 fn copy_to_bytes(&mut self, len: usize) -> Bytes {
562 self.buffer.copy_to_bytes(len)
563 }
564
565 fn remaining(&self) -> usize {
566 self.buffer.remaining()
567 }
568
569 #[inline(always)]
570 fn try_copy_to_slice(&mut self, dst: &mut [u8]) -> Result<(), TryGetError> {
571 self.buffer.try_copy_to_slice(dst)
572 }
573
574 fn chunk(&self) -> &[u8] {
575 self.buffer.chunk()
576 }
577
578 fn advance(&mut self, cnt: usize) {
579 self.buffer.advance(cnt);
580 }
581}
582
583#[cfg(test)]
584mod tests {
585 use super::{super::writer::Writer, *};
586 use crate::{
587 Handle, IoBufs, Runner as _, Storage as _, WriteOptions,
588 buffer::paged::{CacheRef, Sealed},
589 deterministic,
590 };
591 use bytes::Buf as _;
592 use commonware_macros::test_traced;
593 use commonware_utils::{NZU16, NZUsize, channel::oneshot, sync::Mutex};
594 use futures::task::noop_waker_ref;
595 use std::{
596 sync::atomic::{AtomicUsize, Ordering},
597 task::{Context, Poll},
598 };
599
600 type Calls = Arc<Mutex<Vec<(bool, Vec<(u64, usize)>)>>>;
603
604 #[derive(Default)]
606 struct Reads {
607 fail_from: Option<u64>,
609 yields: usize,
611 trailing_error: bool,
613 gated: bool,
615 admission: Option<(Arc<tokio::sync::Semaphore>, Arc<AtomicUsize>)>,
618 fail_once: Option<(u64, Arc<AtomicUsize>)>,
620 }
621
622 type Gate = (u64, oneshot::Sender<()>);
624
625 #[derive(Clone, Default)]
627 struct Gates(Arc<Mutex<Vec<Gate>>>);
628
629 impl Gates {
630 fn release(&self, offset: u64) {
632 let mut gates = self.0.lock();
633 let i = gates
634 .iter()
635 .position(|(o, release)| *o == offset && !release.is_closed())
636 .unwrap();
637 let _ = gates.remove(i).1.send(());
638 }
639
640 fn open(&self) -> Vec<u64> {
642 self.0
643 .lock()
644 .iter()
645 .filter(|(_, release)| !release.is_closed())
646 .map(|(offset, _)| *offset)
647 .collect()
648 }
649 }
650
651 struct GroupingBlob<B> {
653 inner: B,
654 calls: Calls,
655 reads: Reads,
656 gates: Gates,
657 }
658
659 impl<B> GroupingBlob<B> {
660 async fn begin(&self, many: bool, ranges: &[(u64, usize)]) -> Result<(), Error> {
662 self.record(many, ranges)?;
663 for _ in 0..self.reads.yields {
664 crate::reschedule().await;
665 }
666 Ok(())
667 }
668
669 fn record(&self, many: bool, ranges: &[(u64, usize)]) -> Result<(), Error> {
670 self.calls.lock().push((many, ranges.to_vec()));
671 if self
672 .reads
673 .fail_once
674 .as_ref()
675 .is_some_and(|(offset, attempts)| {
676 ranges[0].0 == *offset && attempts.fetch_add(1, Ordering::SeqCst) == 0
677 })
678 {
679 return Err(Error::ReadFailed);
680 }
681 if self
682 .reads
683 .fail_from
684 .is_some_and(|fail| ranges.iter().any(|(offset, _)| *offset >= fail))
685 {
686 return Err(Error::ReadFailed);
687 }
688 Ok(())
689 }
690 }
691
692 impl<B: Blob> Blob for GroupingBlob<B> {
693 async fn read_at_buf(
694 &self,
695 offset: u64,
696 len: usize,
697 bufs: impl Into<IoBufsMut> + Send,
698 options: ReadOptions,
699 ) -> Result<IoBufsMut, Error> {
700 self.begin(false, &[(offset, len)]).await?;
701 self.inner.read_at_buf(offset, len, bufs, options).await
702 }
703
704 async fn read_at(
705 &self,
706 offset: u64,
707 len: usize,
708 options: ReadOptions,
709 ) -> Result<IoBufsMut, Error> {
710 self.begin(false, &[(offset, len)]).await?;
711 self.inner.read_at(offset, len, options).await
712 }
713
714 async fn start_read_many(
718 &self,
719 ranges: &[(u64, usize)],
720 options: ReadOptions,
721 ) -> Handle<Vec<IoBufsMut>> {
722 if let Err(err) = self.record(true, ranges) {
723 return Handle::ready(Err(err));
724 }
725 let bytes = self.inner.start_read_many(ranges, options).await;
726 let offset = ranges[0].0;
727 let gate = self.reads.gated.then(|| {
728 let (release, released) = oneshot::channel();
729 self.gates.0.lock().push((offset, release));
730 released
731 });
732 let admission = self.reads.admission.clone();
733 let yields = self.reads.yields;
734 let trailing_error = self.reads.trailing_error;
735 Handle::from_future(async move {
736 if let Some(released) = gate {
737 let _ = released.await;
738 }
739 let _permit = match &admission {
740 Some((permits, _)) => Some(permits.acquire().await.unwrap()),
741 None => None,
742 };
743 if let Some((_, attempts)) = &admission {
744 let fail = offset == 0 && attempts.fetch_add(1, Ordering::SeqCst) == 0;
745 crate::reschedule().await;
746 crate::reschedule().await;
747 if fail {
748 return Err(Error::ReadFailed);
749 }
750 }
751 for _ in 0..yields {
752 crate::reschedule().await;
753 }
754 let bufs = bytes.await?;
755 if trailing_error {
756 return Err(Error::ReadFailed);
757 }
758 Ok(bufs)
759 })
760 }
761
762 async fn write_at(
763 &self,
764 offset: u64,
765 bufs: impl Into<IoBufs> + Send,
766 options: WriteOptions,
767 ) -> Result<(), Error> {
768 self.inner.write_at(offset, bufs, options).await
769 }
770 async fn resize(&self, len: u64) -> Result<(), Error> {
771 self.inner.resize(len).await
772 }
773 async fn sync(&self) -> Result<(), Error> {
774 self.inner.sync().await
775 }
776 async fn start_sync(&self) -> Handle<()> {
777 self.inner.start_sync().await
778 }
779 }
780
781 const PIPE_PAGE_SIZE: NonZeroU16 = NZU16!(64);
785 const PIPE_PHYSICAL_PAGE: usize = 64 + Checksum::SIZE;
786 const PIPE_BUFFER_PAGES: usize = 64;
787 const PIPE_BUFFER_BYTES: usize = PIPE_BUFFER_PAGES * PIPE_PHYSICAL_PAGE;
788 const PIPE_BATCH_PAGES: usize = PIPE_BUFFER_PAGES / (MAX_READ_AHEAD + 1);
789 const PIPE_BATCH_BYTES: usize = PIPE_BATCH_PAGES * PIPE_PHYSICAL_PAGE;
790 const PIPE_BATCH_LOGICAL: usize = PIPE_BATCH_PAGES * 64;
791 const PIPE_BLOB_BATCHES: usize = MAX_READ_AHEAD + 2;
792 const PIPE_BLOB_BYTES: usize = PIPE_BLOB_BATCHES * PIPE_BATCH_LOGICAL + 10;
793
794 async fn recorded_blob(
798 context: &deterministic::Context,
799 name: &[u8],
800 page_size: NonZeroU16,
801 buffer: usize,
802 len: usize,
803 reads: Reads,
804 ) -> (
805 Sealed<GroupingBlob<<deterministic::Context as crate::Storage>::Blob>>,
806 Calls,
807 Vec<u8>,
808 Gates,
809 ) {
810 let calls: Calls = Arc::new(Mutex::new(Vec::new()));
811 let gates = Gates::default();
812 let (blob, size) = context.open("pipelining", name).await.unwrap();
813 let blob = GroupingBlob {
814 inner: blob,
815 calls: calls.clone(),
816 reads,
817 gates: gates.clone(),
818 };
819 let cache = CacheRef::new(page_size, NZUsize!(4));
820 let writer = Writer::new(blob, size, buffer, cache).await.unwrap();
821 let data: Vec<u8> = (0..len).map(|i| (i % 251) as u8).collect();
822 let (writer, _) = writer.append(&data).await.unwrap();
823 let (_, sealed) = writer.snapshot().await.unwrap();
824 calls.lock().clear();
825 (sealed, calls, data, gates)
826 }
827
828 async fn piped_blob(
830 context: &deterministic::Context,
831 name: &[u8],
832 reads: Reads,
833 ) -> (
834 Sealed<GroupingBlob<<deterministic::Context as crate::Storage>::Blob>>,
835 Calls,
836 Vec<u8>,
837 Gates,
838 ) {
839 recorded_blob(
840 context,
841 name,
842 PIPE_PAGE_SIZE,
843 PIPE_BUFFER_BYTES,
844 PIPE_BLOB_BYTES,
845 reads,
846 )
847 .await
848 }
849
850 fn piped_calls(batches: std::ops::Range<usize>) -> Vec<(bool, Vec<(u64, usize)>)> {
852 batches
853 .map(|i| {
854 (
855 true,
856 vec![((i * PIPE_BATCH_BYTES) as u64, PIPE_BATCH_BYTES)],
857 )
858 })
859 .collect()
860 }
861
862 fn first_fill_calls(batch: usize) -> Vec<(bool, Vec<(u64, usize)>)> {
865 let len = batch * PIPE_PHYSICAL_PAGE;
866 (0..=MAX_READ_AHEAD)
867 .map(|i| (true, vec![((i * len) as u64, len)]))
868 .collect()
869 }
870
871 async fn take<B: Blob>(replay: &mut Replay<B>, expected: &[u8]) {
873 assert!(replay.ensure(expected.len()).await.unwrap());
874 assert_eq!(replay.copy_to_bytes(expected.len()).as_ref(), expected);
875 }
876
877 #[test]
879 fn test_replay_is_sync() {
880 fn assert_sync<T: Sync>() {}
881 assert_sync::<Replay<<deterministic::Context as crate::Storage>::Blob>>();
882 }
883
884 #[test]
886 fn test_failed_read_ahead_is_read_again_on_demand() {
887 deterministic::Runner::default().start(|context| async move {
888 let reads = Reads {
889 fail_once: Some((PIPE_BATCH_BYTES as u64, Arc::new(AtomicUsize::new(0)))),
890 ..Reads::default()
891 };
892 let (sealed, calls, data, _) =
893 piped_blob(&context, b"speculative-failure", reads).await;
894 let mut replay = sealed
895 .replay(NZUsize!(PIPE_BUFFER_BYTES), ReadOptions::default())
896 .unwrap();
897 take(&mut replay, &data[..PIPE_BATCH_LOGICAL]).await;
898 take(&mut replay, &data[PIPE_BATCH_LOGICAL..]).await;
899 assert!(!replay.ensure(1).await.unwrap());
900 assert_eq!(
901 calls
902 .lock()
903 .iter()
904 .filter(|(_, ranges)| ranges[0].0 == PIPE_BATCH_BYTES as u64)
905 .count(),
906 2
907 );
908 });
909 }
910
911 #[rstest::rstest]
914 #[case(1)]
915 #[case(2)]
916 #[case(3)]
917 #[case(4)]
918 fn test_retry_does_not_wait_on_queued_reads(
919 #[case] permits: usize,
920 #[values(false, true)] cancel: bool,
921 ) {
922 use crate::Clock as _;
923 deterministic::Runner::default().start(|context| async move {
924 let admission = Arc::new(tokio::sync::Semaphore::new(permits));
925 let attempts = Arc::new(AtomicUsize::new(0));
926 let reads = Reads {
927 admission: Some((admission, attempts.clone())),
928 ..Reads::default()
929 };
930 let (sealed, _, data, _) = piped_blob(&context, b"permit-retry", reads).await;
931 let mut replay = sealed
932 .replay(NZUsize!(PIPE_BUFFER_BYTES), ReadOptions::default())
933 .unwrap();
934 if cancel {
935 let mut first = Box::pin(replay.ensure(1));
936 assert!(
937 first
938 .as_mut()
939 .poll(&mut Context::from_waker(noop_waker_ref()))
940 .is_pending()
941 );
942 drop(first);
943 } else {
944 assert!(matches!(replay.ensure(1).await, Err(Error::ReadFailed)));
945 }
946 assert!(attempts.load(Ordering::SeqCst) >= 1);
947 commonware_macros::select! {
948 result = replay.ensure(1) => assert!(result.unwrap()),
949 _ = context.sleep(std::time::Duration::from_secs(1)) => {
950 panic!("retry hung with {permits} permits, cancellation {cancel}");
951 },
952 }
953 take(&mut replay, &data[..PIPE_BATCH_LOGICAL]).await;
954 });
955 }
956
957 #[test_traced]
958 fn test_batches_split_the_buffer() {
959 let executor = deterministic::Runner::default();
960 executor.start(|context| async move {
961 let full = MAX_READ_AHEAD + 1;
964 for (pages, batch) in [
965 (1, 1),
966 (full - 1, 1),
967 (full, 1),
968 (full + 1, 2),
969 (2 * full - 1, 2),
970 (2 * full + 1, 3),
971 (PIPE_BUFFER_PAGES, PIPE_BATCH_PAGES),
972 ] {
973 let (sealed, calls, _, _) = recorded_blob(
974 &context,
975 format!("depth-{pages}").as_bytes(),
976 PIPE_PAGE_SIZE,
977 pages * PIPE_PHYSICAL_PAGE,
978 PIPE_BLOB_BYTES,
979 Reads::default(),
980 )
981 .await;
982 let mut replay = sealed
983 .replay(NZUsize!(pages * PIPE_PHYSICAL_PAGE), ReadOptions::default())
984 .unwrap();
985 assert!(replay.ensure(1).await.unwrap());
986 assert_eq!(*calls.lock(), first_fill_calls(batch), "{pages} pages");
987 }
988 });
989 }
990
991 #[test_traced]
992 fn test_fill_keeps_batches_in_flight() {
993 let executor = deterministic::Runner::default();
994 executor.start(|context| async move {
995 let (sealed, calls, data, _) =
996 piped_blob(&context, b"pipeline", Reads::default()).await;
997 let mut replay = sealed
998 .replay(NZUsize!(PIPE_BUFFER_BYTES), ReadOptions::default())
999 .unwrap();
1000
1001 let ahead = MAX_READ_AHEAD;
1004 assert!(replay.ensure(1).await.unwrap());
1005 assert_eq!(*calls.lock(), piped_calls(0..ahead + 1));
1006
1007 take(&mut replay, &data[..PIPE_BATCH_LOGICAL]).await;
1009 assert_eq!(calls.lock().len(), ahead + 1);
1010 take(
1011 &mut replay,
1012 &data[PIPE_BATCH_LOGICAL..2 * PIPE_BATCH_LOGICAL],
1013 )
1014 .await;
1015 assert_eq!(*calls.lock(), piped_calls(0..ahead + 2));
1016 take(
1017 &mut replay,
1018 &data[2 * PIPE_BATCH_LOGICAL..PIPE_BLOB_BATCHES * PIPE_BATCH_LOGICAL],
1019 )
1020 .await;
1021 assert_eq!(*calls.lock(), piped_calls(0..PIPE_BLOB_BATCHES));
1022
1023 take(&mut replay, &data[PIPE_BLOB_BATCHES * PIPE_BATCH_LOGICAL..]).await;
1025 assert!(!replay.ensure(1).await.unwrap());
1026 assert!(replay.is_exhausted());
1027 });
1028 }
1029
1030 #[test_traced]
1031 fn test_seek_discards_the_batches_in_flight() {
1032 let executor = deterministic::Runner::default();
1033 executor.start(|context| async move {
1034 let reads = Reads {
1037 gated: true,
1038 ..Reads::default()
1039 };
1040 let (sealed, calls, data, gates) = piped_blob(&context, b"seek", reads).await;
1041 let mut replay = sealed
1042 .replay(NZUsize!(PIPE_BUFFER_BYTES), ReadOptions::default())
1043 .unwrap();
1044 let ahead = MAX_READ_AHEAD;
1045 {
1046 let mut fill = Box::pin(replay.ensure(1));
1047 assert!(
1048 fill.as_mut()
1049 .poll(&mut Context::from_waker(noop_waker_ref()))
1050 .is_pending()
1051 );
1052 gates.release(0);
1053 assert!(fill.await.unwrap());
1054 }
1055 assert_eq!(calls.lock().len(), ahead + 1);
1056 let in_flight: Vec<u64> = (1..=ahead).map(|i| (i * PIPE_BATCH_BYTES) as u64).collect();
1057 assert_eq!(gates.open(), in_flight);
1058
1059 let batch = PIPE_BLOB_BATCHES - 2;
1062 let offset = batch * PIPE_BATCH_LOGICAL + 5;
1063 replay.seek_to(offset as u64).unwrap();
1064 assert!(gates.open().is_empty());
1065 {
1066 let mut fill = Box::pin(replay.ensure(1));
1067 assert!(
1068 fill.as_mut()
1069 .poll(&mut Context::from_waker(noop_waker_ref()))
1070 .is_pending()
1071 );
1072 for i in batch..PIPE_BLOB_BATCHES {
1073 gates.release((i * PIPE_BATCH_BYTES) as u64);
1074 }
1075 assert!(fill.await.unwrap());
1076 }
1077 take(&mut replay, &data[offset..]).await;
1078 assert!(!replay.ensure(1).await.unwrap());
1079 assert_eq!(
1080 calls.lock()[ahead + 1..],
1081 piped_calls(batch..PIPE_BLOB_BATCHES)
1082 );
1083 });
1084 }
1085
1086 #[test_traced]
1087 fn test_failed_fill_reads_the_same_batch_again() {
1088 let executor = deterministic::Runner::default();
1089 executor.start(|context| async move {
1090 let fail_from = Some(PIPE_BATCH_BYTES as u64);
1091 let reads = Reads {
1092 fail_from,
1093 ..Reads::default()
1094 };
1095 let (sealed, calls, data, _) = piped_blob(&context, b"retry", reads).await;
1096 let mut replay = sealed
1097 .replay(NZUsize!(PIPE_BUFFER_BYTES), ReadOptions::default())
1098 .unwrap();
1099 take(&mut replay, &data[..PIPE_BATCH_LOGICAL]).await;
1100
1101 for _ in 0..2 {
1104 assert!(matches!(replay.ensure(1).await, Err(Error::ReadFailed)));
1105 }
1106 let ahead = MAX_READ_AHEAD;
1107 let mut expected = piped_calls(0..ahead + 2);
1108 expected.extend(piped_calls(1..ahead + 2));
1109 expected.extend(piped_calls(1..ahead + 2));
1110 assert_eq!(*calls.lock(), expected);
1111 });
1112 }
1113
1114 #[test_traced]
1115 fn test_dropped_fill_reads_the_same_batch_again() {
1116 let executor = deterministic::Runner::default();
1117 executor.start(|context| async move {
1118 let reads = Reads {
1120 yields: 2,
1121 ..Reads::default()
1122 };
1123 let (sealed, calls, data, _) = piped_blob(&context, b"dropped", reads).await;
1124 let mut replay = sealed
1125 .replay(NZUsize!(PIPE_BUFFER_BYTES), ReadOptions::default())
1126 .unwrap();
1127
1128 {
1130 let mut fill = Box::pin(replay.ensure(1));
1131 assert!(
1132 fill.as_mut()
1133 .poll(&mut Context::from_waker(noop_waker_ref()))
1134 .is_pending()
1135 );
1136 }
1137 let ahead = MAX_READ_AHEAD;
1138 assert_eq!(*calls.lock(), piped_calls(0..ahead + 1));
1139
1140 take(&mut replay, &data[..2 * PIPE_BATCH_LOGICAL]).await;
1143 let mut expected = piped_calls(0..ahead + 1);
1144 expected.extend(piped_calls(0..ahead + 2));
1145 assert_eq!(*calls.lock(), expected);
1146 });
1147 }
1148
1149 #[test_traced]
1150 fn test_batches_landing_before_and_after_their_wait_are_served() {
1151 let executor = deterministic::Runner::default();
1152 executor.start(|context| async move {
1153 let reads = Reads {
1156 gated: true,
1157 ..Reads::default()
1158 };
1159 let (sealed, calls, data, gates) = piped_blob(&context, b"gated", reads).await;
1160 let mut replay = sealed
1161 .replay(NZUsize!(PIPE_BUFFER_BYTES), ReadOptions::default())
1162 .unwrap();
1163
1164 {
1167 let mut fill = Box::pin(replay.ensure(1));
1168 assert!(
1169 fill.as_mut()
1170 .poll(&mut Context::from_waker(noop_waker_ref()))
1171 .is_pending()
1172 );
1173 assert_eq!(calls.lock().len(), MAX_READ_AHEAD + 1);
1174 gates.release(PIPE_BATCH_BYTES as u64);
1175 gates.release(0);
1176 assert!(fill.await.unwrap());
1177 }
1178 take(&mut replay, &data[..PIPE_BATCH_LOGICAL]).await;
1179
1180 take(
1182 &mut replay,
1183 &data[PIPE_BATCH_LOGICAL..2 * PIPE_BATCH_LOGICAL],
1184 )
1185 .await;
1186 assert_eq!(calls.lock().len(), MAX_READ_AHEAD + 2);
1187
1188 let wakes = Arc::new(Wakes::default());
1191 let waker = futures::task::waker(wakes.clone());
1192 {
1193 let mut fill = Box::pin(replay.ensure(1));
1194 assert!(
1195 fill.as_mut()
1196 .poll(&mut Context::from_waker(&waker))
1197 .is_pending()
1198 );
1199 gates.release(2 * PIPE_BATCH_BYTES as u64);
1200 assert_eq!(wakes.0.load(Ordering::SeqCst), 1);
1201 assert!(matches!(
1202 fill.as_mut().poll(&mut Context::from_waker(&waker)),
1203 Poll::Ready(Ok(true))
1204 ));
1205 }
1206 take(
1207 &mut replay,
1208 &data[2 * PIPE_BATCH_LOGICAL..3 * PIPE_BATCH_LOGICAL],
1209 )
1210 .await;
1211 });
1212 }
1213
1214 #[derive(Default)]
1216 struct Wakes(AtomicUsize);
1217
1218 impl futures::task::ArcWake for Wakes {
1219 fn wake_by_ref(wakes: &Arc<Self>) {
1220 wakes.0.fetch_add(1, Ordering::SeqCst);
1221 }
1222 }
1223
1224 #[test_traced]
1225 fn test_read_error_after_the_batch_fails_the_fill() {
1226 let executor = deterministic::Runner::default();
1227 executor.start(|context| async move {
1228 let reads = Reads {
1230 trailing_error: true,
1231 ..Reads::default()
1232 };
1233 let (sealed, _, _, _) = piped_blob(&context, b"late-error", reads).await;
1234 let mut replay = sealed
1235 .replay(NZUsize!(PIPE_BUFFER_BYTES), ReadOptions::default())
1236 .unwrap();
1237 assert!(matches!(replay.ensure(1).await, Err(Error::ReadFailed)));
1238 });
1239 }
1240
1241 #[test]
1242 fn test_replay_buf_bytes_share_pages() {
1243 let source = bytes::Bytes::from_static(b"abcd............efgh............");
1244 let range = source.as_ptr_range();
1245 let mut replay = ReplayBuf::new(16, 4);
1246 replay.push(
1247 BufferState {
1248 buffer: source,
1249 num_pages: 2,
1250 last_page_len: 4,
1251 },
1252 8,
1253 );
1254
1255 let first = replay.copy_to_bytes(3);
1256 assert_eq!(first.as_ref(), b"abc");
1257 assert!(range.contains(&first.as_ptr()));
1258
1259 let crossing = replay.copy_to_bytes(3);
1260 assert_eq!(crossing.as_ref(), b"def");
1261 assert!(!range.contains(&crossing.as_ptr()));
1262
1263 let last = replay.copy_to_bytes(2);
1264 assert_eq!(last.as_ref(), b"gh");
1265 assert!(range.contains(&last.as_ptr()));
1266 assert_eq!(replay.remaining(), 0);
1267 }
1268
1269 #[test]
1270 fn test_replay_buf_copy_to_slice_page_boundaries() {
1271 let source = bytes::Bytes::from_static(b"abcd............efgh............ijkl............");
1272 let mut replay = ReplayBuf::new(16, 4);
1273 replay.push(
1274 BufferState {
1275 buffer: source,
1276 num_pages: 3,
1277 last_page_len: 4,
1278 },
1279 12,
1280 );
1281
1282 let mut inside = [0u8; 3];
1284 replay.try_copy_to_slice(&mut inside).unwrap();
1285 assert_eq!(&inside, b"abc");
1286 assert_eq!(replay.chunk(), b"d");
1287 assert_eq!(replay.remaining(), 9);
1288
1289 let mut boundary = [0u8; 1];
1291 replay.try_copy_to_slice(&mut boundary).unwrap();
1292 assert_eq!(&boundary, b"d");
1293 assert_eq!(replay.chunk(), b"efgh");
1294 assert_eq!(replay.remaining(), 8);
1295
1296 let mut spanning = [0u8; 6];
1298 replay.try_copy_to_slice(&mut spanning).unwrap();
1299 assert_eq!(&spanning, b"efghij");
1300 assert_eq!(replay.chunk(), b"kl");
1301 assert_eq!(replay.remaining(), 2);
1302
1303 let mut tail = [0u8; 2];
1305 replay.try_copy_to_slice(&mut tail).unwrap();
1306 assert_eq!(&tail, b"kl");
1307 assert_eq!(replay.chunk(), b"");
1308 assert_eq!(replay.remaining(), 0);
1309
1310 let err = replay.try_copy_to_slice(&mut [0u8; 1]).unwrap_err();
1311 assert_eq!(err.requested, 1);
1312 assert_eq!(err.available, 0);
1313 }
1314
1315 const PAGE_SIZE: NonZeroU16 = NZU16!(103);
1316 const BUFFER_PAGES: usize = 2;
1317
1318 #[test_traced("DEBUG")]
1319 fn test_replay_basic() {
1320 let executor = deterministic::Runner::default();
1321 executor.start(|context: deterministic::Context| async move {
1322 let (blob, blob_size) = context.open("test_partition", b"test_blob").await.unwrap();
1323 assert_eq!(blob_size, 0);
1324
1325 let cache_ref = super::super::CacheRef::new(PAGE_SIZE, NZUsize!(BUFFER_PAGES));
1326 let mut append = Writer::new(blob, blob_size, BUFFER_PAGES * 115, cache_ref)
1327 .await
1328 .unwrap();
1329
1330 let data: Vec<u8> = (0u8..=255).cycle().take(300).collect();
1332 (append, _) = append.append(&data).await.unwrap();
1333 append = append.sync().await.unwrap();
1334
1335 let (_, mut replay) = append
1337 .replay(NZUsize!(BUFFER_PAGES), ReadOptions::default())
1338 .await
1339 .unwrap();
1340
1341 replay.ensure(300).await.unwrap();
1343
1344 assert_eq!(replay.remaining(), 300);
1346
1347 let mut collected = Vec::new();
1349 while replay.remaining() > 0 {
1350 let chunk = replay.chunk();
1351 collected.extend_from_slice(chunk);
1352 let len = chunk.len();
1353 replay.advance(len);
1354 }
1355 assert_eq!(collected, data);
1356 });
1357 }
1358
1359 #[test_traced("DEBUG")]
1360 fn test_replay_partial_page() {
1361 let executor = deterministic::Runner::default();
1362 executor.start(|context: deterministic::Context| async move {
1363 let (blob, blob_size) = context.open("test_partition", b"test_blob").await.unwrap();
1364
1365 let cache_ref = super::super::CacheRef::new(PAGE_SIZE, NZUsize!(BUFFER_PAGES));
1366 let mut append = Writer::new(blob, blob_size, BUFFER_PAGES * 115, cache_ref)
1367 .await
1368 .unwrap();
1369
1370 let data: Vec<u8> = (1u8..=(PAGE_SIZE.get() + 10) as u8).collect();
1372 (append, _) = append.append(&data).await.unwrap();
1373 append = append.sync().await.unwrap();
1374
1375 let (_, mut replay) = append
1376 .replay(NZUsize!(BUFFER_PAGES), ReadOptions::default())
1377 .await
1378 .unwrap();
1379
1380 replay.ensure(data.len()).await.unwrap();
1382
1383 assert_eq!(replay.remaining(), data.len());
1384 });
1385 }
1386
1387 #[test_traced("DEBUG")]
1388 fn test_replay_cross_buffer_boundary() {
1389 let executor = deterministic::Runner::default();
1392 executor.start(|context: deterministic::Context| async move {
1393 let (blob, blob_size) = context.open("test_partition", b"test_blob").await.unwrap();
1394 assert_eq!(blob_size, 0);
1395
1396 let cache_ref = super::super::CacheRef::new(PAGE_SIZE, NZUsize!(BUFFER_PAGES));
1397 let mut append = Writer::new(blob, blob_size, BUFFER_PAGES * 115, cache_ref)
1398 .await
1399 .unwrap();
1400
1401 let data: Vec<u8> = (0u8..=255).cycle().take(400).collect();
1403 (append, _) = append.append(&data).await.unwrap();
1404 append = append.sync().await.unwrap();
1405
1406 let (_, mut replay) = append
1410 .replay(NZUsize!(115), ReadOptions::default())
1411 .await
1412 .unwrap();
1413
1414 assert!(replay.ensure(400).await.unwrap());
1416 assert_eq!(replay.remaining(), 400);
1417
1418 let mut collected = Vec::new();
1420 let mut chunks_read = 0;
1421 while replay.remaining() > 0 {
1422 let chunk = replay.chunk();
1423 assert!(
1424 !chunk.is_empty(),
1425 "chunk() returned empty but remaining > 0"
1426 );
1427 collected.extend_from_slice(chunk);
1428 let len = chunk.len();
1429 replay.advance(len);
1430 chunks_read += 1;
1431 }
1432
1433 assert_eq!(collected, data);
1434
1435 assert!(
1438 chunks_read >= 4,
1439 "Expected at least 4 chunks for 4 pages, got {}",
1440 chunks_read
1441 );
1442 });
1443 }
1444
1445 #[test_traced("DEBUG")]
1446 fn test_replay_empty_blob() {
1447 let executor = deterministic::Runner::default();
1450 executor.start(|context: deterministic::Context| async move {
1451 let (blob, blob_size) = context.open("test_partition", b"test_blob").await.unwrap();
1452 assert_eq!(blob_size, 0);
1453
1454 let cache_ref = super::super::CacheRef::new(PAGE_SIZE, NZUsize!(BUFFER_PAGES));
1455 let append = Writer::new(blob, blob_size, BUFFER_PAGES * 115, cache_ref)
1456 .await
1457 .unwrap();
1458
1459 assert_eq!(append.size(), 0);
1461
1462 let (_, mut replay) = append
1464 .replay(NZUsize!(BUFFER_PAGES), ReadOptions::default())
1465 .await
1466 .unwrap();
1467
1468 assert_eq!(replay.remaining(), 0);
1471
1472 assert!(replay.ensure(0).await.unwrap());
1474
1475 assert!(!replay.ensure(1).await.unwrap());
1477
1478 assert!(replay.is_exhausted());
1480
1481 assert!(replay.chunk().is_empty());
1483
1484 assert_eq!(replay.remaining(), 0);
1486 });
1487 }
1488
1489 #[test_traced("DEBUG")]
1490 fn test_replay_seek_to() {
1491 let executor = deterministic::Runner::default();
1492 executor.start(|context: deterministic::Context| async move {
1493 let (blob, blob_size) = context.open("test_partition", b"test_blob").await.unwrap();
1494
1495 let cache_ref = super::super::CacheRef::new(PAGE_SIZE, NZUsize!(BUFFER_PAGES));
1496 let mut append = Writer::new(blob, blob_size, BUFFER_PAGES * 115, cache_ref)
1497 .await
1498 .unwrap();
1499
1500 let data: Vec<u8> = (0u8..=255).cycle().take(300).collect();
1502 (append, _) = append.append(&data).await.unwrap();
1503 append = append.sync().await.unwrap();
1504
1505 let (_, mut replay) = append
1506 .replay(NZUsize!(BUFFER_PAGES), ReadOptions::default())
1507 .await
1508 .unwrap();
1509
1510 replay.seek_to(150).unwrap();
1512 replay.ensure(50).await.unwrap();
1513 assert_eq!(replay.get_u8(), data[150]);
1514
1515 replay.seek_to(0).unwrap();
1517 replay.ensure(1).await.unwrap();
1518 assert_eq!(replay.get_u8(), data[0]);
1519
1520 assert!(replay.seek_to(data.len() as u64 + 1).is_err());
1522
1523 let seek_offset = 150usize;
1525 replay.seek_to(seek_offset as u64).unwrap();
1526 let expected_remaining = data.len() - seek_offset;
1527 let mut collected = Vec::new();
1528 loop {
1529 if !replay.ensure(1).await.unwrap() {
1531 break; }
1533 let chunk = replay.chunk();
1534 if chunk.is_empty() {
1535 break;
1536 }
1537 collected.extend_from_slice(chunk);
1538 let len = chunk.len();
1539 replay.advance(len);
1540 }
1541 assert_eq!(
1542 collected.len(),
1543 expected_remaining,
1544 "After seeking to {}, should read {} bytes but got {}",
1545 seek_offset,
1546 expected_remaining,
1547 collected.len()
1548 );
1549 assert_eq!(collected, &data[seek_offset..]);
1550
1551 for offset in [data.len(), 250, 299] {
1553 replay.seek_to(Widen::widen(offset)).unwrap();
1554 let remaining = data.len() - offset;
1555 assert!(!replay.ensure(remaining + 1).await.unwrap());
1556 assert_eq!(replay.copy_to_bytes(remaining).as_ref(), &data[offset..]);
1557 }
1558 });
1559 }
1560}