From 6be675bf1eb0c510b8ac534672600d23a381685c Mon Sep 17 00:00:00 2001 From: lifegpc Date: Sun, 17 Jul 2022 00:49:18 +0000 Subject: [PATCH] Add task manager --- src/download.rs | 62 +++++++++------- src/ext/any.rs | 6 ++ src/ext/mod.rs | 1 + src/main.rs | 1 + src/task_manager.rs | 167 ++++++++++++++++++++++++++++++++++++++++++++ 5 files changed, 213 insertions(+), 24 deletions(-) create mode 100644 src/ext/any.rs create mode 100644 src/task_manager.rs diff --git a/src/download.rs b/src/download.rs index 45b66af..aa82b7a 100644 --- a/src/download.rs +++ b/src/download.rs @@ -13,6 +13,7 @@ use crate::downloader::DownloaderHelper; use crate::downloader::DownloaderResult; use crate::downloader::LocalFile; use crate::error::PixivDownloaderError; +use crate::ext::any::AsAny; use crate::ext::try_err::TryErr; use crate::fanbox::article::block::FanboxArticleBlock; use crate::fanbox::check::CheckUnknown; @@ -23,6 +24,8 @@ use crate::opthelper::get_helper; use crate::pixiv_link::FanboxPostID; use crate::pixiv_link::PixivID; use crate::pixiv_web::PixivWebClient; +use crate::task_manager::get_progress_bar; +use crate::task_manager::TaskManager; #[cfg(feature = "ugoira")] use crate::ugoira::{convert_ugoira_to_mp4, UgoiraFrames}; use crate::utils::get_file_name_from_url; @@ -33,6 +36,7 @@ use reqwest::IntoUrl; use std::fs::create_dir_all; use std::path::PathBuf; use std::sync::Arc; +use tokio::task::JoinHandle; impl Main { pub async fn download(&mut self) -> i32 { @@ -307,8 +311,8 @@ impl Main { if pages_data.is_some() && helper.download_multiple_images() { let mut np = 0u16; let pages_data = pages_data.as_ref().unwrap(); - let progress_bars = Arc::new(MultiProgress::new()); - let mut tasks = Vec::new(); + let progress_bars = get_progress_bar(); + let tasks = TaskManager::default(); let mut re: Result<(), PixivDownloaderError> = Ok(()); for page in pages_data.members() { let url = page["urls"]["original"].as_str(); @@ -319,23 +323,30 @@ impl Main { ); continue; } - let f = tokio::spawn(Self::download_artwork_link( - url.unwrap().to_owned(), - np, - Some(Arc::clone(&progress_bars)), - Arc::clone(&datas), - Arc::clone(&base), - )); - tasks.push(f); + tasks + .add_task(Self::download_artwork_link( + url.unwrap().to_owned(), + np, + Some(Arc::clone(&progress_bars)), + Arc::clone(&datas), + Arc::clone(&base), + )) + .await; np += 1; } - for task in tasks { - let r = task.await; - let r = match r { - Ok(r) => r, - Err(e) => Err(PixivDownloaderError::from(e)), - }; - concat_pixiv_downloader_error!(re, r); + tasks.join().await; + let tasks = tasks.take_finished_tasks(); + for mut task in tasks { + let t = task.as_any_mut(); + if let Some(task) = t.downcast_mut::>>() + { + let r = task.await; + let r = match r { + Ok(r) => r, + Err(e) => Err(PixivDownloaderError::from(e)), + }; + concat_pixiv_downloader_error!(re, r); + } } return re; } else if pages_data.is_some() { @@ -379,7 +390,6 @@ impl Main { /// * `progress_bars` - Multiple progress bars /// * `base` - The directory of the target pub async fn download_fanbox_file( - &self, dh: DownloaderHelper, progress_bars: Option>, base: Arc, @@ -409,7 +419,6 @@ impl Main { /// * `datas` - The artwork's data /// * `base` - The directory of the target pub async fn download_fanbox_image( - &self, dh: DownloaderHelper, np: u16, progress_bars: Option>, @@ -536,7 +545,7 @@ impl Main { let dh = img .download_original_url()? .try_err(gettext("Can not get original url for image"))?; - self.download_fanbox_image( + Self::download_fanbox_image( dh, np, None, @@ -561,8 +570,7 @@ impl Main { let dh = f .download_url()? .try_err(gettext("Failed to get url of the file."))?; - self.download_fanbox_file(dh, None, Arc::clone(&base)) - .await?; + Self::download_fanbox_file(dh, None, Arc::clone(&base)).await?; } } FanboxPost::Image(img) => { @@ -582,8 +590,14 @@ impl Main { let dh = img .download_original_url()? .try_err(gettext("Can not get original url for image"))?; - self.download_fanbox_image(dh, np, None, Arc::clone(&datas), Arc::clone(&base)) - .await?; + Self::download_fanbox_image( + dh, + np, + None, + Arc::clone(&datas), + Arc::clone(&base), + ) + .await?; np += 1; } } diff --git a/src/ext/any.rs b/src/ext/any.rs new file mode 100644 index 0000000..e9646b1 --- /dev/null +++ b/src/ext/any.rs @@ -0,0 +1,6 @@ +use std::any::Any; + +pub trait AsAny { + fn as_any(&self) -> &T; + fn as_any_mut(&mut self) -> &mut T; +} diff --git a/src/ext/mod.rs b/src/ext/mod.rs index 7588f22..6d83387 100644 --- a/src/ext/mod.rs +++ b/src/ext/mod.rs @@ -1,3 +1,4 @@ +pub mod any; pub mod atomic; pub mod cstr; #[cfg(feature = "flagset")] diff --git a/src/main.rs b/src/main.rs index acd5d01..9d057b8 100644 --- a/src/main.rs +++ b/src/main.rs @@ -43,6 +43,7 @@ mod retry_interval; mod server; mod settings; mod settings_list; +mod task_manager; #[cfg(feature = "ugoira")] mod ugoira; mod utils; diff --git a/src/task_manager.rs b/src/task_manager.rs new file mode 100644 index 0000000..6a12383 --- /dev/null +++ b/src/task_manager.rs @@ -0,0 +1,167 @@ +use crate::ext::any::AsAny; +use crate::ext::replace::ReplaceWith; +use crate::ext::replace::ReplaceWith2; +use crate::ext::rw_lock::GetRwLock; +use futures_util::lock::Mutex; +use indicatif::MultiProgress; +use std::any::Any; +use std::future::Future; +use std::sync::Arc; +use std::sync::RwLock; +use std::time::Duration; +use tokio::task::JoinHandle; + +lazy_static! { + #[doc(hidden)] + static ref TOTAL_DOWNLOAD_TASK_COUNT: Arc> = Arc::new(Mutex::new(0)); + #[doc(hidden)] + static ref PROGRESS_BAR: Arc = Arc::new(MultiProgress::new()); +} + +pub trait IsFinished { + fn is_finished(&self) -> bool; +} + +impl IsFinished for JoinHandle { + fn is_finished(&self) -> bool { + self.is_finished() + } +} + +pub trait IsFinishedAny: IsFinished + Any {} + +impl IsFinishedAny for T where T: IsFinished + Any {} + +impl AsAny for Box { + fn as_any(&self) -> &(dyn Any + Send + Sync + 'static) { + self + } + + fn as_any_mut(&mut self) -> &mut (dyn Any + Send + Sync + 'static) { + self + } +} + +pub trait GetMaxCount { + fn get_max_count(&self) -> usize; +} + +pub struct MaxDownloadTasks { + _unused: [u8; 0], +} + +impl MaxDownloadTasks { + pub fn new() -> Self { + MaxDownloadTasks { _unused: [] } + } +} + +impl GetMaxCount for MaxDownloadTasks { + fn get_max_count(&self) -> usize { + 5 + } +} + +/// Task manager +pub struct TaskManager { + /// Current running task + tasks: RwLock>>, + /// Finished task + finished_tasks: RwLock>>, + /// Total task count + task_count: Arc>, + max_count: Box, +} + +impl TaskManager { + /// Create a new instance + pub fn new(task_count: Arc>, max_count: T) -> Self { + Self { + tasks: RwLock::new(Vec::new()), + finished_tasks: RwLock::new(Vec::new()), + task_count, + max_count: Box::new(max_count), + } + } + + /// Add a new task. + pub async fn add_task(&self, future: F) + where + F: Future + Send + 'static, + F::Output: Send + 'static, + JoinHandle: IsFinishedAny + Send + Sync + 'static, + { + let total_count = self.max_count.get_max_count(); + loop { + { + let mut count = self.task_count.lock().await; + if *count < total_count { + self.tasks + .get_mut() + .push(Box::new(tokio::task::spawn(future))); + count.replace_with(*count + 1); + break; + } + let tasks = self.tasks.replace_with2(Vec::new()); + let mut new_tasks = Vec::new(); + let mut new_count = *count; + for i in tasks { + if i.is_finished() { + self.finished_tasks.get_mut().push(i); + new_count -= 1; + } else { + new_tasks.push(i); + } + } + self.tasks.replace_with2(new_tasks); + count.replace_with(new_count); + } + tokio::time::sleep(Duration::new(0, 10_000_000)).await; + } + } + + /// Wait all tasks finished. + pub async fn join(&self) { + loop { + { + let mut count = self.task_count.lock().await; + let tasks = self.tasks.replace_with2(Vec::new()); + if tasks.len() == 0 { + break; + } + let mut new_tasks = Vec::new(); + let mut new_count = *count; + for i in tasks { + if i.is_finished() { + self.finished_tasks.get_mut().push(i); + new_count -= 1; + } else { + new_tasks.push(i); + } + } + self.tasks.replace_with2(new_tasks); + count.replace_with(new_count); + } + tokio::time::sleep(Duration::new(0, 10_000_000)).await; + } + } + + /// Take all finished tasks + pub fn take_finished_tasks(&self) -> Vec> { + self.finished_tasks.replace_with2(Vec::new()) + } +} + +pub fn get_progress_bar() -> Arc { + Arc::clone(&PROGRESS_BAR) +} + +pub fn get_total_download_task_count() -> Arc> { + Arc::clone(&TOTAL_DOWNLOAD_TASK_COUNT) +} + +impl Default for TaskManager { + fn default() -> Self { + Self::new(get_total_download_task_count(), MaxDownloadTasks::new()) + } +}