bench/
lib.rs

1use core::str;
2use std::{
3    convert::TryInto,
4    net::{IpAddr, Ipv6Addr, SocketAddr},
5    num::ParseIntError,
6    str::FromStr,
7    sync::Arc,
8};
9
10use anyhow::{Context, Result};
11use bytes::Bytes;
12use clap::Parser;
13use quinn::crypto::rustls::QuicClientConfig;
14use rustls::{
15    RootCertStore,
16    pki_types::{CertificateDer, PrivateKeyDer},
17};
18use tokio::runtime::{Builder, Runtime};
19use tracing::trace;
20
21pub mod stats;
22
23pub fn configure_tracing_subscriber() {
24    tracing::subscriber::set_global_default(
25        tracing_subscriber::FmtSubscriber::builder()
26            .with_env_filter(tracing_subscriber::EnvFilter::from_default_env())
27            .finish(),
28    )
29    .unwrap();
30}
31
32/// Creates a server endpoint which runs on the given runtime
33pub fn server_endpoint(
34    rt: &tokio::runtime::Runtime,
35    cert: CertificateDer<'static>,
36    key: PrivateKeyDer<'static>,
37    opt: &Opt,
38) -> (SocketAddr, quinn::Endpoint) {
39    let cert_chain = vec![cert];
40    let mut server_config = quinn::ServerConfig::with_single_cert(cert_chain, key).unwrap();
41    server_config.transport = Arc::new(transport_config(opt));
42
43    let endpoint = {
44        let _guard = rt.enter();
45        quinn::Endpoint::server(
46            server_config,
47            SocketAddr::new(IpAddr::V6(Ipv6Addr::LOCALHOST), 0),
48        )
49        .unwrap()
50    };
51    let server_addr = endpoint.local_addr().unwrap();
52    (server_addr, endpoint)
53}
54
55/// Create a client endpoint and client connection
56pub async fn connect_client(
57    server_addr: SocketAddr,
58    server_cert: CertificateDer<'_>,
59    opt: Opt,
60) -> Result<(quinn::Endpoint, quinn::Connection)> {
61    let endpoint =
62        quinn::Endpoint::client(SocketAddr::new(IpAddr::V6(Ipv6Addr::LOCALHOST), 0)).unwrap();
63
64    let mut roots = RootCertStore::empty();
65    roots.add(server_cert)?;
66
67    let default_provider = rustls::crypto::ring::default_provider();
68    let provider = rustls::crypto::CryptoProvider {
69        cipher_suites: vec![opt.cipher.as_rustls()],
70        ..default_provider
71    };
72
73    let crypto = rustls::ClientConfig::builder_with_provider(provider.into())
74        .with_protocol_versions(&[&rustls::version::TLS13])
75        .unwrap()
76        .with_root_certificates(roots)
77        .with_no_client_auth();
78
79    let mut client_config = quinn::ClientConfig::new(Arc::new(QuicClientConfig::try_from(crypto)?));
80    client_config.transport_config(Arc::new(transport_config(&opt)));
81
82    let connection = endpoint
83        .connect_with(client_config, server_addr, "localhost")
84        .unwrap()
85        .await
86        .context("unable to connect")?;
87    trace!("connected");
88
89    Ok((endpoint, connection))
90}
91
92pub async fn drain_stream(mut stream: quinn::RecvStream, read_unordered: bool) -> Result<usize> {
93    let mut read = 0;
94
95    if read_unordered {
96        let mut stream = stream.into_unordered();
97        while let Some(chunk) = stream.read_chunk(usize::MAX).await? {
98            read += chunk.bytes.len();
99        }
100    } else {
101        // These are 32 buffers, for reading approximately 32kB at once
102        #[rustfmt::skip]
103        let mut bufs = [
104            Bytes::new(), Bytes::new(), Bytes::new(), Bytes::new(),
105            Bytes::new(), Bytes::new(), Bytes::new(), Bytes::new(),
106            Bytes::new(), Bytes::new(), Bytes::new(), Bytes::new(),
107            Bytes::new(), Bytes::new(), Bytes::new(), Bytes::new(),
108            Bytes::new(), Bytes::new(), Bytes::new(), Bytes::new(),
109            Bytes::new(), Bytes::new(), Bytes::new(), Bytes::new(),
110            Bytes::new(), Bytes::new(), Bytes::new(), Bytes::new(),
111            Bytes::new(), Bytes::new(), Bytes::new(), Bytes::new(),
112        ];
113
114        while let Some(n) = stream.read_chunks(&mut bufs[..]).await? {
115            read += bufs.iter().take(n).map(|buf| buf.len()).sum::<usize>();
116        }
117    }
118
119    Ok(read)
120}
121
122pub async fn send_data_on_stream(stream: &mut quinn::SendStream, stream_size: u64) -> Result<()> {
123    const DATA: &[u8] = &[0xAB; 1024 * 1024];
124    let bytes_data = Bytes::from_static(DATA);
125
126    let full_chunks = stream_size / (DATA.len() as u64);
127    let remaining = (stream_size % (DATA.len() as u64)) as usize;
128
129    for _ in 0..full_chunks {
130        stream
131            .write_chunk(bytes_data.clone())
132            .await
133            .context("failed sending data")?;
134    }
135
136    if remaining != 0 {
137        stream
138            .write_chunk(bytes_data.slice(0..remaining))
139            .await
140            .context("failed sending data")?;
141    }
142
143    stream.finish().unwrap();
144    // Wait for stream to close
145    _ = stream.stopped().await;
146
147    Ok(())
148}
149
150pub fn rt() -> Runtime {
151    Builder::new_current_thread().enable_all().build().unwrap()
152}
153
154pub fn transport_config(opt: &Opt) -> quinn::TransportConfig {
155    // High stream windows are chosen because the amount of concurrent streams
156    // is configurable as a parameter.
157    let mut config = quinn::TransportConfig::default();
158    config.max_concurrent_uni_streams(opt.max_streams.try_into().unwrap());
159    config.initial_mtu(opt.initial_mtu);
160
161    let mut acks = quinn::AckFrequencyConfig::default();
162    acks.ack_eliciting_threshold(10u32.into());
163    config.ack_frequency_config(Some(acks));
164
165    config
166}
167
168#[derive(Parser, Debug, Clone, Copy)]
169#[clap(name = "bulk")]
170pub struct Opt {
171    /// The total number of clients which should be created
172    #[clap(long = "clients", short = 'c', default_value = "1")]
173    pub clients: usize,
174    /// The total number of streams which should be created
175    #[clap(long = "streams", short = 'n', default_value = "1")]
176    pub streams: usize,
177    /// The amount of concurrent streams which should be used
178    #[clap(long = "max_streams", short = 'm', default_value = "1")]
179    pub max_streams: usize,
180    /// Number of bytes to transmit from server to client
181    ///
182    /// This can use SI suffixes for sizes. For example, 1M will transfer
183    /// 1MiB, 10G will transfer 10GiB.
184    #[clap(long, default_value = "1G", value_parser = parse_byte_size)]
185    pub download_size: u64,
186    /// Number of bytes to transmit from client to server
187    ///
188    /// This can use SI suffixes for sizes. For example, 1M will transfer
189    /// 1MiB, 10G will transfer 10GiB.
190    #[clap(long, default_value = "0", value_parser = parse_byte_size)]
191    pub upload_size: u64,
192    /// Show connection stats the at the end of the benchmark
193    #[clap(long = "stats")]
194    pub stats: bool,
195    /// Whether to use the unordered read API
196    #[clap(long = "unordered")]
197    pub read_unordered: bool,
198    /// Allows to configure the desired cipher suite
199    ///
200    /// Valid options are: aes128, aes256, chacha20
201    #[clap(long = "cipher", default_value = "aes128")]
202    pub cipher: CipherSuite,
203    /// Starting guess for maximum UDP payload size
204    #[clap(long, default_value = "1200")]
205    pub initial_mtu: u16,
206}
207
208fn parse_byte_size(s: &str) -> Result<u64, ParseIntError> {
209    let s = s.trim();
210
211    let multiplier = match s.chars().last() {
212        Some('T') => 1024 * 1024 * 1024 * 1024,
213        Some('G') => 1024 * 1024 * 1024,
214        Some('M') => 1024 * 1024,
215        Some('k') => 1024,
216        _ => 1,
217    };
218
219    let s = match multiplier {
220        1 => s,
221        _ => &s[..s.len() - 1],
222    };
223
224    Ok(u64::from_str(s)? * multiplier)
225}
226
227#[derive(Debug, PartialEq, Eq, Clone, Copy)]
228pub enum CipherSuite {
229    Aes128,
230    Aes256,
231    Chacha20,
232}
233
234impl CipherSuite {
235    pub fn as_rustls(self) -> rustls::SupportedCipherSuite {
236        use rustls::crypto::ring::cipher_suite;
237        match self {
238            Self::Aes128 => cipher_suite::TLS13_AES_128_GCM_SHA256,
239            Self::Aes256 => cipher_suite::TLS13_AES_256_GCM_SHA384,
240            Self::Chacha20 => cipher_suite::TLS13_CHACHA20_POLY1305_SHA256,
241        }
242    }
243}
244
245impl FromStr for CipherSuite {
246    type Err = anyhow::Error;
247
248    fn from_str(s: &str) -> Result<Self, Self::Err> {
249        match s.to_lowercase().as_str() {
250            "aes128" => Ok(Self::Aes128),
251            "aes256" => Ok(Self::Aes256),
252            "chacha20" => Ok(Self::Chacha20),
253            _ => Err(anyhow::anyhow!("Unknown cipher suite {}", s)),
254        }
255    }
256}
bench/lib.rs

bench/
lib.rs
