use std::collections::{BTreeMap, VecDeque};
use std::io::SeekFrom;
use std::pin::Pin;
use std::sync::Arc;
use bytes::Bytes;
use futures::stream::{self, FuturesUnordered, Stream, StreamExt};
use crate::bmt::DEFAULT_BODY_SIZE;
use crate::chunk::ChunkAddress;
use super::error::{FileError, Result};
use super::frontier::{SubtreeNode, overlapping_children};
use super::joiner::{GenericJoiner, MAX_INTERMEDIATE_IN_FLIGHT};
use super::mode::JoinMode;
use super::tree::{ChunkRange, TreeParams};
use crate::store::{ChunkGet, MaybeSend};
const DEFAULT_LEAF_RETRIES: u32 = 4;
#[cfg(test)]
thread_local! {
static PEAK_LEAF_OCCUPANCY: std::cell::Cell<usize> = const { std::cell::Cell::new(0) };
}
#[cfg(test)]
fn record_occupancy(occupancy: usize) {
PEAK_LEAF_OCCUPANCY.with(|p| p.set(p.get().max(occupancy)));
}
pub struct WindowedReader<G, M: JoinMode, const BODY_SIZE: usize = DEFAULT_BODY_SIZE>
where
G: ChunkGet<BODY_SIZE>,
{
getter: Arc<G>,
subtrees: Vec<SubtreeNode<M>>,
tree: TreeParams<BODY_SIZE>,
span: u64,
concurrency: usize,
position: u64,
window: usize,
#[allow(
dead_code,
reason = "retained reusable joiner state for re-stream-on-seek"
)]
root: ChunkAddress,
}
impl<G, M, const BODY_SIZE: usize> std::fmt::Debug for WindowedReader<G, M, BODY_SIZE>
where
G: ChunkGet<BODY_SIZE>,
M: JoinMode,
{
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
f.debug_struct("WindowedReader")
.field("span", &self.span)
.field("position", &self.position)
.field("concurrency", &self.concurrency)
.finish_non_exhaustive()
}
}
impl<G, M, const BODY_SIZE: usize> GenericJoiner<G, M, BODY_SIZE>
where
G: ChunkGet<BODY_SIZE> + 'static,
M: JoinMode + MaybeSend + Sync,
{
pub fn into_windowed_reader(self, window: usize) -> WindowedReader<G, M, BODY_SIZE> {
WindowedReader {
getter: Arc::clone(self.getter()),
subtrees: self.subtrees().to_vec(),
tree: self.tree(),
span: self.size(),
concurrency: self.concurrency(),
position: self.position(),
root: *self.root(),
window: window.max(1),
}
}
}
impl<G, M, const BODY_SIZE: usize> WindowedReader<G, M, BODY_SIZE>
where
G: ChunkGet<BODY_SIZE> + 'static,
M: JoinMode + MaybeSend + Sync,
{
pub fn stream(&mut self) -> impl Stream<Item = Result<Bytes>> + '_ {
windowed_walk::<G, M, BODY_SIZE>(
Arc::clone(&self.getter),
self.subtrees.clone(),
self.tree,
self.span,
self.concurrency,
self.position,
self.window,
)
}
pub fn seek(&mut self, pos: SeekFrom) -> std::io::Result<u64> {
self.position = super::resolve_seek_position(pos, self.position, self.span)?;
Ok(self.position)
}
#[inline]
pub const fn position(&self) -> u64 {
self.position
}
#[inline]
pub const fn size(&self) -> u64 {
self.span
}
}
#[cfg(not(target_arch = "wasm32"))]
type BoxResolvedFuture<M> = Pin<Box<dyn std::future::Future<Output = Resolved<M>> + Send>>;
#[cfg(target_arch = "wasm32")]
type BoxResolvedFuture<M> = Pin<Box<dyn std::future::Future<Output = Resolved<M>>>>;
struct Pending<M: JoinMode> {
node: SubtreeNode<M>,
retries: u32,
}
enum Resolved<M: JoinMode> {
Leaf(u64, Bytes),
Children(u64, Vec<SubtreeNode<M>>),
Retry(Pending<M>),
Failed(FileError),
}
#[inline]
fn is_leaf<M: JoinMode, const BS: usize>(node: &SubtreeNode<M>) -> bool {
node.span <= BS as u64
}
async fn fetch_one<G, M, const BS: usize>(
getter: &G,
chunk_range: &ChunkRange,
pending: Pending<M>,
) -> Resolved<M>
where
G: ChunkGet<BS>,
M: JoinMode + MaybeSend + Sync,
{
let node = &pending.node;
let body = match super::mode::read_chunk_body::<M, G, BS>(
getter,
&node.addr,
&node.context,
node.span,
)
.await
{
Ok(body) => body,
Err(e) => {
if is_leaf::<M, BS>(node) && pending.retries > 0 {
return Resolved::Retry(Pending {
node: pending.node,
retries: pending.retries - 1,
});
}
return Resolved::Failed(e);
}
};
if is_leaf::<M, BS>(node) {
return Resolved::Leaf(node.byte_offset, body);
}
let parent_offset = node.byte_offset;
match overlapping_children::<M, BS>(&body, node, chunk_range) {
Ok(children) => Resolved::Children(parent_offset, children),
Err(e) => Resolved::Failed(e),
}
}
fn insert_by_offset<M: JoinMode>(queue: &mut VecDeque<Pending<M>>, pending: Pending<M>) {
let pos = queue.partition_point(|p| p.node.byte_offset < pending.node.byte_offset);
queue.insert(pos, pending);
}
fn windowed_walk<G, M, const BODY_SIZE: usize>(
getter: Arc<G>,
subtrees: Vec<SubtreeNode<M>>,
tree: TreeParams<BODY_SIZE>,
span: u64,
concurrency: usize,
start: u64,
window: usize,
) -> impl Stream<Item = Result<Bytes>>
where
G: ChunkGet<BODY_SIZE> + 'static,
M: JoinMode + MaybeSend + Sync,
{
let width = concurrency.max(1);
let window = window.max(1);
let range_start = start.min(span);
let chunk_range = tree.chunks_for_range(range_start, span - range_start);
struct State<G, M: JoinMode> {
getter: Arc<G>,
chunk_range: ChunkRange,
range_start: u64,
width: usize,
window: usize,
next_emit_offset: u64,
leaf_queue: VecDeque<Pending<M>>,
node_queue: VecDeque<Pending<M>>,
pending_intermediate_offsets: BTreeMap<u64, usize>,
intermediate_in_flight: usize,
leaf_in_flight: usize,
in_flight: FuturesUnordered<BoxResolvedFuture<M>>,
buffered: BTreeMap<u64, Bytes>,
}
let mut leaf_queue = VecDeque::new();
let mut node_queue = VecDeque::new();
let mut pending_intermediate_offsets: BTreeMap<u64, usize> = BTreeMap::new();
let range_start_byte = chunk_range.start * BODY_SIZE as u64;
let range_end_byte = chunk_range.end * BODY_SIZE as u64;
for st in subtrees {
if st.byte_offset >= range_end_byte || st.byte_offset + st.span <= range_start_byte {
continue;
}
let pending = Pending {
node: st,
retries: DEFAULT_LEAF_RETRIES,
};
if is_leaf::<M, BODY_SIZE>(&pending.node) {
insert_by_offset(&mut leaf_queue, pending);
} else {
*pending_intermediate_offsets
.entry(pending.node.byte_offset)
.or_insert(0) += 1;
insert_by_offset(&mut node_queue, pending);
}
}
let state = State::<G, M> {
getter,
chunk_range,
range_start,
width,
window,
next_emit_offset: range_start,
leaf_queue,
node_queue,
pending_intermediate_offsets,
intermediate_in_flight: 0,
leaf_in_flight: 0,
in_flight: FuturesUnordered::new(),
buffered: BTreeMap::new(),
};
stream::unfold(state, move |mut state| async move {
loop {
let head_ready = state
.buffered
.first_key_value()
.is_some_and(|(&k, _)| k == state.next_emit_offset);
if head_ready {
let (_, body) = state.buffered.pop_first().expect("head just observed");
state.next_emit_offset += body.len() as u64;
return Some((Ok(body), state));
}
loop {
if state.in_flight.len() >= state.width {
break;
}
let min_pending_intermediate =
state.pending_intermediate_offsets.keys().next().copied();
let leaf_admissible = state.leaf_queue.front().is_some_and(|p| {
min_pending_intermediate.is_none_or(|m| p.node.byte_offset < m)
}) && state.leaf_in_flight + state.buffered.len()
< state.window;
let can_admit_intermediate = state.intermediate_in_flight
< MAX_INTERMEDIATE_IN_FLIGHT
&& !state.node_queue.is_empty();
let admit_intermediate = can_admit_intermediate
&& (!leaf_admissible || state.intermediate_in_flight + 1 < state.width);
let pending = if admit_intermediate {
state.intermediate_in_flight += 1;
state.node_queue.pop_front().expect("node queue non-empty")
} else if leaf_admissible {
state.leaf_in_flight += 1;
state.leaf_queue.pop_front().expect("front admissible")
} else {
break;
};
let getter = Arc::clone(&state.getter);
let range = state.chunk_range;
state.in_flight.push(Box::pin(async move {
fetch_one::<G, M, BODY_SIZE>(&*getter, &range, pending).await
}) as BoxResolvedFuture<M>);
}
let occupancy = state.leaf_in_flight + state.buffered.len();
debug_assert!(
occupancy <= state.window,
"windowed walk exceeded the leaf-body bound"
);
#[cfg(test)]
record_occupancy(occupancy);
let resolved = match state.in_flight.next().await {
Some(r) => r,
None => return None,
};
match resolved {
Resolved::Leaf(leaf_start, body) => {
state.leaf_in_flight -= 1;
let leaf_end = leaf_start + body.len() as u64;
if leaf_end <= state.range_start {
continue;
}
let clip_lo = state.range_start.saturating_sub(leaf_start) as usize;
let offset = leaf_start.max(state.range_start);
state.buffered.insert(offset, body.slice(clip_lo..));
}
Resolved::Children(parent_offset, children) => {
state.intermediate_in_flight -= 1;
if let Some(count) = state.pending_intermediate_offsets.get_mut(&parent_offset)
{
*count -= 1;
if *count == 0 {
state.pending_intermediate_offsets.remove(&parent_offset);
}
}
for child in children {
let pending = Pending {
node: child,
retries: DEFAULT_LEAF_RETRIES,
};
if is_leaf::<M, BODY_SIZE>(&pending.node) {
insert_by_offset(&mut state.leaf_queue, pending);
} else {
*state
.pending_intermediate_offsets
.entry(pending.node.byte_offset)
.or_insert(0) += 1;
insert_by_offset(&mut state.node_queue, pending);
}
}
}
Resolved::Retry(pending) => {
state.leaf_in_flight -= 1;
insert_by_offset(&mut state.leaf_queue, pending);
}
Resolved::Failed(e) => return Some((Err(e), state)),
}
}
})
}
#[cfg(feature = "tokio")]
pub struct WindowedJoinerReader<G, M: JoinMode, const BODY_SIZE: usize = DEFAULT_BODY_SIZE>
where
G: ChunkGet<BODY_SIZE>,
{
reader: WindowedReader<G, M, BODY_SIZE>,
residual: Bytes,
#[allow(clippy::type_complexity)]
stream: Option<Pin<Box<dyn Stream<Item = Result<Bytes>> + Send>>>,
}
#[cfg(feature = "tokio")]
impl<G, M, const BODY_SIZE: usize> std::fmt::Debug for WindowedJoinerReader<G, M, BODY_SIZE>
where
G: ChunkGet<BODY_SIZE>,
M: JoinMode,
{
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
f.debug_struct("WindowedJoinerReader")
.field("reader", &self.reader)
.field("residual_len", &self.residual.len())
.field("has_stream", &self.stream.is_some())
.finish()
}
}
#[cfg(feature = "tokio")]
impl<G, M, const BODY_SIZE: usize> WindowedReader<G, M, BODY_SIZE>
where
G: ChunkGet<BODY_SIZE> + 'static,
M: JoinMode + Send + Sync + 'static,
{
pub fn into_reader(self) -> WindowedJoinerReader<G, M, BODY_SIZE> {
WindowedJoinerReader {
reader: self,
residual: Bytes::new(),
stream: None,
}
}
}
#[cfg(feature = "tokio")]
impl<G: ChunkGet<BODY_SIZE>, M: JoinMode, const BODY_SIZE: usize> Unpin
for WindowedJoinerReader<G, M, BODY_SIZE>
{
}
#[cfg(feature = "tokio")]
impl<G, M, const BODY_SIZE: usize> tokio::io::AsyncRead for WindowedJoinerReader<G, M, BODY_SIZE>
where
G: ChunkGet<BODY_SIZE> + 'static,
M: JoinMode + Send + Sync + 'static,
{
fn poll_read(
self: Pin<&mut Self>,
cx: &mut std::task::Context<'_>,
buf: &mut tokio::io::ReadBuf<'_>,
) -> std::task::Poll<std::io::Result<()>> {
use bytes::Buf;
use std::task::Poll;
let this = self.get_mut();
if !this.residual.is_empty() {
let to_copy = this.residual.len().min(buf.remaining());
buf.put_slice(&this.residual[..to_copy]);
this.residual.advance(to_copy);
return Poll::Ready(Ok(()));
}
if this.stream.is_none() {
let r = &this.reader;
this.stream = Some(Box::pin(windowed_walk::<G, M, BODY_SIZE>(
Arc::clone(&r.getter),
r.subtrees.clone(),
r.tree,
r.span,
r.concurrency,
r.position,
r.window,
)));
}
let stream = this.stream.as_mut().expect("stream just set");
match stream.as_mut().poll_next(cx) {
Poll::Ready(Some(Ok(body))) => {
let to_copy = body.len().min(buf.remaining());
buf.put_slice(&body[..to_copy]);
this.reader.position += body.len() as u64;
if to_copy < body.len() {
this.residual = body.slice(to_copy..);
}
Poll::Ready(Ok(()))
}
Poll::Ready(Some(Err(e))) => Poll::Ready(Err(std::io::Error::other(e))),
Poll::Ready(None) => Poll::Ready(Ok(())),
Poll::Pending => Poll::Pending,
}
}
}
#[cfg(feature = "tokio")]
impl<G, M, const BODY_SIZE: usize> tokio::io::AsyncSeek for WindowedJoinerReader<G, M, BODY_SIZE>
where
G: ChunkGet<BODY_SIZE> + 'static,
M: JoinMode + Send + Sync + 'static,
{
fn start_seek(self: Pin<&mut Self>, pos: SeekFrom) -> std::io::Result<()> {
let this = self.get_mut();
this.stream = None;
this.residual = Bytes::new();
this.reader.seek(pos)?;
Ok(())
}
fn poll_complete(
self: Pin<&mut Self>,
_cx: &mut std::task::Context<'_>,
) -> std::task::Poll<std::io::Result<u64>> {
std::task::Poll::Ready(Ok(self.get_mut().reader.position))
}
}
#[cfg(test)]
mod tests {
use super::*;
use crate::bmt::DEFAULT_BODY_SIZE;
use crate::chunk::AnyChunk;
use crate::file::Joiner;
use crate::file::split;
use futures::executor::block_on;
use std::collections::HashMap;
use std::sync::atomic::{AtomicUsize, Ordering};
fn split_and_store(data: &[u8]) -> (ChunkAddress, HashMap<ChunkAddress, AnyChunk>) {
let (root, store) = split::<DEFAULT_BODY_SIZE>(data).unwrap();
(root, store.into_chunks())
}
async fn drain(
reader: &mut WindowedReader<HashMap<ChunkAddress, AnyChunk>, super::super::mode::PlainMode>,
) -> Vec<u8> {
let stream = reader.stream();
futures::pin_mut!(stream);
let mut out = Vec::new();
while let Some(item) = stream.next().await {
out.extend_from_slice(&item.unwrap());
}
out
}
#[test]
fn stream_from_zero_equals_whole_file() {
let data: Vec<u8> = (0..DEFAULT_BODY_SIZE * 5 + 321)
.map(|i| (i % 256) as u8)
.collect();
let (root, store) = split_and_store(&data);
block_on(async {
let joiner = Joiner::new(store, root).await.unwrap();
let mut reader = joiner.into_windowed_reader(16);
let got = drain(&mut reader).await;
assert_eq!(got, data);
});
}
#[test]
fn seek_then_stream_yields_suffix() {
let data: Vec<u8> = (0..DEFAULT_BODY_SIZE * 5 + 321)
.map(|i| (i % 256) as u8)
.collect();
let bs = DEFAULT_BODY_SIZE as u64;
let cases = [bs + 10, bs, bs * 3, bs / 2, data.len() as u64 - 1];
let (root, store) = split_and_store(&data);
block_on(async {
for k in cases {
let joiner = Joiner::new(store.clone(), root).await.unwrap();
let mut reader = joiner.into_windowed_reader(16);
reader.seek(SeekFrom::Start(k)).unwrap();
let got = drain(&mut reader).await;
assert_eq!(got, &data[k as usize..], "suffix from {k}");
}
});
}
#[test]
fn backward_seek_then_read_is_correct() {
let data: Vec<u8> = (0..DEFAULT_BODY_SIZE * 4 + 17)
.map(|i| (i % 256) as u8)
.collect();
let bs = DEFAULT_BODY_SIZE as u64;
let (root, store) = split_and_store(&data);
block_on(async {
let joiner = Joiner::new(store, root).await.unwrap();
let mut reader = joiner.into_windowed_reader(16);
reader.seek(SeekFrom::Start(bs * 3)).unwrap();
let _ = drain(&mut reader).await;
reader.seek(SeekFrom::Start(bs)).unwrap();
let got = drain(&mut reader).await;
assert_eq!(got, &data[bs as usize..]);
});
}
#[test]
fn width_one_correct() {
let data: Vec<u8> = (0..DEFAULT_BODY_SIZE * 5 + 7)
.map(|i| (i % 256) as u8)
.collect();
let (root, store) = split_and_store(&data);
block_on(async {
let joiner = Joiner::new(store, root).await.unwrap().with_concurrency(1);
let mut reader = joiner.into_windowed_reader(1);
let got = drain(&mut reader).await;
assert_eq!(got, data);
});
}
async fn yield_now() {
struct YieldNow(bool);
impl std::future::Future for YieldNow {
type Output = ();
fn poll(
mut self: Pin<&mut Self>,
cx: &mut std::task::Context<'_>,
) -> std::task::Poll<()> {
if self.0 {
std::task::Poll::Ready(())
} else {
self.0 = true;
cx.waker().wake_by_ref();
std::task::Poll::Pending
}
}
}
YieldNow(false).await
}
#[derive(Clone)]
struct BoundProbe {
chunks: Arc<HashMap<ChunkAddress, AnyChunk>>,
in_flight: Arc<AtomicUsize>,
max_in_flight: Arc<AtomicUsize>,
}
impl ChunkGet<DEFAULT_BODY_SIZE> for BoundProbe {
type Error = crate::store::ChunkStoreError;
async fn get(&self, address: &ChunkAddress) -> std::result::Result<AnyChunk, Self::Error> {
let now = self.in_flight.fetch_add(1, Ordering::SeqCst) + 1;
self.max_in_flight.fetch_max(now, Ordering::SeqCst);
for _ in 0..4 {
yield_now().await;
}
self.in_flight.fetch_sub(1, Ordering::SeqCst);
self.chunks
.get(address)
.cloned()
.ok_or_else(|| crate::store::ChunkStoreError::not_found(address))
}
}
#[test]
fn reorder_buffer_never_exceeds_window() {
let leaves = 40usize;
let window = 6usize;
let data: Vec<u8> = (0..DEFAULT_BODY_SIZE * leaves)
.map(|i| (i % 256) as u8)
.collect();
let (root, store) = split_and_store(&data);
let probe = BoundProbe {
chunks: Arc::new(store),
in_flight: Arc::new(AtomicUsize::new(0)),
max_in_flight: Arc::new(AtomicUsize::new(0)),
};
let max_seen = Arc::clone(&probe.max_in_flight);
block_on(async {
let joiner = Joiner::new(probe, root)
.await
.unwrap()
.with_concurrency(leaves);
let mut reader = joiner.into_windowed_reader(window);
let stream = reader.stream();
futures::pin_mut!(stream);
let mut total = 0usize;
while let Some(item) = stream.next().await {
total += item.unwrap().len();
}
assert_eq!(total, data.len());
});
let peak = max_seen.load(Ordering::SeqCst);
assert!(
peak <= window + 2,
"in-flight peak {peak} exceeds window {window} (+slack)"
);
}
#[derive(Clone)]
struct HeadSlow {
chunks: Arc<HashMap<ChunkAddress, AnyChunk>>,
slow_offset: u64,
released: Arc<AtomicUsize>,
gate: usize,
leaf_offsets: Arc<HashMap<ChunkAddress, u64>>,
}
impl ChunkGet<DEFAULT_BODY_SIZE> for HeadSlow {
type Error = crate::store::ChunkStoreError;
async fn get(&self, address: &ChunkAddress) -> std::result::Result<AnyChunk, Self::Error> {
let is_slow = self.leaf_offsets.get(address) == Some(&self.slow_offset);
if is_slow {
while self.released.load(Ordering::SeqCst) < self.gate {
yield_now().await;
}
for _ in 0..4 {
yield_now().await;
}
} else {
for _ in 0..4 {
yield_now().await;
}
if self.leaf_offsets.contains_key(address) {
self.released.fetch_add(1, Ordering::SeqCst);
}
}
self.chunks
.get(address)
.cloned()
.ok_or_else(|| crate::store::ChunkStoreError::not_found(address))
}
}
fn unique_byte(i: u64) -> u8 {
(i.wrapping_mul(2_654_435_761) >> 11) as u8
}
fn leaf_offsets_of(total: u64) -> HashMap<ChunkAddress, u64> {
use crate::chunk::Chunk;
let mut map = HashMap::new();
let leaves = total.div_ceil(DEFAULT_BODY_SIZE as u64);
for i in 0..leaves {
let off = i * DEFAULT_BODY_SIZE as u64;
let end = (off + DEFAULT_BODY_SIZE as u64).min(total);
let body: Vec<u8> = (off..end).map(unique_byte).collect();
let chunk = crate::chunk::ContentChunk::<DEFAULT_BODY_SIZE>::new(body).unwrap();
map.insert(*chunk.address(), off);
}
map
}
#[test]
fn head_slowest_in_order_and_bounded() {
let leaves = 40usize;
let window = 6usize;
let total = (DEFAULT_BODY_SIZE * leaves) as u64;
let data: Vec<u8> = (0..total).map(unique_byte).collect();
let (root, store) = split_and_store(&data);
PEAK_LEAF_OCCUPANCY.with(|p| p.set(0));
let leaf_offsets = leaf_offsets_of(total);
let getter = HeadSlow {
chunks: Arc::new(store),
slow_offset: 0,
released: Arc::new(AtomicUsize::new(0)),
gate: window - 1,
leaf_offsets: Arc::new(leaf_offsets),
};
block_on(async {
let joiner = Joiner::new(getter, root)
.await
.unwrap()
.with_concurrency(leaves);
let mut reader = joiner.into_windowed_reader(window);
let stream = reader.stream();
futures::pin_mut!(stream);
let mut got = Vec::new();
while let Some(item) = stream.next().await {
got.extend_from_slice(&item.unwrap());
}
assert_eq!(got, data, "head-slowest still yields whole file in order");
});
let peak = PEAK_LEAF_OCCUPANCY.with(std::cell::Cell::get);
assert!(
peak <= window,
"leaf-body occupancy peak {peak} exceeds window {window}"
);
}
const TINY_BODY: usize = 256;
fn tiny_leaf_addresses(data: &[u8]) -> HashMap<ChunkAddress, ()> {
use crate::chunk::Chunk;
let mut set = HashMap::new();
for block in data.chunks(TINY_BODY) {
let chunk = crate::chunk::ContentChunk::<TINY_BODY>::new(block.to_vec()).unwrap();
set.insert(*chunk.address(), ());
}
set
}
fn tiny_deep_data(leaves: usize) -> Vec<u8> {
(0..TINY_BODY * leaves).map(|i| (i % 251) as u8).collect()
}
#[derive(Clone)]
struct TinyOrderProbe {
chunks: Arc<HashMap<ChunkAddress, AnyChunk<TINY_BODY>>>,
leaves: Arc<HashMap<ChunkAddress, ()>>,
kinds: Arc<std::sync::Mutex<Vec<bool>>>,
intermediate_in_flight: Arc<AtomicUsize>,
peak_intermediate: Arc<AtomicUsize>,
}
impl TinyOrderProbe {
fn new(store: HashMap<ChunkAddress, AnyChunk<TINY_BODY>>, data: &[u8]) -> Self {
Self {
chunks: Arc::new(store),
leaves: Arc::new(tiny_leaf_addresses(data)),
kinds: Arc::new(std::sync::Mutex::new(Vec::new())),
intermediate_in_flight: Arc::new(AtomicUsize::new(0)),
peak_intermediate: Arc::new(AtomicUsize::new(0)),
}
}
fn reset(&self) {
self.kinds.lock().unwrap().clear();
self.peak_intermediate.store(0, Ordering::SeqCst);
}
fn intermediates_before_first_leaf(&self) -> usize {
self.kinds
.lock()
.unwrap()
.iter()
.take_while(|is_leaf| !**is_leaf)
.count()
}
fn intermediate_fetches(&self) -> usize {
self.kinds.lock().unwrap().iter().filter(|l| !**l).count()
}
}
impl ChunkGet<TINY_BODY> for TinyOrderProbe {
type Error = crate::store::ChunkStoreError;
async fn get(
&self,
address: &ChunkAddress,
) -> std::result::Result<AnyChunk<TINY_BODY>, Self::Error> {
let is_leaf = self.leaves.contains_key(address);
self.kinds.lock().unwrap().push(is_leaf);
if !is_leaf {
let now = self.intermediate_in_flight.fetch_add(1, Ordering::SeqCst) + 1;
self.peak_intermediate.fetch_max(now, Ordering::SeqCst);
}
for _ in 0..4 {
yield_now().await;
}
if !is_leaf {
self.intermediate_in_flight.fetch_sub(1, Ordering::SeqCst);
}
self.chunks
.get(address)
.cloned()
.ok_or_else(|| crate::store::ChunkStoreError::not_found(address))
}
}
#[test]
fn windowed_first_leaf_before_frontier() {
let data = tiny_deep_data(900);
let (root, store) = split::<TINY_BODY>(&data).unwrap();
let probe = TinyOrderProbe::new(store.into_chunks(), &data);
block_on(async {
let joiner = Joiner::<_, TINY_BODY>::new(probe.clone(), root)
.await
.unwrap();
probe.reset();
let mut reader = joiner.into_windowed_reader(16);
let stream = reader.stream();
futures::pin_mut!(stream);
let mut got = Vec::new();
while let Some(item) = stream.next().await {
got.extend_from_slice(&item.unwrap());
}
assert_eq!(got, data, "deep-tree in-order reassembly is byte-exact");
});
let frontier = probe.intermediate_fetches();
assert!(
frontier >= 40,
"test needs a frontier far larger than the cap, saw {frontier}"
);
let before = probe.intermediates_before_first_leaf();
assert!(
before <= 4 * MAX_INTERMEDIATE_IN_FLIGHT,
"first leaf fetched after {before} intermediates (frontier {frontier}); \
expected a short descent, not the whole frontier"
);
let peak = probe.peak_intermediate.load(Ordering::SeqCst);
assert!(
peak <= MAX_INTERMEDIATE_IN_FLIGHT,
"intermediate in-flight peak {peak} exceeds cap {MAX_INTERMEDIATE_IN_FLIGHT}"
);
}
#[test]
fn windowed_deep_frontier_in_order_and_bounded() {
let leaves = 2000usize;
let window = 5usize;
let data = tiny_deep_data(leaves);
let (root, store) = split::<TINY_BODY>(&data).unwrap();
let probe = TinyOrderProbe::new(store.into_chunks(), &data);
PEAK_LEAF_OCCUPANCY.with(|p| p.set(0));
block_on(async {
let joiner = Joiner::<_, TINY_BODY>::new(probe.clone(), root)
.await
.unwrap()
.with_concurrency(leaves);
probe.reset();
let mut reader = joiner.into_windowed_reader(window);
let stream = reader.stream();
futures::pin_mut!(stream);
let mut got = Vec::new();
while let Some(item) = stream.next().await {
got.extend_from_slice(&item.unwrap());
}
assert_eq!(got, data, "deep frontier still yields whole file in order");
});
let occupancy = PEAK_LEAF_OCCUPANCY.with(std::cell::Cell::get);
assert!(
occupancy <= window,
"leaf-body occupancy peak {occupancy} exceeds window {window}"
);
let peak = probe.peak_intermediate.load(Ordering::SeqCst);
assert!(
peak <= MAX_INTERMEDIATE_IN_FLIGHT,
"intermediate in-flight peak {peak} exceeds cap {MAX_INTERMEDIATE_IN_FLIGHT}"
);
}
#[cfg(feature = "encryption")]
#[test]
fn head_slowest_encrypted_in_order_and_bounded() {
use crate::file::EncryptedJoiner;
use crate::file::split_encrypted;
let leaves = 30usize;
let window = 5usize;
let total = (DEFAULT_BODY_SIZE * leaves) as u64;
let data: Vec<u8> = (0..total).map(unique_byte).collect();
let (root_ref, store) = split_encrypted::<DEFAULT_BODY_SIZE>(&data).unwrap();
let store = store.into_chunks();
PEAK_LEAF_OCCUPANCY.with(|p| p.set(0));
let getter = UniformSlow {
chunks: Arc::new(store),
};
block_on(async {
let joiner = EncryptedJoiner::new(getter, root_ref)
.await
.unwrap()
.with_concurrency(leaves);
let mut reader = joiner.into_windowed_reader(window);
let stream = reader.stream();
futures::pin_mut!(stream);
let mut out = Vec::new();
while let Some(item) = stream.next().await {
out.extend_from_slice(&item.unwrap());
}
assert_eq!(out, data, "encrypted short-body slide stays in order");
});
let peak = PEAK_LEAF_OCCUPANCY.with(std::cell::Cell::get);
assert!(
peak <= window,
"encrypted leaf-body occupancy peak {peak} exceeds window {window}"
);
}
#[cfg(feature = "encryption")]
#[derive(Clone)]
struct UniformSlow {
chunks: Arc<HashMap<ChunkAddress, AnyChunk>>,
}
#[cfg(feature = "encryption")]
impl ChunkGet<DEFAULT_BODY_SIZE> for UniformSlow {
type Error = crate::store::ChunkStoreError;
async fn get(&self, address: &ChunkAddress) -> std::result::Result<AnyChunk, Self::Error> {
for _ in 0..4 {
yield_now().await;
}
self.chunks
.get(address)
.cloned()
.ok_or_else(|| crate::store::ChunkStoreError::not_found(address))
}
}
#[cfg(feature = "encryption")]
mod encrypted {
use super::*;
use crate::file::EncryptedJoiner;
use crate::file::split_encrypted;
async fn drain_enc(
reader: &mut WindowedReader<
HashMap<ChunkAddress, AnyChunk>,
crate::file::mode::EncryptedMode,
>,
) -> Vec<u8> {
let stream = reader.stream();
futures::pin_mut!(stream);
let mut out = Vec::new();
while let Some(item) = stream.next().await {
out.extend_from_slice(&item.unwrap());
}
out
}
#[test]
fn encrypted_stream_and_seek() {
let data: Vec<u8> = (0..DEFAULT_BODY_SIZE * 5 + 321)
.map(|i| (i % 256) as u8)
.collect();
let bs = DEFAULT_BODY_SIZE as u64;
let (root_ref, store) = split_encrypted::<DEFAULT_BODY_SIZE>(&data).unwrap();
let store = store.into_chunks();
block_on(async {
let joiner = EncryptedJoiner::new(store.clone(), root_ref.clone())
.await
.unwrap();
let mut reader = joiner.into_windowed_reader(16);
assert_eq!(drain_enc(&mut reader).await, data);
for k in [bs + 10, bs, bs * 3] {
let joiner = EncryptedJoiner::new(store.clone(), root_ref.clone())
.await
.unwrap();
let mut reader = joiner.into_windowed_reader(16);
reader.seek(SeekFrom::Start(k)).unwrap();
assert_eq!(drain_enc(&mut reader).await, &data[k as usize..]);
}
});
}
}
#[cfg(feature = "tokio")]
#[tokio::test]
async fn reader_seek_and_read() {
use tokio::io::{AsyncReadExt, AsyncSeekExt};
let data: Vec<u8> = (0..DEFAULT_BODY_SIZE * 3 + 50)
.map(|i| (i % 256) as u8)
.collect();
let (root, store) = split_and_store(&data);
let joiner = Joiner::new(store, root).await.unwrap();
let mut reader = joiner.into_windowed_reader(16).into_reader();
let mut all = Vec::new();
reader.read_to_end(&mut all).await.unwrap();
assert_eq!(all, data);
reader
.seek(SeekFrom::Start(DEFAULT_BODY_SIZE as u64))
.await
.unwrap();
let mut tail = Vec::new();
reader.read_to_end(&mut tail).await.unwrap();
assert_eq!(tail, &data[DEFAULT_BODY_SIZE..]);
}
}