lifegpc/pixiv_downloader
1
0
Fork 0
mirror of https://github.com/lifegpc/pixiv_downloader.git synced 2026-07-08 01:32:41 +08:00
Code Issues Packages Projects Releases Wiki Activity

Add task manager

Browse Source
This commit is contained in:
lifegpc
2022-07-17 00:49:18 +00:00
committed by GitHub
parent 8ff1520d63
commit 6be675bf1e
5 changed files with 213 additions and 24 deletions
Download Patch File Download Diff File Expand all files Collapse all files

62
src/download.rs
View File

@@ -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::<JoinHandle<Result<(), PixivDownloaderError>>>()
{
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<Arc<MultiProgress>>,
base: Arc<PathBuf>,
@@ -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<Arc<MultiProgress>>,
@@ -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;
}
}

6
src/ext/any.rs Normal file
View File

@@ -0,0 +1,6 @@
use std::any::Any;
pub trait AsAny<T: Any + ?Sized> {
fn as_any(&self) -> &T;
fn as_any_mut(&mut self) -> &mut T;
}

1
src/ext/mod.rs
View File

@@ -1,3 +1,4 @@
pub mod any;
pub mod atomic;
pub mod cstr;
#[cfg(feature = "flagset")]

1
src/main.rs
View File

@@ -43,6 +43,7 @@ mod retry_interval;
mod server;
mod settings;
mod settings_list;
mod task_manager;
#[cfg(feature = "ugoira")]
mod ugoira;
mod utils;

167
src/task_manager.rs Normal file
View File

@@ -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<Mutex<usize>> = Arc::new(Mutex::new(0));
#[doc(hidden)]
static ref PROGRESS_BAR: Arc<MultiProgress> = Arc::new(MultiProgress::new());
}
pub trait IsFinished {
fn is_finished(&self) -> bool;
}
impl<T> IsFinished for JoinHandle<T> {
fn is_finished(&self) -> bool {
self.is_finished()
}
}
pub trait IsFinishedAny: IsFinished + Any {}
impl<T> IsFinishedAny for T where T: IsFinished + Any {}
impl AsAny<dyn Any + Send + Sync + 'static> for Box<dyn IsFinishedAny + Send + Sync> {
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<Vec<Box<dyn IsFinishedAny + Send + Sync>>>,
/// Finished task
finished_tasks: RwLock<Vec<Box<dyn IsFinishedAny + Send + Sync>>>,
/// Total task count
task_count: Arc<Mutex<usize>>,
max_count: Box<dyn GetMaxCount>,
}
impl TaskManager {
/// Create a new instance
pub fn new<T: GetMaxCount + 'static>(task_count: Arc<Mutex<usize>>, 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<F>(&self, future: F)
where
F: Future + Send + 'static,
F::Output: Send + 'static,
JoinHandle<F::Output>: 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<Box<dyn IsFinishedAny + Send + Sync>> {
self.finished_tasks.replace_with2(Vec::new())
}
}
pub fn get_progress_bar() -> Arc<MultiProgress> {
Arc::clone(&PROGRESS_BAR)
}
pub fn get_total_download_task_count() -> Arc<Mutex<usize>> {
Arc::clone(&TOTAL_DOWNLOAD_TASK_COUNT)
}
impl Default for TaskManager {
fn default() -> Self {
Self::new(get_total_download_task_count(), MaxDownloadTasks::new())
}
}