Composite Scenarios Execution - Setup parallely, spam parallely (and in stages)

Cargo.toml

@@ -1,7 +1,8 @@
 [workspace]
 members = [
     "crates/bundle_provider",
-    "crates/cli/",
+    "crates/cli/", 
+    "crates/composefile/",
     "crates/core/",
     "crates/engine_provider",
     "crates/sqlite_db/",
@@ -20,6 +21,7 @@ repository = "https://github.com/flashbots/contender"
 
 [workspace.dependencies]
 contender_core = { path = "crates/core/" }
+contender_composefile = { path = "crates/composefile/" }
 contender_sqlite = { path = "crates/sqlite_db/" }
 contender_testfile = { path = "crates/testfile/" }
 contender_bundle_provider = { path = "crates/bundle_provider/" }

contender-compose.yml

@@ -0,0 +1,48 @@
+setup:
+  simpler:
+    testfile: ./scenarios/simpler.toml
+    min_balance: 12.1
+    call_forkchoice: true
+    optimism: false
+    auth_rpc_url: "https://rpc.chain.com"
+    jwt_secret: SECRET
+  uniV2:
+    testfile: ./scenarios/uniV2.toml
+    rpc_url: http://localhost:8545 # Optional
+    min_balance: "11"
+    env:
+      - Key1=Valu1
+      - Key2=Valu2
+    tx_type: eip1559  # Optional
+
+spam:
+  stages:
+    warmup:
+      - testfile: ./scenarios/simpler.toml
+        tps: 3
+        duration: 3
+        loop: 2
+        seed: "123"
+      - testfile: ./scenarios/simpler.toml
+        tps: 1
+        duration: 3
+        loop: 2
+      - testfile: ./scenarios/simpler.toml
+        tps: 4
+        duration: 3
+        loop: 2
+
+        call_forkchoice: true
+        optimism: false
+        auth_rpc_url: "https://rpc.chain.com"
+        jwt_secret: SECRET
+    medium:
+      - testfile: ./scenarios/uniV2.toml
+        tps: 10
+        min_balance: 20
+        loop: 1
+        duration: 2
+      - testfile: ./scenarios/simpler.toml
+        tps: 14
+        duration: 3
+        loop: 2

crates/cli/Cargo.toml

@@ -18,6 +18,7 @@ contender_core = { workspace = true }
 contender_sqlite = { workspace = true }
 contender_testfile = { workspace = true }
 contender_engine_provider = { workspace = true }
+contender_composefile = { workspace = true }
 
 ansi_term = { workspace = true }
 serde = { workspace = true }
@@ -42,7 +43,9 @@ serde_json = { workspace = true }
 tracing = { workspace = true }
 tracing-subscriber = { workspace = true }
 webbrowser = { workspace = true }
+yaml-rust2 = "0.10.2"
 op-alloy-network = { workspace = true }
+hashlink = "0.10.0"
 
 [dev-dependencies]
 tempfile = "3.15.0"

crates/cli/src/commands/composite.rs

@@ -0,0 +1,113 @@
+use std::{fmt::Debug, sync::Arc};
+
+use contender_composefile::composefile::{ComposeFile, CompositeSpamConfiguration};
+use tokio::{sync::Mutex, task};
+use tracing::{error, info};
+
+use crate::commands::{setup, spam, SetupCommandArgs, SpamCommandArgs};
+
+#[derive(Debug, clap::Args)]
+pub struct CompositeScenarioArgs {
+    pub filename: Option<String>,
+    pub private_keys: Option<Vec<String>>,
+}
+
+pub async fn composite(
+    db: &(impl contender_core::db::DbOps + Clone + Send + Sync + 'static),
+    args: CompositeScenarioArgs,
+) -> Result<(), Box<dyn std::error::Error>> {
+    let compose_file_name = match args.filename {
+        Some(filepath) => filepath,
+        None => String::from("./contender-compose.yml"),
+    };
+
+    let compose_file = ComposeFile::init_from_path(compose_file_name)?;
+    let sharable_db = Arc::new(Mutex::new(db.clone()));
+    let sharable_private_keys = Arc::new(Mutex::new(args.private_keys));
+
+    let setup_scenarios = compose_file.get_setup_config()?;
+    let setup_tasks: Vec<_> = setup_scenarios
+        .into_iter()
+        .enumerate()
+        .map(|(index, scenario)| {
+            let db_clone = sharable_db.clone();
+            let scenario_config = scenario.clone();
+            let private_keys = sharable_private_keys.clone();
+            let setup_command_args = SetupCommandArgs::from_json(scenario_config.config);
+
+            task::spawn(async move {
+                let setup_command = setup_command_args
+                    .await
+                    .expect("msg")
+                    .with_private_keys(private_keys.lock().await.clone());
+                match setup(&*db_clone.lock().await, setup_command).await {
+                    Ok(_) => info!(
+                        "Setup [{index}] - {}: completed successfully",
+                        &scenario_config.name
+                    ),
+                    Err(err) => error!(
+                        "Setup [{index}] - {} failed: {err:?}",
+                        &scenario_config.name
+                    ),
+                };
+                //setup(&*db_clone.lock().await, scenario_config.config).await.map_err(|e| Err("s".into()))
+            })
+        })
+        .collect();
+
+    futures::future::join_all(setup_tasks).await;
+
+    info!("================================================================================================= Done Composite run for setup =================================================================================================");
+
+    let spam_scenarios = compose_file.get_spam_config()?;
+    for scenario in spam_scenarios {
+        let CompositeSpamConfiguration {
+            stage_name,
+            spam_configs,
+        } = scenario;
+        info!("================================================================================================= Running stage: {stage_name:?} =================================================================================================");
+
+        let mut spam_tasks = Vec::new();
+        let sharable_stage_name_object = Arc::new(Mutex::new(stage_name.clone()));
+        for (spam_scenario_index, spam_command) in spam_configs.into_iter().enumerate() {
+            info!("Starting scenario [{spam_scenario_index:?}]");
+            let db_clone = sharable_db.clone();
+            let private_keys = sharable_private_keys.clone();
+            let private_keys_clone = private_keys.clone().lock().await.clone();
+            let task = task::spawn(async move {
+                let spam_call = async || -> Result<(), Box<dyn std::error::Error>> {
+                    let spam_command_args = SpamCommandArgs::from_json(spam_command)
+                        .await?
+                        .with_private_keys(private_keys_clone);
+
+                    let mut test_scenario = spam_command_args
+                        .init_scenario(&*db_clone.lock().await)
+                        .await?;
+                    spam(
+                        &*db_clone.lock().await,
+                        &spam_command_args,
+                        &mut test_scenario,
+                    )
+                    .await?;
+                    Ok(())
+                };
+
+                match spam_call().await {
+                    Ok(()) => {
+                        info!("Successful: Scenario [{spam_scenario_index:?}]");
+                    }
+                    Err(err) => {
+                        error!("Error occured while Scenario [{spam_scenario_index:?}]: {err:?}")
+                    }
+                };
+            });
+            spam_tasks.push(task);
+        }
+
+        for task in spam_tasks {
+            task.await?;
+        }
+        info!("================================================================================================= Done Composite run for spam - Stage [{:?}] =================================================================================================", &*sharable_stage_name_object.clone().lock().await);
+    }
+    Ok(())
+}

crates/cli/src/commands/contender_subcommand.rs

@@ -9,6 +9,27 @@ use super::spam::SpamCliArgs;
 
 #[derive(Debug, Subcommand)]
 pub enum ContenderSubcommand {
+    #[command(
+        name = "compose",
+        about = "Composite scenario execution, setup and spam multiple contracts from a single YML file."
+    )]
+    Composite {
+        #[arg(
+            short,
+            long,
+            long_help = "File path for composite scenario file",
+            default_value = "./contender-compose.yml"
+        )]
+        filename: Option<String>,
+        #[arg(
+            short,
+            long = "priv-key",
+            long_help = "Add private keys to wallet map. Used to fund agent accounts or sign transactions.
+May be specified multiple times."
+        )]
+        private_keys: Option<Vec<String>>,
+    },
+
     #[command(name = "admin", about = "Admin commands")]
     Admin {
         #[command(subcommand)]

crates/cli/src/commands/mod.rs

@@ -1,5 +1,6 @@
 pub mod admin;
 pub mod common;
+pub mod composite;
 mod contender_subcommand;
 pub mod db;
 mod report;
@@ -9,6 +10,7 @@ mod spamd;
 
 use clap::Parser;
 
+pub use composite::composite;
 pub use contender_subcommand::{ContenderSubcommand, DbCommand};
 pub use report::report;
 pub use setup::{setup, SetupCliArgs, SetupCommandArgs};