Merge pull request #413 from AOSC-Dev/add-download-threads-arg

feat: add `--download-threads` / `-t` arg to control download threads

Author: eatradish

i18n/en-US/oma.ftl

@@ -280,3 +280,4 @@ path-not-exist = Path { $path } does not exist.
 not-allow-oma-pending-piped = oma cannot run with piped output in interactive mode.
 space-warn = You seem to be running out of space. Consider using { $cmd } to purge the package cache, which will free up { $size } in storage space; Please also consider cleaning up your rarely used software or personal data.
 space-warn-with-zero = You seem to be running out of space, Please consider cleaning up your rarely used software or personal data.
+wrong-thread-count = The specified number of threads { $count } is not allowed (min: 1, max: 255).

i18n/zh-CN/oma.ftl

@@ -269,3 +269,4 @@ path-not-exist = 路径 { $path } 不存在。
 not-allow-oma-pending-piped = 不允许搭配管道使用 oma 的交互模式界面。
 space-warn = 检测到您的存储空间告急，您可以使用 { $cmd } 清除软件安装包缓存，可释放 { $size } 存储空间；也请考虑删除不必要的软件包或个人数据。
 space-warn-with-zero = 检测到您的存储空间告急，请考虑删除不必要的软件包或个人数据。
+wrong-thread-count = 您指定的线程数 { $count } 不合法（最小：1，最大：255）。

oma-fetch/src/download.rs

@@ -8,7 +8,7 @@ use std::{
 };
 
 use async_compression::futures::bufread::{BzDecoder, GzipDecoder, XzDecoder, ZstdDecoder};
-use bon::Builder;
+use bon::bon;
 use futures::{AsyncRead, TryStreamExt, io::BufReader};
 use reqwest::{
     Client, Method, RequestBuilder,
@@ -29,22 +29,16 @@ use crate::{DownloadEntry, DownloadSourceType};
 const READ_FILE_BUFSIZE: usize = 65536;
 const DOWNLOAD_BUFSIZE: usize = 8192;
 
-#[derive(Snafu, Debug)]
-#[snafu(display("source list is empty"))]
-pub struct EmptySource {
-    file_name: String,
+#[derive(Debug, Snafu)]
+pub enum BuilderError {
+    #[snafu(display("Download task {file_name} sources is empty"))]
+    EmptySource { file_name: String },
+    #[snafu(display("Not allow set illegal download threads: {count}"))]
+    IllegalDownloadThread { count: usize },
 }
 
-#[derive(Builder)]
 pub(crate) struct SingleDownloader<'a> {
     client: &'a Client,
-    #[builder(with = |entry: &'a DownloadEntry| -> Result<_, EmptySource> {
-        if entry.source.is_empty() {
-            return Err(EmptySource { file_name: entry.filename.to_string() });
-        } else {
-            return Ok(entry);
-        }
-    })]
     pub entry: &'a DownloadEntry,
     progress: (usize, usize),
     retry_times: usize,
@@ -99,7 +93,39 @@ pub enum SingleDownloadError {
     ChecksumMismatch,
 }
 
