From 960288e350d12ff9867fb1e6770dd5f7900b94d2 Mon Sep 17 00:00:00 2001 From: lifegpc Date: Mon, 13 Jun 2022 04:31:59 +0000 Subject: [PATCH] Update --- src/downloader/downloader.rs | 66 ++++++++++++++++++++++++++++++++---- src/downloader/tasks.rs | 53 +++++++++++++++++++++++++++-- 2 files changed, 109 insertions(+), 10 deletions(-) diff --git a/src/downloader/downloader.rs b/src/downloader/downloader.rs index ef497ad..b44b571 100644 --- a/src/downloader/downloader.rs +++ b/src/downloader/downloader.rs @@ -4,9 +4,10 @@ use super::enums::DownloaderResult; use super::enums::DownloaderStatus; use super::error::DownloaderError; use super::local_file::LocalFile; -use super::tasks::create_download_tasks_simple; +use super::tasks::check_tasks; use crate::ext::atomic::AtomicQuick; use crate::ext::io::ClearFile; +use crate::ext::replace::ReplaceWith2; use crate::ext::rw_lock::GetRwLock; use crate::ext::try_err::TryErr; use crate::utils::ask_need_overwrite; @@ -14,6 +15,8 @@ use crate::webclient::WebClient; use crate::webclient::ToHeaders; use reqwest::IntoUrl; use std::collections::HashMap; +#[cfg(test)] +use std::fs::create_dir; use std::fs::remove_file; use std::io::Seek; use std::io::SeekFrom; @@ -42,7 +45,7 @@ pub struct DownloaderInternal { /// The status of the downloader status: RwLock, /// All tasks - tasks: RwLock>>>, + pub tasks: RwLock>>>, /// Whether to enable mulitple thread mode multi: AtomicBool, } @@ -142,6 +145,12 @@ impl DownloaderInternal { *self.status.get_ref() == DownloaderStatus::Created } + #[inline] + /// Returns true if the downloader is downloading now. + pub fn is_downloading(&self) -> bool { + *self.status.get_ref() == DownloaderStatus::Downloading + } + #[inline] /// Returns true if is multiple thread mode. pub fn is_multi_threads(&self) -> bool { @@ -161,6 +170,18 @@ impl DownloaderInternal { } } + #[inline] + /// Set the status to [DownloaderStatus::Downloading] and returns the current value + pub fn set_downloading(&self) -> DownloaderStatus { + self.status.replace_with2(DownloaderStatus::Downloading) + } + + #[inline] + /// Set the status to [DownloaderStatus::Downloaded] and returns the current value + pub fn set_downloaded(&self) -> DownloaderStatus { + self.status.replace_with2(DownloaderStatus::Downloaded) + } + /// Write datas to the file. /// * `data` - Data pub fn write(&self, data: &[u8]) -> Result<(), DownloaderError> { @@ -176,6 +197,8 @@ impl DownloaderInternal { pub struct Downloader { /// internal type downloader: Arc>, + /// The task to check status. + task: RwLock>>>, } impl Downloader { @@ -187,7 +210,7 @@ impl Downloader { pub fn new + ?Sized>(url: U, headers: H, path: Option<&P>, overwrite: Option) -> Result, DownloaderError> { Ok(match DownloaderInternal::::new(url, headers, path, overwrite)? { DownloaderResult::Ok(d) => { - DownloaderResult::Ok(Self { downloader: Arc::new(d) }) + DownloaderResult::Ok(Self { downloader: Arc::new(d), task: RwLock::new(None) }) } DownloaderResult::Canceled => { DownloaderResult::Canceled } }) @@ -213,13 +236,42 @@ impl Downloader { if !self.is_created() { return self.downloader.get_status(); } - if !self.is_multi_threads() { - let task = tokio::spawn(create_download_tasks_simple(Arc::clone(&self.downloader))); - self.downloader.add_task(task); - } + self.downloader.set_downloading(); + self.task.get_mut().replace(tokio::spawn(check_tasks(Arc::clone(&self.downloader)))); self.downloader.get_status() } + pub async fn join(&self) -> Result<(), DownloaderError> { + match self.task.get_mut().deref_mut() { + Some(v) => { + let re = v.await; + re.unwrap() + } + None => { Ok(()) } + } + } + define_downloader_fn!(is_created, bool, "Returns true if the downloader is created just now."); + define_downloader_fn!(is_downloading, bool, "Returns true if the downloader is downloading now."); define_downloader_fn!(is_multi_threads, bool, "Returns true if is multiple thread mode."); } + +#[tokio::test(flavor = "multi_thread", worker_threads = 4)] +async fn test_downloader() { + let p = Path::new("./test"); + if !p.exists() { + let re = create_dir("./test"); + assert!(re.is_ok() || p.exists()); + } + let url = "https://i.pximg.net/img-original/img/2022/06/12/23/49/43/99014872_p0.png"; + let pb = p.join("99014872_p0.png"); + let downloader = Downloader::::new(url, json::object!{"referer": "https://www.pixiv.net/"}, Some(&pb), Some(true)).unwrap(); + match downloader { + DownloaderResult::Ok(v) => { + assert_eq!(v.is_created(), true); + v.download(); + assert!(v.join().await.is_ok()); + } + DownloaderResult::Canceled => { panic!("This should not happened.") } + } +} diff --git a/src/downloader/tasks.rs b/src/downloader/tasks.rs index 2657567..9140a30 100644 --- a/src/downloader/tasks.rs +++ b/src/downloader/tasks.rs @@ -1,4 +1,5 @@ use crate::ext::io::ClearFile; +use crate::ext::rw_lock::GetRwLock; use crate::ext::try_err::TryErr; use crate::gettext; use super::error::DownloaderError; @@ -6,11 +7,13 @@ use super::downloader::DownloaderInternal; use futures_util::StreamExt; use http_content_range::ContentRange; use reqwest::Response; +use spin_on::spin_on; use std::ops::Deref; use std::io::Seek; use std::io::SeekFrom; use std::io::Write; use std::sync::Arc; +use std::time::Duration; /// Create a download tasks in simple thread mode. pub async fn create_download_tasks_simple(d: Arc>) -> Result<(), DownloaderError> { @@ -40,10 +43,10 @@ pub async fn create_download_tasks_simple(d: Arc { match data { Ok(data) => { - d.pd.inc(data.len() as u64)?; + if !is_multi { + d.pd.inc(data.len() as u64)?; + } d.write(&data)?; } Err(e) => { @@ -136,3 +141,45 @@ pub async fn handle_download(d: Arc(d: Arc>) -> Result<(), DownloaderError> { + if !d.is_multi_threads() { + let task = tokio::spawn(create_download_tasks_simple(Arc::clone(&d))); + d.add_task(task); + } + loop { + tokio::time::sleep(Duration::new(0, 10_000_000)).await; + let mut need_break = false; + { + let mut tasks = d.tasks.get_mut(); + tasks.retain_mut(|task| { + if task.is_finished() { + let re = spin_on(task).unwrap(); + match re { + Ok(_) => { + if !d.is_multi_threads() { + d.set_downloaded(); + need_break = true; + } + } + Err(_) => { + let task = tokio::spawn(create_download_tasks_simple(Arc::clone(&d))); + d.add_task(task); + } + } + false + } else { + true + } + }); + } + if !d.is_multi_threads() && d.tasks.get_ref().len() == 0 { + need_break = true; + } + if need_break { + break; + } + } + Ok(()) +}