feat(rt): read default worker thread from env

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/runtime/builder.rs

@@ -6,6 +6,10 @@ use std::fmt;
 use std::io;
 use std::time::Duration;
 
+/// This key is used to specify the default worker threads for multi-thread runtime.
+#[cfg(all(feature = "rt-multi-thread", not(tokio_wasi)))]
+const ENV_WORKER_THREADS: &str = "TOKIO_WORKER_THREADS";
+
 /// Builds Tokio Runtime with custom configuration values.
 ///
 /// Methods can be chained in order to set the configuration values. The
@@ -181,6 +185,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)))]
@@ -235,8 +240,9 @@ 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,
+            worker_threads: Self::default_worker_threads(kind),
 
             max_blocking_threads: 512,
 
@@ -302,6 +308,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.
@@ -915,6 +923,33 @@ impl Builder {
             blocking_pool,
         ))
     }
+
+    #[cfg(any(not(feature = "rt-multi-thread"), tokio_wasi))]
+    fn default_worker_threads(_: Kind) -> Option<usize> {
+        None
+    }
+
+    #[cfg(all(feature = "rt-multi-thread", not(tokio_wasi)))]
+    fn default_worker_threads(kind: Kind) -> Option<usize> {
+        match kind {
+            // Always return None if using current thread
+            Kind::CurrentThread => return None,
+            Kind::MultiThread => {}
+        };
+        match std::env::var(ENV_WORKER_THREADS) {
+            Ok(s) => {
+                let n: usize = 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);
+                Some(n)
+            }
+            Err(_) => None,
+        }
+    }
 }
 
 cfg_io_driver! {