Always consider RUSTUP_IO_THREADS as input for thread count

Author: djc

src/diskio/mod.rs

@@ -63,7 +63,7 @@ use std::sync::mpsc::Receiver;
 use std::time::{Duration, Instant};
 use std::{fmt::Debug, fs::OpenOptions};
 
-use anyhow::{Context, Result};
+use anyhow::Result;
 
 use crate::process::Process;
 use crate::utils;
@@ -447,15 +447,9 @@ pub(crate) fn get_executor<'a>(
     notify_handler: Option<&'a dyn Fn(Notification<'_>)>,
     ram_budget: usize,
     process: &Process,
-) -> Result<Box<dyn Executor + 'a>> {
+) -> anyhow::Result<Box<dyn Executor + 'a>> {
     // If this gets lots of use, consider exposing via the config file.
-    let thread_count = match process.var("RUSTUP_IO_THREADS") {
-        Err(_) => utils::io_thread_count(),
-        Ok(n) => n
-            .parse::<usize>()
-            .context("invalid value in RUSTUP_IO_THREADS. Must be a natural number")?,
-    };
-    Ok(match thread_count {
+    Ok(match utils::io_thread_count(process)? {
         0 | 1 => Box::new(immediate::ImmediateUnpacker::new()),
         n => Box::new(threaded::Threaded::new(notify_handler, n, ram_budget)),
     })

src/install.rs

@@ -38,11 +38,11 @@ pub(crate) enum InstallMethod<'a> {
 impl InstallMethod<'_> {
     // Install a toolchain
     #[tracing::instrument(level = "trace", err(level = "trace"), skip_all)]
-    pub(crate) async fn install(&self) -> Result<UpdateStatus> {
+    pub(crate) async fn install(&self) -> anyhow::Result<UpdateStatus> {
         // Initialize rayon for use by the remove_dir_all crate limiting the number of threads.
         // This will error if rayon is already initialized but it's fine to ignore that.
         let _ = rayon::ThreadPoolBuilder::new()
-            .num_threads(utils::io_thread_count())
+            .num_threads(utils::io_thread_count(self.cfg().process)?)
             .build_global();
         let nh = &self.cfg().notify_handler;
         match self {

src/utils/mod.rs

@@ -7,6 +7,7 @@ use std::io::{self, BufReader, Write};
 use std::ops::{BitAnd, BitAndAssign};
 use std::path::{Path, PathBuf};
 use std::process::ExitStatus;
+use std::str::FromStr;
 use std::thread;
 
 use anyhow::{Context, Result, anyhow, bail};
@@ -27,14 +28,24 @@ pub(crate) mod notify;
 pub mod raw;
 pub(crate) mod units;
 
-pub fn io_thread_count() -> usize {
-    match thread::available_parallelism() {
+pub fn io_thread_count(process: &Process) -> anyhow::Result<usize> {
+    if let Ok(n) = process.var("RUSTUP_IO_THREADS") {
+        let threads = usize::from_str(&n).context(
+            "invalid value in RUSTUP_IO_THREADS -- must be a natural number greater than zero",
+        )?;
+        match threads {
+            0 => bail!("RUSTUP_IO_THREADS must be a natural number greater than zero"),
+            _ => return Ok(threads),
+        }
+    };
+
+    Ok(match thread::available_parallelism() {
         // Don't spawn more than 8 I/O threads unless the user tells us to.
         // Feel free to increase this value if it improves performance.
         Ok(threads) => Ord::min(threads.get(), 8),
         // Unknown for target platform or no permission to query.
         Err(_) => 1,
-    }
+    })
 }
 
 #[must_use]