rustwave 0.1.0

use rustwave::{Client, Server};
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 server = Server::bind(addr);

    let server_received = Arc::new(AtomicU64::new(0));
    let server_replied = Arc::new(AtomicU64::new(0));

    {
        let server_received = server_received.clone();
        let server_replied = server_replied.clone();

        server.on("ping", move |data, ctx| {
            let server_received = server_received.clone();
            let server_replied = server_replied.clone();
            async move {
                server_received.fetch_add(1, Ordering::Relaxed);
                let text = data.parse::<String>().unwrap_or_default();
                if ctx.emit("pong", text).await.is_ok() {
                    server_replied.fetch_add(1, Ordering::Relaxed);
                }
            }
        });
    }

    tokio::spawn(async move {
        let _ = server.run().await;
    });

    sleep(Duration::from_millis(300)).await;

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

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

    println!(
        "Starting 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 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(_) => 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!();
    println!("========== BENCHMARK RESULTS ==========");
    println!("Clients spawned:            {}", NUM_CLIENTS);
    println!("Messages per client:        {}", MESSAGES_PER_CLIENT);
    println!("Expected total messages:    {}", expected_total);
    println!("----------------------------------------");
    println!(
        "Client -> sent messages:    {}",
        client_sent.load(Ordering::Relaxed)
    );
    println!(
        "Server -> received:         {}",
        server_received.load(Ordering::Relaxed)
    );
    println!(
        "Server -> replies sent:     {}",
        server_replied.load(Ordering::Relaxed)
    );
    println!(
        "Client -> replies received: {}",
        client_received.load(Ordering::Relaxed)
    );
    println!("----------------------------------------");
    println!("Time to finish sending:     {:?}", sending_done_at);
    println!("Total time (incl. replies): {:?}", total_elapsed);
    println!(
        "Throughput:                 {:.0} msgs/sec",
        expected_total as f64 / total_elapsed.as_secs_f64()
    );
    println!("=========================================");
}