use crate::{
DownloadResult, Event, ProgressEntry, Puller, PullerError, Pusher, multi::TokioExecutor,
};
use bytes::Bytes;
use core::time::Duration;
use crossfire::{mpmc, spsc};
use futures::TryStreamExt;
use std::sync::Arc;
use tokio_util::sync::CancellationToken;
#[derive(Debug, Clone, Copy)]
pub struct DownloadOptions {
pub retry_gap: Duration,
pub push_queue_cap: usize,
}
#[allow(clippy::too_many_lines)]
pub fn download_single<R: Puller, W: Pusher>(
mut puller: R,
mut pusher: W,
options: DownloadOptions,
) -> DownloadResult<TokioExecutor<R, W::Error>, R::Error, W::Error> {
const ID: usize = 0;
let token = CancellationToken::new();
let (tx, event_chain) = mpmc::unbounded_async();
pusher.set_listener({
let tx = tx.clone();
Box::new(move |p| {
let _ = tx.send(Event::PushProgress(p));
})
});
let (tx_push, rx_push) =
spsc::bounded_async_blocking::<(ProgressEntry, Bytes)>(options.push_queue_cap);
let push_thread = Arc::new(std::sync::OnceLock::new());
let push_handle = tokio::task::spawn_blocking({
let push_thread = push_thread.clone();
let token = token.clone();
let tx = tx.clone();
move || {
let _ = push_thread.set(std::thread::current());
while let Ok((mut spin, mut data)) = rx_push.recv() {
loop {
if token.is_cancelled() {
return;
}
let _ = tx.send(Event::Pushing(ID, spin.clone()));
let len_before_push = data.len();
match pusher.push(&spin, data) {
Ok(()) => break,
Err((err, bytes)) => {
let _ = tx.send(Event::PushError(ID, spin.clone(), err));
let written = len_before_push.saturating_sub(bytes.len());
data = bytes;
spin.start += written as u64;
}
}
std::thread::park_timeout(options.retry_gap);
}
}
loop {
if token.is_cancelled() {
break;
}
let _ = tx.send(Event::Flushing);
match pusher.flush() {
Ok(()) => break,
Err(err) => {
let _ = tx.send(Event::FlushError(err));
}
}
std::thread::park_timeout(options.retry_gap);
}
}
});
let pull_handle = tokio::spawn(async move {
'redownload: loop {
let _ = tx.send(Event::Pulling(ID));
let mut downloaded: u64 = 0;
let mut stream = loop {
match puller.pull(None).await {
Ok(t) => break t,
Err((e, retry_gap)) => {
let _ = tx.send(Event::PullError(ID, e));
tokio::time::sleep(retry_gap.unwrap_or(options.retry_gap)).await;
}
}
};
loop {
match stream.try_next().await {
Ok(Some(chunk)) => {
if chunk.is_empty() {
continue;
}
let len = chunk.len() as u64;
let span = downloaded..(downloaded + len);
let _ = tx.send(Event::PullProgress(ID, span.clone()));
let _ = tx_push.send((span, chunk)).await;
downloaded += len;
}
Ok(None) => break 'redownload,
Err((e, retry_gap)) => {
let is_irrecoverable = e.is_irrecoverable();
let _ = tx.send(Event::PullError(ID, e));
tokio::time::sleep(retry_gap.unwrap_or(options.retry_gap)).await;
if is_irrecoverable {
continue 'redownload;
}
}
}
}
}
let _ = tx.send(Event::Finished(ID));
});
tokio::spawn({
let token = token.clone();
async move {
tokio::select! {
_ = push_handle => {},
() = token.cancelled() => {
pull_handle.abort();
if let Some(t) = push_thread.get() {
t.unpark();
}
}
}
}
});
DownloadResult::new(event_chain, None, token)
}
#[cfg(test)]
mod tests {
#![allow(clippy::cast_possible_truncation)]
use super::*;
use crate::BufWriterPusher;
use crate::{
MemPusher, Merge, ProgressEntry,
mock::{MockPuller, build_mock_data},
};
use futures::stream;
use std::{dbg, vec};
use tokio::time::{sleep, timeout};
use vec::Vec;
#[tokio::test]
async fn test_sequential_download() {
let mock_data = build_mock_data(3 * 1024);
let puller = MockPuller::new(&mock_data);
let pusher = MemPusher::with_capacity(mock_data.len());
let receive = pusher.receive.clone();
#[allow(clippy::single_range_in_vec_init)]
let download_chunks = [0..mock_data.len() as u64];
let result = download_single(
puller,
pusher,
DownloadOptions {
retry_gap: Duration::from_secs(1),
push_queue_cap: 1024,
},
);
let mut pull_progress: Vec<ProgressEntry> = Vec::new();
let mut push_progress: Vec<ProgressEntry> = Vec::new();
while let Ok(e) = result.event_chain().recv().await {
match e {
Event::PullProgress(_, p) => pull_progress.merge_progress(p),
Event::PushProgress(p) => push_progress.merge_progress(p),
_ => {}
}
}
dbg!(&pull_progress);
dbg!(&push_progress);
assert_eq!(pull_progress, download_chunks);
assert_eq!(push_progress, download_chunks);
assert_eq!(&**receive.lock(), mock_data);
}
#[tokio::test]
async fn test_sequential_download_abort_discards() {
let mock_data = build_mock_data(3 * 1024);
let puller = MockPuller::new(&mock_data);
let pusher = MemPusher::with_capacity(mock_data.len());
let receive = pusher.receive.clone();
let result = download_single(
puller,
pusher,
DownloadOptions {
retry_gap: Duration::from_secs(1),
push_queue_cap: 1024,
},
);
result.abort();
assert!(result.is_aborted());
tokio::time::timeout(Duration::from_secs(10), async {
while result.event_chain().recv().await.is_ok() {}
})
.await
.expect("event loop hung after abort");
let written = receive.lock().len();
assert!(
written <= mock_data.len(),
"abort must not write beyond the source"
);
}
#[derive(Debug, Clone)]
struct SlowMockPuller {
data: Arc<[u8]>,
delay: Duration,
}
impl Puller for SlowMockPuller {
type Error = std::convert::Infallible;
#[allow(clippy::cast_possible_truncation)]
fn pull(
&mut self,
range: Option<&ProgressEntry>,
) -> impl Future<
Output = crate::PullResult<impl crate::PullStream<Self::Error>, Self::Error>,
> + Send {
type PullItem = crate::PullResult<Bytes, std::convert::Infallible>;
let owned: Vec<u8> = match range {
Some(r) => self.data[r.start as usize..r.end as usize].to_vec(),
None => self.data.to_vec(),
};
let delay = self.delay;
async move {
sleep(delay).await;
let items: Vec<PullItem> = owned
.chunks(2)
.map(|c| Ok(Bytes::from(c.to_vec())))
.collect();
Ok(stream::iter(items))
}
}
}
#[tokio::test(flavor = "multi_thread")]
async fn test_sequential_download_abort_discards_buffered() {
let mock_data = build_mock_data(64 * 1024);
let puller = SlowMockPuller {
data: Arc::from(mock_data.as_slice()),
delay: Duration::from_millis(50),
};
let inner = MemPusher::with_capacity(mock_data.len());
let receive = inner.receive.clone();
let pusher = BufWriterPusher::new(inner, mock_data.len() + 1);
let result = download_single(
puller,
pusher,
DownloadOptions {
retry_gap: Duration::from_secs(1),
push_queue_cap: 1024,
},
);
let mut aborted = false;
while let Ok(e) = result.event_chain().recv().await {
if matches!(e, Event::Pushing(_, _)) {
result.abort();
assert!(result.is_aborted());
aborted = true;
break;
}
}
assert!(aborted, "expected a Pushing event before aborting");
timeout(Duration::from_secs(10), async {
while result.event_chain().recv().await.is_ok() {}
})
.await
.expect("event loop hung after abort");
let written = receive.lock().len();
assert_eq!(
written, 0,
"abort must discard buffered bytes, not write them to the sink"
);
}
#[cfg(feature = "file")]
#[tokio::test(flavor = "multi_thread")]
async fn test_sequential_download_abort_discards_file() {
use std::io::Read;
let mock_data = build_mock_data(64 * 1024);
let puller = SlowMockPuller {
data: Arc::from(mock_data.as_slice()),
delay: Duration::from_millis(50),
};
let tmp = tempfile::NamedTempFile::new().unwrap();
let path = tmp.path().to_path_buf();
let file = tokio::fs::File::from(tmp.reopen().unwrap());
let inner = crate::StdFilePusher::new(file, mock_data.len() as u64, false)
.await
.unwrap();
let pusher = BufWriterPusher::new(inner, mock_data.len() + 1);
let result = download_single(
puller,
pusher,
DownloadOptions {
retry_gap: Duration::from_secs(1),
push_queue_cap: 1024,
},
);
let mut aborted = false;
while let Ok(e) = result.event_chain().recv().await {
if matches!(e, Event::Pushing(_, _)) {
result.abort();
assert!(result.is_aborted());
aborted = true;
break;
}
}
assert!(aborted, "expected a Pushing event before aborting");
timeout(Duration::from_secs(10), async {
while result.event_chain().recv().await.is_ok() {}
})
.await
.expect("event loop hung after abort");
let mut f = std::fs::File::open(&path).unwrap();
let mut buf = Vec::new();
f.read_to_end(&mut buf).unwrap();
assert_eq!(buf.len(), mock_data.len(), "file should remain pre-sized");
assert!(
buf.iter().all(|&b| b == 0),
"abort must not write buffered bytes to the file"
);
}
use parking_lot::Mutex;
use std::sync::Arc;
use std::sync::atomic::{AtomicBool, Ordering};
#[derive(Debug)]
struct FatalErr;
impl std::fmt::Display for FatalErr {
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
f.write_str("fatal")
}
}
impl std::error::Error for FatalErr {}
impl crate::PullerError for FatalErr {
fn is_irrecoverable(&self) -> bool {
true
}
}
#[derive(Debug)]
struct RecoverableErr;
impl std::fmt::Display for RecoverableErr {
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
f.write_str("recoverable")
}
}
impl std::error::Error for RecoverableErr {}
impl crate::PullerError for RecoverableErr {
fn is_irrecoverable(&self) -> bool {
false
}
}
struct FlakySink {
fail_push: Arc<AtomicBool>,
fail_flush: Arc<AtomicBool>,
receive: Arc<Mutex<Vec<u8>>>,
listener: Option<crate::ProgressListener>,
}
impl FlakySink {
fn new() -> Self {
Self {
fail_push: Arc::new(AtomicBool::new(false)),
fail_flush: Arc::new(AtomicBool::new(false)),
receive: Arc::new(Mutex::new(Vec::new())),
listener: None,
}
}
}
impl crate::Pusher for FlakySink {
type Error = std::io::Error;
fn set_listener(&mut self, cb: crate::ProgressListener) {
self.listener = Some(cb);
}
fn push(
&mut self,
range: &crate::ProgressEntry,
bytes: Bytes,
) -> Result<(), (Self::Error, Bytes)> {
if self.fail_push.swap(false, Ordering::SeqCst) {
return Err((std::io::Error::other("push"), bytes));
}
let mut g = self.receive.lock();
if range.start as usize == g.len() {
g.extend_from_slice(&bytes);
} else {
if g.len() < range.end as usize {
g.resize(range.end as usize, 0);
}
g[range.start as usize..range.end as usize].copy_from_slice(&bytes);
}
drop(g);
if let Some(l) = &mut self.listener {
l(range.clone());
}
Ok(())
}
fn flush(&mut self) -> Result<(), Self::Error> {
if self.fail_flush.swap(false, Ordering::SeqCst) {
Err(std::io::Error::other("flush"))
} else {
Ok(())
}
}
}
#[derive(Debug, Clone)]
struct PullErrOncePuller {
data: Arc<[u8]>,
failed: Arc<AtomicBool>,
}
impl crate::Puller for PullErrOncePuller {
type Error = RecoverableErr;
fn pull(
&mut self,
range: Option<&crate::ProgressEntry>,
) -> impl Future<
Output = crate::PullResult<impl crate::PullStream<Self::Error>, Self::Error>,
> + Send {
if !self.failed.swap(true, Ordering::SeqCst) {
return std::future::ready(Err((RecoverableErr, Some(Duration::ZERO))));
}
let data = match range {
Some(r) => &self.data[r.start as usize..r.end as usize],
None => &self.data,
};
let items: Vec<crate::PullResult<Bytes, RecoverableErr>> = data
.chunks(2)
.map(|c| Ok(Bytes::copy_from_slice(c)))
.collect();
std::future::ready(Ok(stream::iter(items)))
}
}
#[derive(Debug, Clone)]
struct StreamErrOncePuller {
data: Arc<[u8]>,
failed: Arc<AtomicBool>,
}
impl crate::Puller for StreamErrOncePuller {
type Error = FatalErr;
fn pull(
&mut self,
range: Option<&crate::ProgressEntry>,
) -> impl Future<
Output = crate::PullResult<impl crate::PullStream<Self::Error>, Self::Error>,
> + Send {
if !self.failed.swap(true, Ordering::SeqCst) {
let items: Vec<crate::PullResult<Bytes, FatalErr>> =
vec![Err((FatalErr, Some(Duration::ZERO)))];
return std::future::ready(Ok(stream::iter(items)));
}
let data = match range {
Some(r) => &self.data[r.start as usize..r.end as usize],
None => &self.data,
};
let items: Vec<crate::PullResult<Bytes, FatalErr>> = data
.chunks(2)
.map(|c| Ok(Bytes::copy_from_slice(c)))
.collect();
std::future::ready(Ok(stream::iter(items)))
}
}
#[derive(Debug, Clone)]
struct RecoverableStreamErrOncePuller {
data: Arc<[u8]>,
failed: Arc<AtomicBool>,
}
impl crate::Puller for RecoverableStreamErrOncePuller {
type Error = RecoverableErr;
fn pull(
&mut self,
range: Option<&crate::ProgressEntry>,
) -> impl Future<
Output = crate::PullResult<impl crate::PullStream<Self::Error>, Self::Error>,
> + Send {
let first = !self.failed.swap(true, Ordering::SeqCst);
let data = match range {
Some(r) => &self.data[r.start as usize..r.end as usize],
None => &self.data,
};
let mut items: Vec<crate::PullResult<Bytes, RecoverableErr>> = data
.chunks(2)
.map(|c| Ok(Bytes::copy_from_slice(c)))
.collect();
if first {
let mut with_err = vec![Err((RecoverableErr, Some(Duration::ZERO)))];
with_err.append(&mut items);
return std::future::ready(Ok(stream::iter(with_err)));
}
std::future::ready(Ok(stream::iter(items)))
}
}
#[tokio::test]
async fn test_single_push_error_retries() {
let mock_data = build_mock_data(3 * 1024);
let puller = MockPuller::new(&mock_data);
let sink = FlakySink::new();
sink.fail_push.store(true, Ordering::SeqCst);
let receive = sink.receive.clone();
let result = download_single(
puller,
sink,
DownloadOptions {
retry_gap: Duration::ZERO,
push_queue_cap: 1024,
},
);
while result.event_chain().recv().await.is_ok() {}
assert_eq!(&**receive.lock(), mock_data);
}
#[tokio::test]
async fn test_single_flush_error_retries() {
let mock_data = build_mock_data(3 * 1024);
let puller = MockPuller::new(&mock_data);
let sink = FlakySink::new();
sink.fail_flush.store(true, Ordering::SeqCst);
let receive = sink.receive.clone();
let result = download_single(
puller,
sink,
DownloadOptions {
retry_gap: Duration::ZERO,
push_queue_cap: 1024,
},
);
while result.event_chain().recv().await.is_ok() {}
assert_eq!(&**receive.lock(), mock_data);
}
#[tokio::test]
async fn test_single_pull_error_retries() {
let mock_data = build_mock_data(3 * 1024);
let puller = PullErrOncePuller {
data: Arc::from(mock_data.as_slice()),
failed: Arc::new(AtomicBool::new(false)),
};
let pusher = MemPusher::with_capacity(mock_data.len());
let receive = pusher.receive.clone();
let result = download_single(
puller,
pusher,
DownloadOptions {
retry_gap: Duration::ZERO,
push_queue_cap: 1024,
},
);
while result.event_chain().recv().await.is_ok() {}
assert_eq!(&**receive.lock(), mock_data);
}
#[tokio::test]
async fn test_single_stream_error_irrecoverable_retries() {
let mock_data = build_mock_data(3 * 1024);
let puller = StreamErrOncePuller {
data: Arc::from(mock_data.as_slice()),
failed: Arc::new(AtomicBool::new(false)),
};
let pusher = MemPusher::with_capacity(mock_data.len());
let receive = pusher.receive.clone();
let result = download_single(
puller,
pusher,
DownloadOptions {
retry_gap: Duration::ZERO,
push_queue_cap: 1024,
},
);
while result.event_chain().recv().await.is_ok() {}
assert_eq!(&**receive.lock(), mock_data);
}
#[tokio::test]
async fn test_single_stream_error_recoverable_retries() {
let mock_data = build_mock_data(3 * 1024);
let puller = RecoverableStreamErrOncePuller {
data: Arc::from(mock_data.as_slice()),
failed: Arc::new(AtomicBool::new(false)),
};
let pusher = MemPusher::with_capacity(mock_data.len());
let receive = pusher.receive.clone();
let result = download_single(
puller,
pusher,
DownloadOptions {
retry_gap: Duration::ZERO,
push_queue_cap: 1024,
},
);
while result.event_chain().recv().await.is_ok() {}
assert_eq!(&**receive.lock(), mock_data);
}
#[tokio::test]
async fn puller_and_error_coverage() {
assert_eq!(format!("{FatalErr}"), "fatal");
assert_eq!(format!("{RecoverableErr}"), "recoverable");
let mut pull_err = PullErrOncePuller {
data: Arc::from(b"abcdef".as_slice()),
failed: Arc::new(AtomicBool::new(false)),
};
let _ = pull_err.pull(Some(&(0..2u64))).await; let _ = pull_err.pull(Some(&(0..2u64))).await; let _ = pull_err.pull(None).await;
let mut stream_err = StreamErrOncePuller {
data: Arc::from(b"abcdef".as_slice()),
failed: Arc::new(AtomicBool::new(false)),
};
let _ = stream_err.pull(Some(&(0..2u64))).await;
let _ = stream_err.pull(Some(&(0..2u64))).await; let _ = stream_err.pull(None).await;
let mut rec_stream_err = RecoverableStreamErrOncePuller {
data: Arc::from(b"abcdef".as_slice()),
failed: Arc::new(AtomicBool::new(false)),
};
let _ = rec_stream_err.pull(Some(&(0..2u64))).await; let _ = rec_stream_err.pull(Some(&(0..2u64))).await; let _ = rec_stream_err.pull(None).await; }
#[test]
fn flaky_sink_noncontiguous_write_rebuffers() {
let mut sink = FlakySink::new();
sink.push(&(5..8u64), Bytes::from_static(b"xyz")).unwrap();
assert_eq!(&**sink.receive.lock(), b"\0\0\0\0\0xyz");
}
#[tokio::test]
async fn test_slow_mock_puller_some_range() {
let mut p = SlowMockPuller {
data: Arc::from(b"hello world".as_slice()),
delay: Duration::ZERO,
};
assert!(p.pull(Some(&(0..5))).await.is_ok());
}
#[tokio::test]
async fn test_sequential_download_empty_file() {
let mock_data: Vec<u8> = Vec::new();
let puller = MockPuller::new(&mock_data);
let pusher = MemPusher::with_capacity(0);
let receive = pusher.receive.clone();
let result = download_single(
puller,
pusher,
DownloadOptions {
retry_gap: Duration::from_secs(1),
push_queue_cap: 1024,
},
);
while result.event_chain().recv().await.is_ok() {}
timeout(Duration::from_secs(10), async {
while result.event_chain().recv().await.is_ok() {}
})
.await
.expect("event loop hung on empty file");
assert_eq!(receive.lock().len(), 0);
}
#[derive(Debug, Clone)]
struct PullErrNoGapPuller {
data: Arc<[u8]>,
failed: Arc<AtomicBool>,
}
impl crate::Puller for PullErrNoGapPuller {
type Error = RecoverableErr;
fn pull(
&mut self,
range: Option<&crate::ProgressEntry>,
) -> impl Future<
Output = crate::PullResult<impl crate::PullStream<Self::Error>, Self::Error>,
> + Send {
if !self.failed.swap(true, Ordering::SeqCst) {
return std::future::ready(Err((RecoverableErr, None)));
}
let data = match range {
Some(r) => &self.data[r.start as usize..r.end as usize],
None => &self.data,
};
let items: Vec<crate::PullResult<Bytes, RecoverableErr>> = data
.chunks(2)
.map(|c| Ok(Bytes::copy_from_slice(c)))
.collect();
std::future::ready(Ok(stream::iter(items)))
}
}
#[tokio::test]
async fn test_single_pull_error_without_retry_gap_uses_options_default() {
let mock_data = build_mock_data(3 * 1024);
let puller = PullErrNoGapPuller {
data: Arc::from(mock_data.as_slice()),
failed: Arc::new(AtomicBool::new(false)),
};
let pusher = MemPusher::with_capacity(mock_data.len());
let receive = pusher.receive.clone();
let result = download_single(
puller,
pusher,
DownloadOptions {
retry_gap: Duration::ZERO,
push_queue_cap: 1024,
},
);
while result.event_chain().recv().await.is_ok() {}
assert_eq!(&**receive.lock(), mock_data);
}
use crate::{PullResult, PullStream};
use std::pin::Pin;
use std::task::{Context, Poll};
struct PendingStream;
impl futures::Stream for PendingStream {
type Item = Result<Bytes, (std::convert::Infallible, Option<Duration>)>;
fn poll_next(self: Pin<&mut Self>, _cx: &mut Context<'_>) -> Poll<Option<Self::Item>> {
Poll::Pending
}
}
impl Unpin for PendingStream {}
#[derive(Clone)]
struct StallPuller;
impl Puller for StallPuller {
type Error = std::convert::Infallible;
fn pull(
&mut self,
_range: Option<&ProgressEntry>,
) -> impl Future<Output = PullResult<impl PullStream<Self::Error>, Self::Error>> + Send
{
std::future::ready(Ok(PendingStream))
}
}
#[tokio::test]
#[ignore = "regression baseline: download_single has no pull_timeout, a stalled body hangs forever; enable after adding pull_timeout to DownloadOptions and surfacing a PullTimeout/error"]
async fn test_single_stall_body_hangs_without_timeout() {
let puller = StallPuller;
let pusher = MemPusher::with_capacity(0);
let result = download_single(
puller,
pusher,
DownloadOptions {
retry_gap: Duration::from_secs(1),
push_queue_cap: 1024,
},
);
let drained = tokio::time::timeout(Duration::from_secs(3), async {
while result.event_chain().recv().await.is_ok() {}
})
.await;
assert!(
drained.is_ok(),
"download_single must not hang forever on a stalled body; a pull_timeout should surface a PullTimeout or error and end the session"
);
}
}