feat: add global -j, --threads (foundry-rs#9367)

* feat: add global -j, --threads

* Update crates/cli/src/opts/global.rs

* fix tests after comment update

---------

Co-authored-by: zerosnacks <[email protected]>
Co-authored-by: zerosnacks <[email protected]>

Author: 2 people

Cargo.lock

@@ -4,7 +4,6 @@ use eyre::{Result, WrapErr};
 use rand::{rngs::StdRng, RngCore, SeedableRng};
 use regex::RegexSetBuilder;
 use std::{
-    num::NonZeroUsize,
     sync::{
         atomic::{AtomicBool, Ordering},
         Arc,
@@ -73,9 +72,9 @@ pub struct Create2Args {
     #[arg(alias = "ch", long, value_name = "HASH", required_unless_present = "init_code")]
     init_code_hash: Option<String>,
 
-    /// Number of threads to use. Defaults to and caps at the number of logical cores.
-    #[arg(short, long)]
-    jobs: Option<NonZeroUsize>,
+    /// Number of threads to use. Specifying 0 defaults to the number of logical cores.
+    #[arg(global = true, long, short = 'j', visible_alias = "jobs")]
+    threads: Option<usize>,
 
     /// Address of the caller. Used for the first 20 bytes of the salt.
     #[arg(long, value_name = "ADDRESS")]
@@ -107,7 +106,7 @@ impl Create2Args {
             salt,
             init_code,
             init_code_hash,
-            jobs,
+            threads,
             caller,
             seed,
             no_random,
@@ -168,8 +167,8 @@ impl Create2Args {
         let regex = RegexSetBuilder::new(regexs).case_insensitive(!case_sensitive).build()?;
 
         let mut n_threads = std::thread::available_parallelism().map_or(1, |n| n.get());
-        if let Some(jobs) = jobs {
-            n_threads = n_threads.min(jobs.get());
+        if let Some(threads) = threads {
+            n_threads = n_threads.min(threads);
         }
         if cfg!(test) {
             n_threads = n_threads.min(2);
@@ -433,8 +432,14 @@ mod tests {
 
     #[test]
     fn j0() {
-        let e =
-            Create2Args::try_parse_from(["foundry-cli", "--starts-with=00", "-j0"]).unwrap_err();
-        let _ = e.print();
+        let args = Create2Args::try_parse_from([
+            "foundry-cli",
+            "--starts-with=00",
+            "--init-code-hash",
+            &B256::ZERO.to_string(),
+            "-j0",
+        ])
+        .unwrap();
+        assert_eq!(args.threads, Some(0));
     }
 }

crates/cast/bin/cmd/create2.rs

@@ -28,6 +28,11 @@ Options:
   -h, --help
           Print help (see a summary with '-h')
 
+  -j, --threads <THREADS>
+          Number of threads to use. Specifying 0 defaults to the number of logical cores
+          
+          [aliases: jobs]
+
   -V, --version
           Print version
 

crates/cast/tests/cli/main.rs

@@ -37,15 +37,16 @@ dotenvy = "0.15"
 eyre.workspace = true
 futures.workspace = true
 indicatif = "0.17"
+rayon.workspace = true
 regex = { workspace = true, default-features = false }
+serde_json.workspace = true
 serde.workspace = true
 strsim = "0.11"
 strum = { workspace = true, features = ["derive"] }
 tokio = { workspace = true, features = ["macros"] }
 tracing-subscriber = { workspace = true, features = ["registry", "env-filter"] }
 tracing.workspace = true
 yansi.workspace = true
-serde_json.workspace = true
 
 tracing-tracy = { version = "0.11", optional = true }
 

crates/cli/Cargo.toml

@@ -3,7 +3,7 @@ use foundry_common::shell::{ColorChoice, OutputFormat, OutputMode, Shell, Verbos
 use serde::{Deserialize, Serialize};
 
 /// Global options.
-#[derive(Clone, Copy, Debug, Default, Serialize, Deserialize, Parser)]
+#[derive(Clone, Debug, Default, Serialize, Deserialize, Parser)]
 pub struct GlobalOpts {
     /// Verbosity level of the log messages.
     ///
@@ -16,39 +16,47 @@ pub struct GlobalOpts {
     /// - 3 (-vvv): Print execution traces for failing tests.
     /// - 4 (-vvvv): Print execution traces for all tests, and setup traces for failing tests.
     /// - 5 (-vvvvv): Print execution and setup traces for all tests.
-    #[clap(short, long, global = true, verbatim_doc_comment, conflicts_with = "quiet", action = ArgAction::Count, help_heading = "Display options")]
-    pub verbosity: Verbosity,
+    #[arg(help_heading = "Display options", global = true, short, long, verbatim_doc_comment, conflicts_with = "quiet", action = ArgAction::Count)]
+    verbosity: Verbosity,
 
     /// Do not print log messages.
-    #[clap(short, long, global = true, alias = "silent", help_heading = "Display options")]
+    #[arg(help_heading = "Display options", global = true, short, long, alias = "silent")]
     quiet: bool,
 
     /// Format log messages as JSON.
-    #[clap(
-        long,
-        global = true,
-        alias = "format-json",
-        conflicts_with_all = &["quiet", "color"],
-        help_heading = "Display options"
-    )]
+    #[arg(help_heading = "Display options", global = true, long, alias = "format-json", conflicts_with_all = &["quiet", "color"])]
     json: bool,
 
     /// The color of the log messages.
-    #[clap(long, global = true, value_enum, help_heading = "Display options")]
+    #[arg(help_heading = "Display options", global = true, long, value_enum)]
     color: Option<ColorChoice>,
+
+    /// Number of threads to use. Specifying 0 defaults to the number of logical cores.
+    #[arg(global = true, long, short = 'j', visible_alias = "jobs")]
+    threads: Option<usize>,
 }
 
 impl GlobalOpts {
     /// Initialize the global options.
-    pub fn init(self) -> eyre::Result<()> {
+    pub fn init(&self) -> eyre::Result<()> {
         // Set the global shell.
         self.shell().set();
 
+        // Initialize the thread pool only if `threads` was requested to avoid unnecessary overhead.
+        if self.threads.is_some() {
+            self.force_init_thread_pool()?;
+        }
+
         Ok(())
     }
 
+    /// Initialize the global thread pool.
+    pub fn force_init_thread_pool(&self) -> eyre::Result<()> {
+        init_thread_pool(self.threads.unwrap_or(0))
+    }
+
     /// Create a new shell instance.
-    pub fn shell(self) -> Shell {
+    pub fn shell(&self) -> Shell {
         let mode = match self.quiet {
             true => OutputMode::Quiet,
             false => OutputMode::Normal,
@@ -62,3 +70,12 @@ impl GlobalOpts {
         Shell::new_with(format, mode, color, self.verbosity)
     }
 }
+
+/// Initialize the global thread pool.
+pub fn init_thread_pool(threads: usize) -> eyre::Result<()> {
+    rayon::ThreadPoolBuilder::new()
+        .thread_name(|i| format!("foundry-{i}"))
+        .num_threads(threads)
+        .build_global()?;
+    Ok(())
+}

crates/cli/src/opts/global.rs

@@ -115,7 +115,7 @@ no_match_contract = "Bar"
 match_path = "*/Foo*"
 no_match_path = "*/Bar*"
 no_match_coverage = "Baz"
-# Number of threads to use. Not set or zero specifies the number of logical cores.
+# Number of threads to use. Specifying 0 defaults to the number of logical cores.
 threads = 0
 # whether to show test execution progress
 show_progress = true

crates/config/README.md

@@ -65,7 +65,7 @@ foundry_config::merge_impl_figment_convert!(TestArgs, opts, evm_opts);
 #[derive(Clone, Debug, Parser)]
 #[command(next_help_heading = "Test options")]
 pub struct TestArgs {
-    /// Include the global options.
+    // Include global options for users of this struct.
     #[command(flatten)]
     pub global: GlobalOpts,
 
@@ -149,11 +149,6 @@ pub struct TestArgs {
     #[arg(long)]
     pub fuzz_input_file: Option<String>,
 
-    /// Max concurrent threads to use.
-    /// Default value is the number of available CPUs.
-    #[arg(long, short = 'j', visible_alias = "jobs")]
-    pub threads: Option<usize>,
-
     /// Show test execution progress.
     #[arg(long)]
     pub show_progress: bool,
@@ -276,13 +271,6 @@ impl TestArgs {
         // Merge all configs.
         let (mut config, mut evm_opts) = self.load_config_and_evm_opts_emit_warnings()?;
 
-        // Set number of max threads to execute tests.
-        // If not specified then the number of threads determined by rayon will be used.
-        if let Some(test_threads) = config.threads {
-            trace!(target: "forge::test", "execute tests with {} max threads", test_threads);
-            rayon::ThreadPoolBuilder::new().num_threads(test_threads).build_global()?;
-        }
-
         // Explicitly enable isolation for gas reports for more correct gas accounting.
         if self.gas_report {
             evm_opts.isolate = true;
@@ -898,10 +886,6 @@ impl Provider for TestArgs {
             dict.insert("show_progress".to_string(), true.into());
         }
 
-        if let Some(threads) = self.threads {
-            dict.insert("threads".to_string(), threads.into());
-        }
-
         Ok(Map::from([(Config::selected_profile(), dict)]))
     }
 }

crates/forge/bin/cmd/test/mod.rs

@@ -33,6 +33,11 @@ Options:
   -h, --help
           Print help (see a summary with '-h')
 
+  -j, --threads <THREADS>
+          Number of threads to use. Specifying 0 defaults to the number of logical cores
+          
+          [aliases: jobs]
+
   -V, --version
           Print version
 

crates/forge/tests/cli/cmd.rs

@@ -78,7 +78,7 @@ foundry_config::merge_impl_figment_convert!(ScriptArgs, opts, evm_opts);
 /// CLI arguments for `forge script`.
 #[derive(Clone, Debug, Default, Parser)]
 pub struct ScriptArgs {
-    /// Include the global options.
+    // Include global options for users of this struct.
     #[command(flatten)]
     pub global: GlobalOpts,