load_generator.rs

// Copyright (c) Mysten Labs, Inc.
// SPDX-License-Identifier: Apache-2.0

use std::{net::SocketAddr, path::PathBuf};

use anyhow::{anyhow, Context};
use clap::Parser;
use remora::{
    //executor::sui::{check_logs_for_shared_object, import_from_files},
    client::load_generator::{default_metrics_address, LoadGenerator},
    config::{BenchmarkParameters, ImportExport, ValidatorConfig},
    executor::api::Executor,
    executor::fake::FakeExecutor,
    executor::sui::SuiExecutor,
    executor::tpcc::TpccExecutor,
};

#[derive(Parser, Debug)]
#[clap(rename_all = "kebab-case")]
#[command(author, version, about = "Remora load generator", long_about = None)]
struct Args {
    /// The path to the validator configuration.
    #[clap(long, value_name = "FILE")]
    validator_config: PathBuf,
    /// The path to the configuration for the benchmark.
    #[clap(long, value_name = "FILE")]
    benchmark_config: Option<PathBuf>,
    /// The address to expose metrics on.
    #[clap(long, value_name = "ADDRESS", default_value_t = default_metrics_address())]
    metrics_address: SocketAddr,
}

/// The main function for the load generator.
#[tokio::main]
async fn main() -> anyhow::Result<()> {
    let args = Args::parse();
    let validator_config =
        ValidatorConfig::load(&args.validator_config).context("Failed to load validator config")?;
    let benchmark_config = match args.benchmark_config {
        Some(path) => BenchmarkParameters::load(path).context("Failed to load benchmark config")?,
        None => BenchmarkParameters::default(),
    };
    let metrics_address = args.metrics_address;

    tracing::info!("Load generator exposing metrics on {metrics_address}");
    tracing_subscriber::fmt::try_init().map_err(|e| anyhow!("{e}"))?;
    let _registry = mysten_metrics::start_prometheus_server(metrics_address);

    // Create genesis and generate transactions.
    let primary_address = validator_config.client_server_address;

    // Initialize based on executor type
    if benchmark_config.workload.is_fake() {
        let load_generator =
            LoadGenerator::<FakeExecutor>::new(benchmark_config.clone(), primary_address);
        run_load_generator(load_generator).await?;
    } else if benchmark_config.workload.is_tpcc() {
        let load_generator =
            LoadGenerator::<TpccExecutor>::new(benchmark_config.clone(), primary_address);
        run_load_generator(load_generator).await?;
    } else {
        let load_generator =
            LoadGenerator::<SuiExecutor>::new(benchmark_config.clone(), primary_address);
        run_load_generator(load_generator).await?;
    }

    Ok(())
}

async fn run_load_generator<E>(mut load_generator: LoadGenerator<E>) -> anyhow::Result<()>
where
    E: Executor + Send + Sync + 'static,
    <E as Executor>::Transaction: Send + Sync,
{
    let transactions = load_generator.initialize().await;

    // Submit transactions to the server.
    load_generator.run(transactions).await;

    Ok(())
}