From 682e93df93292a47e7d6d04bd9e4922c5bbae5b4 Mon Sep 17 00:00:00 2001
From: Pure White <wu.purewhite@gmail.com>
Date: Wed, 21 Dec 2022 22:43:35 +0800
Subject: [PATCH] rt: read environment variable for worker thread count (#4250)

---
 tokio/src/lib.rs             | 15 +++++++++------
 tokio/src/loom/std/mod.rs    | 22 +++++++++++++++++++++-
 tokio/src/runtime/builder.rs |  4 ++++
 3 files changed, 34 insertions(+), 7 deletions(-)

diff --git a/tokio/src/lib.rs b/tokio/src/lib.rs
index e745fe9995f..05767d017bc 100644
--- a/tokio/src/lib.rs
+++ b/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`].
diff --git a/tokio/src/loom/std/mod.rs b/tokio/src/loom/std/mod.rs
index f0fcd46d23d..6bd1ad93dcf 100644
--- a/tokio/src/loom/std/mod.rs
+++ b/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"))]
diff --git a/tokio/src/runtime/builder.rs b/tokio/src/runtime/builder.rs
index 64cf403aaf7..ea0df2e3b4c 100644
--- a/tokio/src/runtime/builder.rs
+++ b/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.