bench/
bench.rs

1use rustwave::{Client, Server};
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    let server = Server::bind(addr);
16
17    let server_received = Arc::new(AtomicU64::new(0));
18    let server_replied = Arc::new(AtomicU64::new(0));
19
20    {
21        let server_received = server_received.clone();
22        let server_replied = server_replied.clone();
23
24        server.on("ping", move |data, ctx| {
25            let server_received = server_received.clone();
26            let server_replied = server_replied.clone();
27            async move {
28                server_received.fetch_add(1, Ordering::Relaxed);
29                let text = data.parse::<String>().unwrap_or_default();
30                if ctx.emit("pong", text).await.is_ok() {
31                    server_replied.fetch_add(1, Ordering::Relaxed);
32                }
33            }
34        });
35    }
36
37    tokio::spawn(async move {
38        let _ = server.run().await;
39    });
40
41    sleep(Duration::from_millis(300)).await;
42
43    let client_sent = Arc::new(AtomicU64::new(0));
44    let client_received = Arc::new(AtomicU64::new(0));
45
46    let connect_gate = Arc::new(Semaphore::new(MAX_CONCURRENT_CONNECTS));
47
48    println!(
49        "Starting benchmark: {} clients x {} messages = {} total messages expected",
50        NUM_CLIENTS,
51        MESSAGES_PER_CLIENT,
52        NUM_CLIENTS * MESSAGES_PER_CLIENT
53    );
54
55    let start = Instant::now();
56
57    let mut handles = Vec::with_capacity(NUM_CLIENTS);
58
59    for _ in 0..NUM_CLIENTS {
60        let addr = addr.to_string();
61        let client_sent = client_sent.clone();
62        let client_received = client_received.clone();
63        let connect_gate = connect_gate.clone();
64
65        let handle = tokio::spawn(async move {
66            let permit = connect_gate.acquire_owned().await.unwrap();
67
68            let client = match Client::connect(&format!("ws://{}", addr)).await {
69                Ok(c) => c,
70                Err(_) => return None,
71            };
72
73            drop(permit);
74
75            {
76                let client_received = client_received.clone();
77                client.on("pong", move |_data| {
78                    let client_received = client_received.clone();
79                    async move {
80                        client_received.fetch_add(1, Ordering::Relaxed);
81                    }
82                });
83            }
84
85            for i in 0..MESSAGES_PER_CLIENT {
86                if client.emit("ping", format!("msg-{}", i)).await.is_ok() {
87                    client_sent.fetch_add(1, Ordering::Relaxed);
88                }
89            }
90
91            Some(client)
92        });
93
94        handles.push(handle);
95    }
96
97    let mut clients_kept_alive = Vec::with_capacity(NUM_CLIENTS);
98    for handle in handles {
99        if let Ok(Some(client)) = handle.await {
100            clients_kept_alive.push(client);
101        }
102    }
103
104    let sending_done_at = start.elapsed();
105
106    let expected_total = (NUM_CLIENTS * MESSAGES_PER_CLIENT) as u64;
107
108    let wait_timeout = Duration::from_secs(30);
109    let wait_start = Instant::now();
110
111    loop {
112        let received_now = client_received.load(Ordering::Relaxed);
113        if received_now >= expected_total || wait_start.elapsed() > wait_timeout {
114            break;
115        }
116        sleep(Duration::from_millis(50)).await;
117    }
118
119    let total_elapsed = start.elapsed();
120
121    println!();
122    println!("========== BENCHMARK RESULTS ==========");
123    println!("Clients spawned:            {}", NUM_CLIENTS);
124    println!("Messages per client:        {}", MESSAGES_PER_CLIENT);
125    println!("Expected total messages:    {}", expected_total);
126    println!("----------------------------------------");
127    println!(
128        "Client -> sent messages:    {}",
129        client_sent.load(Ordering::Relaxed)
130    );
131    println!(
132        "Server -> received:         {}",
133        server_received.load(Ordering::Relaxed)
134    );
135    println!(
136        "Server -> replies sent:     {}",
137        server_replied.load(Ordering::Relaxed)
138    );
139    println!(
140        "Client -> replies received: {}",
141        client_received.load(Ordering::Relaxed)
142    );
143    println!("----------------------------------------");
144    println!("Time to finish sending:     {:?}", sending_done_at);
145    println!("Total time (incl. replies): {:?}", total_elapsed);
146    println!(
147        "Throughput:                 {:.0} msgs/sec",
148        expected_total as f64 / total_elapsed.as_secs_f64()
149    );
150    println!("=========================================");
151}
bench/bench.rs

bench/
bench.rs
