diff --git a/Cargo.lock b/Cargo.lock index 8f48fcd..d5f1a70 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -1074,6 +1074,12 @@ version = "6.1.0" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "21326818e99cfe6ce1e524c2a805c189a99b5ae555a35d19f9a284b427d86afa" +[[package]] +name = "parse-size" +version = "1.0.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "944553dd59c802559559161f9816429058b869003836120e262e8caec061b7ae" + [[package]] name = "peeking_take_while" version = "0.1.2" @@ -1166,6 +1172,7 @@ dependencies = [ "lazy_static", "link-cplusplus", "modular-bitfield", + "parse-size", "proc_macros", "regex", "reqwest", diff --git a/Cargo.toml b/Cargo.toml index 1487f71..c70ad92 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -23,6 +23,7 @@ int-enum = "0.4" json = "0.12" lazy_static = "1.4" modular-bitfield = "0.11" +parse-size = "1" proc_macros = { path = "proc_macros" } regex = "1" reqwest = { version = "0.11", features = ["brotli", "deflate", "gzip", "rustls-tls", "socks", "stream"] } diff --git a/src/downloader/downloader.rs b/src/downloader/downloader.rs index b8e9668..8c02aab 100644 --- a/src/downloader/downloader.rs +++ b/src/downloader/downloader.rs @@ -33,6 +33,7 @@ use std::ops::DerefMut; use std::path::Path; use std::sync::atomic::AtomicBool; use std::sync::atomic::AtomicI64; +use std::sync::atomic::AtomicU32; use std::sync::atomic::AtomicU64; use std::sync::Arc; use std::sync::RwLock; @@ -86,6 +87,8 @@ pub struct DownloaderInternal { @@ -167,6 +170,7 @@ impl DownloaderInternal { retry_count: AtomicI64::new(0), max_part_retry_count: AtomicI64::new(3), max_threads: AtomicU64::new(8), + part_size: AtomicU32::new(0x10000), })) } } @@ -255,6 +259,16 @@ impl DownloaderIn self.error.get_mut().take() } + #[inline] + /// Return the size of the each part + pub fn get_part_size(&self) -> u32 { + if self.pd.is_downloading() { + self.pd.get_part_size() + } else { + self.part_size.qload() + } + } + /// Increase the retry count and return the duration should waited. /// If [None] is returned, should stop retry. pub fn get_retry_duration(&self) -> Option { @@ -356,6 +370,12 @@ impl DownloaderIn self.error.get_mut().replace(err); } + #[inline] + /// Set the size of the each part when downloading file + pub fn set_part_size(&self, part_size: u32) { + self.part_size.qsave(part_size) + } + #[inline] /// Sets the length of the progress bar pub fn set_progress_bar_length(&self, length: u64) { @@ -483,6 +503,18 @@ impl Do if !self.is_created() { return self.downloader.get_status(); } + if !self.downloader.is_downloading() && self.is_multi_threads() { + match self + .downloader + .pd + .set_part_size(self.downloader.get_part_size()) + { + Ok(_) => {} + Err(e) => { + println!("{}", e); + } + } + } self.downloader.set_downloading(); tokio::spawn(check_tasks(Arc::clone(&self.downloader))); self.downloader.get_status() @@ -556,6 +588,7 @@ impl Do self.enable_multiple_download() } self.set_max_threads(helper.max_threads()); + self.set_part_size(helper.part_size()); } #[inline] @@ -578,6 +611,12 @@ impl Do self.downloader.set_max_threads(max_threads) } + #[inline] + /// Set the size of the each part when downloading file + pub fn set_part_size(&self, part_size: u32) { + self.downloader.set_part_size(part_size) + } + #[inline] /// Set the retry interval. pub fn set_retry_interval(&self, retry_interval: NonTailList) { diff --git a/src/downloader/pd_file/file.rs b/src/downloader/pd_file/file.rs index 01dcdb2..728ef68 100644 --- a/src/downloader/pd_file/file.rs +++ b/src/downloader/pd_file/file.rs @@ -46,6 +46,8 @@ const TYPE_OFFSET: SeekFrom = SeekFrom::Start(11); const FILE_SIZE_OFFSET: SeekFrom = SeekFrom::Start(12); /// The offset of the downloaded_file_size in pd file const DOWNLOADED_FILE_SIZE_OFFSET: SeekFrom = SeekFrom::Start(20); +/// The offset of the part_size in pd file +const PART_SIZE_OFFSET: SeekFrom = SeekFrom::Start(28); #[derive(Debug)] /// The pd file @@ -185,6 +187,12 @@ impl PdFile { } } + #[inline] + /// Return the size of the each part. + pub fn get_part_size(&self) -> u32 { + self.part_size.qload() + } + /// Increase the downloaded file size. /// * `size` - The file size want to added. /// @@ -441,6 +449,20 @@ impl PdFile { Ok(()) } + /// Set the size of the each part. Ignored in single thread mode. + pub fn set_part_size(&self, part_size: u32) -> Result<(), PdFileError> { + self.part_size.qstore(part_size); + 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(PART_SIZE_OFFSET)?; + f.write_le_u32(part_size)?; + self.need_saved.qstore(false); + } + Ok(()) + } + /// Write all data to the file. pub fn write(&self) -> Result<(), PdFileError> { let mut f = self.file.get_mut(); diff --git a/src/main.rs b/src/main.rs index a5ec5b7..7590843 100644 --- a/src/main.rs +++ b/src/main.rs @@ -17,6 +17,7 @@ extern crate lazy_static; #[cfg(all(feature = "link-cplusplus", target_env = "gnu"))] extern crate link_cplusplus; extern crate modular_bitfield; +extern crate parse_size; extern crate proc_macros; extern crate regex; extern crate reqwest; diff --git a/src/opt/mod.rs b/src/opt/mod.rs index 12b4b24..d23532c 100644 --- a/src/opt/mod.rs +++ b/src/opt/mod.rs @@ -1 +1,2 @@ +pub mod size; pub mod use_progress_bar; diff --git a/src/opt/size.rs b/src/opt/size.rs new file mode 100644 index 0000000..b0dfee7 --- /dev/null +++ b/src/opt/size.rs @@ -0,0 +1,27 @@ +use json::JsonValue; +use std::convert::TryFrom; + +/// Parse file size +pub fn parse_size(obj: &JsonValue) -> Option { + match obj.as_u64() { + Some(s) => Some(s), + None => match obj.as_str() { + Some(s) => match parse_size::parse_size(s) { + Ok(s) => Some(s), + Err(_) => None, + }, + None => None, + }, + } +} + +/// Parse file size as [i32] +pub fn parse_u32_size(obj: &JsonValue) -> Option { + match parse_size(obj) { + Some(s) => match u32::try_from(s) { + Ok(s) => Some(s), + Err(_) => None, + }, + None => None, + } +} diff --git a/src/opthelper.rs b/src/opthelper.rs index 5d1f02c..d3e677a 100644 --- a/src/opthelper.rs +++ b/src/opthelper.rs @@ -5,6 +5,7 @@ use crate::ext::rw_lock::GetRwLock; use crate::ext::use_or_not::ToBool; use crate::ext::use_or_not::UseOrNot; use crate::list::NonTailList; +use crate::opt::size::parse_u32_size; use crate::opt::use_progress_bar::UseProgressBar; use crate::opts::CommandOpts; use crate::retry_interval::parse_retry_interval_from_json; @@ -139,6 +140,21 @@ impl OptHelper { false } + /// Return the size of the each part when downloading file. + pub fn part_size(&self) -> u32 { + match self.opt.get_ref().part_size { + Some(r) => { + return r; + } + None => {} + } + let re = self.settings.get_ref().get("part-size"); + if re.is_some() { + return parse_u32_size(re.as_ref().unwrap()).unwrap(); + } + 65536 + } + 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 38fa18b..0018fc0 100644 --- a/src/opts.rs +++ b/src/opts.rs @@ -9,8 +9,10 @@ use crate::utils::check_file_exists; use crate::utils::get_exe_path_else_current; use getopts::HasArg; use getopts::Options; +use std::convert::TryFrom; use std::env; use std::num::ParseIntError; +use std::num::TryFromIntError; use std::str::FromStr; use std::time::Duration; @@ -82,6 +84,8 @@ pub struct CommandOpts { pub download_part_retry: Option, /// The maximun threads when downloading file. pub max_threads: Option, + /// The size of the each part when downloading file. + pub part_size: Option, } impl CommandOpts { @@ -107,6 +111,7 @@ impl CommandOpts { multiple_threads_download: None, download_part_retry: None, max_threads: None, + part_size: None, } } @@ -152,6 +157,15 @@ pub fn print_usage(prog: &str, opts: &Options) { println!("{}", opts.usage(brief.as_str())); } +/// Error when parsing size +#[derive(Debug, derive_more::Display, derive_more::From)] +pub enum ParseSizeError { + /// Failed to parse size. + ParseSize(parse_size::Error), + /// The size is too big. + Overflow(TryFromIntError), +} + /// Prase [bool] from string pub fn parse_bool>(s: Option) -> Result, String> { let tmp = match s { @@ -189,6 +203,17 @@ pub fn parse_i64>(s: Option) -> Result, ParseIntErr } } +/// Prase size as [u32] from string +pub fn parse_u32_size>(s: Option) -> Result, ParseSizeError> { + match s { + Some(s) => { + let s = parse_size::parse_size(s.as_ref())?; + Ok(Some(u32::try_from(s)?)) + } + None => Ok(None), + } +} + /// Prase [u64] from string pub fn parse_u64>(s: Option) -> Result, ParseIntError> { match s { @@ -328,6 +353,12 @@ pub fn parse_cmd() -> Option { gettext("The maximun threads when downloading file."), "COUNT", ); + opts.optopt( + "k", + "part-size", + gettext("The size of the each part when downloading file."), + "SIZE", + ); let result = match opts.parse(&argv[1..]) { Ok(m) => m, Err(err) => { @@ -526,6 +557,15 @@ pub fn parse_cmd() -> Option { return None; } } + match parse_u32_size(result.opt_str("part-size")) { + Ok(r) => { + re.as_mut().unwrap().part_size = r; + } + Err(e) => { + println!("{} {}", gettext("Failed to parse part size:"), e); + return None; + } + } re } diff --git a/src/settings_list.rs b/src/settings_list.rs index da4d926..fe924bd 100644 --- a/src/settings_list.rs +++ b/src/settings_list.rs @@ -5,6 +5,7 @@ use crate::gettext; use crate::retry_interval::check_retry_interval; use crate::settings::SettingDes; use crate::settings::JsonValueType; +use crate::opt::size::parse_u32_size; use json::JsonValue; pub fn get_settings_list() -> Vec { @@ -26,6 +27,7 @@ pub fn get_settings_list() -> Vec { 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(), SettingDes::new("max-threads", gettext("The maximun threads when downloading file."), JsonValueType::Number, Some(check_u64)).unwrap(), + SettingDes::new("part-size", gettext("The size of the each part when downloading file."), JsonValueType::Number, Some(check_parse_size_u32)).unwrap(), ] } @@ -39,6 +41,11 @@ fn check_u64(obj: &JsonValue) -> bool { r.is_some() } +#[inline] +fn check_parse_size_u32(obj: &JsonValue) -> bool { + parse_u32_size(obj).is_some() +} + fn check_nonempty_str(obj: &JsonValue) -> bool { let r = obj.as_str(); r.is_some() && r.unwrap().len() != 0