Merge pull request #61 from flashbots/refactor-cli

Refactor cli

crates/cli/src/commands/mod.rs

@@ -0,0 +1,25 @@
+mod report;
+mod run;
+mod setup;
+mod spam;
+mod types;
+
+use clap::Parser;
+
+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 {
+    #[command(subcommand)]
+    pub command: ContenderSubcommand,
+}
+
+impl ContenderCli {
+    pub fn parse_args() -> Self {
+        Self::parse()
+    }
+}

crates/cli/src/commands/report.rs

@@ -0,0 +1,43 @@
+use contender_core::db::DbOps;
+use csv::WriterBuilder;
+
+use crate::util::write_run_txs;
+
+pub fn report(
+    db: &(impl DbOps + Clone + Send + Sync + 'static),
+    id: Option<u64>,
+    out_file: Option<String>,
+) -> Result<(), Box<dyn std::error::Error>> {
+    let num_runs = db.num_runs()?;
+    let id = if let Some(id) = id {
+        if id == 0 || id > num_runs {
+            panic!("Invalid run ID: {}", id);
+        }
+        id
+    } else {
+        if num_runs == 0 {
+            panic!("No runs to report");
+        }
+        // get latest run
+        println!("No run ID provided. Using latest run ID: {}", num_runs);
+        num_runs
+    };
+    let txs = db.get_run_txs(id)?;
+    println!("found {} txs", txs.len());
+    println!(
+        "Exporting report for run ID {:?} to out_file {:?}",
+        id, out_file
+    );
+
+    if let Some(out_file) = out_file {
+        let mut writer = WriterBuilder::new().has_headers(true).from_path(out_file)?;
+        write_run_txs(&mut writer, &txs)?;
+    } else {
+        let mut writer = WriterBuilder::new()
+            .has_headers(true)
+            .from_writer(std::io::stdout());
+        write_run_txs(&mut writer, &txs)?; // TODO: write a macro that lets us generalize the writer param to write_run_txs, then refactor this duplication
+    };
+
+    Ok(())
+}

crates/cli/src/commands/run.rs

@@ -0,0 +1,121 @@
+use std::{env, str::FromStr, sync::Arc};
+
+use alloy::{
+    eips::BlockId,
+    network::AnyNetwork,
+    providers::{Provider, ProviderBuilder},
+    rpc::types::BlockTransactionsKind,
+    transports::http::reqwest::Url,
+};
+use contender_core::{
+    agent_controller::AgentStore,
+    db::DbOps,
+    error::ContenderError,
+    generator::RandSeed,
+    spammer::{LogCallback, Spammer, TimedSpammer},
+    test_scenario::TestScenario,
+};
+use contender_testfile::TestConfig;
+
+use crate::{
+    default_scenarios::{BuiltinScenario, BuiltinScenarioConfig},
+    util::{check_private_keys, get_signers_with_defaults, prompt_cli},
+};
+
+pub async fn run(
+    db: &(impl DbOps + Clone + Send + Sync + 'static),
+    scenario: BuiltinScenario,
+    rpc_url: String,
+    private_key: Option<String>,
+    interval: usize,
+    duration: usize,
+    txs_per_duration: usize,
+) -> Result<(), Box<dyn std::error::Error>> {
+    let user_signers = get_signers_with_defaults(private_key.map(|s| vec![s]));
+    let admin_signer = &user_signers[0];
+    let rand_seed = RandSeed::default();
+    let provider = ProviderBuilder::new()
+        .network::<AnyNetwork>()
+        .on_http(Url::parse(&rpc_url).expect("Invalid RPC URL"));
+    let block_gas_limit = provider
+        .get_block(BlockId::latest(), BlockTransactionsKind::Hashes)
+        .await?
+        .map(|b| b.header.gas_limit)
+        .ok_or(ContenderError::SetupError(
+            "failed getting gas limit from block",
+            None,
+        ))?;
+
+    let fill_percent = env::var("C_FILL_PERCENT")
+        .map(|s| u16::from_str(&s).expect("invalid u16: fill_percent"))
+        .unwrap_or(100u16);
+
+    let scenario_config = match scenario {
+        BuiltinScenario::FillBlock => BuiltinScenarioConfig::fill_block(
+            block_gas_limit,
+            txs_per_duration as u64,
+            admin_signer.address(),
+            fill_percent,
+        ),
+    };
+    let testconfig: TestConfig = scenario_config.into();
+    check_private_keys(&testconfig, &user_signers);
+
+    let rpc_url = Url::parse(&rpc_url).expect("Invalid RPC URL");
+    let mut scenario = TestScenario::new(
+        testconfig,
+        db.clone().into(),
+        rpc_url.to_owned(),
+        None,
+        rand_seed,
+        &user_signers,
+        AgentStore::default(),
+    )
+    .await?;
+
+    let contract_name = "SpamMe";
+    let contract_result = db.get_named_tx(contract_name)?;
+    let do_deploy_contracts = if contract_result.is_some() {
+        let input = prompt_cli(format!(
+            "{} deployment already detected. Re-deploy? [y/N]",
+            contract_name
+        ));
+        input.to_lowercase() == "y"
+    } else {
+        true
+    };
+
+    if do_deploy_contracts {
+        println!("deploying contracts...");
+        scenario.deploy_contracts().await?;
+    }
+
+    println!("running setup...");
+    scenario.run_setup().await?;
+
+    let wait_duration = std::time::Duration::from_secs(interval as u64);
+    let spammer = TimedSpammer::new(wait_duration);
+    let timestamp = std::time::SystemTime::now()
+        .duration_since(std::time::UNIX_EPOCH)
+        .expect("Time went backwards")
+        .as_millis();
+    let run_id = db.insert_run(timestamp as u64, duration * txs_per_duration)?;
+    let callback = LogCallback::new(Arc::new(
+        ProviderBuilder::new()
+            .network::<AnyNetwork>()
+            .on_http(rpc_url),
+    ));
+
+    println!("starting spammer...");
+    spammer
+        .spam_rpc(
+            &mut scenario,
+            txs_per_duration,
+            duration,
+            Some(run_id),
+            callback.into(),
+        )
+        .await?;
+
+    Ok(())
+}

crates/cli/src/commands/setup.rs

@@ -0,0 +1,64 @@
+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_testfile::TestConfig;
+use std::str::FromStr;
+
+use crate::util::{
+    check_private_keys_fns, find_insufficient_balance_addrs, get_signers_with_defaults,
+};
+
+pub async fn setup(
+    db: &(impl contender_core::db::DbOps + Clone + Send + Sync + 'static),
+    testfile: impl AsRef<str>,
+    rpc_url: impl AsRef<str>,
+    private_keys: Option<Vec<String>>,
+    min_balance: String,
+) -> 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 testconfig: TestConfig = TestConfig::from_file(testfile.as_ref())?;
+    let min_balance = parse_ether(&min_balance)?;
+
+    let user_signers = private_keys
+        .as_ref()
+        .unwrap_or(&vec![])
+        .iter()
+        .map(|key| PrivateKeySigner::from_str(key).expect("invalid private key"))
+        .collect::<Vec<PrivateKeySigner>>();
+    let signers = get_signers_with_defaults(private_keys);
+    check_private_keys_fns(
+        &testconfig.setup.to_owned().unwrap_or_default(),
+        signers.as_slice(),
+    );
+    let broke_accounts = find_insufficient_balance_addrs(
+        &user_signers.iter().map(|s| s.address()).collect::<Vec<_>>(),
+        min_balance,
+        &rpc_client,
+    )
+    .await?;
+    if !broke_accounts.is_empty() {
+        panic!("Some accounts do not have sufficient balance");
+    }
+
+    let mut scenario = TestScenario::new(
+        testconfig.to_owned(),
+        db.clone().into(),
+        url,
+        None,
+        RandSeed::new(),
+        &signers,
+        Default::default(),
+    )
+    .await?;
+
+    scenario.deploy_contracts().await?;
+    scenario.run_setup().await?;
+    // TODO: catch failures and prompt user to retry specific steps
+
+    Ok(())
+}