bench/
lib.rs

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

bench/
lib.rs
