Add response simple rate

Author: RicoGit

Cargo.lock

@@ -1,15 +1,15 @@
 [package]
 name = "ws_stress"
-version = "0.1.0"
+version = "0.1.1"
 edition = "2021"
 authors = ["konstantin Solovev"]
 
 [dependencies]
 clap = { version = "4", features = ["derive"] }
 tokio = { version = "1.17", features = ["full"] }
 tokio-tungstenite = { version = "0.19.0", features = ["rustls-tls-native-roots"]}
-#tokio-rustls = { version = "0.23.4", default-features = false, features = ["logging"] }
 futures = "0.3"
 human_bytes = { version = "0.4.2", default-features = false }
 hashbrown = "0.14"
+rand = "0.8"
 

Cargo.toml

@@ -1,3 +1,64 @@
 
+# Simple stress test tool for websockets
 
-cargo run -- --addr wss://0.0.0.0:8080 -m 10000 -c 1 '{"foo":"bar"}'
+```
+Usage: ws_stress [OPTIONS] --addr <ADDR> [MESSAGE]
+
+Arguments:
+  [MESSAGE]  Message to send [default: "hello world!"]
+
+Options:
+  -a, --addr <ADDR>
+          The websocket address to connect to, e.g. ws://0.0.0.0:8080
+  -c, --connections <CONNECTIONS>
+          How many concurrent sockets to open [default: 1]
+  -m, --messages <MESSAGES>
+          How many messages send per 1 socket [default: 1]
+  -e, --log-errors
+          Collect and print errors in the end of the tests
+  -r, --resp-sample-rate <RESP_SAMPLE_RATE>
+          Show responses from the server with the specified sample rate [0-1] [default: 0.0]
+  -h, --help
+          Print help
+  -V, --version
+          Print version
+
+```
+                      
+## Examples
+
+### Simple run: 1 connection send 1 default message
+
+```bash
+./ws_stress --addr ws://0.0.0.0:8080
+```
+    Benchmarking: 1 messages, 1 connections
+    Errors: {}
+    Benchmark results:
+    Ok: 1
+    Err: 0
+    Total: 1
+    Elapsed time: 0.00 s
+    Connections: 1
+    Msg len: 12 B
+    Requests per second: 11810.98 rps
+    Throughput: 1.13 Mbit/sec
+
+### Stress test: 1 connection sends 100000 specifies json messages
+    
+```bash
+./ws_stress --addr ws://0.0.0.0:8080 '{"traceId":"2644795231141946144","correlationId":121212121212,"payload":{"type":"Foo","data":[{"time":"2023-01-01T10:28:31.236453182Z","amount":0.01},{"time":"2023-02-02T10:28:31.236453182Z","amount":0.0051}]}}' \
+-m 100000 -c 1 -e
+```
+    Benchmarking: 100000 messages, 1 connections
+    Errors: {}
+    Benchmark results:
+    Ok: 100000
+    Err: 0
+    Total: 100000
+    Elapsed time: 0.55 s
+    Connections: 1
+    Msg len: 511 B
+    Requests per second: 183013.38 rps
+    Throughput: 748.16 Mbit/sec
+       

README.MD

@@ -3,22 +3,35 @@ use clap::Parser;
 #[derive(Parser, Debug)]
 #[command(author, version, about, long_about = None)]
 pub struct Args {
-    /// The websocket address to connect to
+    /// The websocket address to connect to, e.g. ws://0.0.0.0:8080
     #[arg(short, long)]
     pub addr: String,
 
-    /// How many concurrent sockets to use
+    /// How many concurrent sockets to open
     #[arg(short, long, default_value = "1")]
     pub connections: usize,
 
-    /// How many messages send per socket
+    /// How many messages send per 1 socket
     #[arg(short, long, default_value = "1")]
     pub messages: usize,
 
     /// Message to send
     #[arg(default_value = "hello world!")]
     pub message: String,
 
-    #[arg(short, long, default_value = "false")]
+    /// Collect and print errors in the end of the tests
+    #[arg(short = 'e', long, default_value = "false")]
     pub log_errors: bool,
+
+    /// Show responses from the server with the specified sample rate [0-1]
+    #[arg(short = 'r', long, default_value = "0.0", value_parser = parse_sample_rate)]
+    pub resp_sample_rate: f64,
+}
+
+fn parse_sample_rate(src: &str) -> Result<f64, String> {
+    let rate = src.parse::<f64>().map_err(|err| err.to_string())?;
+    if !(0.0..=1.0).contains(&rate) {
+        return Err(format!("Sample rate must be in range [0-1], got {}", rate));
+    }
+    Ok(rate)
 }

src/args.rs

@@ -9,10 +9,13 @@ use std::{
 };
 
 use clap::Parser;
-use futures::SinkExt;
+use futures::stream::SplitStream;
+use futures::{SinkExt, StreamExt};
 use human_bytes::human_bytes;
-use tokio_tungstenite::connect_async;
+use rand::Rng;
+use tokio::net::TcpStream;
 use tokio_tungstenite::tungstenite::Message::Text;
+use tokio_tungstenite::{connect_async, MaybeTlsStream, WebSocketStream};
 
 #[tokio::main]
 async fn main() {
@@ -28,14 +31,18 @@ async fn main() {
         let msg = args.message.clone();
         let barrier = barrier.clone();
         let task = tokio::spawn(async move {
-            let mut client_tungstenite = connect_async(addr).await.unwrap().0;
+            let client_tungstenite = connect_async(addr).await.expect("Failed to connect").0;
+            let (mut tx, rx) = client_tungstenite.split();
+
+            read_responses(rx, args.resp_sample_rate);
+
             barrier.wait().await;
 
             // Async, quite fast, with feedback (send_all is faster but we can't count
             // messages)
             let mut map = hashbrown::HashMap::new();
             for _ in 0..args.messages {
-                match client_tungstenite.send(Text(msg.clone())).await {
+                match tx.feed(Text(msg.clone())).await {
                     Ok(_) => {
                         ok_counter.fetch_add(1, Ordering::Relaxed);
                     }
@@ -52,7 +59,7 @@ async fn main() {
                     }
                 }
             }
-
+            tx.flush().await.expect("flush failed");
             println!("Errors: {:?}", map);
         });
         tasks.push(task);
@@ -92,3 +99,24 @@ async fn main() {
 fn string_to_static_str(s: &str) -> &'static str {
     Box::leak(s.to_string().into_boxed_str())
 }
+
+fn read_responses(
+    mut rx: SplitStream<WebSocketStream<MaybeTlsStream<TcpStream>>>,
+    sample_rate: f64,
+) {
+    if sample_rate == 0.0 {
+        return;
+    }
+
+    // read all messages from rx and write in into stdout
+    tokio::spawn(async move {
+        while let Some(msg) = rx.next().await {
+            if sample_rate == 1.0 || rand::thread_rng().gen_bool(sample_rate) {
+                let received = msg.expect("Failed to receive a message");
+                if received.is_text() {
+                    println!("> {}", received.into_text().unwrap());
+                }
+            };
+        }
+    });
+}