use std::sync::Arc;
use std::sync::atomic::{AtomicU32, Ordering};
use std::time::{Duration, Instant};
use tokio::sync::{RwLock, broadcast, mpsc, watch};
use tracing::{error, info};
use crate::event_bus::{EventBus, EventEnvelope};
use crate::event_consumer::EventConsumer;
use crate::event_wrapper::EventWrapper;
use gromnie_client::client::Client;
use gromnie_client::transport::ClientTransport;
pub use gromnie_client::config::ClientConfig;
pub type TransportFactory = Arc<dyn Fn(&ClientConfig) -> Box<dyn ClientTransport> + Send + Sync>;
pub(crate) async fn create_client_from_config(
config: &ClientConfig,
raw_event_tx: mpsc::Sender<gromnie_events::ClientEvent>,
transport_factory: Option<&TransportFactory>,
) -> (
Client,
mpsc::UnboundedSender<gromnie_events::SimpleClientAction>,
) {
if let Some(factory) = transport_factory {
Client::new_with_transport(
config.id,
config.address.clone(),
config.account_name.clone(),
config.password.clone(),
config.character_name.clone(),
raw_event_tx,
config.reconnect,
factory(config),
)
.await
} else {
Client::new_with_reconnect(
config.id,
config.address.clone(),
config.account_name.clone(),
config.password.clone(),
config.character_name.clone(),
raw_event_tx,
config.reconnect,
)
.await
}
}
#[derive(Clone, Debug)]
pub enum RunConfig {
Single {
client: ClientConfig,
},
Multi {
server_address: String,
num_clients: u32,
spawn_interval_ms: u64,
shared_event_bus: bool,
},
}
impl RunConfig {
pub fn single(client: ClientConfig) -> Self {
Self::Single { client }
}
pub fn multi(
server_address: String,
num_clients: u32,
spawn_interval_ms: u64,
shared_event_bus: bool,
) -> Self {
Self::Multi {
server_address,
num_clients,
spawn_interval_ms,
shared_event_bus,
}
}
pub fn multi_default(server_address: String, num_clients: u32) -> Self {
Self::Multi {
server_address,
num_clients,
spawn_interval_ms: 1000,
shared_event_bus: false,
}
}
}
#[derive(Clone, Debug)]
pub struct MultiClientConfig {
pub server_address: String,
pub num_clients: u32,
pub spawn_interval_ms: u64,
pub shared_event_bus: bool,
}
impl MultiClientConfig {
pub fn new(server_address: String, num_clients: u32) -> Self {
Self {
server_address,
num_clients,
spawn_interval_ms: 1000,
shared_event_bus: false,
}
}
pub fn with_spawn_interval(mut self, interval_ms: u64) -> Self {
self.spawn_interval_ms = interval_ms;
self
}
pub fn with_shared_event_bus(mut self, shared: bool) -> Self {
self.shared_event_bus = shared;
self
}
}
#[derive(Default)]
pub struct MultiClientStats {
pub attempted: AtomicU32,
pub spawned: AtomicU32,
pub authenticated: AtomicU32,
pub character_created: AtomicU32,
pub logged_in: AtomicU32,
pub errors: AtomicU32,
pub task_failures: AtomicU32,
}
impl MultiClientStats {
pub fn print(&self, elapsed_secs: u64) {
let attempted = self.attempted.load(Ordering::SeqCst);
let spawned = self.spawned.load(Ordering::SeqCst);
let auth = self.authenticated.load(Ordering::SeqCst);
let login = self.logged_in.load(Ordering::SeqCst);
let char_created = self.character_created.load(Ordering::SeqCst);
let errors = self.errors.load(Ordering::SeqCst);
let task_failures = self.task_failures.load(Ordering::SeqCst);
info!(
"[Stats @ {}s] Attempted: {} | Spawned: {} | Auth: {} | LoggedIn: {} | CharCreated: {} | Errors: {} | TaskFailures: {}",
elapsed_secs, attempted, spawned, auth, login, char_created, errors, task_failures
);
}
pub fn print_final(&self, total_time: Duration) {
let attempted = self.attempted.load(Ordering::SeqCst);
let spawned = self.spawned.load(Ordering::SeqCst);
let auth = self.authenticated.load(Ordering::SeqCst);
let login = self.logged_in.load(Ordering::SeqCst);
let errors = self.errors.load(Ordering::SeqCst);
let task_failures = self.task_failures.load(Ordering::SeqCst);
info!("========================================");
info!("Multi-client run complete");
info!("Total time: {:.2}s", total_time.as_secs_f64());
info!("========================================");
info!("Client Launch:");
info!(" Attempted: {}", attempted);
info!(" Spawned: {}", spawned);
info!(" Task failures: {}", task_failures);
info!(
" Failed to spawn: {}",
attempted.saturating_sub(spawned + task_failures)
);
info!("========================================");
info!("Client Progress:");
info!(" Authenticated: {}", auth);
info!(" Logged in: {}", login);
info!(" Auth/login errors: {}", errors);
info!("========================================");
if attempted > 0 {
info!(
"Success rate: {:.1}% ({} / {} attempted)",
(login as f64 / attempted as f64) * 100.0,
login,
attempted
);
}
info!("========================================");
}
}
pub trait MultiClientConsumerFactory: Send + Sync {
fn create_consumer(
&self,
client_id: u32,
client_config: &ClientConfig,
action_tx: mpsc::UnboundedSender<gromnie_events::SimpleClientAction>,
) -> Box<dyn EventConsumer>;
}
pub struct FnConsumerFactory<F>
where
F: Fn(
u32,
&ClientConfig,
mpsc::UnboundedSender<gromnie_events::SimpleClientAction>,
) -> Box<dyn EventConsumer>
+ Send
+ Sync,
{
f: F,
}
impl<F> FnConsumerFactory<F>
where
F: Fn(
u32,
&ClientConfig,
mpsc::UnboundedSender<gromnie_events::SimpleClientAction>,
) -> Box<dyn EventConsumer>
+ Send
+ Sync,
{
pub fn new(f: F) -> Self {
Self { f }
}
}
impl<F> MultiClientConsumerFactory for FnConsumerFactory<F>
where
F: Fn(
u32,
&ClientConfig,
mpsc::UnboundedSender<gromnie_events::SimpleClientAction>,
) -> Box<dyn EventConsumer>
+ Send
+ Sync,
{
fn create_consumer(
&self,
client_id: u32,
client_config: &ClientConfig,
action_tx: mpsc::UnboundedSender<gromnie_events::SimpleClientAction>,
) -> Box<dyn EventConsumer> {
(self.f)(client_id, client_config, action_tx)
}
}
#[derive(Clone)]
pub struct EventBusManager {
pub(crate) event_bus: Arc<EventBus>,
}
impl EventBusManager {
pub fn new(capacity: usize) -> Self {
let (event_bus, _) = EventBus::new(capacity);
Self {
event_bus: Arc::new(event_bus),
}
}
pub fn create_sender(&self, client_id: u32) -> crate::event_bus::EventSender {
self.event_bus.create_sender(client_id)
}
pub fn subscribe(&self) -> broadcast::Receiver<EventEnvelope> {
self.event_bus.subscribe()
}
pub fn subscriber_count(&self) -> usize {
self.event_bus.subscriber_count()
}
}
pub async fn run_client<C, F>(
config: ClientConfig,
event_bus_manager: Arc<EventBusManager>,
event_consumer_factory: F,
shutdown_rx: Option<tokio::sync::watch::Receiver<bool>>,
) where
C: EventConsumer,
F: FnOnce(mpsc::UnboundedSender<gromnie_events::SimpleClientAction>) -> C,
{
let (raw_event_tx, raw_event_rx) = mpsc::channel::<gromnie_events::ClientEvent>(256);
let event_wrapper = EventWrapper::new(config.id, event_bus_manager.event_bus.clone());
tokio::spawn(async move {
event_wrapper.run(raw_event_rx).await;
});
let event_rx = event_bus_manager.subscribe();
let (client, action_tx) = create_client_from_config(&config, raw_event_tx, None).await;
let client = Arc::new(RwLock::new(client));
let event_consumer = event_consumer_factory(action_tx);
run_client_internal(client, event_rx, Box::new(event_consumer), shutdown_rx).await;
}
pub async fn run_client_with_consumers<F>(
config: ClientConfig,
event_bus_manager: Arc<EventBusManager>,
consumers_factory: F,
shutdown_rx: Option<tokio::sync::watch::Receiver<bool>>,
) where
F: FnOnce(
mpsc::UnboundedSender<gromnie_events::SimpleClientAction>,
) -> Vec<Box<dyn EventConsumer>>,
{
use crate::event_bus::{EventEnvelope, EventSource, EventType, SystemEvent};
let (raw_event_tx, raw_event_rx) = mpsc::channel::<gromnie_events::ClientEvent>(256);
let event_wrapper = EventWrapper::new(config.id, event_bus_manager.event_bus.clone());
tokio::spawn(async move {
event_wrapper.run(raw_event_rx).await;
});
let (client, action_tx) = create_client_from_config(&config, raw_event_tx, None).await;
let client = Arc::new(RwLock::new(client));
let mut consumers = consumers_factory(action_tx);
let shutdown_sender = event_bus_manager.create_sender(config.id);
let mut consumer_tasks = Vec::new();
for (idx, mut consumer) in consumers.drain(..).enumerate() {
let mut consumer_rx = event_bus_manager.subscribe();
let handle = tokio::spawn(async move {
info!(target: "events", "Event consumer {} started", idx);
loop {
match consumer_rx.recv().await {
Ok(envelope) => {
if matches!(&envelope.event, EventType::System(SystemEvent::Shutdown)) {
info!(target: "events", "Event consumer {} received shutdown signal", idx);
break;
}
consumer.handle_event(envelope);
}
Err(tokio::sync::broadcast::error::RecvError::Lagged(skipped)) => {
error!(target: "events", "Consumer {} lagged, {} messages were skipped", idx, skipped);
}
Err(tokio::sync::broadcast::error::RecvError::Closed) => {
info!(target: "events", "Event channel closed for consumer {}", idx);
break;
}
}
}
info!(target: "events", "Event consumer {} stopped", idx);
});
consumer_tasks.push(handle);
}
run_client_loop(client, shutdown_rx).await;
info!(target: "events", "Sending shutdown signal to consumers");
let shutdown_event =
EventEnvelope::system_event(SystemEvent::Shutdown, config.id, 0, EventSource::System);
shutdown_sender.publish(shutdown_event);
tokio::time::sleep(tokio::time::Duration::from_millis(100)).await;
info!(target: "events", "Waiting for {} consumer tasks to finish", consumer_tasks.len());
for task in consumer_tasks {
let _ = task.await;
}
}
pub async fn run_client_with_action_channel<C, F>(
config: ClientConfig,
event_bus_manager: Arc<EventBusManager>,
event_consumer_factory: F,
action_tx_sender: mpsc::UnboundedSender<
mpsc::UnboundedSender<gromnie_events::SimpleClientAction>,
>,
shutdown_rx: tokio::sync::watch::Receiver<bool>,
) where
C: EventConsumer,
F: FnOnce(mpsc::UnboundedSender<gromnie_events::SimpleClientAction>) -> C,
{
let (raw_event_tx, raw_event_rx) = mpsc::channel::<gromnie_events::ClientEvent>(256);
let event_wrapper = EventWrapper::new(config.id, event_bus_manager.event_bus.clone());
tokio::spawn(async move {
event_wrapper.run(raw_event_rx).await;
});
let event_rx = event_bus_manager.subscribe();
let (client, action_tx) = create_client_from_config(&config, raw_event_tx, None).await;
let client = Arc::new(RwLock::new(client));
let _ = action_tx_sender.send(action_tx.clone());
let event_consumer = event_consumer_factory(action_tx);
run_client_internal(
client,
event_rx,
Box::new(event_consumer),
Some(shutdown_rx),
)
.await;
}
pub(crate) async fn run_client_internal(
client: Arc<RwLock<Client>>,
mut event_rx: broadcast::Receiver<EventEnvelope>,
mut event_consumer: Box<dyn EventConsumer>,
shutdown_rx: Option<tokio::sync::watch::Receiver<bool>>,
) {
let event_task = tokio::spawn(async move {
info!(target: "events", "Event handler task started");
loop {
match event_rx.recv().await {
Ok(envelope) => {
event_consumer.handle_event(envelope);
}
Err(tokio::sync::broadcast::error::RecvError::Lagged(skipped)) => {
error!(target: "events", "Event receiver lagged, {} messages were skipped", skipped);
}
Err(tokio::sync::broadcast::error::RecvError::Closed) => {
info!(target: "events", "Event channel closed");
break;
}
}
}
info!(target: "events", "Event handler task stopped");
});
run_client_loop(client, shutdown_rx).await;
info!(target: "events", "Waiting for event handler task to finish");
let timeout = tokio::time::Duration::from_secs(1);
match tokio::time::timeout(timeout, event_task).await {
Ok(Ok(())) => {
info!(target: "events", "Event handler task finished gracefully");
}
Ok(Err(e)) => {
error!(target: "events", "Event handler task panicked: {}", e);
}
Err(_) => {
info!(target: "events", "Event handler task did not finish within timeout, continuing shutdown");
}
}
}
async fn run_client_loop(
client: Arc<RwLock<Client>>,
mut shutdown_rx: Option<tokio::sync::watch::Receiver<bool>>,
) {
use tracing::error;
let handle = gromnie_client::client::spawn_client_loop(
client,
std::time::Duration::from_millis(gromnie_client::client::UI_DELAY_MS),
)
.await;
let shutdown_tx = handle.shutdown_sender();
match shutdown_rx.take() {
Some(mut rx) => {
tokio::spawn(async move {
if rx.changed().await.is_ok() {
info!("External shutdown signal received");
let _ = shutdown_tx.send(true);
}
});
}
None => {
tokio::spawn(async move {
if tokio::signal::ctrl_c().await.is_ok() {
info!("Received Ctrl+C, shutting down gracefully...");
let _ = shutdown_tx.send(true);
}
});
}
}
match handle.wait_for_exit().await {
Some(reason) => info!(target: "net", "Client network loop exited: {reason:?}"),
None => error!(target: "net", "Client network loop stopped without reporting a reason"),
}
}
pub fn create_event_bus_manager() -> EventBusManager {
EventBusManager::new(100)
}
pub async fn run_multi_client<G>(
config: MultiClientConfig,
consumer_factory: Arc<dyn MultiClientConsumerFactory>,
client_config_generator: G,
shutdown_rx: Option<watch::Receiver<bool>>,
transport_factory: Option<TransportFactory>,
) -> Arc<MultiClientStats>
where
G: Fn(u32) -> ClientConfig + Send + Sync + 'static,
{
let stats = Arc::new(MultiClientStats::default());
let (local_shutdown_tx, local_shutdown_rx) = watch::channel(false);
let shutdown_rx = shutdown_rx.unwrap_or(local_shutdown_rx);
let shared_event_bus = if config.shared_event_bus {
Some(Arc::new(EventBusManager::new(100)))
} else {
None
};
let mut join_handles = vec![];
let _start_time = Instant::now();
info!(
"Starting multi-client run: {} clients to {}",
config.num_clients, config.server_address
);
info!(
"Rate limiting: {} ms between connections",
config.spawn_interval_ms
);
for client_id in 0..config.num_clients {
stats.attempted.fetch_add(1, Ordering::SeqCst);
let client_config = client_config_generator(client_id);
let event_bus_manager = shared_event_bus
.clone()
.unwrap_or_else(|| Arc::new(EventBusManager::new(100)));
let consumer_factory = consumer_factory.clone();
let stats = stats.clone();
let mut shutdown_rx = shutdown_rx.clone();
let spawn_interval = config.spawn_interval_ms;
let transport_factory = transport_factory.clone();
let handle = tokio::spawn(async move {
let sleep_duration = Duration::from_millis(client_id as u64 * spawn_interval);
tokio::select! {
_ = tokio::time::sleep(sleep_duration) => {
}
_ = shutdown_rx.changed() => {
return;
}
}
let (raw_event_tx, raw_event_rx) = mpsc::channel::<gromnie_events::ClientEvent>(256);
let event_wrapper =
EventWrapper::new(client_config.id, event_bus_manager.event_bus.clone());
tokio::spawn(async move {
event_wrapper.run(raw_event_rx).await;
});
let event_rx = event_bus_manager.subscribe();
let (client, action_tx) =
create_client_from_config(&client_config, raw_event_tx, transport_factory.as_ref())
.await;
let client = Arc::new(RwLock::new(client));
stats.spawned.fetch_add(1, Ordering::SeqCst);
let event_consumer =
consumer_factory.create_consumer(client_config.id, &client_config, action_tx);
run_client_internal(client, event_rx, event_consumer, Some(shutdown_rx)).await;
});
join_handles.push(handle);
}
info!("All clients spawned, waiting for events...");
match tokio::signal::ctrl_c().await {
Ok(()) => {
info!("Received Ctrl+C, shutting down all clients...");
info!("Press Ctrl+C again to force quit");
let _ = local_shutdown_tx.send(true);
}
Err(e) => {
error!("Failed to listen for Ctrl+C: {}", e);
}
}
let shutdown_timeout = Duration::from_secs(5);
let wait_future = async {
for (idx, handle) in join_handles.into_iter().enumerate() {
match handle.await {
Ok(_) => {
}
Err(e) => {
stats.task_failures.fetch_add(1, Ordering::SeqCst);
if e.is_panic() {
error!("[Client {}] Task panicked: {:?}", idx, e);
} else {
error!("[Client {}] Task was cancelled", idx);
}
}
}
}
};
tokio::select! {
_ = wait_future => {
info!("All clients shut down gracefully");
}
_ = tokio::time::sleep(shutdown_timeout) => {
info!("Shutdown timeout reached, forcing exit");
}
_ = tokio::signal::ctrl_c() => {
info!("Second Ctrl+C received, forcing immediate exit");
}
}
stats
}
pub enum RunResult {
Single,
Multi(Arc<MultiClientStats>),
}
pub trait ConsumerBuilder: Send + Sync {
fn build(
&self,
client_id: u32,
client_config: &ClientConfig,
action_tx: mpsc::UnboundedSender<gromnie_events::SimpleClientAction>,
) -> Box<dyn EventConsumer>;
}
type ConsumerFactoryFn = dyn Fn(
u32,
&ClientConfig,
mpsc::UnboundedSender<gromnie_events::SimpleClientAction>,
) -> Box<dyn EventConsumer>
+ Send
+ Sync;
pub struct FnConsumerBuilder {
f: Box<ConsumerFactoryFn>,
}
impl FnConsumerBuilder {
pub fn new<F>(f: F) -> Self
where
F: Fn(
u32,
&ClientConfig,
mpsc::UnboundedSender<gromnie_events::SimpleClientAction>,
) -> Box<dyn EventConsumer>
+ Send
+ Sync
+ 'static,
{
Self { f: Box::new(f) }
}
}
impl ConsumerBuilder for FnConsumerBuilder {
fn build(
&self,
client_id: u32,
client_config: &ClientConfig,
action_tx: mpsc::UnboundedSender<gromnie_events::SimpleClientAction>,
) -> Box<dyn EventConsumer> {
(self.f)(client_id, client_config, action_tx)
}
}
pub async fn run<B, G>(
config: RunConfig,
consumer_builder: B,
client_config_generator: Option<G>,
shutdown_rx: Option<watch::Receiver<bool>>,
) -> RunResult
where
B: ConsumerBuilder + 'static,
G: Fn(u32) -> ClientConfig + Send + Sync + 'static,
{
match config {
RunConfig::Single { client } => {
let event_bus_manager = Arc::new(EventBusManager::new(100));
let (raw_event_tx, raw_event_rx) = mpsc::channel::<gromnie_events::ClientEvent>(256);
let event_wrapper = EventWrapper::new(client.id, event_bus_manager.event_bus.clone());
tokio::spawn(async move {
event_wrapper.run(raw_event_rx).await;
});
let event_rx = event_bus_manager.subscribe();
let (client_obj, action_tx) =
create_client_from_config(&client, raw_event_tx, None).await;
let client_obj = Arc::new(RwLock::new(client_obj));
let event_consumer = consumer_builder.build(client.id, &client, action_tx);
run_client_internal(client_obj, event_rx, event_consumer, shutdown_rx).await;
RunResult::Single
}
RunConfig::Multi {
server_address,
num_clients,
spawn_interval_ms,
shared_event_bus,
} => {
let multi_config = MultiClientConfig {
server_address,
num_clients,
spawn_interval_ms,
shared_event_bus,
};
struct Adapter<T> {
inner: T,
}
impl<T: ConsumerBuilder> MultiClientConsumerFactory for Adapter<T> {
fn create_consumer(
&self,
client_id: u32,
client_config: &ClientConfig,
action_tx: mpsc::UnboundedSender<gromnie_events::SimpleClientAction>,
) -> Box<dyn EventConsumer> {
self.inner.build(client_id, client_config, action_tx)
}
}
let factory = Adapter {
inner: consumer_builder,
};
let stats = match client_config_generator {
Some(config_gen) => {
run_multi_client(
multi_config,
Arc::new(factory),
config_gen,
shutdown_rx,
None,
)
.await
}
None => {
let default_gen = |id| {
ClientConfig::new(
id,
"localhost:9000".to_string(),
format!("client_{}", id),
format!("client_{}", id),
)
.with_reconnect(Default::default())
};
run_multi_client(
multi_config,
Arc::new(factory),
default_gen,
shutdown_rx,
None,
)
.await
}
};
RunResult::Multi(stats)
}
}
}
#[cfg(test)]
mod tests {
use super::*;
use gromnie_client::client::Client;
use gromnie_client::transport::{ClientTransport, TransportChannel, TransportFuture};
use std::net::SocketAddr;
use std::sync::Arc;
use tokio::sync::{RwLock, mpsc, watch};
struct PendingTransport;
impl ClientTransport for PendingTransport {
fn send<'a>(
&'a mut self,
_server: &'a gromnie_client::client::ServerInfo,
_channel: TransportChannel,
_bytes: Vec<u8>,
) -> TransportFuture<'a, ()> {
Box::pin(async { Ok(()) })
}
fn recv<'a>(&'a mut self, _buf: &'a mut [u8]) -> TransportFuture<'a, (usize, SocketAddr)> {
Box::pin(std::future::pending())
}
}
#[tokio::test]
async fn runner_does_not_stop_the_client_immediately() {
let (raw_event_tx, _raw_event_rx) = mpsc::channel(8);
let (client, _action_tx) = Client::new_with_transport(
0,
"127.0.0.1:9000".to_string(),
"account".to_string(),
"password".to_string(),
None,
raw_event_tx,
false,
Box::new(PendingTransport),
)
.await;
let client = Arc::new(RwLock::new(client));
let (shutdown_tx, shutdown_rx) = watch::channel(false);
let runner = tokio::spawn(run_client_loop(client, Some(shutdown_rx)));
tokio::time::sleep(tokio::time::Duration::from_millis(
gromnie_client::client::UI_DELAY_MS + 500,
))
.await;
assert!(
!runner.is_finished(),
"runner stopped without a shutdown signal"
);
shutdown_tx.send(true).expect("shutdown receiver");
tokio::time::timeout(tokio::time::Duration::from_secs(2), runner)
.await
.expect("runner did not stop after the shutdown signal")
.expect("runner task panicked");
}
}