rt: read environment variable for worker thread count (#4250)

tokio/src/lib.rs

@@ -174,12 +174,15 @@
 //! swapping the currently running task on each thread. However, this kind of
 //! swapping can only happen at `.await` points, so code that spends a long time
 //! without reaching an `.await` will prevent other tasks from running. To
-//! combat this, Tokio provides two kinds of threads: Core threads and blocking
-//! threads. The core threads are where all asynchronous code runs, and Tokio
-//! will by default spawn one for each CPU core. The blocking threads are
-//! spawned on demand, can be used to run blocking code that would otherwise
-//! block other tasks from running and are kept alive when not used for a certain
-//! amount of time which can be configured with [`thread_keep_alive`].
+//! combat this, Tokio provides two kinds of threads: Core threads and blocking threads.
+//!
+//! The core threads are where all asynchronous code runs, and Tokio will by default
+//! spawn one for each CPU core. You can use the environment variable `TOKIO_WORKER_THREADS`
+//! to override the default value.
+//!
+//! The blocking threads are spawned on demand, can be used to run blocking code
+//! that would otherwise block other tasks from running and are kept alive when
+//! not used for a certain amount of time which can be configured with [`thread_keep_alive`].
 //! Since it is not possible for Tokio to swap out blocking tasks, like it
 //! can do with asynchronous code, the upper limit on the number of blocking
 //! threads is very large. These limits can be configured on the [`Builder`].

tokio/src/loom/std/mod.rs

@@ -81,7 +81,27 @@ pub(crate) mod sync {
 pub(crate) mod sys {
     #[cfg(feature = "rt-multi-thread")]
     pub(crate) fn num_cpus() -> usize {
-        usize::max(1, num_cpus::get())
+        const ENV_WORKER_THREADS: &str = "TOKIO_WORKER_THREADS";
+
+        match std::env::var(ENV_WORKER_THREADS) {
+            Ok(s) => {
+                let n = s.parse().unwrap_or_else(|e| {
+                    panic!(
+                        "\"{}\" must be usize, error: {}, value: {}",
+                        ENV_WORKER_THREADS, e, s
+                    )
+                });
+                assert!(n > 0, "\"{}\" cannot be set to 0", ENV_WORKER_THREADS);
+                n
+            }
+            Err(std::env::VarError::NotPresent) => usize::max(1, num_cpus::get()),
+            Err(std::env::VarError::NotUnicode(e)) => {
+                panic!(
+                    "\"{}\" must be valid unicode, error: {:?}",
+                    ENV_WORKER_THREADS, e
+                )
+            }
+        }
     }
 
     #[cfg(not(feature = "rt-multi-thread"))]

tokio/src/runtime/builder.rs

@@ -182,6 +182,7 @@ cfg_unstable! {
 
 pub(crate) type ThreadNameFn = std::sync::Arc<dyn Fn() -> String + Send + Sync + 'static>;
 
+#[derive(Clone, Copy)]
 pub(crate) enum Kind {
     CurrentThread,
     #[cfg(all(feature = "rt-multi-thread", not(tokio_wasi)))]
@@ -237,6 +238,7 @@ impl Builder {
             // The clock starts not-paused
             start_paused: false,
 
+            // Read from environment variable first in multi-threaded mode.
             // Default to lazy auto-detection (one thread per CPU core)
             worker_threads: None,
 
@@ -304,6 +306,8 @@ impl Builder {
     /// This can be any number above 0 though it is advised to keep this value
     /// on the smaller side.
     ///
+    /// This will override the value read from environment variable `TOKIO_WORKER_THREADS`.
+    ///
     /// # Default
     ///
     /// The default value is the number of cores available to the system.