Skip to main content

fast_down_ffi/
download.rs

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
41/// This function can be cancelled by dropping the returned Future.
42///
43/// # Errors
44///
45/// - Returns an error if proxy configuration fails.
46/// - Returns an error if the maximum number of retries is reached.
47pub 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    /// Cannot be cancelled by dropping the Future — doing so would leave the
89    /// written content incomplete.
90    /// Multiple calls will resume from where the last call left off.
91    ///
92    /// # Errors
93    ///
94    /// - Returns `Err` if the task is already running.
95    /// - Returns `Err` if proxy configuration fails.
96    /// - Returns `Err` if the download fails unexpectedly.
97    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    /// Cannot be cancelled by dropping the Future — doing so would leave the
181    /// written content incomplete.
182    /// The pusher is determined by [`crate::WriteMethod`].
183    ///
184    /// # Errors
185    ///
186    /// Returns `Err` if opening or creating the file fails. Otherwise see
187    /// [`crate::DownloadTask::start_with_pusher`].
188    #[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    /// Cannot be cancelled by dropping the Future — doing so would leave the
212    /// written content incomplete.
213    ///
214    /// # Errors
215    ///
216    /// Same as [`crate::DownloadTask::start_with_pusher`].
217    #[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")]
228/// # Errors
229///
230/// Returns an error if opening or creating the file fails.
231pub 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}