1use std::collections::{BTreeMap, VecDeque};
12use std::io::SeekFrom;
13use std::pin::Pin;
14use std::sync::Arc;
15
16use bytes::Bytes;
17use futures::stream::{self, FuturesUnordered, Stream, StreamExt};
18
19use crate::bmt::DEFAULT_BODY_SIZE;
20use crate::chunk::ChunkAddress;
21
22use super::error::{FileError, Result};
23use super::frontier::{SubtreeNode, overlapping_children};
24use super::joiner::GenericJoiner;
25use super::mode::JoinMode;
26use super::tree::{ChunkRange, TreeParams};
27use crate::store::{ChunkGet, MaybeSend};
28
29const DEFAULT_LEAF_RETRIES: u32 = 4;
32
33#[cfg(test)]
34thread_local! {
35 static PEAK_LEAF_OCCUPANCY: std::cell::Cell<usize> = const { std::cell::Cell::new(0) };
40}
41
42#[cfg(test)]
44fn record_occupancy(occupancy: usize) {
45 PEAK_LEAF_OCCUPANCY.with(|p| p.set(p.get().max(occupancy)));
46}
47
48pub struct WindowedReader<G, M: JoinMode, const BODY_SIZE: usize = DEFAULT_BODY_SIZE>
56where
57 G: ChunkGet<BODY_SIZE>,
58{
59 getter: Arc<G>,
60 subtrees: Vec<SubtreeNode<M>>,
61 tree: TreeParams<BODY_SIZE>,
62 span: u64,
63 concurrency: usize,
64 position: u64,
65 window: usize,
66 #[allow(
67 dead_code,
68 reason = "retained reusable joiner state for re-stream-on-seek"
69 )]
70 root: ChunkAddress,
71}
72
73impl<G, M, const BODY_SIZE: usize> std::fmt::Debug for WindowedReader<G, M, BODY_SIZE>
74where
75 G: ChunkGet<BODY_SIZE>,
76 M: JoinMode,
77{
78 fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
79 f.debug_struct("WindowedReader")
80 .field("span", &self.span)
81 .field("position", &self.position)
82 .field("concurrency", &self.concurrency)
83 .finish_non_exhaustive()
84 }
85}
86
87impl<G, M, const BODY_SIZE: usize> GenericJoiner<G, M, BODY_SIZE>
88where
89 G: ChunkGet<BODY_SIZE> + 'static,
90 M: JoinMode + MaybeSend + Sync,
91{
92 pub fn into_windowed_reader(self, window: usize) -> WindowedReader<G, M, BODY_SIZE> {
98 WindowedReader {
99 getter: Arc::clone(self.getter()),
100 subtrees: self.subtrees().to_vec(),
101 tree: self.tree(),
102 span: self.size(),
103 concurrency: self.concurrency(),
104 position: self.position(),
105 root: *self.root(),
106 window: window.max(1),
107 }
108 }
109}
110
111impl<G, M, const BODY_SIZE: usize> WindowedReader<G, M, BODY_SIZE>
112where
113 G: ChunkGet<BODY_SIZE> + 'static,
114 M: JoinMode + MaybeSend + Sync,
115{
116 pub fn stream(&mut self) -> impl Stream<Item = Result<Bytes>> + '_ {
121 windowed_walk::<G, M, BODY_SIZE>(
122 Arc::clone(&self.getter),
123 self.subtrees.clone(),
124 self.tree,
125 self.span,
126 self.concurrency,
127 self.position,
128 self.window,
129 )
130 }
131
132 pub fn seek(&mut self, pos: SeekFrom) -> std::io::Result<u64> {
135 self.position = super::resolve_seek_position(pos, self.position, self.span)?;
136 Ok(self.position)
137 }
138
139 #[inline]
141 pub const fn position(&self) -> u64 {
142 self.position
143 }
144
145 #[inline]
147 pub const fn size(&self) -> u64 {
148 self.span
149 }
150}
151
152#[cfg(not(target_arch = "wasm32"))]
154type BoxResolvedFuture<M> = Pin<Box<dyn std::future::Future<Output = Resolved<M>> + Send>>;
155#[cfg(target_arch = "wasm32")]
156type BoxResolvedFuture<M> = Pin<Box<dyn std::future::Future<Output = Resolved<M>>>>;
157
158struct Pending<M: JoinMode> {
161 node: SubtreeNode<M>,
162 retries: u32,
163}
164
165enum Resolved<M: JoinMode> {
167 Leaf(u64, Bytes),
169 Children(Vec<SubtreeNode<M>>),
171 Retry(Pending<M>),
173 Failed(FileError),
175}
176
177#[inline]
179fn is_leaf<M: JoinMode, const BS: usize>(node: &SubtreeNode<M>) -> bool {
180 node.span <= BS as u64
181}
182
183async fn fetch_one<G, M, const BS: usize>(
186 getter: &G,
187 chunk_range: &ChunkRange,
188 pending: Pending<M>,
189) -> Resolved<M>
190where
191 G: ChunkGet<BS>,
192 M: JoinMode + MaybeSend + Sync,
193{
194 let node = &pending.node;
195 let body = match super::mode::read_chunk_body::<M, G, BS>(
196 getter,
197 &node.addr,
198 &node.context,
199 node.span,
200 )
201 .await
202 {
203 Ok(body) => body,
204 Err(e) => {
205 if is_leaf::<M, BS>(node) && pending.retries > 0 {
206 return Resolved::Retry(Pending {
207 node: pending.node,
208 retries: pending.retries - 1,
209 });
210 }
211 return Resolved::Failed(e);
212 }
213 };
214
215 if is_leaf::<M, BS>(node) {
216 return Resolved::Leaf(node.byte_offset, body);
217 }
218 match overlapping_children::<M, BS>(&body, node, chunk_range) {
219 Ok(children) => Resolved::Children(children),
220 Err(e) => Resolved::Failed(e),
221 }
222}
223
224fn windowed_walk<G, M, const BODY_SIZE: usize>(
241 getter: Arc<G>,
242 subtrees: Vec<SubtreeNode<M>>,
243 tree: TreeParams<BODY_SIZE>,
244 span: u64,
245 concurrency: usize,
246 start: u64,
247 window: usize,
248) -> impl Stream<Item = Result<Bytes>>
249where
250 G: ChunkGet<BODY_SIZE> + 'static,
251 M: JoinMode + MaybeSend + Sync,
252{
253 let width = concurrency.max(1);
254 let window = window.max(1);
255
256 let range_start = start.min(span);
257 let chunk_range = tree.chunks_for_range(range_start, span - range_start);
258
259 struct State<G, M: JoinMode> {
260 getter: Arc<G>,
261 chunk_range: ChunkRange,
262 range_start: u64,
263 width: usize,
264 window: usize,
265 next_emit_offset: u64,
266 leaf_queue: VecDeque<Pending<M>>,
269 node_queue: VecDeque<Pending<M>>,
271 leaf_in_flight: usize,
273 in_flight: FuturesUnordered<BoxResolvedFuture<M>>,
274 buffered: BTreeMap<u64, Bytes>,
275 }
276
277 let mut leaf_queue = VecDeque::new();
280 let mut node_queue = VecDeque::new();
281 let range_start_byte = chunk_range.start * BODY_SIZE as u64;
282 let range_end_byte = chunk_range.end * BODY_SIZE as u64;
283 for st in subtrees {
284 if st.byte_offset >= range_end_byte || st.byte_offset + st.span <= range_start_byte {
285 continue;
286 }
287 let pending = Pending {
288 node: st,
289 retries: DEFAULT_LEAF_RETRIES,
290 };
291 if is_leaf::<M, BODY_SIZE>(&pending.node) {
292 leaf_queue.push_back(pending);
293 } else {
294 node_queue.push_back(pending);
295 }
296 }
297
298 let state = State::<G, M> {
299 getter,
300 chunk_range,
301 range_start,
302 width,
303 window,
304 next_emit_offset: range_start,
305 leaf_queue,
306 node_queue,
307 leaf_in_flight: 0,
308 in_flight: FuturesUnordered::new(),
309 buffered: BTreeMap::new(),
310 };
311
312 stream::unfold(state, move |mut state| async move {
313 loop {
314 let head_ready = state
318 .buffered
319 .first_key_value()
320 .is_some_and(|(&k, _)| k == state.next_emit_offset);
321 if head_ready {
322 let (_, body) = state.buffered.pop_first().expect("head just observed");
323 state.next_emit_offset += body.len() as u64;
324 return Some((Ok(body), state));
325 }
326
327 loop {
333 if state.in_flight.len() >= state.width {
334 break;
335 }
336 if let Some(pending) = state.node_queue.pop_front() {
337 let getter = Arc::clone(&state.getter);
338 let range = state.chunk_range;
339 state.in_flight.push(Box::pin(async move {
340 fetch_one::<G, M, BODY_SIZE>(&*getter, &range, pending).await
341 }) as BoxResolvedFuture<M>);
342 continue;
343 }
344 if state.leaf_in_flight + state.buffered.len() >= state.window
345 || state.leaf_queue.is_empty()
346 {
347 break;
348 }
349 let pending = state
350 .leaf_queue
351 .pop_front()
352 .expect("non-empty just observed");
353 state.leaf_in_flight += 1;
354 let getter = Arc::clone(&state.getter);
355 let range = state.chunk_range;
356 state.in_flight.push(Box::pin(async move {
357 fetch_one::<G, M, BODY_SIZE>(&*getter, &range, pending).await
358 }) as BoxResolvedFuture<M>);
359 }
360
361 let occupancy = state.leaf_in_flight + state.buffered.len();
363 debug_assert!(
364 occupancy <= state.window,
365 "windowed walk exceeded the leaf-body bound"
366 );
367 #[cfg(test)]
368 record_occupancy(occupancy);
369
370 let resolved = match state.in_flight.next().await {
372 Some(r) => r,
373 None => return None,
374 };
375 match resolved {
376 Resolved::Leaf(leaf_start, body) => {
377 state.leaf_in_flight -= 1;
378 let leaf_end = leaf_start + body.len() as u64;
382 if leaf_end <= state.range_start {
383 continue;
384 }
385 let clip_lo = state.range_start.saturating_sub(leaf_start) as usize;
386 let offset = leaf_start.max(state.range_start);
387 state.buffered.insert(offset, body.slice(clip_lo..));
388 }
389 Resolved::Children(children) => {
390 for child in children {
391 let pending = Pending {
392 node: child,
393 retries: DEFAULT_LEAF_RETRIES,
394 };
395 if is_leaf::<M, BODY_SIZE>(&pending.node) {
396 let pos = state
399 .leaf_queue
400 .partition_point(|p| p.node.byte_offset < pending.node.byte_offset);
401 state.leaf_queue.insert(pos, pending);
402 } else {
403 state.node_queue.push_back(pending);
404 }
405 }
406 }
407 Resolved::Retry(pending) => {
408 state.leaf_in_flight -= 1;
411 let pos = state
412 .leaf_queue
413 .partition_point(|p| p.node.byte_offset < pending.node.byte_offset);
414 state.leaf_queue.insert(pos, pending);
415 }
416 Resolved::Failed(e) => return Some((Err(e), state)),
417 }
418 }
419 })
420}
421
422#[cfg(feature = "tokio")]
425pub struct WindowedJoinerReader<G, M: JoinMode, const BODY_SIZE: usize = DEFAULT_BODY_SIZE>
426where
427 G: ChunkGet<BODY_SIZE>,
428{
429 reader: WindowedReader<G, M, BODY_SIZE>,
430 residual: Bytes,
431 #[allow(clippy::type_complexity)]
432 stream: Option<Pin<Box<dyn Stream<Item = Result<Bytes>> + Send>>>,
433}
434
435#[cfg(feature = "tokio")]
436impl<G, M, const BODY_SIZE: usize> std::fmt::Debug for WindowedJoinerReader<G, M, BODY_SIZE>
437where
438 G: ChunkGet<BODY_SIZE>,
439 M: JoinMode,
440{
441 fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
442 f.debug_struct("WindowedJoinerReader")
443 .field("reader", &self.reader)
444 .field("residual_len", &self.residual.len())
445 .field("has_stream", &self.stream.is_some())
446 .finish()
447 }
448}
449
450#[cfg(feature = "tokio")]
451impl<G, M, const BODY_SIZE: usize> WindowedReader<G, M, BODY_SIZE>
452where
453 G: ChunkGet<BODY_SIZE> + 'static,
454 M: JoinMode + Send + Sync + 'static,
455{
456 pub fn into_reader(self) -> WindowedJoinerReader<G, M, BODY_SIZE> {
458 WindowedJoinerReader {
459 reader: self,
460 residual: Bytes::new(),
461 stream: None,
462 }
463 }
464}
465
466#[cfg(feature = "tokio")]
467impl<G: ChunkGet<BODY_SIZE>, M: JoinMode, const BODY_SIZE: usize> Unpin
468 for WindowedJoinerReader<G, M, BODY_SIZE>
469{
470}
471
472#[cfg(feature = "tokio")]
473impl<G, M, const BODY_SIZE: usize> tokio::io::AsyncRead for WindowedJoinerReader<G, M, BODY_SIZE>
474where
475 G: ChunkGet<BODY_SIZE> + 'static,
476 M: JoinMode + Send + Sync + 'static,
477{
478 fn poll_read(
479 self: Pin<&mut Self>,
480 cx: &mut std::task::Context<'_>,
481 buf: &mut tokio::io::ReadBuf<'_>,
482 ) -> std::task::Poll<std::io::Result<()>> {
483 use bytes::Buf;
484 use std::task::Poll;
485
486 let this = self.get_mut();
487
488 if !this.residual.is_empty() {
489 let to_copy = this.residual.len().min(buf.remaining());
490 buf.put_slice(&this.residual[..to_copy]);
491 this.residual.advance(to_copy);
492 return Poll::Ready(Ok(()));
493 }
494
495 if this.stream.is_none() {
496 let r = &this.reader;
500 this.stream = Some(Box::pin(windowed_walk::<G, M, BODY_SIZE>(
501 Arc::clone(&r.getter),
502 r.subtrees.clone(),
503 r.tree,
504 r.span,
505 r.concurrency,
506 r.position,
507 r.window,
508 )));
509 }
510
511 let stream = this.stream.as_mut().expect("stream just set");
512 match stream.as_mut().poll_next(cx) {
513 Poll::Ready(Some(Ok(body))) => {
514 let to_copy = body.len().min(buf.remaining());
515 buf.put_slice(&body[..to_copy]);
516 this.reader.position += body.len() as u64;
517 if to_copy < body.len() {
518 this.residual = body.slice(to_copy..);
519 }
520 Poll::Ready(Ok(()))
521 }
522 Poll::Ready(Some(Err(e))) => Poll::Ready(Err(std::io::Error::other(e))),
523 Poll::Ready(None) => Poll::Ready(Ok(())),
524 Poll::Pending => Poll::Pending,
525 }
526 }
527}
528
529#[cfg(feature = "tokio")]
530impl<G, M, const BODY_SIZE: usize> tokio::io::AsyncSeek for WindowedJoinerReader<G, M, BODY_SIZE>
531where
532 G: ChunkGet<BODY_SIZE> + 'static,
533 M: JoinMode + Send + Sync + 'static,
534{
535 fn start_seek(self: Pin<&mut Self>, pos: SeekFrom) -> std::io::Result<()> {
536 let this = self.get_mut();
537 this.stream = None;
538 this.residual = Bytes::new();
539 this.reader.seek(pos)?;
540 Ok(())
541 }
542
543 fn poll_complete(
544 self: Pin<&mut Self>,
545 _cx: &mut std::task::Context<'_>,
546 ) -> std::task::Poll<std::io::Result<u64>> {
547 std::task::Poll::Ready(Ok(self.get_mut().reader.position))
548 }
549}
550
551#[cfg(test)]
552mod tests {
553 use super::*;
554 use crate::bmt::DEFAULT_BODY_SIZE;
555 use crate::chunk::AnyChunk;
556 use crate::file::Joiner;
557 use crate::file::split;
558 use futures::executor::block_on;
559 use std::collections::HashMap;
560 use std::sync::atomic::{AtomicUsize, Ordering};
561
562 fn split_and_store(data: &[u8]) -> (ChunkAddress, HashMap<ChunkAddress, AnyChunk>) {
563 let (root, store) = split::<DEFAULT_BODY_SIZE>(data).unwrap();
564 (root, store.into_chunks())
565 }
566
567 async fn drain(
568 reader: &mut WindowedReader<HashMap<ChunkAddress, AnyChunk>, super::super::mode::PlainMode>,
569 ) -> Vec<u8> {
570 let stream = reader.stream();
571 futures::pin_mut!(stream);
572 let mut out = Vec::new();
573 while let Some(item) = stream.next().await {
574 out.extend_from_slice(&item.unwrap());
575 }
576 out
577 }
578
579 #[test]
580 fn stream_from_zero_equals_whole_file() {
581 let data: Vec<u8> = (0..DEFAULT_BODY_SIZE * 5 + 321)
582 .map(|i| (i % 256) as u8)
583 .collect();
584 let (root, store) = split_and_store(&data);
585 block_on(async {
586 let joiner = Joiner::new(store, root).await.unwrap();
587 let mut reader = joiner.into_windowed_reader(16);
588 let got = drain(&mut reader).await;
589 assert_eq!(got, data);
590 });
591 }
592
593 #[test]
594 fn seek_then_stream_yields_suffix() {
595 let data: Vec<u8> = (0..DEFAULT_BODY_SIZE * 5 + 321)
596 .map(|i| (i % 256) as u8)
597 .collect();
598 let bs = DEFAULT_BODY_SIZE as u64;
599 let cases = [bs + 10, bs, bs * 3, bs / 2, data.len() as u64 - 1];
600 let (root, store) = split_and_store(&data);
601 block_on(async {
602 for k in cases {
603 let joiner = Joiner::new(store.clone(), root).await.unwrap();
604 let mut reader = joiner.into_windowed_reader(16);
605 reader.seek(SeekFrom::Start(k)).unwrap();
606 let got = drain(&mut reader).await;
607 assert_eq!(got, &data[k as usize..], "suffix from {k}");
608 }
609 });
610 }
611
612 #[test]
613 fn backward_seek_then_read_is_correct() {
614 let data: Vec<u8> = (0..DEFAULT_BODY_SIZE * 4 + 17)
615 .map(|i| (i % 256) as u8)
616 .collect();
617 let bs = DEFAULT_BODY_SIZE as u64;
618 let (root, store) = split_and_store(&data);
619 block_on(async {
620 let joiner = Joiner::new(store, root).await.unwrap();
621 let mut reader = joiner.into_windowed_reader(16);
622 reader.seek(SeekFrom::Start(bs * 3)).unwrap();
623 let _ = drain(&mut reader).await;
624 reader.seek(SeekFrom::Start(bs)).unwrap();
625 let got = drain(&mut reader).await;
626 assert_eq!(got, &data[bs as usize..]);
627 });
628 }
629
630 #[test]
631 fn width_one_correct() {
632 let data: Vec<u8> = (0..DEFAULT_BODY_SIZE * 5 + 7)
633 .map(|i| (i % 256) as u8)
634 .collect();
635 let (root, store) = split_and_store(&data);
636 block_on(async {
637 let joiner = Joiner::new(store, root).await.unwrap().with_concurrency(1);
638 let mut reader = joiner.into_windowed_reader(1);
639 let got = drain(&mut reader).await;
640 assert_eq!(got, data);
641 });
642 }
643
644 async fn yield_now() {
647 struct YieldNow(bool);
648 impl std::future::Future for YieldNow {
649 type Output = ();
650 fn poll(
651 mut self: Pin<&mut Self>,
652 cx: &mut std::task::Context<'_>,
653 ) -> std::task::Poll<()> {
654 if self.0 {
655 std::task::Poll::Ready(())
656 } else {
657 self.0 = true;
658 cx.waker().wake_by_ref();
659 std::task::Poll::Pending
660 }
661 }
662 }
663 YieldNow(false).await
664 }
665
666 #[derive(Clone)]
669 struct BoundProbe {
670 chunks: Arc<HashMap<ChunkAddress, AnyChunk>>,
671 in_flight: Arc<AtomicUsize>,
672 max_in_flight: Arc<AtomicUsize>,
673 }
674
675 impl ChunkGet<DEFAULT_BODY_SIZE> for BoundProbe {
676 type Error = crate::store::ChunkStoreError;
677
678 async fn get(&self, address: &ChunkAddress) -> std::result::Result<AnyChunk, Self::Error> {
679 let now = self.in_flight.fetch_add(1, Ordering::SeqCst) + 1;
680 self.max_in_flight.fetch_max(now, Ordering::SeqCst);
681 for _ in 0..4 {
684 yield_now().await;
685 }
686 self.in_flight.fetch_sub(1, Ordering::SeqCst);
687 self.chunks
688 .get(address)
689 .cloned()
690 .ok_or_else(|| crate::store::ChunkStoreError::not_found(address))
691 }
692 }
693
694 #[test]
695 fn reorder_buffer_never_exceeds_window() {
696 let leaves = 40usize;
699 let window = 6usize;
700 let data: Vec<u8> = (0..DEFAULT_BODY_SIZE * leaves)
701 .map(|i| (i % 256) as u8)
702 .collect();
703 let (root, store) = split_and_store(&data);
704
705 let probe = BoundProbe {
706 chunks: Arc::new(store),
707 in_flight: Arc::new(AtomicUsize::new(0)),
708 max_in_flight: Arc::new(AtomicUsize::new(0)),
709 };
710 let max_seen = Arc::clone(&probe.max_in_flight);
711
712 block_on(async {
713 let joiner = Joiner::new(probe, root)
714 .await
715 .unwrap()
716 .with_concurrency(leaves);
717 let mut reader = joiner.into_windowed_reader(window);
718 let stream = reader.stream();
719 futures::pin_mut!(stream);
720 let mut total = 0usize;
721 while let Some(item) = stream.next().await {
722 total += item.unwrap().len();
723 }
724 assert_eq!(total, data.len());
725 });
726
727 let peak = max_seen.load(Ordering::SeqCst);
732 assert!(
733 peak <= window + 2,
734 "in-flight peak {peak} exceeds window {window} (+slack)"
735 );
736 }
737
738 #[derive(Clone)]
745 struct HeadSlow {
746 chunks: Arc<HashMap<ChunkAddress, AnyChunk>>,
747 slow_offset: u64,
749 released: Arc<AtomicUsize>,
751 gate: usize,
753 leaf_offsets: Arc<HashMap<ChunkAddress, u64>>,
755 }
756
757 impl ChunkGet<DEFAULT_BODY_SIZE> for HeadSlow {
758 type Error = crate::store::ChunkStoreError;
759
760 async fn get(&self, address: &ChunkAddress) -> std::result::Result<AnyChunk, Self::Error> {
761 let is_slow = self.leaf_offsets.get(address) == Some(&self.slow_offset);
762 if is_slow {
763 while self.released.load(Ordering::SeqCst) < self.gate {
766 yield_now().await;
767 }
768 for _ in 0..4 {
769 yield_now().await;
770 }
771 } else {
772 for _ in 0..4 {
773 yield_now().await;
774 }
775 if self.leaf_offsets.contains_key(address) {
776 self.released.fetch_add(1, Ordering::SeqCst);
777 }
778 }
779
780 self.chunks
781 .get(address)
782 .cloned()
783 .ok_or_else(|| crate::store::ChunkStoreError::not_found(address))
784 }
785 }
786
787 fn unique_byte(i: u64) -> u8 {
792 (i.wrapping_mul(2_654_435_761) >> 11) as u8
793 }
794
795 fn leaf_offsets_of(total: u64) -> HashMap<ChunkAddress, u64> {
799 use crate::chunk::Chunk;
800 let mut map = HashMap::new();
801 let leaves = total.div_ceil(DEFAULT_BODY_SIZE as u64);
802 for i in 0..leaves {
803 let off = i * DEFAULT_BODY_SIZE as u64;
804 let end = (off + DEFAULT_BODY_SIZE as u64).min(total);
805 let body: Vec<u8> = (off..end).map(unique_byte).collect();
806 let chunk = crate::chunk::ContentChunk::<DEFAULT_BODY_SIZE>::new(body).unwrap();
807 map.insert(*chunk.address(), off);
808 }
809 map
810 }
811
812 #[test]
813 fn head_slowest_in_order_and_bounded() {
814 let leaves = 40usize;
824 let window = 6usize;
825 let total = (DEFAULT_BODY_SIZE * leaves) as u64;
826 let data: Vec<u8> = (0..total).map(unique_byte).collect();
827 let (root, store) = split_and_store(&data);
828
829 PEAK_LEAF_OCCUPANCY.with(|p| p.set(0));
830 let leaf_offsets = leaf_offsets_of(total);
831 let getter = HeadSlow {
832 chunks: Arc::new(store),
833 slow_offset: 0,
834 released: Arc::new(AtomicUsize::new(0)),
835 gate: window - 1,
836 leaf_offsets: Arc::new(leaf_offsets),
837 };
838
839 block_on(async {
840 let joiner = Joiner::new(getter, root)
841 .await
842 .unwrap()
843 .with_concurrency(leaves);
844 let mut reader = joiner.into_windowed_reader(window);
845 let stream = reader.stream();
846 futures::pin_mut!(stream);
847 let mut got = Vec::new();
848 while let Some(item) = stream.next().await {
849 got.extend_from_slice(&item.unwrap());
850 }
851 assert_eq!(got, data, "head-slowest still yields whole file in order");
852 });
853
854 let peak = PEAK_LEAF_OCCUPANCY.with(std::cell::Cell::get);
855 assert!(
856 peak <= window,
857 "leaf-body occupancy peak {peak} exceeds window {window}"
858 );
859 }
860
861 #[cfg(feature = "encryption")]
862 #[test]
863 fn head_slowest_encrypted_in_order_and_bounded() {
864 use crate::file::EncryptedJoiner;
871 use crate::file::split_encrypted;
872
873 let leaves = 30usize;
874 let window = 5usize;
875 let total = (DEFAULT_BODY_SIZE * leaves) as u64;
876 let data: Vec<u8> = (0..total).map(unique_byte).collect();
877 let (root_ref, store) = split_encrypted::<DEFAULT_BODY_SIZE>(&data).unwrap();
878 let store = store.into_chunks();
879
880 PEAK_LEAF_OCCUPANCY.with(|p| p.set(0));
881 let getter = UniformSlow {
882 chunks: Arc::new(store),
883 };
884
885 block_on(async {
886 let joiner = EncryptedJoiner::new(getter, root_ref)
887 .await
888 .unwrap()
889 .with_concurrency(leaves);
890 let mut reader = joiner.into_windowed_reader(window);
891 let stream = reader.stream();
892 futures::pin_mut!(stream);
893 let mut out = Vec::new();
894 while let Some(item) = stream.next().await {
895 out.extend_from_slice(&item.unwrap());
896 }
897 assert_eq!(out, data, "encrypted short-body slide stays in order");
898 });
899
900 let peak = PEAK_LEAF_OCCUPANCY.with(std::cell::Cell::get);
901 assert!(
902 peak <= window,
903 "encrypted leaf-body occupancy peak {peak} exceeds window {window}"
904 );
905 }
906
907 #[cfg(feature = "encryption")]
912 #[derive(Clone)]
913 struct UniformSlow {
914 chunks: Arc<HashMap<ChunkAddress, AnyChunk>>,
915 }
916
917 #[cfg(feature = "encryption")]
918 impl ChunkGet<DEFAULT_BODY_SIZE> for UniformSlow {
919 type Error = crate::store::ChunkStoreError;
920
921 async fn get(&self, address: &ChunkAddress) -> std::result::Result<AnyChunk, Self::Error> {
922 for _ in 0..4 {
923 yield_now().await;
924 }
925 self.chunks
926 .get(address)
927 .cloned()
928 .ok_or_else(|| crate::store::ChunkStoreError::not_found(address))
929 }
930 }
931
932 #[cfg(feature = "encryption")]
933 mod encrypted {
934 use super::*;
935 use crate::file::EncryptedJoiner;
936 use crate::file::split_encrypted;
937
938 async fn drain_enc(
939 reader: &mut WindowedReader<
940 HashMap<ChunkAddress, AnyChunk>,
941 crate::file::mode::EncryptedMode,
942 >,
943 ) -> Vec<u8> {
944 let stream = reader.stream();
945 futures::pin_mut!(stream);
946 let mut out = Vec::new();
947 while let Some(item) = stream.next().await {
948 out.extend_from_slice(&item.unwrap());
949 }
950 out
951 }
952
953 #[test]
954 fn encrypted_stream_and_seek() {
955 let data: Vec<u8> = (0..DEFAULT_BODY_SIZE * 5 + 321)
956 .map(|i| (i % 256) as u8)
957 .collect();
958 let bs = DEFAULT_BODY_SIZE as u64;
959 let (root_ref, store) = split_encrypted::<DEFAULT_BODY_SIZE>(&data).unwrap();
960 let store = store.into_chunks();
961 block_on(async {
962 let joiner = EncryptedJoiner::new(store.clone(), root_ref.clone())
964 .await
965 .unwrap();
966 let mut reader = joiner.into_windowed_reader(16);
967 assert_eq!(drain_enc(&mut reader).await, data);
968
969 for k in [bs + 10, bs, bs * 3] {
971 let joiner = EncryptedJoiner::new(store.clone(), root_ref.clone())
972 .await
973 .unwrap();
974 let mut reader = joiner.into_windowed_reader(16);
975 reader.seek(SeekFrom::Start(k)).unwrap();
976 assert_eq!(drain_enc(&mut reader).await, &data[k as usize..]);
977 }
978 });
979 }
980 }
981
982 #[cfg(feature = "tokio")]
983 #[tokio::test]
984 async fn reader_seek_and_read() {
985 use tokio::io::{AsyncReadExt, AsyncSeekExt};
986
987 let data: Vec<u8> = (0..DEFAULT_BODY_SIZE * 3 + 50)
988 .map(|i| (i % 256) as u8)
989 .collect();
990 let (root, store) = split_and_store(&data);
991 let joiner = Joiner::new(store, root).await.unwrap();
992 let mut reader = joiner.into_windowed_reader(16).into_reader();
993
994 let mut all = Vec::new();
995 reader.read_to_end(&mut all).await.unwrap();
996 assert_eq!(all, data);
997
998 reader
999 .seek(SeekFrom::Start(DEFAULT_BODY_SIZE as u64))
1000 .await
1001 .unwrap();
1002 let mut tail = Vec::new();
1003 reader.read_to_end(&mut tail).await.unwrap();
1004 assert_eq!(tail, &data[DEFAULT_BODY_SIZE..]);
1005 }
1006}