From d7c6929ae624865e9045d803b7f5a896e8bba912 Mon Sep 17 00:00:00 2001
From: lifegpc <g1710431395@gmail.com>
Date: Sun, 19 Jun 2022 15:29:21 +0000
Subject: [PATCH] Update

---
 src/downloader/pd_file/file.rs        | 26 ++++++++++++++++++++++++++
 src/downloader/pd_file/part_status.rs |  8 ++++++++
 src/downloader/tasks.rs               | 11 ++++++++---
 3 files changed, 42 insertions(+), 3 deletions(-)

diff --git a/src/downloader/pd_file/file.rs b/src/downloader/pd_file/file.rs
index 4194a85..9762d79 100644
--- a/src/downloader/pd_file/file.rs
+++ b/src/downloader/pd_file/file.rs
@@ -77,6 +77,8 @@ pub struct PdFile {
     mem_only: AtomicBool,
     /// The status of the each part.
     part_datas: RwLock<Vec<Arc<PdFilePartStatus>>>,
+    /// The offset of the first part data
+    part_data_offset: AtomicU64,
 }
 
 impl PdFile {
@@ -95,6 +97,7 @@ impl PdFile {
             part_size: AtomicU32::new(0),
             mem_only: AtomicBool::new(true),
             part_datas: RwLock::new(Vec::new()),
+            part_data_offset: AtomicU64::new(32),
         }
     }
 
@@ -403,6 +406,7 @@ impl PdFile {
             part_size: AtomicU32::new(part_size),
             mem_only: AtomicBool::new(false),
             part_datas: RwLock::new(part_datas),
+            part_data_offset: AtomicU64::new(32 + (file_name_len as u64)),
         })
     }
 
@@ -480,6 +484,7 @@ impl PdFile {
             Err(gettext("File name should not be empty."))?
         } else {
             self.file_name.get_mut().replace(String::from(fname));
+            self.part_data_offset.qstore(32 + (fname.len() as u64));
             if !self.is_mem_only() {
                 self.need_saved.qstore(true);
                 self.reopen()?;
@@ -530,6 +535,27 @@ impl PdFile {
         Ok(())
     }
 
+    /// Update part data in file
+    /// * `index` - The index of the part data.
+    pub fn update_part_data(&self, index: usize) -> Result<(), PdFileError> {
+        let part_datas = self.part_datas.get_ref();
+        if index >= part_datas.len() {
+            return Ok(());
+        }
+        let part = &part_datas[index];
+        if !self.is_mem_only() {
+            self.need_saved.qstore(true);
+            let mut f = self.file.get_mut();
+            let mut f = f.as_mut().unwrap();
+            f.seek(SeekFrom::Start(
+                self.part_data_offset.qload() + 2 * (index as u64),
+            ))?;
+            part.write_to(&mut f)?;
+            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/downloader/pd_file/part_status.rs b/src/downloader/pd_file/part_status.rs
index 45bc811..e2d98e6 100644
--- a/src/downloader/pd_file/part_status.rs
+++ b/src/downloader/pd_file/part_status.rs
@@ -126,6 +126,14 @@ impl PdFilePartStatus {
         }
     }
 
+    #[inline]
+    /// Increase the size of the part
+    /// * `size` - The size to increase
+    pub fn inc(&self, size: u32) -> Result<(), PdFileError> {
+        let s = self.downloaded_size() + size;
+        self.set_downloaded_size(s)
+    }
+
     #[inline]
     /// Returns true if the download is waited
     pub fn is_waited(&self) -> bool {
diff --git a/src/downloader/tasks.rs b/src/downloader/tasks.rs
index 5c50752..4d31040 100644
--- a/src/downloader/tasks.rs
+++ b/src/downloader/tasks.rs
@@ -141,7 +141,7 @@ pub async fn create_download_tasks_simple<
             gettext("Downloading \"<loc>\".").replace("<loc>", d.get_file_name().as_str()),
         );
     }
-    handle_download(d, result, None).await
+    handle_download(d, result, None, None).await
 }
 
 /// Do first job when download in multiple mode.
@@ -212,7 +212,7 @@ pub async fn create_download_tasks_multi<
     if status.as_u16() != 206 {
         return Err(DownloaderError::from(status));
     }
-    handle_download(d, result, Some(pd)).await
+    handle_download(d, result, Some(pd), Some(index)).await
 }
 
 /// Handle download process
@@ -220,6 +220,7 @@ pub async fn handle_download<T: Seek + Write + Send + Sync + ClearFile + GetTarg
     d: Arc<DownloaderInternal<T>>,
     re: Response,
     pd: Option<Arc<PdFilePartStatus>>,
+    index: Option<usize>,
 ) -> Result<(), DownloaderError> {
     let mut stream = re.bytes_stream();
     let is_multi = d.is_multi_threads();
@@ -237,12 +238,16 @@ pub async fn handle_download<T: Seek + Write + Send + Sync + ClearFile + GetTarg
                             if d.enabled_progress_bar() {
                                 d.inc_progress_bar(len);
                             }
+                            d.write(&data)?;
                         } else {
                             if !d.is_multi_threads() {
                                 return Ok(());
                             }
+                            let len = data.len() as u32;
+                            pd.as_ref().unwrap().inc(len)?;
+                            d.pd.inc(len as u64)?;
+                            d.pd.update_part_data(index.unwrap())?;
                         }
-                        d.write(&data)?;
                     }
                     Err(e) => {
                         if !is_multi {