diff --git a/src/downloader/downloader.rs b/src/downloader/downloader.rs index 2247d3a..1a5ae4c 100644 --- a/src/downloader/downloader.rs +++ b/src/downloader/downloader.rs @@ -81,6 +81,8 @@ pub struct DownloaderInternal>, /// Current retry count retry_count: AtomicI64, + /// The maximum retry count for each part. + pub max_part_retry_count: AtomicI64, } impl DownloaderInternal { @@ -160,6 +162,7 @@ impl DownloaderInternal { max_retry_count: AtomicI64::new(3), retry_interval: RwLock::new(l), retry_count: AtomicI64::new(0), + max_part_retry_count: AtomicI64::new(3), })) } } @@ -186,6 +189,24 @@ impl DownloaderIn self.progress.get_mut().take(); } + /// Enable multiple download + pub fn enable_multiple_download(&self) { + self.multi.qstore(true); + if !self.is_multi_threads() { + println!( + "{}", + gettext("Warning: This file will still use single thread mode to download.") + ); + } else { + match self.pd.enable_multi() { + Ok(_) => {} + Err(e) => { + println!("{}", e); + } + } + } + } + /// Enable the progress bar /// * `style` - The style of the progress bar /// * `mults` - The instance of [MultiProgress] if multiple progress bars are needed. @@ -474,6 +495,12 @@ impl Do self.downloader.disable_progress_bar() } + #[inline] + /// Enable multiple download + pub fn enable_multiple_download(&self) { + self.downloader.enable_multiple_download() + } + #[inline] /// Enable the progress bar /// * `style` - The style of the progress bar @@ -511,10 +538,25 @@ impl Do None => {} } self.set_retry_interval(helper.download_retry_interval()); + match helper.download_part_retry() { + Some(u) => self.set_max_part_retry_count(u), + None => {} + } + if helper.multiple_threads_download() { + self.enable_multiple_download() + } } #[inline] - /// Set the maximum retry count. -1 means always. + /// Set the maximum retry count for each part. < 0 means always. + pub fn set_max_part_retry_count(&self, max_part_retry_count: i64) { + self.downloader + .max_part_retry_count + .qstore(max_part_retry_count) + } + + #[inline] + /// Set the maximum retry count. < 0 means always. pub fn set_max_retry_count(&self, max_retry_count: i64) { self.downloader.set_max_retry_count(max_retry_count) } @@ -611,6 +653,7 @@ async fn test_failed_downloader() { DownloaderResult::Ok(v) => { v.set_retry_interval(retry_interval); v.set_max_retry_count(1); + v.set_max_part_retry_count(1); assert_eq!(v.is_created(), true); v.disable_progress_bar(); v.download(); diff --git a/src/downloader/pd_file/file.rs b/src/downloader/pd_file/file.rs index 83ff728..01dcdb2 100644 --- a/src/downloader/pd_file/file.rs +++ b/src/downloader/pd_file/file.rs @@ -40,6 +40,8 @@ lazy_static! { /// The offset of the status in pd file const STATUS_OFFSET: SeekFrom = SeekFrom::Start(10); +/// The offset of the type of the pd file +const TYPE_OFFSET: SeekFrom = SeekFrom::Start(11); /// The offset of the file_size in pd file const FILE_SIZE_OFFSET: SeekFrom = SeekFrom::Start(12); /// The offset of the downloaded_file_size in pd file @@ -133,6 +135,22 @@ impl PdFile { Ok(()) } + /// Enable multiple thread support + pub fn enable_multi(&self) -> Result<(), PdFileError> { + if !self.is_downloading() { + self.ftype.replace_with2(PdFileType::MultiThread); + if !self.is_mem_only() { + self.need_saved.qstore(true); + let mut f = self.file.get_mut(); + let f = f.as_mut().unwrap(); + f.seek(TYPE_OFFSET)?; + f.write_le_u8(self.ftype.get_ref().int_value())?; + self.need_saved.qstore(false); + } + } + Ok(()) + } + /// Force close the file pub fn force_close(&self) { let mut f = self.file.get_mut(); diff --git a/src/downloader/pd_file/part_status.rs b/src/downloader/pd_file/part_status.rs index 65f8b93..45bc811 100644 --- a/src/downloader/pd_file/part_status.rs +++ b/src/downloader/pd_file/part_status.rs @@ -1,5 +1,6 @@ use crate::downloader::pd_file::enums::PdFilePartStatus as PdFilePartStatus2; use crate::downloader::pd_file::error::PdFileError; +use crate::ext::atomic::AtomicQuick; use crate::ext::rw_lock::GetRwLock; use crate::gettext; use int_enum::IntEnum; @@ -9,6 +10,7 @@ use std::fmt::Debug; use std::fmt::Display; use std::io::Read; use std::io::Write; +use std::sync::atomic::AtomicI64; use std::sync::RwLock; /// The data is out of bounds. @@ -61,7 +63,10 @@ impl Debug for PdFilePartStatusInternal { #[derive(Debug)] /// The status of the each part in pd file pub struct PdFilePartStatus { + /// The internal type status: RwLock, + /// Retry count + _retry_count: AtomicI64, } impl PdFilePartStatus { @@ -74,6 +79,7 @@ impl PdFilePartStatus { status.set_downloaded_size(0); Self { status: RwLock::new(status), + _retry_count: AtomicI64::new(0), } } @@ -83,6 +89,12 @@ impl PdFilePartStatus { self.status.get_ref().downloaded_size() } + #[inline] + /// Returns the current retry count. + pub fn retry_count(&self) -> i64 { + self._retry_count.qload() + } + #[inline] /// Returns the status of this part pub fn status(&self) -> PdFilePartStatus2 { @@ -97,6 +109,12 @@ impl PdFilePartStatus { } } + #[inline] + /// Set the new retry count. + pub fn set_retry_count(&self, count: i64) { + self._retry_count.qstore(count) + } + /// Set the status of this part pub fn set_status(&self, status: PdFilePartStatus2) -> Result<(), PdFileError> { match self.status.get_mut().set_status_checked(status) { @@ -152,6 +170,7 @@ impl PdFilePartStatus { status.set_downloaded_size(ds); Ok(Self { status: RwLock::new(status), + _retry_count: AtomicI64::new(0), }) } @@ -200,4 +219,7 @@ fn test_part_status() { status.set_downloaded_size(0x323133).unwrap(); assert_eq!(status.downloaded_size(), 0x323133); assert_eq!(status.to_bytes(), vec![179, 196, 200, 0]); + assert_eq!(status.retry_count(), 0); + status.set_retry_count(3); + assert_eq!(status.retry_count(), 3); } diff --git a/src/opthelper.rs b/src/opthelper.rs index 0d692d8..6ac1711 100644 --- a/src/opthelper.rs +++ b/src/opthelper.rs @@ -45,6 +45,18 @@ impl OptHelper { } } + /// Return max retry count of each part when downloading in multiple thread mode. + pub fn download_part_retry(&self) -> Option { + if self.opt.get_ref().download_part_retry.is_some() { + return Some(self.opt.get_ref().download_part_retry.unwrap()); + } + let re = self.settings.get_ref().get("download-part-retry"); + if re.is_some() { + return Some(re.unwrap().as_i64().unwrap()); + } + None + } + /// Return retry count when downloading failed. pub fn download_retry(&self) -> Option { if self.opt.get_ref().download_retry.is_some() { @@ -90,6 +102,28 @@ impl OptHelper { } } + /// Return whether to enable multiple threads download. + pub fn multiple_threads_download(&self) -> bool { + match self.opt.get_ref().multiple_threads_download { + Some(r) => { + return r; + } + None => {} + } + if self + .settings + .get_ref() + .have_bool("multiple-threads-download") + { + return self + .settings + .get_ref() + .get_bool("multiple-threads-download") + .unwrap(); + } + false + } + pub fn update(&self, opt: CommandOpts, settings: SettingStore) { if settings.have("author-name-filters") { self._author_name_filters.replace_with2( diff --git a/src/opts.rs b/src/opts.rs index 50b2d82..3e3e2ea 100644 --- a/src/opts.rs +++ b/src/opts.rs @@ -76,6 +76,10 @@ pub struct CommandOpts { pub download_retry: Option, /// Retry interval when downloading files. pub download_retry_interval: Option>, + /// Whether to enable multiple threads download. + pub multiple_threads_download: Option, + /// Max retry count of each part when downloading in multiple thread mode. + pub download_part_retry: Option, } impl CommandOpts { @@ -98,6 +102,8 @@ impl CommandOpts { download_multiple_images: None, download_retry: None, download_retry_interval: None, + multiple_threads_download: None, + download_part_retry: None, } } @@ -280,6 +286,26 @@ pub fn parse_cmd() -> Option { gettext("The interval (in seconds) between two retries when downloading files."), "LIST", ); + opts.opt( + "", + "multiple-threads-download", + format!( + "{} ({} {})", + gettext("Whether to enable multiple threads download."), + gettext("Default:"), + "yes" + ) + .as_str(), + "yes/no", + HasArg::Maybe, + getopts::Occur::Optional, + ); + opts.optopt( + "", + "download-part-retry", + gettext("Max retry count of each part when downloading in multiple thread mode."), + "COUNT", + ); let result = match opts.parse(&argv[1..]) { Ok(m) => m, Err(err) => { @@ -447,6 +473,28 @@ pub fn parse_cmd() -> Option { } re.as_mut().unwrap().download_retry_interval = Some(r.unwrap()); } + match parse_i64(result.opt_str("download-part-retry")) { + Ok(r) => { + re.as_mut().unwrap().download_part_retry = r; + } + Err(e) => { + println!("{} {}", gettext("Failed to parse retry count:"), e); + return None; + } + } + match parse_optional_opt(&result, "multiple-threads-download", true, parse_bool) { + Ok(b) => re.as_mut().unwrap().multiple_threads_download = b, + Err(e) => { + println!( + "{} {}", + gettext("Failed to parse :") + .replace("", "multiple-threads-download") + .as_str(), + e + ); + return None; + } + } re } diff --git a/src/settings_list.rs b/src/settings_list.rs index 4e0eb74..296dedf 100644 --- a/src/settings_list.rs +++ b/src/settings_list.rs @@ -22,7 +22,9 @@ pub fn get_settings_list() -> Vec { SettingDes::new("use-progress-bar", gettext("Whether to enable progress bar."), JsonValueType::Multiple, Some(check_user_or_not)).unwrap(), SettingDes::new("download-multiple-images", gettext("Download multiple images at the same time."), JsonValueType::Boolean, None).unwrap(), SettingDes::new("download-retry", gettext("Max retry count if download failed."), JsonValueType::Number, Some(check_i64)).unwrap(), - SettingDes::new("download-retry-interval", gettext("The interval (in seconds) between two retries when downloading files."), JsonValueType::Multiple, Some(check_retry_interval)).unwrap() + SettingDes::new("download-retry-interval", gettext("The interval (in seconds) between two retries when downloading files."), JsonValueType::Multiple, Some(check_retry_interval)).unwrap(), + SettingDes::new("multiple-threads-download", gettext("Whether to enable multiple threads download."), JsonValueType::Boolean, None).unwrap(), + SettingDes::new("download-part-retry", gettext("Max retry count of each part when downloading in multiple thread mode."), JsonValueType::Number, Some(check_i64)).unwrap(), ] }