Skip to main content

nectar_primitives/file/
windowed.rs

1//! Bounded seek-and-play reader over the async joiner.
2//!
3//! [`WindowedReader`] seeks to a byte offset then streams forward in file order
4//! with peak memory bounded by `window` leaf bodies rather than the file size.
5//! The window slides with the read cursor: at most `window` leaf fetches plus
6//! buffered leaves are held at once, and the lowest-offset (head) leaf is always
7//! fetched first, so emission makes progress and the bound holds regardless of
8//! resolve order. Reseeking drops the window and repositions against the
9//! retained frontier, so no intermediate chunk is re-fetched.
10
11use 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
29/// Number of times a failed leaf fetch is re-enqueued before the walk surfaces
30/// the error.
31const DEFAULT_LEAF_RETRIES: u32 = 4;
32
33#[cfg(test)]
34thread_local! {
35    /// Test hook: peak observed `leaf_in_flight + buffered.len()` across a walk
36    /// on this thread, so a test can assert the true window bound on leaf bodies
37    /// held at once. Thread-local because `block_on` runs each test's walk on its
38    /// own thread, so the parallel test runner never interleaves the counter.
39    static PEAK_LEAF_OCCUPANCY: std::cell::Cell<usize> = const { std::cell::Cell::new(0) };
40}
41
42/// Record a sampled occupancy, keeping the running peak for this thread.
43#[cfg(test)]
44fn record_occupancy(occupancy: usize) {
45    PEAK_LEAF_OCCUPANCY.with(|p| p.set(p.get().max(occupancy)));
46}
47
48/// Seek-and-play reader: forward in-order delivery over a window that slides
49/// with the read cursor, bounded to `window` in-flight plus buffered leaves.
50///
51/// Retains the joiner's reusable state (getter, pre-expanded frontier, tree
52/// params) so each [`stream`](Self::stream) rebuilds a fresh window-bounded walk
53/// from `[position, size)` without re-expanding the frontier. [`seek`](Self::seek)
54/// only repositions; it never re-fetches intermediates.
55pub 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    /// Seek-and-play reader over a window that slides with the read cursor. Peak
93    /// memory is `window` leaf bodies: in-flight plus buffered leaves never
94    /// exceed `window`. `window` is clamped to at least 1; a value of at least
95    /// `2 * concurrency` keeps the pool full even while the head leaf is the
96    /// straggler.
97    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    /// In-order leaf bodies from the current position to EOF. Each item is one
117    /// contiguous run; reassembly is pure concatenation. Peak memory is `window`
118    /// leaf bodies: the window slides with the emit cursor, so only leaves within
119    /// `window` spans ahead of it are ever fetched or held.
120    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    /// Reposition: drop the window and restart the walk from `pos`. Cheap; reuses
133    /// the frontier, no intermediate re-fetch.
134    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    /// Current read position.
140    #[inline]
141    pub const fn position(&self) -> u64 {
142        self.position
143    }
144
145    /// Total file size.
146    #[inline]
147    pub const fn size(&self) -> u64 {
148        self.span
149    }
150}
151
152/// Cfg-gated boxed resolved-future alias: `+ Send` on native, unbounded on wasm.
153#[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
158/// One unit of pending work: a tree node plus its remaining retry budget. The
159/// budget only decrements on a failed leaf fetch.
160struct Pending<M: JoinMode> {
161    node: SubtreeNode<M>,
162    retries: u32,
163}
164
165/// What a worker future resolves to once its chunk lands.
166enum Resolved<M: JoinMode> {
167    /// A leaf: its absolute byte offset and decoded body.
168    Leaf(u64, Bytes),
169    /// An intermediate: its overlapping children to re-queue.
170    Children(Vec<SubtreeNode<M>>),
171    /// A leaf fetch failed with retries left: re-queue this node.
172    Retry(Pending<M>),
173    /// A fetch failed terminally (retries exhausted or intermediate error).
174    Failed(FileError),
175}
176
177/// Is this node a leaf?
178#[inline]
179fn is_leaf<M: JoinMode, const BS: usize>(node: &SubtreeNode<M>) -> bool {
180    node.span <= BS as u64
181}
182
183/// Fetch one node: a leaf yields its body, an intermediate yields its children.
184/// A leaf error consumes one retry, then re-queues or fails.
185async 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
224/// Window-bounded in-order walk from `start` to EOF.
225///
226/// A single walk integrating fetch and reorder: intermediate nodes are always
227/// descended (few and small), but a leaf is moved from the pending queue into
228/// the in-flight pool only while in-flight plus buffered leaves are below
229/// `window`, and always the lowest-offset pending leaf first. Resolved leaves
230/// land in a `BTreeMap` reorder buffer keyed by absolute offset; the head is
231/// emitted when its offset equals `next_emit_offset`, which then advances by
232/// `body.len()` (never a fixed stride, so shorter encrypted leaves stay aligned)
233/// and frees a slot, sliding the window forward.
234///
235/// The head leaf is always the lowest-offset pending leaf, so it is admitted as
236/// soon as a slot is free and emission always makes progress: no deadlock. The
237/// count gate enforces the invariant that in-flight leaf fetches plus buffered
238/// leaves never exceed `window`, so peak memory is `window` leaf bodies
239/// regardless of resolve order, leaf span, or file size.
240fn 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        /// Leaves not yet admitted to the pool, kept in offset order so the
267        /// lowest-offset (head) leaf is always admitted first.
268        leaf_queue: VecDeque<Pending<M>>,
269        /// Intermediate nodes pending descent; always admissible.
270        node_queue: VecDeque<Pending<M>>,
271        /// Count of in-flight leaf fetches (excludes intermediates).
272        leaf_in_flight: usize,
273        in_flight: FuturesUnordered<BoxResolvedFuture<M>>,
274        buffered: BTreeMap<u64, Bytes>,
275    }
276
277    // Seed the queues with the subtrees overlapping `[range_start, span)`,
278    // separating leaves (offset-gated) from intermediates (always admissible).
279    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            // Emit the head leaf as soon as it is buffered, advancing the cursor
315            // and sliding the window so the refill below admits newly-in-range
316            // leaves on the next turn.
317            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            // Refill the pool: intermediates always; a leaf only while a leaf slot
328            // is free, that is while in-flight plus buffered leaves are below
329            // `window`, and always the lowest-offset (head) leaf first. The head
330            // is therefore admitted as soon as a slot frees, so progress is
331            // guaranteed and at most `window` leaf bodies are ever held.
332            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            // In-flight leaf fetches plus buffered leaves never exceed `window`.
362            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            // Pool empty and nothing left to admit: the file is drained.
371            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                    // Clip the first partial leaf at the read position; offsets
379                    // stay absolute, so a boundary leaf buffers only its in-range
380                    // tail keyed at the read position.
381                    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                            // Keep the leaf queue ordered by offset so the front
397                            // is always the lowest, which the window gate tests.
398                            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                    // Re-enqueue in offset order so the gate still sees the lowest
409                    // leaf at the front and the slot it held is freed for refill.
410                    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/// Native `AsyncRead` + `AsyncSeek` shim, draining [`WindowedReader::stream`]
423/// into a residual buffer.
424#[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    /// Wrap as a tokio `AsyncRead` + `AsyncSeek` reader. Native-only.
457    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            // Rebuild an owned `'static` stream from the retained joiner parts,
497            // so the boxed stream does not borrow `this`. Cloning the parts is the
498            // same cost as `stream()`; no frontier re-expansion.
499            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    /// A self-waking yield: returns `Pending` once and immediately schedules a
645    /// re-poll, so it makes progress under `futures::executor::block_on`.
646    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    /// A getter that records peak concurrent fetches so a test can assert the
667    /// reorder buffer never admits more than `window` leaves at once.
668    #[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            // Hold the fetch open across several self-waking yields so
682            // concurrently admitted leaves overlap here before any resolves.
683            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        // Many leaves, small window: leaf fetches overlap, so without the bound
697        // the in-flight peak would climb to the leaf count. The window caps it.
698        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        // The reorder buffer admits at most `window` leaves, so concurrent leaf
728        // fetches (plus the intermediate node fetches that seed children) never
729        // climb far above the window. Allow a small slack for in-flight
730        // intermediates, but assert it is nowhere near the leaf count.
731        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    /// A getter that delays the leaf at `slow_offset` until `gate` other leaves
739    /// have resolved. This is the adversarial resolve order the original reorder
740    /// buffer mishandled: with the head leaf landing last, the old buffer kept
741    /// fetching and buffering later leaves, so resident leaf bodies (in-flight
742    /// plus buffered) climbed to roughly twice the window. The sliding window
743    /// refuses to fetch past a free slot, holding the total at the window.
744    #[derive(Clone)]
745    struct HeadSlow {
746        chunks: Arc<HashMap<ChunkAddress, AnyChunk>>,
747        /// File offset of the head leaf to delay.
748        slow_offset: u64,
749        /// Number of leaves released so far (used to gate the slow leaf).
750        released: Arc<AtomicUsize>,
751        /// How many other leaves must release before the slow leaf may.
752        gate: usize,
753        /// Address-to-offset map so the getter can recognise the head leaf.
754        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                // Park until the gate of other leaves has released, then a few
764                // more yields so a buffer would overflow if it were unbounded.
765                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    /// Distinct byte per file position, so every leaf body (and thus its
788    /// content address) is unique. A `(i % 256)` fill would make all full leaves
789    /// identical, collapsing the address-to-offset map and defeating the
790    /// per-leaf delay below.
791    fn unique_byte(i: u64) -> u8 {
792        (i.wrapping_mul(2_654_435_761) >> 11) as u8
793    }
794
795    /// Map each leaf chunk address to its file offset. Leaf bodies are
796    /// content-addressed, so rebuilding each leaf from `unique_byte` reproduces
797    /// the stored address.
798    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        // The head leaf at `position` resolves last among the leaves in its
815        // window. Under the old buffer this leaks: with the head missing, the
816        // walk kept `window` leaf fetches in flight while also buffering the
817        // resolved stragglers, so resident leaf bodies reached roughly twice the
818        // window. The sliding window keeps in-flight plus buffered at `window` and
819        // delivery in order. A naive run at `gate = window` instead would deadlock
820        // the correct walk (the head can never release), so the gate is the most a
821        // non-deadlocking adversary can demand, and the old buffer still overruns
822        // it.
823        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        // Encrypted leaf bodies are shorter than BODY_SIZE, so the slide must
865        // advance by body.len(), not the stride. Encrypted leaf addresses hash
866        // the ciphertext and are not reconstructible from plaintext here, so the
867        // head cannot be singled out; a uniform delay across all fetches still
868        // overlaps concurrently admitted leaves, and the occupancy hook proves a
869        // correct walk holds at most `window` leaf bodies on short bodies.
870        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    /// A getter that holds every fetch open across several self-waking yields so
908    /// concurrently admitted leaves overlap. Used where addresses cannot be
909    /// mapped to offsets (encrypted leaves), so the head cannot be singled out;
910    /// the occupancy hook still proves the window bound.
911    #[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                // whole file
963                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                // mid-leaf and leaf-aligned suffixes
970                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}