fix(diskio): fall back to single-threaded unpacking when ram_budget < 512MB to avoid OOM on memory-constrained systems

Rustup OOMs on systems with <1GB RAM because it spawns multiple I/O threads for unpacking regardless of available memory. The Threaded executor uses far more memory than its buffer pool budget tracks. This PR adds a 512MB ram_budget threshold below which rustup falls back to single-threaded unpacking, which peaks at ~110MB.

Fixes #3125

Author: cachebag

src/bin/rustup-init.rs

@@ -40,7 +40,7 @@ fn main() -> Result<ExitCode> {
     let process = Process::os();
     let runtime = tokio::runtime::Builder::new_multi_thread()
         .enable_all()
-        .worker_threads(process.io_thread_count()?)
+        .worker_threads(process.io_thread_count()?.count())
         .build()
         .unwrap();
 

src/diskio/mod.rs

@@ -62,6 +62,7 @@ use std::{fmt::Debug, fs::OpenOptions};
 use anyhow::Result;
 use tracing::{error, trace, warn};
 
+use crate::process::IoThreadCount;
 use crate::utils::units::Size;
 
 pub(crate) mod immediate;
@@ -442,18 +443,45 @@ pub(crate) fn create_dir<P: AsRef<Path>>(path: P) -> io::Result<()> {
     std::fs::create_dir(path)
 }
 
+/// The Threaded executor uses substantially more memory than its ram_budget
+/// accounts for (pool overhead, sharded_slab metadata, multiple in-flight
+/// operations, thread stacks, allocator fragmentation). On systems where the
+/// ram_budget is under this threshold, fall back to single-threaded unpacking
+/// which peaks at ~110MB.
+/// See https://github.com/rust-lang/rustup/issues/3125
+const LOW_MEMORY_THRESHOLD: usize = 512 * 1024 * 1024;
+
 /// Get the executor for disk IO.
+// If this gets lots of use, consider exposing via the config file.
 pub(crate) fn get_executor<'a>(
     ram_budget: usize,
-    io_thread_count: usize,
+    thread_count: IoThreadCount,
 ) -> Box<dyn Executor + 'a> {
-    // If this gets lots of use, consider exposing via the config file.
-    match io_thread_count {
+    let threads = effective_thread_count(ram_budget, thread_count);
+    match threads {
         0 | 1 => Box::new(immediate::ImmediateUnpacker::new()),
         n => Box::new(threaded::Threaded::new(n, ram_budget)),
     }
 }
 
+fn effective_thread_count(ram_budget: usize, thread_count: IoThreadCount) -> usize {
+    match thread_count {
+        IoThreadCount::UserSpecified(n) => n,
+        IoThreadCount::Default(n) if n <= 1 => n,
+        IoThreadCount::Default(n) if ram_budget < LOW_MEMORY_THRESHOLD => {
+            warn!(
+                "Using single-threaded unpacking due to low memory \
+                 (ram budget: {} < {} threshold). \
+                 Set RUSTUP_IO_THREADS to override.",
+                Size::new(ram_budget),
+                Size::new(LOW_MEMORY_THRESHOLD)
+            );
+            1
+        }
+        IoThreadCount::Default(n) => n,
+    }
+}
+
 pub(crate) fn unpack_ram(io_chunk_size: usize, budget: Option<usize>) -> usize {
     const RAM_ALLOWANCE_FOR_RUSTUP_AND_BUFFERS: usize = 200 * 1024 * 1024;
     let minimum_ram = io_chunk_size * 2;

src/diskio/test.rs

@@ -163,3 +163,34 @@ fn test_complete_file_immediate() {
 fn test_complete_file_threaded() {
     test_complete_file("2").unwrap()
 }
+
+#[test]
+fn test_effective_thread_count() {
+    use super::{LOW_MEMORY_THRESHOLD, effective_thread_count};
+    use crate::process::IoThreadCount::{Default, UserSpecified};
+
+    // Already single-threaded: no change regardless of budget
+    assert_eq!(effective_thread_count(32 * 1024 * 1024, Default(1)), 1);
+    assert_eq!(effective_thread_count(32 * 1024 * 1024, Default(0)), 0);
+
+    // Below threshold: forced to single-threaded
+    assert_eq!(effective_thread_count(256 * 1024 * 1024, Default(8)), 1);
+    assert_eq!(
+        effective_thread_count(LOW_MEMORY_THRESHOLD - 1, Default(4)),
+        1
+    );
+
+    // At or above threshold: thread count unchanged
+    assert_eq!(effective_thread_count(LOW_MEMORY_THRESHOLD, Default(4)), 4);
+    assert_eq!(effective_thread_count(1024 * 1024 * 1024, Default(8)), 8);
+
+    // User-specified threads are always respected
+    assert_eq!(
+        effective_thread_count(32 * 1024 * 1024, UserSpecified(4)),
+        4
+    );
+    assert_eq!(
+        effective_thread_count(LOW_MEMORY_THRESHOLD - 1, UserSpecified(8)),
+        8
+    );
+}

src/install.rs

@@ -41,7 +41,7 @@ impl InstallMethod<'_, '_> {
         // 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(self.cfg().process.io_thread_count()?)
+            .num_threads(self.cfg().process.io_thread_count()?.count())
             .build_global();
         match &self {
             InstallMethod::Copy { .. }

src/process.rs

@@ -31,6 +31,19 @@ mod file_source;
 mod terminal_source;
 pub use terminal_source::ColorableTerminal;
 
+pub enum IoThreadCount {
+    Default(usize),
+    UserSpecified(usize),
+}
+
+impl IoThreadCount {
+    pub fn count(&self) -> usize {
+        match self {
+            Self::Default(n) | Self::UserSpecified(n) => *n,
+        }
+    }
+}
+
 /// Allows concrete types for the process abstraction.
 #[derive(Clone, Debug)]
 pub enum Process {
@@ -69,23 +82,23 @@ impl Process {
         home::env::rustup_home_with_env(self).context("failed to determine rustup home dir")
     }
 
-    pub fn io_thread_count(&self) -> Result<usize> {
+    pub fn io_thread_count(&self) -> Result<IoThreadCount> {
         if let Ok(n) = self.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),
+                _ => return Ok(IoThreadCount::UserSpecified(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),
+            Ok(threads) => IoThreadCount::Default(Ord::min(threads.get(), 8)),
             // Unknown for target platform or no permission to query.
-            Err(_) => 1,
+            Err(_) => IoThreadCount::Default(1),
         })
     }