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 noq::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: &Runtime,
35    cert: CertificateDer<'static>,
36    key: PrivateKeyDer<'static>,
37    opt: &Opt,
38) -> (SocketAddr, noq::Endpoint) {
39    let cert_chain = vec![cert];
40    let mut server_config = noq::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        noq::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<(noq::Endpoint, noq::Connection)> {
61    let endpoint =
62        noq::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 = noq::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: noq::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_many_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 noq::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_type: RuntimeType) -> Runtime {
151    match runtime_type {
152        RuntimeType::Tokio => {
153            let counter = std::sync::atomic::AtomicUsize::new(0);
154            Builder::new_multi_thread()
155                .thread_name_fn(move || {
156                    format!(
157                        "tokio-runtime-{}",
158                        counter.fetch_add(1, std::sync::atomic::Ordering::Relaxed)
159                    )
160                })
161                .enable_all()
162                .build()
163                .unwrap()
164        }
165        RuntimeType::TokioCurrentThread => {
166            Builder::new_current_thread().enable_all().build().unwrap()
167        }
168    }
169}
170
171pub fn transport_config(opt: &Opt) -> noq::TransportConfig {
172    // High stream windows are chosen because the amount of concurrent streams
173    // is configurable as a parameter.
174    let mut config = noq::TransportConfig::default();
175    config.max_concurrent_uni_streams(opt.max_streams.try_into().unwrap());
176    config.initial_mtu(opt.initial_mtu);
177
178    let mut acks = noq::AckFrequencyConfig::default();
179    acks.ack_eliciting_threshold(10u32.into());
180    config.ack_frequency_config(Some(acks));
181
182    config
183}
184
185#[derive(Parser, Debug, Clone, Copy)]
186#[clap(name = "bulk")]
187pub struct Opt {
188    /// The total number of clients which should be created
189    #[clap(long = "clients", short = 'c', default_value = "1")]
190    pub clients: usize,
191    /// The total number of streams which should be created
192    #[clap(long = "streams", short = 'n', default_value = "1")]
193    pub streams: usize,
194    /// The amount of concurrent streams which should be used
195    #[clap(long = "max_streams", short = 'm', default_value = "1")]
196    pub max_streams: usize,
197    /// Number of bytes to transmit from server to client
198    ///
199    /// This can use SI suffixes for sizes. For example, 1M will transfer
200    /// 1MiB, 10G will transfer 10GiB.
201    #[clap(long, default_value = "1G", value_parser = parse_byte_size)]
202    pub download_size: u64,
203    /// Number of bytes to transmit from client to server
204    ///
205    /// This can use SI suffixes for sizes. For example, 1M will transfer
206    /// 1MiB, 10G will transfer 10GiB.
207    #[clap(long, default_value = "0", value_parser = parse_byte_size)]
208    pub upload_size: u64,
209    /// Show connection stats the at the end of the benchmark
210    #[clap(long = "stats")]
211    pub stats: bool,
212    /// Whether to use the unordered read API
213    #[clap(long = "unordered")]
214    pub read_unordered: bool,
215    /// Allows to configure the desired cipher suite
216    ///
217    /// Valid options are: aes128, aes256, chacha20
218    #[clap(long = "cipher", default_value = "aes128")]
219    pub cipher: CipherSuite,
220    /// Starting guess for maximum UDP payload size
221    #[clap(long, default_value = "1200")]
222    pub initial_mtu: u16,
223    /// The runtime type to use
224    #[clap(long, default_value = "tokio")]
225    pub runtime_type: RuntimeType,
226}
227
228#[derive(Debug, Clone, Copy)]
229pub enum RuntimeType {
230    Tokio,
231    TokioCurrentThread,
232}
233
234impl FromStr for RuntimeType {
235    type Err = anyhow::Error;
236
237    fn from_str(s: &str) -> Result<Self, Self::Err> {
238        match s.to_lowercase().as_str() {
239            "tokio" => Ok(Self::Tokio),
240            "tokio-current-thread" => Ok(Self::TokioCurrentThread),
241            _ => Err(anyhow::anyhow!("Unknown runtime type {}", s)),
242        }
243    }
244}
245
246fn parse_byte_size(s: &str) -> Result<u64, ParseIntError> {
247    let s = s.trim();
248
249    let multiplier = match s.chars().last() {
250        Some('T') => 1024 * 1024 * 1024 * 1024,
251        Some('G') => 1024 * 1024 * 1024,
252        Some('M') => 1024 * 1024,
253        Some('k') => 1024,
254        _ => 1,
255    };
256
257    let s = match multiplier {
258        1 => s,
259        _ => &s[..s.len() - 1],
260    };
261
262    Ok(u64::from_str(s)? * multiplier)
263}
264
265#[derive(Debug, PartialEq, Eq, Clone, Copy)]
266pub enum CipherSuite {
267    Aes128,
268    Aes256,
269    Chacha20,
270}
271
272impl CipherSuite {
273    pub fn as_rustls(self) -> rustls::SupportedCipherSuite {
274        use rustls::crypto::ring::cipher_suite;
275        match self {
276            Self::Aes128 => cipher_suite::TLS13_AES_128_GCM_SHA256,
277            Self::Aes256 => cipher_suite::TLS13_AES_256_GCM_SHA384,
278            Self::Chacha20 => cipher_suite::TLS13_CHACHA20_POLY1305_SHA256,
279        }
280    }
281}
282
283impl FromStr for CipherSuite {
284    type Err = anyhow::Error;
285
286    fn from_str(s: &str) -> Result<Self, Self::Err> {
287        match s.to_lowercase().as_str() {
288            "aes128" => Ok(Self::Aes128),
289            "aes256" => Ok(Self::Aes256),
290            "chacha20" => Ok(Self::Chacha20),
291            _ => Err(anyhow::anyhow!("Unknown cipher suite {}", s)),
292        }
293    }
294}
bench/lib.rs

bench/
lib.rs
