1extern crate alloc;
2use super::macros::poll_ok;
3use crate::{DownloadResult, Event, ProgressEntry, RandPuller, RandPusher, Total, WorkerId};
4use alloc::{sync::Arc, vec::Vec};
5use bytes::Bytes;
6use core::{
7 num::NonZeroUsize,
8 sync::atomic::{AtomicUsize, Ordering},
9 time::Duration,
10};
11use fast_steal::{Executor, Handle, Task, TaskList};
12use futures::TryStreamExt;
13use tokio::task::AbortHandle;
14
15#[derive(Debug, Clone)]
16pub struct DownloadOptions {
17 pub download_chunks: Vec<ProgressEntry>,
18 pub concurrent: NonZeroUsize,
19 pub retry_gap: Duration,
20 pub push_queue_cap: usize,
21 pub min_chunk_size: u64,
22}
23
24pub async fn download_multi<R, W>(
25 puller: R,
26 mut pusher: W,
27 options: DownloadOptions,
28) -> DownloadResult<R::Error, W::Error>
29where
30 R: RandPuller + 'static + Sync,
31 W: RandPusher + 'static,
32{
33 let (tx, event_chain) = kanal::unbounded_async();
34 let (tx_push, rx_push) =
35 kanal::bounded_async::<(WorkerId, ProgressEntry, Bytes)>(options.push_queue_cap);
36 let tx_clone = tx.clone();
37 let push_handle = tokio::spawn(async move {
38 while let Ok((id, spin, data)) = rx_push.recv().await {
39 poll_ok!(
40 {},
41 pusher.push(spin.clone(), data.clone()).await,
42 id @ tx_clone => PushError,
43 options.retry_gap
44 );
45 tx_clone.send(Event::PushProgress(id, spin)).await.unwrap();
46 }
47 poll_ok!(
48 {},
49 pusher.flush().await,
50 tx_clone => FlushError,
51 options.retry_gap
52 );
53 });
54 let executor: TokioExecutor<R, W> = TokioExecutor {
55 tx,
56 tx_push,
57 puller,
58 retry_gap: options.retry_gap,
59 id: Arc::new(AtomicUsize::new(0)),
60 min_chunk_size: options.min_chunk_size,
61 };
62 let task_list = TaskList::run(
63 options.concurrent.get(),
64 options.min_chunk_size,
65 &options.download_chunks[..],
66 executor,
67 );
68 DownloadResult::new(
69 event_chain,
70 push_handle,
71 &task_list
72 .handles()
73 .iter()
74 .map(|h| h.0.clone())
75 .collect::<Arc<[_]>>(),
76 )
77}
78
79#[derive(Clone)]
80pub struct TokioHandle(AbortHandle);
81impl Handle for TokioHandle {
82 type Output = ();
83 fn abort(&mut self) -> Self::Output {
84 self.0.abort();
85 }
86}
87pub struct TokioExecutor<R, W>
88where
89 R: RandPuller + 'static,
90 W: RandPusher + 'static,
91{
92 tx: kanal::AsyncSender<Event<R::Error, W::Error>>,
93 tx_push: kanal::AsyncSender<(WorkerId, ProgressEntry, Bytes)>,
94 puller: R,
95 retry_gap: Duration,
96 id: Arc<AtomicUsize>,
97 min_chunk_size: u64,
98}
99impl<R, W> Executor for TokioExecutor<R, W>
100where
101 R: RandPuller + 'static + Sync,
102 W: RandPusher + 'static,
103{
104 type Handle = TokioHandle;
105 fn execute(self: Arc<Self>, task: Arc<Task>, task_list: Arc<TaskList<Self>>) -> Self::Handle {
106 let id = self.id.fetch_add(1, Ordering::SeqCst);
107 let handle = tokio::spawn(async move {
108 'steal_task: loop {
109 let mut start = task.start();
110 if start >= task.end() {
111 if task_list.steal(&task, 2 * self.min_chunk_size) {
112 continue;
113 }
114 break;
115 }
116 self.tx.send(Event::Pulling(id)).await.unwrap();
117 let download_range = start..task.end();
118 let mut puller = self.puller.clone();
119 let mut stream = puller.pull(&download_range);
120 loop {
121 match stream.try_next().await {
122 Ok(Some(mut chunk)) => {
123 let len = chunk.len() as u64;
124 task.fetch_add_start(len);
125 let range_start = start;
126 start += len;
127 let range_end = start.min(task.end());
128 if range_start >= range_end {
129 continue 'steal_task;
130 }
131 let span = range_start..range_end;
132 let len = span.total() as usize;
133 self.tx
134 .send(Event::PullProgress(id, span.clone()))
135 .await
136 .unwrap();
137 self.tx_push
138 .send((id, span, chunk.split_to(len)))
139 .await
140 .unwrap();
141 }
142 Ok(None) => break,
143 Err(e) => {
144 self.tx.send(Event::PullError(id, e)).await.unwrap();
145 tokio::time::sleep(self.retry_gap).await;
146 }
147 }
148 }
149 }
150 self.tx.send(Event::Finished(id)).await.unwrap();
151 task_list.remove(&task);
152 });
153 TokioHandle(handle.abort_handle())
154 }
155}
156
157#[cfg(test)]
158mod tests {
159 extern crate std;
160 use super::*;
161 use crate::{
162 MergeProgress, ProgressEntry,
163 core::mock::{MockRandPuller, MockRandPusher, build_mock_data},
164 };
165 use alloc::vec;
166 use std::dbg;
167
168 #[tokio::test]
169 async fn test_concurrent_download() {
170 let mock_data = build_mock_data(3 * 1024);
171 let puller = MockRandPuller::new(&mock_data);
172 let pusher = MockRandPusher::new(&mock_data);
173 #[allow(clippy::single_range_in_vec_init)]
174 let download_chunks = vec![0..mock_data.len() as u64];
175 let result = download_multi(
176 puller,
177 pusher.clone(),
178 DownloadOptions {
179 concurrent: NonZeroUsize::new(32).unwrap(),
180 retry_gap: Duration::from_secs(1),
181 push_queue_cap: 1024,
182 download_chunks: download_chunks.clone(),
183 min_chunk_size: 1,
184 },
185 )
186 .await;
187
188 let mut pull_progress: Vec<ProgressEntry> = Vec::new();
189 let mut push_progress: Vec<ProgressEntry> = Vec::new();
190 let mut pull_ids = [false; 32];
191 let mut push_ids = [false; 32];
192 while let Ok(e) = result.event_chain.recv().await {
193 match e {
194 Event::PullProgress(id, p) => {
195 pull_ids[id] = true;
196 pull_progress.merge_progress(p);
197 }
198 Event::PushProgress(id, p) => {
199 push_ids[id] = true;
200 push_progress.merge_progress(p);
201 }
202 _ => {}
203 }
204 }
205 dbg!(&pull_progress);
206 dbg!(&push_progress);
207 assert_eq!(pull_progress, download_chunks);
208 assert_eq!(push_progress, download_chunks);
209 assert_eq!(pull_ids, [true; 32]);
210 assert_eq!(push_ids, [true; 32]);
211
212 result.join().await.unwrap();
213 pusher.assert().await;
214 }
215}