diff --git a/src/downloader/pd_file/file.rs b/src/downloader/pd_file/file.rs index 9c8f53a..4194a85 100644 --- a/src/downloader/pd_file/file.rs +++ b/src/downloader/pd_file/file.rs @@ -197,6 +197,25 @@ impl PdFile { self.file_size.qload() } + /// Get the next waited part data + /// * `data` - The part data if found + /// + /// Returns the index of the part data + pub fn get_next_waited_part_data( + &self, + data: &mut Option>, + ) -> Option { + let datas = self.part_datas.get_ref(); + let index = 0; + for d in datas.iter() { + if d.is_waited() { + data.replace(Arc::clone(d)); + return Some(index); + } + } + None + } + /// Return status data of a part /// * `index` - The part index pub fn get_part_data(&self, index: usize) -> Option> { @@ -233,6 +252,27 @@ impl PdFile { Ok(downloaded_size) } + /// Initialize all part datas + pub fn initialize_part_datas(&self) -> Result<(), PdFileError> { + if self.is_multi_threads() && self.is_downloading() { + let file_size = self.file_size.qload(); + let part_size = self.part_size.qload(); + let mut part_counts = (file_size + (part_size as u64) - 1) / (part_size as u64); + let mut parts = Vec::new(); + while part_counts > 0 { + parts.push(Arc::new(PdFilePartStatus::new())); + part_counts -= 1; + } + self.part_datas.replace_with2(parts); + if !self.is_mem_only() { + self.need_saved.qstore(true); + self.write()?; + self.need_saved.qstore(false); + } + } + Ok(()) + } + #[inline] /// Returns true if the download is completed. pub fn is_completed(&self) -> bool { diff --git a/src/downloader/tasks.rs b/src/downloader/tasks.rs index c74e1d4..5c50752 100644 --- a/src/downloader/tasks.rs +++ b/src/downloader/tasks.rs @@ -1,6 +1,8 @@ use super::downloader::DownloaderInternal; use super::downloader::GetTargetFileName; use super::error::DownloaderError; +use super::pd_file::PdFilePartStatus; +use crate::ext::atomic::AtomicQuick; use crate::ext::io::ClearFile; use crate::ext::replace::ReplaceWith2; use crate::ext::rw_lock::GetRwLock; @@ -139,7 +141,7 @@ pub async fn create_download_tasks_simple< gettext("Downloading \"\".").replace("", d.get_file_name().as_str()), ); } - handle_download(d, result).await + handle_download(d, result, None).await } /// Do first job when download in multiple mode. @@ -171,13 +173,53 @@ pub async fn create_download_tasks_multi_first< ))); } } + match d.pd.initialize_part_datas() { + Ok(_) => {} + Err(e) => { + println!("{}", e); + } + } Ok(()) } +/// Create a new download task in multiple thread mode. +pub async fn create_download_tasks_multi< + T: Seek + Write + Send + Sync + ClearFile + GetTargetFileName, +>( + d: Arc>, + pd: Arc, + index: usize, +) -> Result<(), DownloaderError> { + let part_size = d.get_part_size() as u64; + let file_size = d.pd.get_file_size(); + let start = part_size * (index as u64); + let end = std::cmp::min(start + part_size - 1, file_size); + let mut headers = d.headers.deref().clone(); + headers.insert(String::from("Range"), format!("{}-{}", start, end)); + let result = d + .client + .get(d.url.deref().clone(), headers) + .await + .try_err(gettext("Failed to get url."))?; + let status = result.status(); + if status == 200 || status == 416 { + d.fallback_to_simp(); + d.tasks.replace_with2(Vec::new()); + return Err(DownloaderError::from(gettext( + "Warning: The server seems does not support range.", + ))); + } + if status.as_u16() != 206 { + return Err(DownloaderError::from(status)); + } + handle_download(d, result, Some(pd)).await +} + /// Handle download process pub async fn handle_download( d: Arc>, re: Response, + pd: Option>, ) -> Result<(), DownloaderError> { let mut stream = re.bytes_stream(); let is_multi = d.is_multi_threads(); @@ -195,6 +237,10 @@ pub async fn handle_download( + d: &Arc>, +) -> Result<(), DownloaderError> { + let mut needed_size = (d.max_threads.qload() as usize) - d.tasks.get_ref().len(); + while needed_size > 0 { + check_dropped!(d); + let mut data = None; + let index = d.pd.get_next_waited_part_data(&mut data); + match index { + Some(index) => { + let task = tokio::spawn(create_download_tasks_multi( + Arc::clone(d), + data.unwrap(), + index, + )); + d.add_task(task); + } + None => { + return Ok(()); + } + } + needed_size -= 1; + } + Ok(()) +} + /// Check tasks are completed or not. And create new tasks if needed. pub async fn check_tasks< T: Seek + Write + Send + Sync + ClearFile + GetTargetFileName + 'static, @@ -260,6 +346,8 @@ pub async fn check_tasks< if d.pd.is_started() { let task = tokio::spawn(create_download_tasks_multi_first(Arc::clone(&d))); d.add_task(task); + } else { + add_new_multi_tasks(&d).await?; } } loop { @@ -326,6 +414,10 @@ pub async fn check_tasks< } let task = tokio::spawn(create_download_tasks_multi_first(Arc::clone(&d))); d.add_task(task); + } else { + if d.tasks.get_ref().len() < (d.max_threads.qload() as usize) { + add_new_multi_tasks(&d).await?; + } } } if need_break {