rustwave 0.1.0

use rustwave::Client;
use std::sync::Arc;
use std::sync::atomic::{AtomicU64, Ordering};
use std::time::{Duration, Instant};
use tokio::sync::Semaphore;
use tokio::time::sleep;

const NUM_CLIENTS: usize = 10_000;
const MESSAGES_PER_CLIENT: usize = 100;
const MAX_CONCURRENT_CONNECTS: usize = 1_000;

#[tokio::main]
async fn main() {
    let addr = "127.0.0.1:9500";

    let client_sent = Arc::new(AtomicU64::new(0));
    let client_received = Arc::new(AtomicU64::new(0));
    let client_connect_failed = Arc::new(AtomicU64::new(0));

    let connect_gate = Arc::new(Semaphore::new(MAX_CONCURRENT_CONNECTS));

    println!(
        "Starting client benchmark: {} clients x {} messages = {} total messages expected",
        NUM_CLIENTS,
        MESSAGES_PER_CLIENT,
        NUM_CLIENTS * MESSAGES_PER_CLIENT
    );

    let start = Instant::now();

    let mut handles = Vec::with_capacity(NUM_CLIENTS);

    for _ in 0..NUM_CLIENTS {
        let addr = addr.to_string();
        let client_sent = client_sent.clone();
        let client_received = client_received.clone();
        let client_connect_failed = client_connect_failed.clone();
        let connect_gate = connect_gate.clone();

        let handle = tokio::spawn(async move {
            let permit = connect_gate.acquire_owned().await.unwrap();

            let client = match Client::connect(&format!("ws://{}", addr)).await {
                Ok(c) => c,
                Err(_) => {
                    client_connect_failed.fetch_add(1, Ordering::Relaxed);
                    return None;
                }
            };

            drop(permit);

            {
                let client_received = client_received.clone();
                client.on("pong", move |_data| {
                    let client_received = client_received.clone();
                    async move {
                        client_received.fetch_add(1, Ordering::Relaxed);
                    }
                });
            }

            for i in 0..MESSAGES_PER_CLIENT {
                if client.emit("ping", format!("msg-{}", i)).await.is_ok() {
                    client_sent.fetch_add(1, Ordering::Relaxed);
                }
            }

            Some(client)
        });

        handles.push(handle);
    }

    let mut clients_kept_alive = Vec::with_capacity(NUM_CLIENTS);
    for handle in handles {
        if let Ok(Some(client)) = handle.await {
            clients_kept_alive.push(client);
        }
    }

    let sending_done_at = start.elapsed();

    let expected_total = (NUM_CLIENTS * MESSAGES_PER_CLIENT) as u64;

    let wait_timeout = Duration::from_secs(30);
    let wait_start = Instant::now();

    loop {
        let received_now = client_received.load(Ordering::Relaxed);
        if received_now >= expected_total || wait_start.elapsed() > wait_timeout {
            break;
        }
        sleep(Duration::from_millis(50)).await;
    }

    let total_elapsed = start.elapsed();

    println!("========== CLIENT-SIDE STATS ==========");
    println!("Clients:                    {}", NUM_CLIENTS);
    println!(
        "Connect failures:           {}",
        client_connect_failed.load(Ordering::Relaxed)
    );
    println!("Messages per client:        {}", MESSAGES_PER_CLIENT);
    println!("Expected total messages:    {}", expected_total);
    println!("------------------------------------------");
    println!(
        "Sent messages:               {}",
        client_sent.load(Ordering::Relaxed)
    );
    println!(
        "Replies received:            {}",
        client_received.load(Ordering::Relaxed)
    );
    println!("------------------------------------------");
    println!("Time to finish sending:     {:?}", sending_done_at);
    println!("Total time (incl. replies): {:?}", total_elapsed);
    println!(
        "Client-side throughput:     {:.0} msgs/sec",
        expected_total as f64 / total_elapsed.as_secs_f64()
    );
    println!("==========================================");
}