Skip to main content

client_banch/
client_banch.rs

1use rustwave::Client;
2use std::sync::Arc;
3use std::sync::atomic::{AtomicU64, Ordering};
4use std::time::{Duration, Instant};
5use tokio::sync::Semaphore;
6use tokio::time::sleep;
7
8const NUM_CLIENTS: usize = 10_000;
9const MESSAGES_PER_CLIENT: usize = 100;
10const MAX_CONCURRENT_CONNECTS: usize = 1_000;
11
12#[tokio::main]
13async fn main() {
14    let addr = "127.0.0.1:9500";
15
16    let client_sent = Arc::new(AtomicU64::new(0));
17    let client_received = Arc::new(AtomicU64::new(0));
18    let client_connect_failed = Arc::new(AtomicU64::new(0));
19
20    let connect_gate = Arc::new(Semaphore::new(MAX_CONCURRENT_CONNECTS));
21
22    println!(
23        "Starting client benchmark: {} clients x {} messages = {} total messages expected",
24        NUM_CLIENTS,
25        MESSAGES_PER_CLIENT,
26        NUM_CLIENTS * MESSAGES_PER_CLIENT
27    );
28
29    let start = Instant::now();
30
31    let mut handles = Vec::with_capacity(NUM_CLIENTS);
32
33    for _ in 0..NUM_CLIENTS {
34        let addr = addr.to_string();
35        let client_sent = client_sent.clone();
36        let client_received = client_received.clone();
37        let client_connect_failed = client_connect_failed.clone();
38        let connect_gate = connect_gate.clone();
39
40        let handle = tokio::spawn(async move {
41            let permit = connect_gate.acquire_owned().await.unwrap();
42
43            let client = match Client::connect(&format!("ws://{}", addr)).await {
44                Ok(c) => c,
45                Err(_) => {
46                    client_connect_failed.fetch_add(1, Ordering::Relaxed);
47                    return None;
48                }
49            };
50
51            drop(permit);
52
53            {
54                let client_received = client_received.clone();
55                client.on("pong", move |_data| {
56                    let client_received = client_received.clone();
57                    async move {
58                        client_received.fetch_add(1, Ordering::Relaxed);
59                    }
60                });
61            }
62
63            for i in 0..MESSAGES_PER_CLIENT {
64                if client.emit("ping", format!("msg-{}", i)).await.is_ok() {
65                    client_sent.fetch_add(1, Ordering::Relaxed);
66                }
67            }
68
69            Some(client)
70        });
71
72        handles.push(handle);
73    }
74
75    let mut clients_kept_alive = Vec::with_capacity(NUM_CLIENTS);
76    for handle in handles {
77        if let Ok(Some(client)) = handle.await {
78            clients_kept_alive.push(client);
79        }
80    }
81
82    let sending_done_at = start.elapsed();
83
84    let expected_total = (NUM_CLIENTS * MESSAGES_PER_CLIENT) as u64;
85
86    let wait_timeout = Duration::from_secs(30);
87    let wait_start = Instant::now();
88
89    loop {
90        let received_now = client_received.load(Ordering::Relaxed);
91        if received_now >= expected_total || wait_start.elapsed() > wait_timeout {
92            break;
93        }
94        sleep(Duration::from_millis(50)).await;
95    }
96
97    let total_elapsed = start.elapsed();
98
99    println!("========== CLIENT-SIDE STATS ==========");
100    println!("Clients:                    {}", NUM_CLIENTS);
101    println!(
102        "Connect failures:           {}",
103        client_connect_failed.load(Ordering::Relaxed)
104    );
105    println!("Messages per client:        {}", MESSAGES_PER_CLIENT);
106    println!("Expected total messages:    {}", expected_total);
107    println!("------------------------------------------");
108    println!(
109        "Sent messages:               {}",
110        client_sent.load(Ordering::Relaxed)
111    );
112    println!(
113        "Replies received:            {}",
114        client_received.load(Ordering::Relaxed)
115    );
116    println!("------------------------------------------");
117    println!("Time to finish sending:     {:?}", sending_done_at);
118    println!("Total time (incl. replies): {:?}", total_elapsed);
119    println!(
120        "Client-side throughput:     {:.0} msgs/sec",
121        expected_total as f64 / total_elapsed.as_secs_f64()
122    );
123    println!("==========================================");
124}