use clap::Parser;
use std::sync::Arc;
use std::time::{Duration, Instant};
use tracing::error;
use gromnie_runner::{
AutoLoginConsumer, ClientConfig, ClientNaming, CompositeConsumer, FnConsumerBuilder,
MultiClientStats, RunConfig, RunResult, StatsConsumer, logging,
};
#[derive(Parser)]
#[command(name = "load-tester")]
#[command(about = "Load testing tool for AC server")]
pub struct Args {
#[command(subcommand)]
command: Option<Command>,
#[command(flatten)]
run_args: RunArgs,
}
#[derive(Parser, Debug)]
pub enum Command {
Naming {
client_id: u32,
},
Run(RunArgs),
}
#[derive(Parser, Debug, Clone, Default)]
pub struct RunArgs {
#[arg(short, long, default_value = "5")]
clients: u32,
#[arg(long, default_value = "localhost")]
host: String,
#[arg(short, long, default_value = "9000")]
port: u16,
#[arg(short, long, default_value = "1000")]
rate_limit: u64,
#[arg(short, long)]
verbose: bool,
#[arg(long, default_value = "5")]
stats_interval: u64,
}
#[tokio::main]
async fn main() {
let args = Args::parse();
if let Some(Command::Naming { client_id }) = args.command {
let naming = ClientNaming::new(client_id);
println!("Client ID: {}", client_id);
println!("Account: {}", naming.account_name());
println!("Password: {}", naming.password());
println!("Character: {}", naming.character_name());
return;
}
let _log_guard = logging::init_logging("load-tester").expect("failed to initialize logging");
let run_args = if let Some(Command::Run(args)) = args.command {
args
} else {
args.run_args
};
if run_args.rate_limit == 0 {
error!("Rate limit must be a non-zero value (milliseconds)");
std::process::exit(1);
}
let server_address = format!("{}:{}", run_args.host, run_args.port);
let config = RunConfig::multi(
server_address,
run_args.clients,
run_args.rate_limit,
false, );
let stats = Arc::new(MultiClientStats::default());
let stats_for_builder = stats.clone();
let verbose = run_args.verbose;
let consumer_builder = FnConsumerBuilder::new(move |client_id, client_config, action_tx| {
let character_name = format!("{}-A", client_config.account_name);
let stats_consumer = Box::new(
StatsConsumer::new(client_id, stats_for_builder.clone()).with_verbose(verbose),
);
let auto_login_consumer = Box::new(
AutoLoginConsumer::new(client_id, character_name, action_tx).with_verbose(verbose),
);
Box::new(CompositeConsumer::new(vec![
stats_consumer,
auto_login_consumer,
]))
});
let start_time = Instant::now();
let stats_for_task = stats.clone();
let stats_interval = run_args.stats_interval;
let stats_handle = tokio::spawn(async move {
loop {
tokio::time::sleep(Duration::from_secs(stats_interval)).await;
stats_for_task.print(start_time.elapsed().as_secs());
}
});
let host = run_args.host.clone();
let port = run_args.port;
let result = gromnie_runner::run(
config,
consumer_builder,
Some(move |client_id| {
let naming = ClientNaming::new(client_id);
ClientConfig::new(
client_id,
format!("{}:{}", host, port),
naming.account_name(),
naming.password(),
)
}),
None, )
.await;
stats_handle.abort();
if let RunResult::Multi(final_stats) = result {
final_stats.print_final(start_time.elapsed());
}
}