Repository: tokio-rs/axum
File: examples/websockets/src/client.rs (Rust)

Code before:
//! Note that this is not currently optimized for performance, especially around
//! stdout mutex management. Rather it's intended to show an example of working with a
//! websocket server.
//!

//boilerplate
use clap::Parser;
use std::borrow::Cow;
use std::time::Instant;

//we need tungstenite for websocket impl (same library as what axum is using)
use tokio::task::JoinHandle;
use tokio_tungstenite::tungstenite::protocol::frame::coding::CloseCode;
use tokio_tungstenite::tungstenite::protocol::CloseFrame;
use tokio_tungstenite::{connect_async, tungstenite::protocol::Message};
use url::Url;
//stream splitting
use futures_util::{SinkExt, StreamExt};

#[derive(Parser)]
struct Args {
    #[arg(short, long, default_value_t = String::from("ws://127.0.0.1:3000/ws"))]
    url: String,
    #[arg(short, long, default_value_t = 2)]
    number_of_clients: usize,
}

#[tokio::main]
async fn main() {
    let args = Args::parse();
    let url = Url::parse(&args.url).expect("Invalid URL supplied!");
    let start_time = Instant::now();
    //spawn a whole bunch of clients
    let clients: Vec<JoinHandle<_>> = (0..args.number_of_clients)
        .into_iter()
        .map(|cli| {
            let uurl = url.clone();
            tokio::spawn(async move { spawn_client(uurl, cli).await })
        })
        .collect();

    //wait for our clients to exit
    futures::future::join_all(clients).await;
    let end_time = Instant::now();
    println!("Total time taken {:#?}.", end_time - start_time);
}

//creates a client connected to a given url. quietly exits on failure.
async fn spawn_client(url: Url, who: usize) {
    let ws_stream = match connect_async(url).await {
        Ok((stream, response)) => {

Review comment (suggestion):
I don't think we need all this argument parsing stuff, its not important to what the example is trying to show.

Could just hardcode the url and have `let number_of_clients = 10;` variable.

Code after:
//! Note that this is not currently optimized for performance, especially around
//! stdout mutex management. Rather it's intended to show an example of working with axum's
//! websocket server and how the client-side and server-side code can be quite similar.
//!

use futures_util::stream::FuturesUnordered;
use futures_util::{SinkExt, StreamExt};
use std::borrow::Cow;
use std::ops::ControlFlow;
use std::time::Instant;

// we will use tungstenite for websocket client impl (same library as what axum is using)
use tokio_tungstenite::{
    connect_async,
    tungstenite::protocol::{frame::coding::CloseCode, CloseFrame, Message},
};

const N_CLIENTS: usize = 2; //set to desired number
const SERVER: &'static str = "ws://127.0.0.1:3000/ws";

#[tokio::main]
async fn main() {
    let start_time = Instant::now();
    //spawn several clients that will concurrently talk to the server
    let mut clients = (0..N_CLIENTS)
        .into_iter()
        .map(|cli| tokio::spawn(spawn_client(cli)))
        .collect::<FuturesUnordered<_>>();

    //wait for all our clients to exit
    while clients.next().await.is_some() {}

    let end_time = Instant::now();

    //total time should be the same no matter how many clients we spawn
    println!(
        "Total time taken {:#?} with {N_CLIENTS} concurrent clients, should be about 6.45 seconds.",
        end_time - start_time
    );
}

//creates a client. quietly exits on failure.
async fn spawn_client(who: usize) {
    let ws_stream = match connect_async(SERVER).await {
        Ok((stream, response)) => {
            println!("Handshake for client {} has been completed", who);
            // This will be the HTTP response, same as with server this is the last moment we
            // can still access HTTP stuff.
            println!("Server response was {:?}", response);
            stream
        }