1use crate::{Config, Error, Event, Tx};
2use fast_down::{
3 BoxPusher, Merge, UrlInfo,
4 fast_puller::{FastDownPuller, FastDownPullerOptions, build_client},
5 http::Prefetch,
6 invert,
7 multi::{self, download_multi},
8 single::{self, download_single},
9};
10use parking_lot::Mutex;
11use reqwest::{Response, header::HeaderMap};
12use std::{
13 collections::HashMap,
14 sync::{
15 Arc,
16 atomic::{AtomicBool, Ordering},
17 },
18};
19use tokio_util::sync::CancellationToken;
20use url::Url;
21
22#[derive(Debug)]
23pub struct DownloadTask {
24 pub info: UrlInfo,
25 pub config: Config,
26 pub resp: Option<Arc<Mutex<Option<Response>>>>,
27 pub tx: Tx,
28 pub is_running: AtomicBool,
29}
30
31#[must_use]
32fn parse_headers(headers: &HashMap<String, String>) -> Arc<HeaderMap> {
33 headers
34 .iter()
35 .map(|(k, v)| (k.parse(), v.parse()))
36 .filter_map(|(k, v)| k.ok().zip(v.ok()))
37 .collect::<HeaderMap>()
38 .into()
39}
40
41pub async fn prefetch(url: Url, config: Config, tx: Tx) -> Result<DownloadTask, Error> {
48 let headers = parse_headers(&config.headers);
49 let local_addr: Arc<[_]> = config.local_address.clone().into();
50 let client = build_client(
51 &headers,
52 config.proxy.as_deref(),
53 config.accept_invalid_certs,
54 config.accept_invalid_hostnames,
55 local_addr.first().copied(),
56 )?;
57 let mut retry_count = 0;
58 let (info, resp) = loop {
59 match client.prefetch(url.clone()).await {
60 Ok(t) => break t,
61 Err((e, t)) => {
62 let _ = tx.send(Event::PrefetchError(format!("{e:?}")));
63 retry_count += 1;
64 if retry_count >= config.retry_times {
65 return Err(Error::PrefetchTimeout(e));
66 }
67 tokio::time::sleep(t.unwrap_or(config.retry_gap)).await;
68 }
69 }
70 };
71 Ok(DownloadTask {
72 config,
73 resp: Some(Arc::new(Mutex::new(Some(resp)))),
74 tx,
75 info,
76 is_running: AtomicBool::new(false),
77 })
78}
79
80struct RunGuard<'a>(&'a AtomicBool);
81impl Drop for RunGuard<'_> {
82 fn drop(&mut self) {
83 self.0.store(false, Ordering::Release);
84 }
85}
86
87impl DownloadTask {
88 pub async fn start_with_pusher(
98 &self,
99 pusher: BoxPusher,
100 cancel_token: CancellationToken,
101 ) -> Result<(), Error> {
102 if self
103 .is_running
104 .compare_exchange(false, true, Ordering::Acquire, Ordering::Relaxed)
105 .is_err()
106 {
107 return Err(Error::AlreadyRunning);
108 }
109 let _guard = RunGuard(&self.is_running);
110 let progress = self.config.downloaded_chunk.clone();
111 let puller = FastDownPuller::new(FastDownPullerOptions {
112 url: self.info.final_url.clone(),
113 headers: parse_headers(&self.config.headers),
114 proxy: self.config.proxy.as_deref(),
115 available_ips: self.config.local_address.clone().into(),
116 accept_invalid_certs: self.config.accept_invalid_certs,
117 accept_invalid_hostnames: self.config.accept_invalid_hostnames,
118 file_id: self.info.file_id.clone(),
119 resp: self.resp.clone(),
120 })?;
121 let threads = if self.info.fast_download {
122 self.config.threads.max(1)
123 } else {
124 1
125 };
126 let result = if self.info.fast_download {
127 download_multi(
128 puller,
129 pusher,
130 multi::DownloadOptions {
131 download_chunks: invert(
132 progress.lock().iter().cloned(),
133 self.info.size,
134 self.config.chunk_window,
135 ),
136 retry_gap: self.config.retry_gap,
137 concurrent: threads,
138 pull_timeout: self.config.pull_timeout,
139 push_queue_cap: self.config.write_queue_cap,
140 min_chunk_size: self.config.min_chunk_size,
141 max_speculative: self.config.max_speculative,
142 },
143 )
144 } else {
145 download_single(
146 puller,
147 pusher,
148 single::DownloadOptions {
149 retry_gap: self.config.retry_gap,
150 push_queue_cap: self.config.write_queue_cap,
151 },
152 )
153 };
154 let mut cancelled = false;
155 loop {
156 tokio::select! {
157 () = cancel_token.cancelled(), if !cancelled => {
158 result.abort();
159 cancelled = true;
160 },
161 e = result.event_chain.recv() => match e {
162 Ok(e) => {
163 let _ = self.tx.send((&e).into());
164 if let fast_down::Event::PushProgress(_, range) = e {
165 let mut p = progress.lock();
166 if range.start == 0 && !self.info.fast_download {
167 p.clear();
168 }
169 p.merge_progress(range);
170 }
171 },
172 Err(_) => break,
173 }
174 }
175 }
176 result.join().await?;
177 Ok(())
178 }
179
180 #[cfg(feature = "file")]
189 pub async fn start(
190 &self,
191 save_path: std::path::PathBuf,
192 cancel_token: CancellationToken,
193 ) -> Result<(), Error> {
194 let pusher = get_pusher(
195 &self.info,
196 self.config.write_method.clone(),
197 self.config.sync_all,
198 self.config.cache_high_watermark,
199 self.config.cache_low_watermark,
200 self.config.write_buffer_size,
201 &save_path,
202 );
203 let pusher = tokio::select! {
204 () = cancel_token.cancelled() => return Ok(()),
205 pusher = pusher => pusher.map_err(Error::Io)?,
206 };
207 self.start_with_pusher(pusher, cancel_token).await
208 }
209
210 #[allow(clippy::missing_panics_doc)]
211 #[cfg(feature = "mem")]
218 pub async fn start_in_memory(&self, cancel_token: CancellationToken) -> Result<Vec<u8>, Error> {
219 #[allow(clippy::cast_possible_truncation)]
220 let pusher = fast_down::mem::MemPusher::with_capacity(self.info.size as usize);
221 self.start_with_pusher(BoxPusher::new(pusher.clone()), cancel_token)
222 .await?;
223 Ok(Arc::try_unwrap(pusher.receive).unwrap().into_inner())
224 }
225}
226
227#[cfg(feature = "file")]
228pub async fn get_pusher(
232 info: &UrlInfo,
233 write_method: crate::WriteMethod,
234 sync_all: bool,
235 high_watermark: usize,
236 low_watermark: usize,
237 buffer_size: usize,
238 save_path: &std::path::Path,
239) -> Result<BoxPusher, String> {
240 let file = tokio::fs::OpenOptions::new()
241 .create(true)
242 .write(true)
243 .read(true)
244 .truncate(false)
245 .open(&save_path)
246 .await
247 .map_err(|e| format!("{e:?}"))?;
248 #[cfg(target_pointer_width = "64")]
249 if info.fast_download && write_method == crate::WriteMethod::Mmap {
250 use fast_down::file::MmapFilePusher;
251 let pusher = BoxPusher::new(
252 MmapFilePusher::new(file, info.size, sync_all)
253 .await
254 .map_err(|e| format!("{e:?}"))?,
255 );
256 return Ok(pusher);
257 }
258 let pusher = fast_down::file::CacheFilePusher::new(
259 file,
260 info.size,
261 sync_all,
262 high_watermark,
263 low_watermark,
264 buffer_size,
265 )
266 .await
267 .map_err(|e| format!("{e:?}"))?;
268 Ok(BoxPusher::new(pusher))
269}