From 26850df87e686cb022ce2ad7587e661893f620c3 Mon Sep 17 00:00:00 2001
From: lifegpc <g1710431395@gmail.com>
Date: Sun, 17 Jul 2022 01:46:50 +0000
Subject: [PATCH] add new option max-download-tasks check existing task before
 push new task

---
 src/opthelper.rs     | 17 +++++++++++++++++
 src/opts.rs          | 40 ++++++++++++++++++++++++++++++++++++++++
 src/settings_list.rs |  6 ++++++
 src/task_manager.rs  | 17 +++++++++--------
 4 files changed, 72 insertions(+), 8 deletions(-)

diff --git a/src/opthelper.rs b/src/opthelper.rs
index f18793d..c8ec535 100644
--- a/src/opthelper.rs
+++ b/src/opthelper.rs
@@ -128,6 +128,23 @@ impl OptHelper {
         }
     }
 
+    /// Return the maximun number of tasks to download simultaneously.
+    pub fn max_download_tasks(&self) -> usize {
+        match self.opt.get_ref().max_download_tasks {
+            Some(r) => {
+                return r;
+            }
+            None => {}
+        }
+        match self.settings.get_ref().get("max-download-tasks") {
+            Some(re) => {
+                return re.as_usize().unwrap();
+            }
+            None => {}
+        }
+        5
+    }
+
     /// Return the maximun threads when downloading file.
     pub fn max_threads(&self) -> u64 {
         match self.opt.get_ref().max_threads {
diff --git a/src/opts.rs b/src/opts.rs
index e8a3cf9..1ceddb4 100644
--- a/src/opts.rs
+++ b/src/opts.rs
@@ -92,6 +92,8 @@ pub struct CommandOpts {
     #[cfg(feature = "server")]
     /// Server listen address
     pub server: Option<SocketAddr>,
+    /// Maximun number of tasks to download simultaneously
+    pub max_download_tasks: Option<usize>,
 }
 
 impl CommandOpts {
@@ -120,6 +122,7 @@ impl CommandOpts {
             part_size: None,
             #[cfg(feature = "server")]
             server: None,
+            max_download_tasks: None,
         }
     }
 
@@ -260,6 +263,18 @@ pub fn parse_u64<T: AsRef<str>>(s: Option<T>) -> Result<Option<u64>, ParseIntErr
     }
 }
 
+pub fn parse_nonempty_usize<T: AsRef<str>>(s: Option<T>) -> Result<Option<usize>, ParseIntError> {
+    match s {
+        Some(s) => {
+            let s = s.as_ref();
+            let s = s.trim();
+            let c = s.parse::<std::num::NonZeroUsize>()?;
+            Ok(Some(c.get()))
+        }
+        None => Ok(None),
+    }
+}
+
 /// Parse optional option
 /// * `opts` - The result of options. See [getopts::Matches].
 /// * `key` - The key of the option.
@@ -392,6 +407,20 @@ pub fn parse_cmd() -> Option<CommandOpts> {
         gettext("The size of the each part when downloading file."),
         "SIZE",
     );
+    opts.opt(
+        "",
+        "max-download-tasks",
+        format!(
+            "{} ({} {})",
+            gettext("The maximun number of tasks to download simultaneously."),
+            gettext("Default:"),
+            "5"
+        )
+        .as_str(),
+        "COUNT",
+        HasArg::Maybe,
+        getopts::Occur::Optional,
+    );
     let result = match opts.parse(&argv[1..]) {
         Ok(m) => m,
         Err(err) => {
@@ -605,6 +634,17 @@ pub fn parse_cmd() -> Option<CommandOpts> {
             return None;
         }
     }
+    match parse_optional_opt(&result, "max-download-tasks", 5, parse_nonempty_usize) {
+        Ok(r) => re.as_mut().unwrap().max_download_tasks = r,
+        Err(e) => {
+            println!(
+                "{} {}",
+                gettext("Failed to parse <opt>:").replace("<opt>", "max-download-tasks"),
+                e
+            );
+            return None;
+        }
+    }
     re
 }
 
diff --git a/src/settings_list.rs b/src/settings_list.rs
index ecf6922..5207867 100644
--- a/src/settings_list.rs
+++ b/src/settings_list.rs
@@ -40,6 +40,7 @@ pub fn get_settings_list() -> Vec<SettingDes> {
         SettingDes::new("server", gettext("Server address."), JsonValueType::Str, Some(check_socket_addr)).unwrap(),
         #[cfg(feature = "server")]
         SettingDes::new("cors-entries", gettext("The domains allowed to send CORS requests."), JsonValueType::Array, Some(check_cors_entries)).unwrap(),
+        SettingDes::new("max-download-tasks", gettext("The maximun number of tasks to download simultaneously."), JsonValueType::Number, Some(check_nozero_usize)).unwrap(),
     ]
 }
 
@@ -70,6 +71,11 @@ fn check_socket_addr(obj: &JsonValue) -> bool {
     }
 }
 
+fn check_nozero_usize(obj: &JsonValue) -> bool {
+    let r = obj.as_usize();
+    r.is_some() && r.unwrap() > 0
+}
+
 fn check_u64(obj: &JsonValue) -> bool {
     let r = obj.as_u64();
     r.is_some()
diff --git a/src/task_manager.rs b/src/task_manager.rs
index 6a12383..a9de463 100644
--- a/src/task_manager.rs
+++ b/src/task_manager.rs
@@ -2,6 +2,7 @@ use crate::ext::any::AsAny;
 use crate::ext::replace::ReplaceWith;
 use crate::ext::replace::ReplaceWith2;
 use crate::ext::rw_lock::GetRwLock;
+use crate::opthelper::get_helper;
 use futures_util::lock::Mutex;
 use indicatif::MultiProgress;
 use std::any::Any;
@@ -58,7 +59,7 @@ impl MaxDownloadTasks {
 
 impl GetMaxCount for MaxDownloadTasks {
     fn get_max_count(&self) -> usize {
-        5
+        get_helper().max_download_tasks()
     }
 }
 
@@ -95,13 +96,6 @@ impl TaskManager {
         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;
@@ -115,6 +109,13 @@ impl TaskManager {
                 }
                 self.tasks.replace_with2(new_tasks);
                 count.replace_with(new_count);
+                if *count < total_count {
+                    self.tasks
+                        .get_mut()
+                        .push(Box::new(tokio::task::spawn(future)));
+                    count.replace_with(*count + 1);
+                    break;
+                }
             }
             tokio::time::sleep(Duration::new(0, 10_000_000)).await;
         }