Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Advanced setup #65

Merged
merged 11 commits into from
Dec 19, 2024
Original file line number Diff line number Diff line change
Expand Up @@ -47,13 +47,14 @@ Requires --priv-key to be set for each 'from' address in the given testfile.",
)]
duration: Option<usize>,

/// The seed to use for generating spam transactions. If not provided, one is generated.
/// The seed to use for generating spam transactions & accounts.
#[arg(
short,
long,
long_help = "The seed to use for generating spam transactions"
long_help = "The seed to use for generating spam transactions",
default_value = "0xffffffffffffffffffffffffffffffff13131313131313131313131313131313"
)]
seed: Option<String>,
seed: String,

/// The private keys to use for blockwise spamming.
/// Required if `txs_per_block` is set.
Expand Down Expand Up @@ -109,6 +110,23 @@ May be specified multiple times."
default_value = "1.0"
)]
min_balance: String,

/// The seed used to generate pool accounts.
#[arg(
short,
long,
long_help = "The seed used to generate pool accounts.",
default_value = "0xffffffffffffffffffffffffffffffff13131313131313131313131313131313"
)]
seed: String,

/// The number of signers to generate for each pool.
#[arg(
short,
long = "signers-per-pool",
long_help = "The number of signers to generate for each pool."
)]
num_signers_per_pool: Option<usize>,
},

#[command(
Expand Down
4 changes: 2 additions & 2 deletions crates/cli/src/commands/mod.rs
Original file line number Diff line number Diff line change
@@ -1,16 +1,16 @@
mod contender_subcommand;
mod report;
mod run;
mod setup;
mod spam;
mod types;

use clap::Parser;

pub use contender_subcommand::ContenderSubcommand;
pub use report::report;
pub use run::run;
pub use setup::setup;
pub use spam::{spam, SpamCommandArgs};
pub use types::ContenderSubcommand;

