fast_pull/core/
single.rs

1extern crate alloc;
2use super::macros::poll_ok;
3use crate::{DownloadResult, Event, ProgressEntry, SeqPuller, SeqPusher};
4use bytes::Bytes;
5use core::time::Duration;
6use fast_steal::{Executor, Handle};
7use futures::TryStreamExt;
8
9#[derive(Debug, Clone)]
10pub struct DownloadOptions {
11    pub retry_gap: Duration,
12    pub push_queue_cap: usize,
13}
14
15#[derive(Clone)]
16pub struct EmptyHandle;
17impl Handle for EmptyHandle {
18    type Output = ();
19    fn abort(&mut self) -> Self::Output {}
20}
21
22#[derive(Debug, Clone)]
23pub struct EmptyExecutor;
24impl Executor for EmptyExecutor {
25    type Handle = EmptyHandle;
26    fn execute(
27        self: alloc::sync::Arc<Self>,
28        _: alloc::sync::Arc<fast_steal::Task>,
29        _: alloc::sync::Arc<fast_steal::TaskList<Self>>,
30    ) -> Self::Handle {
31        EmptyHandle
32    }
33}
34
35pub async fn download_single<R, W>(
36    mut puller: R,
37    mut pusher: W,
38    options: DownloadOptions,
39) -> DownloadResult<EmptyExecutor, R::Error, W::Error>
40where
41    R: SeqPuller + 'static,
42    W: SeqPusher + 'static,
43{
44    let (tx, event_chain) = kanal::unbounded_async();
45    let (tx_push, rx_push) = kanal::bounded_async::<(ProgressEntry, Bytes)>(options.push_queue_cap);
46    let tx_clone = tx.clone();
47    const ID: usize = 0;
48    let push_handle = tokio::spawn(async move {
49        while let Ok((spin, data)) = rx_push.recv().await {
50            poll_ok!(
51                pusher.push(&data).await,
52                ID @ tx_clone => PushError,
53                options.retry_gap
54            );
55            tx_clone.send(Event::PushProgress(ID, spin)).await.unwrap();
56        }
57        poll_ok!(
58            pusher.flush().await,
59            tx_clone => FlushError,
60            options.retry_gap
61        );
62    });
63    let handle = tokio::spawn(async move {
64        tx.send(Event::Pulling(ID)).await.unwrap();
65        let mut downloaded: u64 = 0;
66        let mut stream = puller.pull();
67        loop {
68            match stream.try_next().await {
69                Ok(Some(chunk)) => {
70                    let len = chunk.len() as u64;
71                    let span = downloaded..(downloaded + len);
72                    tx.send(Event::PullProgress(ID, span.clone()))
73                        .await
74                        .unwrap();
75                    tx_push.send((span, chunk)).await.unwrap();
76                    downloaded += len;
77                }
78                Ok(None) => break,
79                Err(e) => {
80                    tx.send(Event::PullError(ID, e)).await.unwrap();
81                    tokio::time::sleep(options.retry_gap).await;
82                }
83            }
84        }
85        tx.send(Event::Finished(ID)).await.unwrap();
86    });
87    DownloadResult::new(event_chain, push_handle, &[handle.abort_handle()], None)
88}
89
90#[cfg(test)]
91mod tests {
92    extern crate std;
93    use super::*;
94    use crate::{
95        MergeProgress,
96        mem::MemPusher,
97        mock::{MockPuller, build_mock_data},
98    };
99    use alloc::vec;
100    use std::dbg;
101    use vec::Vec;
102
103    #[tokio::test]
104    async fn test_sequential_download() {
105        let mock_data = build_mock_data(3 * 1024);
106        let puller = MockPuller::new(&mock_data);
107        let pusher = MemPusher::with_capacity(mock_data.len());
108        #[allow(clippy::single_range_in_vec_init)]
109        let download_chunks = vec![0..mock_data.len() as u64];
110        let result = download_single(
111            puller,
112            pusher.clone(),
113            DownloadOptions {
114                retry_gap: Duration::from_secs(1),
115                push_queue_cap: 1024,
116            },
117        )
118        .await;
119
120        let mut pull_progress: Vec<ProgressEntry> = Vec::new();
121        let mut push_progress: Vec<ProgressEntry> = Vec::new();
122        while let Ok(e) = result.event_chain.recv().await {
123            match e {
124                Event::PullProgress(_, p) => {
125                    pull_progress.merge_progress(p);
126                }
127                Event::PushProgress(_, p) => {
128                    push_progress.merge_progress(p);
129                }
130                _ => {}
131            }
132        }
133        dbg!(&pull_progress);
134        dbg!(&push_progress);
135        assert_eq!(pull_progress, download_chunks);
136        assert_eq!(push_progress, download_chunks);
137
138        result.join().await.unwrap();
139        assert_eq!(&**pusher.receive.lock(), mock_data);
140    }
141}
fast_pull/core/single.rs

fast_pull/core/
single.rs
