3
\$\begingroup\$

I wrote the following HTTP server stress tester in Rust. My actual goal was to write an efficient template for writing load testers, so I can swap out the inner loop to do e.g. requests to a Redis cluster. Looking for advice on two things:

  1. Idiomatic usage of Rust?
  2. Any optimization advice for Tokio i.e. if I'm joining on the tasks in a dumb way?

Here's the code:

use clap::Parser;
use futures::future;
use hyper::{Body, Client, Method, Request, StatusCode};
use std::sync::atomic::{AtomicU64, Ordering};
use std::time::Duration;

static REQUEST_COUNT: AtomicU64 = AtomicU64::new(0);
static ERROR_COUNT: AtomicU64 = AtomicU64::new(0);

#[derive(Parser, Debug, Clone)]
struct Args {
    #[arg(short, long, default_value = "http://127.0.0.1:8080")]
    address: String,

    #[arg(short, long, default_value_t = 1000)]
    connections: u32,
}

#[tokio::main]
async fn main() {
    let args = Args::parse();
    let client = Client::new();

    future::join_all((0..(args.connections + 1)).map(|i| {
        let client = client.clone();
        let address = args.address.clone();
        if i == 0 {
            tokio::spawn(async move {
                loop {
                    tokio::time::sleep(Duration::from_secs(1)).await;
                    println!(
                        "Made {} requests and got {} errors",
                        REQUEST_COUNT.load(Ordering::SeqCst),
                        ERROR_COUNT.load(Ordering::SeqCst)
                    );
                }
            })
        } else {
            tokio::spawn(async move {
                loop {
                    let request = Request::builder()
                        .method(Method::GET)
                        .uri(address.clone())
                        .body(Body::from(""))
                        .unwrap();
                    let res = client.request(request).await;
                    match res {
                        Ok(res) => {
                            if res.status() == StatusCode::OK {
                                REQUEST_COUNT.fetch_add(1, Ordering::SeqCst);
                                continue;
                            }
                        }
                        Err(e) => println!("Got error: {:?}", e),
                    }
                    ERROR_COUNT.fetch_add(1, Ordering::SeqCst);
                }
            })
        }
    }))
    .await;
}
\$\endgroup\$
3
  • \$\begingroup\$ "REQUEST_COUNT: AtomicU64" N GETs will cost N memory barriers. Maybe that's OK. If a spawned task did K GETs, then we could see barrier operations go down by a factor of K. It would be helpful to publish example timings of no-op, to get a feel for the overhead imposed by this measurement routine. I didn't notice any TCP connection caching. This becomes even more important for https:// URLs, as SSL handshake requires some roundtrips. \$\endgroup\$
    – J_H
    Mar 8 at 20:13
  • \$\begingroup\$ The fetch_adds are not visible in a flame graph as they are negligible compared to the I/O, so I didn't even bother to put acquire/release semantics there or use per connection accumulators that I would sum up once a second in the task that prints current status. \$\endgroup\$
    – eof
    Mar 9 at 6:53
  • \$\begingroup\$ I didn't have time yet to look at the number of TCP connections being created and if it matches the request count. This is somewhat specific to the particular HTTP client and not the template itself. Redis clients do pipelining automatically under the hood over the same socket. \$\endgroup\$
    – eof
    Mar 9 at 6:55

1 Answer 1

2
\$\begingroup\$

Firstly, I'd suggest using some functions. Right now your logic is all in main. But the logic of main could be simpler and easier to follow if the output logic and server testing loop were split out into their own functions. I also find that async fn works just a bit more intuitively then async blocks.

The way that you spin up n + 1 iterations of the loop and then do something completely different for the first iteration is awkward. Instead, you could use join_all to create a future for the sequence, and join! to combine that with a future for the display loop.

But actually, you don't need to wait for tokio::spawn, as it runs the future in the background whether or not you wait for it. So just spawn your display loop and don't worry about waiting for it to complete.

You actually probably don't need tokio::spawn for the connections either. In principle, using tokio::spawn allows the different loops to run on different physical cores. However, from some quicks tests and looking at the code, it seems that hyper::Client spawn background tasks anyways, and so there's little benefit to spawning yourself.

Rather than creating distinct loops, it probably make sense to use a stream and process the elements concurrently. Something like:

futures::stream::iter(std::iter::repeat(()))
    .for_each_concurrent(args.connections, |_| hit_server(&client, &args.address))
    .await;

In this case hit_server makes a single request to the server via client. for_each_concurrent will take care of ensuring that there always args.connections connections in process.

Here's my reworking of your example:

use clap::Parser;
use futures::StreamExt;
use futures::future;
use hyper::{Body, Client, Method, Request, StatusCode};
use std::sync::atomic::{AtomicU64, Ordering};
use std::time::Duration;
use futures::future::Either;

static REQUEST_COUNT: AtomicU64 = AtomicU64::new(0);
static ERROR_COUNT: AtomicU64 = AtomicU64::new(0);

#[derive(Parser, Debug, Clone)]
struct Args {
    #[arg(short, long, default_value = "http://127.0.0.1:8080")]
    address: String,

    #[arg(short, long, default_value_t = 1000)]
    connections: usize,
}

async fn show_stats() {
    loop {
        tokio::time::sleep(Duration::from_secs(1)).await;
        println!(
            "Made {} requests and got {} errors",
            REQUEST_COUNT.load(Ordering::SeqCst),
            ERROR_COUNT.load(Ordering::SeqCst)
        );
    }
}

async fn hit_server(client: &Client<hyper::client::HttpConnector>, address: &String) {
    let request = Request::builder()
        .method(Method::GET)
        .uri(address.clone())
        .body(Body::from(""))
        .unwrap();
    let res = client.request(request).await;
    match res {
        Ok(res) => {
            if res.status() == StatusCode::OK {
                REQUEST_COUNT.fetch_add(1, Ordering::SeqCst);
                return;
            }
        }
        Err(e) => println!("Got error: {:?}", e),
    }
    ERROR_COUNT.fetch_add(1, Ordering::SeqCst);
}

#[tokio::main]
async fn main() {
    let args = Args::parse();
    let client = Client::new();

    tokio::spawn(show_stats());

    futures::stream::iter(std::iter::repeat(()))
        .for_each_concurrent(args.connections, |_| hit_server(&client, &args.address))
        .await;
}
\$\endgroup\$

Your Answer

By clicking “Post Your Answer”, you agree to our terms of service and acknowledge that you have read and understand our privacy policy and code of conduct.

Not the answer you're looking for? Browse other questions tagged or ask your own question.