#[derive(Parser, Debug)]
pub struct ContenderCli {
Expand Down
48 changes: 40 additions & 8 deletions crates/cli/src/commands/setup.rs
Original file line number Diff line number Diff line change
Expand Up @@ -2,12 +2,17 @@ use alloy::{
network::AnyNetwork, primitives::utils::parse_ether, providers::ProviderBuilder,
signers::local::PrivateKeySigner, transports::http::reqwest::Url,
};
use contender_core::{generator::RandSeed, test_scenario::TestScenario};
use contender_core::{
agent_controller::{AgentStore, SignerStore},
generator::RandSeed,
test_scenario::TestScenario,
};
use contender_testfile::TestConfig;
use std::str::FromStr;

use crate::util::{
check_private_keys_fns, find_insufficient_balance_addrs, get_signers_with_defaults,
check_private_keys_fns, find_insufficient_balance_addrs, fund_accounts, get_setup_pools,
get_signers_with_defaults,
};

pub async fn setup(
Expand All @@ -16,11 +21,14 @@ pub async fn setup(
rpc_url: impl AsRef<str>,
private_keys: Option<Vec<String>>,
min_balance: String,
seed: RandSeed,
signers_per_period: usize,
) -> Result<(), Box<dyn std::error::Error>> {
let url = Url::parse(rpc_url.as_ref()).expect("Invalid RPC URL");
let rpc_client = ProviderBuilder::new()
.network::<AnyNetwork>()
.on_http(url.to_owned());
let eth_client = ProviderBuilder::new().on_http(url.to_owned());
let testconfig: TestConfig = TestConfig::from_file(testfile.as_ref())?;
let min_balance = parse_ether(&min_balance)?;

Expand All @@ -30,10 +38,10 @@ pub async fn setup(
.iter()
.map(|key| PrivateKeySigner::from_str(key).expect("invalid private key"))
.collect::<Vec<PrivateKeySigner>>();
let signers = get_signers_with_defaults(private_keys);
let default_signers = get_signers_with_defaults(private_keys);
check_private_keys_fns(
&testconfig.setup.to_owned().unwrap_or_default(),
signers.as_slice(),
&default_signers,
);
let broke_accounts = find_insufficient_balance_addrs(
&user_signers.iter().map(|s| s.address()).collect::<Vec<_>>(),
Expand All @@ -42,17 +50,41 @@ pub async fn setup(
)
.await?;
if !broke_accounts.is_empty() {
panic!("Some accounts do not have sufficient balance");
panic!("Insufficient balance in provided user account(s)");
}

let mut agents = AgentStore::new();
let from_pools = get_setup_pools(&testconfig);
let mut all_signers = vec![];
all_signers.extend_from_slice(&user_signers);

for from_pool in &from_pools {
if agents.has_agent(from_pool) {
continue;
}

let agent = SignerStore::new_random(signers_per_period, &seed, from_pool);
all_signers.extend_from_slice(&agent.signers);
agents.add_agent(from_pool, agent);
}

fund_accounts(
&rpc_client,
&eth_client,
min_balance,
&all_signers,
&default_signers[0],
)
.await?;

let mut scenario = TestScenario::new(
testconfig.to_owned(),
db.clone().into(),
url,
None,
RandSeed::new(),
&signers,
Default::default(),
seed,
&default_signers,
agents,
)
.await?;

Expand Down
11 changes: 4 additions & 7 deletions crates/cli/src/commands/spam.rs
Original file line number Diff line number Diff line change
Expand Up @@ -14,7 +14,7 @@ use contender_core::{
use contender_testfile::TestConfig;

use crate::util::{
check_private_keys, fund_accounts, get_from_pools, get_signers_with_defaults,
check_private_keys, fund_accounts, get_signers_with_defaults, get_spam_pools,
spam_callback_default, SpamCallbackType,
};

Expand All @@ -26,7 +26,7 @@ pub struct SpamCommandArgs {
pub txs_per_block: Option<usize>,
pub txs_per_second: Option<usize>,
pub duration: Option<usize>,
pub seed: Option<String>,
pub seed: String,
pub private_keys: Option<Vec<String>>,
pub disable_reports: bool,
pub min_balance: String,
Expand All @@ -37,10 +37,7 @@ pub async fn spam(
args: SpamCommandArgs,
) -> Result<(), Box<dyn std::error::Error>> {
let testconfig = TestConfig::from_file(&args.testfile)?;
let rand_seed = args
.seed
.map(|s| RandSeed::seed_from_str(s.as_ref()))
.unwrap_or_default();
let rand_seed = RandSeed::seed_from_str(&args.seed);
let url = Url::parse(&args.rpc_url).expect("Invalid RPC URL");
let rpc_client = ProviderBuilder::new()
.network::<AnyNetwork>()
Expand All @@ -57,7 +54,7 @@ pub async fn spam(
.expect("No spam function calls found in testfile");

// distill all from_pool arguments from the spam requests
let from_pools = get_from_pools(&testconfig);
let from_pools = get_spam_pools(&testconfig);

let mut agents = AgentStore::new();
let signers_per_period = args
Expand Down
26 changes: 19 additions & 7 deletions crates/cli/src/main.rs
Original file line number Diff line number Diff line change
Expand Up @@ -5,7 +5,7 @@ mod util;
use std::sync::LazyLock;

use commands::{ContenderCli, ContenderSubcommand, SpamCommandArgs};
use contender_core::db::DbOps;
use contender_core::{db::DbOps, generator::RandSeed};
use contender_sqlite::SqliteDb;

static DB: LazyLock<SqliteDb> = std::sync::LazyLock::new(|| {
Expand All @@ -16,14 +16,28 @@ static DB: LazyLock<SqliteDb> = std::sync::LazyLock::new(|| {
async fn main() -> Result<(), Box<dyn std::error::Error>> {
let args = ContenderCli::parse_args();
let _ = DB.create_tables(); // ignore error; tables already exist
let db = DB.clone();

match args.command {
ContenderSubcommand::Setup {
testfile,
rpc_url,
private_keys,
min_balance,
} => commands::setup(&DB.clone(), testfile, rpc_url, private_keys, min_balance).await?,
seed,
num_signers_per_pool,
} => {
commands::setup(
&db,
testfile,
rpc_url,
private_keys,
min_balance,
RandSeed::seed_from_str(&seed),
num_signers_per_pool.unwrap_or(1),
)
.await?
}

ContenderSubcommand::Spam {
testfile,
Expand All @@ -38,7 +52,7 @@ async fn main() -> Result<(), Box<dyn std::error::Error>> {
min_balance,
} => {
commands::spam(
&DB.clone(),
&db,
SpamCommandArgs {
testfile,
rpc_url,
Expand All @@ -55,9 +69,7 @@ async fn main() -> Result<(), Box<dyn std::error::Error>> {
.await?
}

ContenderSubcommand::Report { id, out_file } => {
commands::report(&DB.clone(), id, out_file)?
}
ContenderSubcommand::Report { id, out_file } => commands::report(&db, id, out_file)?,

ContenderSubcommand::Run {
scenario,
Expand All @@ -68,7 +80,7 @@ async fn main() -> Result<(), Box<dyn std::error::Error>> {
txs_per_duration,
} => {
commands::run(
&DB.clone(),
&db,
scenario,
rpc_url,
private_key,
Expand Down
12 changes: 11 additions & 1 deletion crates/cli/src/util.rs
Original file line number Diff line number Diff line change
Expand Up @@ -33,7 +33,17 @@ pub const DEFAULT_PRV_KEYS: [&str; 10] = [
"0x2a871d0798f97d79848a013d4936a73bf4cc922c825d33c1cf7073dff6d409c6",
];

pub fn get_from_pools(testconfig: &TestConfig) -> Vec<String> {
pub fn get_setup_pools(testconfig: &TestConfig) -> Vec<String> {
testconfig
.setup
.to_owned()
.unwrap_or_default()
.into_iter()
.filter_map(|s| s.from_pool)
.collect()
}

pub fn get_spam_pools(testconfig: &TestConfig) -> Vec<String> {
let mut from_pools = vec![];
let spam = testconfig
.spam
Expand Down
58 changes: 36 additions & 22 deletions crates/core/src/generator/mod.rs
Original file line number Diff line number Diff line change
Expand Up @@ -203,27 +203,41 @@ where
}
PlanType::Setup(on_setup_step) => {
let setup_steps = conf.get_setup_steps()?;
for step in setup_steps.iter() {
// lookup placeholders in DB & update map before templating
templater.find_fncall_placeholders(step, db, &mut placeholder_map)?;

// setup tx with template values
let tx = NamedTxRequest::new(
templater.template_function_call(
&self.make_strict_call(step, 0)?,
&placeholder_map,
)?,
None,
step.kind.to_owned(),
);

let handle = on_setup_step(tx.to_owned())?;
if let Some(handle) = handle {
handle.await.map_err(|e| {
ContenderError::with_err(e, "join error; callback crashed")
})?;
let agents = self.get_agent_store();
let num_accts = agents
.all_agents()
.next()
.map(|(_, store)| store.signers.len())
.unwrap_or(1);

for i in 0..(num_accts) {
for step in setup_steps.iter() {
if i > 0 && step.from_pool.is_none() {
// only loop on from_pool steps; single-account steps can't be repeated
continue;
}
// lookup placeholders in DB & update map before templating
templater.find_fncall_placeholders(step, db, &mut placeholder_map)?;

// setup tx with template values
let tx = NamedTxRequest::new(
templater.template_function_call(
&self.make_strict_call(step, i)?, // 'from' address injected here
&placeholder_map,
)?,
None,
step.kind.to_owned(),
);

let handle = on_setup_step(tx.to_owned())?;
if let Some(handle) = handle {
handle.await.map_err(|e| {
ContenderError::with_err(e, "join error; callback crashed")
})?;
}
txs.push(tx.into());
}
txs.push(tx.into());
}
}
PlanType::Spam(num_txs, on_spam_setup) => {
Expand Down Expand Up @@ -281,7 +295,7 @@ where
for step in spam_steps.iter() {
// converts a FunctionCallDefinition to a NamedTxRequest (filling in fuzzable args),
// returns a callback handle and the processed tx request
let process_tx = |req| {
let prepare_tx = |req| {
let args = get_fuzzed_args(req, &canonical_fuzz_map, i);
let fuzz_tx_value = get_fuzzed_tx_value(req, &canonical_fuzz_map, i);
let mut req = req.to_owned();
Expand All @@ -304,7 +318,7 @@ where

match step {
SpamRequest::Tx(req) => {
let (handle, tx) = process_tx(req)?;
let (handle, tx) = prepare_tx(req)?;
if let Some(handle) = handle {
handle.await.map_err(|e| {
ContenderError::with_err(e, "error from callback")
Expand All @@ -315,7 +329,7 @@ where
SpamRequest::Bundle(req) => {
let mut bundle_txs = vec![];
for tx in req.txs.iter() {
let (handle, txr) = process_tx(tx)?;
let (handle, txr) = prepare_tx(tx)?;
if let Some(handle) = handle {
handle.await.map_err(|e| {
ContenderError::with_err(e, "error from callback")
Expand Down
Loading
Loading