Skip to content

Commit

Permalink
feat: add global -j, --threads
Browse files Browse the repository at this point in the history
  • Loading branch information
DaniPopes committed Nov 20, 2024
1 parent 622f922 commit 9eb3651
Show file tree
Hide file tree
Showing 8 changed files with 61 additions and 43 deletions.
1 change: 1 addition & 0 deletions Cargo.lock

Some generated files are not rendered by default. Learn more about how customized files appear on GitHub.

25 changes: 15 additions & 10 deletions crates/cast/bin/cmd/create2.rs
Original file line number Diff line number Diff line change
Expand Up @@ -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,
Expand Down Expand Up @@ -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. Zero specifies 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")]
Expand Down Expand Up @@ -107,7 +106,7 @@ impl Create2Args {
salt,
init_code,
init_code_hash,
jobs,
threads,
caller,
seed,
no_random,
Expand Down Expand Up @@ -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);
Expand Down Expand Up @@ -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));
}
}
5 changes: 5 additions & 0 deletions crates/cast/tests/cli/main.rs
Original file line number Diff line number Diff line change
Expand Up @@ -28,6 +28,11 @@ Options:
-h, --help
Print help (see a summary with '-h')
-j, --threads <THREADS>
Number of threads to use. Zero specifies the number of logical cores
[aliases: jobs]
-V, --version
Print version
Expand Down
3 changes: 2 additions & 1 deletion crates/cli/Cargo.toml
Original file line number Diff line number Diff line change
Expand Up @@ -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 }

Expand Down
45 changes: 31 additions & 14 deletions crates/cli/src/opts/global.rs
Original file line number Diff line number Diff line change
Expand Up @@ -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.
///
Expand All @@ -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. Zero specifies 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,
Expand All @@ -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(())
}
18 changes: 1 addition & 17 deletions crates/forge/bin/cmd/test/mod.rs
Original file line number Diff line number Diff line change
Expand Up @@ -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,

Expand Down Expand Up @@ -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,
Expand Down Expand Up @@ -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;
Expand Down Expand Up @@ -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)]))
}
}
Expand Down
5 changes: 5 additions & 0 deletions crates/forge/tests/cli/cmd.rs
Original file line number Diff line number Diff line change
Expand Up @@ -33,6 +33,11 @@ Options:
-h, --help
Print help (see a summary with '-h')
-j, --threads <THREADS>
Number of threads to use. Zero specifies the number of logical cores
[aliases: jobs]
-V, --version
Print version
Expand Down
2 changes: 1 addition & 1 deletion crates/script/src/lib.rs
Original file line number Diff line number Diff line change
Expand Up @@ -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,

Expand Down

0 comments on commit 9eb3651

Please sign in to comment.