-impl SingleDownloader<'_> {
+#[bon]
+impl<'a> SingleDownloader<'a> {
+    #[builder]
+    pub(crate) fn new(
+        client: &'a Client,
+        entry: &'a DownloadEntry,
+        progress: (usize, usize),
+        retry_times: usize,
+        msg: Option<String>,
+        download_list_index: usize,
+        file_type: CompressFile,
+        set_permission: Option<u32>,
+        timeout: Duration,
+    ) -> Result<SingleDownloader<'a>, BuilderError> {
+        if entry.source.is_empty() {
+            return Err(BuilderError::EmptySource {
+                file_name: entry.filename.to_string(),
+            });
+        }
+
+        Ok(Self {
+            client,
+            entry,
+            progress,
+            retry_times,
+            msg,
+            download_list_index,
+            file_type,
+            set_permission,
+            timeout,
+        })
+    }
+
     pub(crate) async fn try_download(self, callback: &impl AsyncFn(Event)) -> DownloadResult {
         let mut sources = self.entry.source.clone();
         assert!(!sources.is_empty());

oma-fetch/src/lib.rs

@@ -2,14 +2,14 @@ use std::{cmp::Ordering, path::PathBuf, time::Duration};
 
 use bon::{Builder, builder};
 use checksum::Checksum;
-use download::{EmptySource, SingleDownloader, SuccessSummary};
+use download::{BuilderError, SingleDownloader, SuccessSummary};
 use futures::StreamExt;
 
 use reqwest::{Client, Method, RequestBuilder};
 use tracing::debug;
 
 pub mod checksum;
-mod download;
+pub mod download;
 pub use crate::download::SingleDownloadError;
 
 pub use reqwest;
@@ -206,7 +206,13 @@ impl DownloadManager<'_> {
     pub async fn start_download(
         &self,
         callback: impl AsyncFn(Event),
-    ) -> Result<Summary, EmptySource> {
+    ) -> Result<Summary, BuilderError> {
+        if self.threads == 0 || self.threads > 255 {
+            return Err(BuilderError::IllegalDownloadThread {
+                count: self.threads,
+            });
+        }
+
         let mut tasks = Vec::new();
         let mut list = vec![];
         for (i, c) in self.download_list.iter().enumerate() {
@@ -215,13 +221,13 @@ impl DownloadManager<'_> {
                 .client(self.client)
                 .maybe_msg(msg)
                 .download_list_index(i)
-                .entry(c)?
+                .entry(c)
                 .progress((i + 1, self.download_list.len()))
                 .retry_times(self.retry_times)
                 .file_type(c.file_type)
                 .maybe_set_permission(self.set_permission)
                 .timeout(self.timeout)
-                .build();
+                .build()?;
 
             list.push(single);
         }

oma-refresh/src/db.rs

@@ -25,6 +25,7 @@ use oma_apt_sources_lists::SourcesListError;
 use oma_fetch::{
     CompressFile, DownloadEntry, DownloadManager, DownloadSource, DownloadSourceType,
     checksum::{Checksum, ChecksumError},
+    download::BuilderError,
     reqwest::{
         Client, Response,
         header::{CONTENT_LENGTH, HeaderValue},
@@ -100,6 +101,10 @@ pub enum RefreshError {
     SourceListsEmpty,
     #[error("Failed to operate file: {0}")]
     OperateFile(PathBuf, std::io::Error),
+    #[error("thread count is not illegal: {0}")]
+    WrongThreadCount(usize),
+    #[error("Failed to build download manager")]
+    DownloadManagerBuilderError(BuilderError),
 }
 
 type Result<T> = std::result::Result<T, RefreshError>;
@@ -189,6 +194,10 @@ pub enum Event {
 
 impl<'a> OmaRefresh<'a> {
     pub async fn start(mut self, callback: impl AsyncFn(Event)) -> Result<()> {
+        if self.threads == 0 || self.threads > 255 {
+            return Err(RefreshError::WrongThreadCount(self.threads));
+        }
+
         let arch = dpkg_arch(&self.source)?;
         let sourcelist = scan_sources_lists(&self.source, &arch, self.apt_config, &callback)
             .await
@@ -272,7 +281,7 @@ impl<'a> OmaRefresh<'a> {
                 callback(Event::DownloadEvent(event)).await;
             })
             .await
-            .unwrap();
+            .map_err(RefreshError::DownloadManagerBuilderError)?;
 
         if !res.is_download_success() {
             return Err(RefreshError::DownloadFailed(None));

src/error.rs

@@ -8,6 +8,7 @@ use oma_console::writer::{Writeln, Writer};
 use oma_contents::OmaContentsError;
 use oma_fetch::SingleDownloadError;
 use oma_fetch::checksum::ChecksumError;
+use oma_fetch::download::BuilderError;
 use oma_history::HistoryError;
 
 #[cfg(feature = "aosc")]
@@ -437,6 +438,20 @@ impl From<RefreshError> for OutputError {
                 description: fl!("failed-to-operate-path", p = path.display().to_string()),
                 source: Some(Box::new(error)),
             },
+            RefreshError::WrongThreadCount(count) => Self {
+                description: fl!("wrong-thread-count", count = count),
+                source: None,
+            },
+            RefreshError::DownloadManagerBuilderError(builder_error) => match builder_error {
+                BuilderError::EmptySource { file_name } => Self {
+                    description: format!("BUG: task {} should is not empty", file_name),
+                    source: None,
+                },
+                BuilderError::IllegalDownloadThread { count } => Self {
+                    description: fl!("wrong-thread-count", count = count),
+                    source: None,
+                },
+            },
         }
     }
 }

src/main.rs

@@ -124,6 +124,9 @@ pub struct GlobalOptions {
     #[arg(long, global = true, env = "OMA_NO_BELL", value_parser = FalseyValueParser::new()
 )]
     no_bell: bool,
+    /// Setup download threads (default as 4)
+    #[arg(long, short = 't', global = true, env = "OMA_DOWNLOAD_THREADS")]
+    download_threads: Option<usize>,
 }
 
 fn main() {

src/subcommand/download.rs

@@ -29,6 +29,9 @@ pub struct Download {
     /// Run oma in "dry-run" mode. Useful for testing changes and operations without making changes to the system
     #[arg(from_global)]
     dry_run: bool,
+    /// Setup download threads (default as 4)
+    #[arg(from_global)]
+    download_threads: Option<usize>,
 }
 
 impl CliExecuter for Download {
@@ -37,6 +40,7 @@ impl CliExecuter for Download {
             packages,
             path,
             dry_run,
+            download_threads,
         } = self;
 
         let path = path.canonicalize().map_err(|e| OutputError {
@@ -73,7 +77,7 @@ impl CliExecuter for Download {
             &HTTP_CLIENT,
             pkgs,
             DownloadConfig {
-                network_thread: Some(config.network_thread()),
+                network_thread: Some(download_threads.unwrap_or_else(|| config.network_thread())),
                 download_dir: Some(&path),
                 auth: auth_config("/").as_ref(),
             },

src/subcommand/fix_broken.rs

@@ -44,6 +44,9 @@ pub struct FixBroken {
     /// Set apt options
     #[arg(from_global)]
     apt_options: Vec<String>,
+    /// Setup download threads (default as 4)
+    #[arg(from_global)]
+    download_threads: Option<usize>,
 }
 
 impl CliExecuter for FixBroken {
@@ -62,6 +65,7 @@ impl CliExecuter for FixBroken {
             sysroot,
             apt_options,
             no_fix_dpkg_status,
+            download_threads,
         } = self;
 
         let mut _fds = None;
@@ -97,7 +101,7 @@ impl CliExecuter for FixBroken {
             .autoremove(autoremove)
             .remove_config(remove_config)
             .maybe_auth_config(auth_config)
-            .network_thread(config.network_thread())
+            .network_thread(download_threads.unwrap_or_else(|| config.network_thread()))
             .build()
             .run()
     }

src/subcommand/history.rs

@@ -103,6 +103,9 @@ pub struct Undo {
     /// Set apt options
     #[arg(from_global)]
     apt_options: Vec<String>,
+    /// Setup download threads (default as 4)
+    #[arg(from_global)]
+    download_threads: Option<usize>,
 }
 
 impl CliExecuter for Undo {
@@ -122,6 +125,7 @@ impl CliExecuter for Undo {
             sysroot,
             apt_options,
             no_fix_dpkg_status,
+            download_threads,
         } = self;
 
         let _fds = if !no_check_dbus && !config.no_check_dbus() && !dry_run {
@@ -243,7 +247,7 @@ impl CliExecuter for Undo {
             .yes(false)
             .remove_config(remove_config)
             .autoremove(autoremove)
-            .network_thread(config.network_thread())
+            .network_thread(download_threads.unwrap_or_else(|| config.network_thread()))
             .maybe_auth_config(auth_config)
             .build()
             .run()?;

src/subcommand/install.rs

@@ -92,6 +92,9 @@ pub struct Install {
     /// Remove package(s) also remove configuration file(s), like apt purge
     #[arg(long, visible_alias = "purge")]
     remove_config: bool,
+    /// Setup download threads (default as 4)
+    #[arg(from_global)]
+    download_threads: Option<usize>,
 }
 
 impl CliExecuter for Install {
@@ -119,6 +122,7 @@ impl CliExecuter for Install {
             autoremove,
             remove_config,
             no_fix_dpkg_status,
+            download_threads,
         } = self;
 
         if !dry_run {
@@ -144,7 +148,7 @@ impl CliExecuter for Install {
                 .client(&HTTP_CLIENT)
                 .dry_run(dry_run)
                 .no_progress(no_progress)
-                .network_thread(config.network_thread())
+                .network_thread(download_threads.unwrap_or_else(|| config.network_thread()))
                 .sysroot(&sysroot)
                 .config(&apt_config)
                 .maybe_auth_config(auth_config);
@@ -219,7 +223,7 @@ impl CliExecuter for Install {
             .yes(yes)
             .remove_config(remove_config)
             .autoremove(autoremove)
-            .network_thread(config.network_thread())
+            .network_thread(download_threads.unwrap_or_else(|| config.network_thread()))
             .maybe_auth_config(auth_config)
             .fix_dpkg_status(!no_fix_dpkg_status)
             .build()