perf/
client.rs

1#[cfg(feature = "json-output")]
2use std::path::{Path, PathBuf};
3use std::{
4    net::{IpAddr, Ipv4Addr, Ipv6Addr, SocketAddr},
5    sync::Arc,
6    time::{Duration, Instant},
7};
8
9use anyhow::{Context, Result};
10use bytes::Bytes;
11use clap::Parser;
12use noq::{TokioRuntime, crypto::rustls::QuicClientConfig};
13use rustls::pki_types::{CertificateDer, ServerName, UnixTime};
14use tokio::sync::Semaphore;
15use tracing::{debug, error, info};
16
17use crate::{
18    CommonOpt, PERF_CIPHER_SUITES,
19    noprotection::NoProtectionClientConfig,
20    parse_byte_size,
21    stats::{OpenStreamStats, Stats},
22};
23
24/// Connects to a QUIC perf server and maintains a specified pattern of requests until interrupted
25#[derive(Parser)]
26#[clap(name = "client")]
27pub struct Opt {
28    /// Host to connect to
29    #[clap(default_value = "localhost:4433")]
30    host: String,
31    /// Override DNS resolution for host
32    #[clap(long)]
33    ip: Option<IpAddr>,
34    /// Specify the local socket address
35    #[clap(long)]
36    local_addr: Option<SocketAddr>,
37    /// Number of unidirectional requests to maintain concurrently
38    #[clap(long, default_value = "0")]
39    uni_requests: u64,
40    /// Number of bidirectional requests to maintain concurrently
41    #[clap(long, default_value = "1")]
42    bi_requests: u64,
43    /// Number of bytes to request
44    ///
45    /// This can use SI suffixes for sizes. For example, 1M will transfer
46    /// 1MiB, 10G will transfer 10GiB.
47    #[clap(long, default_value = "1M", value_parser = parse_byte_size)]
48    download_size: u64,
49    /// Number of bytes to transmit, in addition to the request header
50    ///
51    /// This can use SI suffixes for sizes. For example, 1M will transfer
52    /// 1MiB, 10G will transfer 10GiB.
53    #[clap(long, default_value = "1M", value_parser = parse_byte_size)]
54    upload_size: u64,
55    /// The time to run in seconds
56    #[clap(long, default_value = "60")]
57    duration: u64,
58    /// The interval in seconds at which stats are reported
59    #[clap(long, default_value = "1")]
60    interval: u64,
61    /// File path to output JSON statistics to. If the file is '-', stdout will be used
62    #[cfg(feature = "json-output")]
63    #[clap(long)]
64    json: Option<PathBuf>,
65    /// Common options
66    #[command(flatten)]
67    common: CommonOpt,
68}
69
70pub async fn run(opt: Opt) -> Result<()> {
71    let mut host_parts = opt.host.split(':');
72    let host_name = host_parts.next().unwrap();
73    let host_port = host_parts
74        .next()
75        .map_or(Ok(443), |x| x.parse())
76        .context("parsing port")?;
77    let addr = match opt.ip {
78        None => tokio::net::lookup_host(&opt.host)
79            .await
80            .context("resolving host")?
81            .next()
82            .unwrap(),
83        Some(ip) => SocketAddr::new(ip, host_port),
84    };
85
86    info!("connecting to {} at {}", host_name, addr);
87
88    let bind_addr = opt.local_addr.unwrap_or_else(|| {
89        let unspec = if addr.is_ipv4() {
90            Ipv4Addr::UNSPECIFIED.into()
91        } else {
92            Ipv6Addr::UNSPECIFIED.into()
93        };
94        SocketAddr::new(unspec, 0)
95    });
96
97    info!("local addr {:?}", bind_addr);
98
99    let socket = opt.common.bind_socket(bind_addr)?;
100
101    let mut endpoint_cfg = noq::EndpointConfig::default();
102    endpoint_cfg.max_udp_payload_size(opt.common.max_udp_payload_size)?;
103
104    let endpoint = noq::Endpoint::new(endpoint_cfg, None, socket, Arc::new(TokioRuntime))?;
105
106    let default_provider = rustls::crypto::ring::default_provider();
107    let provider = Arc::new(rustls::crypto::CryptoProvider {
108        cipher_suites: PERF_CIPHER_SUITES.into(),
109        ..default_provider
110    });
111
112    let mut crypto = rustls::ClientConfig::builder_with_provider(provider.clone())
113        .with_protocol_versions(&[&rustls::version::TLS13])
114        .unwrap()
115        .dangerous()
116        .with_custom_certificate_verifier(SkipServerVerification::new(provider))
117        .with_no_client_auth();
118    crypto.alpn_protocols = vec![b"perf".to_vec()];
119
120    if opt.common.keylog {
121        crypto.key_log = Arc::new(rustls::KeyLogFile::new());
122    }
123
124    let transport = opt.common.build_transport_config(
125        #[cfg(feature = "qlog")]
126        "perf-client",
127    )?;
128
129    let crypto = QuicClientConfig::try_from(crypto)?;
130    let mut config = noq::ClientConfig::new(match opt.common.no_protection {
131        true => Arc::new(NoProtectionClientConfig::new(crypto)),
132        false => Arc::new(crypto),
133    });
134    config.transport_config(Arc::new(transport));
135
136    let stream_stats = OpenStreamStats::default();
137
138    let connection = endpoint
139        .connect_with(config, addr, host_name)?
140        .await
141        .context("connecting")?;
142
143    info!("established");
144
145    let drive_fut = async {
146        tokio::try_join!(
147            drive_uni(
148                connection.clone(),
149                stream_stats.clone(),
150                opt.uni_requests,
151                opt.upload_size,
152                opt.download_size
153            ),
154            drive_bi(
155                connection.clone(),
156                stream_stats.clone(),
157                opt.bi_requests,
158                opt.upload_size,
159                opt.download_size
160            )
161        )
162    };
163
164    let mut stats = Stats::default();
165
166    let stats_fut = async {
167        let interval_duration = Duration::from_secs(opt.interval);
168
169        #[cfg(feature = "json-output")]
170        let allow_table_output = opt.json.clone().is_none_or(|path| path != Path::new("-"));
171        #[cfg(not(feature = "json-output"))]
172        let allow_table_output = true;
173
174        loop {
175            let start = Instant::now();
176            tokio::time::sleep(interval_duration).await;
177            {
178                stats.on_interval(start, &stream_stats);
179
180                if allow_table_output {
181                    stats.print();
182                    if opt.common.conn_stats {
183                        println!("{:?}\n", connection.stats());
184                    }
185                }
186            }
187        }
188    };
189
190    tokio::select! {
191        _ = drive_fut => {}
192        _ = stats_fut => {}
193        _ = tokio::signal::ctrl_c() => {
194            info!("shutting down");
195            connection.close(0u32.into(), b"interrupted");
196        }
197        // Add a small duration so the final interval can be reported
198        _ = tokio::time::sleep(Duration::from_secs(opt.duration) + Duration::from_millis(200)) => {
199            info!("shutting down");
200            connection.close(0u32.into(), b"done");
201        }
202    }
203
204    endpoint.wait_all_draining().await;
205
206    #[cfg(feature = "json-output")]
207    if let Some(path) = opt.json {
208        stats.print_json(path.as_path())?;
209    }
210
211    Ok(())
212}
213
214async fn drain_stream(
215    mut stream: noq::RecvStream,
216    download: u64,
217    stream_stats: OpenStreamStats,
218) -> Result<()> {
219    if download == 0 {
220        return Ok(());
221    }
222
223    #[rustfmt::skip]
224    let mut bufs = [
225        Bytes::new(), Bytes::new(), Bytes::new(), Bytes::new(),
226        Bytes::new(), Bytes::new(), Bytes::new(), Bytes::new(),
227        Bytes::new(), Bytes::new(), Bytes::new(), Bytes::new(),
228        Bytes::new(), Bytes::new(), Bytes::new(), Bytes::new(),
229        Bytes::new(), Bytes::new(), Bytes::new(), Bytes::new(),
230        Bytes::new(), Bytes::new(), Bytes::new(), Bytes::new(),
231        Bytes::new(), Bytes::new(), Bytes::new(), Bytes::new(),
232        Bytes::new(), Bytes::new(), Bytes::new(), Bytes::new(),
233    ];
234    let download_start = Instant::now();
235    let recv_stream_stats = stream_stats.new_receiver(&stream, download);
236
237    let mut first_byte = true;
238
239    while let Some(size) = stream.read_many_chunks(&mut bufs[..]).await? {
240        if first_byte {
241            recv_stream_stats.on_first_byte(download_start.elapsed());
242            first_byte = false;
243        }
244        let bytes_received = bufs[..size].iter().map(|b| b.len()).sum();
245        recv_stream_stats.on_bytes(bytes_received);
246    }
247
248    if first_byte {
249        recv_stream_stats.on_first_byte(download_start.elapsed());
250    }
251    recv_stream_stats.finish(download_start.elapsed());
252
253    debug!("response finished on {}", stream.id());
254    Ok(())
255}
256
257async fn drive_uni(
258    connection: noq::Connection,
259    stream_stats: OpenStreamStats,
260    concurrency: u64,
261    upload: u64,
262    download: u64,
263) -> Result<()> {
264    if concurrency == 0 {
265        return Ok(());
266    }
267
268    let sem = Arc::new(Semaphore::new(concurrency as usize));
269
270    loop {
271        let permit = sem.clone().acquire_owned().await.unwrap();
272        let send = connection.open_uni().await?;
273        let stream_stats = stream_stats.clone();
274
275        debug!("sending request on {}", send.id());
276        let connection = connection.clone();
277        tokio::spawn(async move {
278            if let Err(e) = request_uni(send, connection, upload, download, stream_stats).await {
279                error!("sending request failed: {:#}", e);
280            }
281
282            drop(permit);
283        });
284    }
285}
286
287async fn request_uni(
288    send: noq::SendStream,
289    conn: noq::Connection,
290    upload: u64,
291    download: u64,
292    stream_stats: OpenStreamStats,
293) -> Result<()> {
294    request(send, upload, download, stream_stats.clone()).await?;
295    let recv = conn.accept_uni().await?;
296    drain_stream(recv, download, stream_stats).await?;
297    Ok(())
298}
299
300async fn request(
301    mut send: noq::SendStream,
302    mut upload: u64,
303    download: u64,
304    stream_stats: OpenStreamStats,
305) -> Result<()> {
306    let upload_start = Instant::now();
307    send.write_all(&download.to_be_bytes()).await?;
308    if upload == 0 {
309        send.finish().unwrap();
310        return Ok(());
311    }
312
313    let send_stream_stats = stream_stats.new_sender(&send, upload);
314
315    static DATA: [u8; 1024 * 1024] = [42; 1024 * 1024];
316    while upload > 0 {
317        let chunk_len = upload.min(DATA.len() as u64);
318        send.write_chunk(Bytes::from_static(&DATA[..chunk_len as usize]))
319            .await
320            .context("sending response")?;
321        send_stream_stats.on_bytes(chunk_len as usize);
322        upload -= chunk_len;
323    }
324    send.finish().unwrap();
325    // Wait for stream to close
326    _ = send.stopped().await;
327    send_stream_stats.finish(upload_start.elapsed());
328
329    debug!("upload finished on {}", send.id());
330    Ok(())
331}
332
333async fn drive_bi(
334    connection: noq::Connection,
335    stream_stats: OpenStreamStats,
336    concurrency: u64,
337    upload: u64,
338    download: u64,
339) -> Result<()> {
340    if concurrency == 0 {
341        return Ok(());
342    }
343
344    let sem = Arc::new(Semaphore::new(concurrency as usize));
345
346    loop {
347        let permit = sem.clone().acquire_owned().await.unwrap();
348        let (send, recv) = connection.open_bi().await?;
349        let stream_stats = stream_stats.clone();
350
351        debug!("sending request on {}", send.id());
352        tokio::spawn(async move {
353            if let Err(e) = request_bi(send, recv, upload, download, stream_stats).await {
354                error!("request failed: {:#}", e);
355            }
356
357            drop(permit);
358        });
359    }
360}
361
362async fn request_bi(
363    send: noq::SendStream,
364    recv: noq::RecvStream,
365    upload: u64,
366    download: u64,
367    stream_stats: OpenStreamStats,
368) -> Result<()> {
369    request(send, upload, download, stream_stats.clone()).await?;
370    drain_stream(recv, download, stream_stats).await?;
371    Ok(())
372}
373
374#[derive(Debug)]
375struct SkipServerVerification(Arc<rustls::crypto::CryptoProvider>);
376
377impl SkipServerVerification {
378    fn new(provider: Arc<rustls::crypto::CryptoProvider>) -> Arc<Self> {
379        Arc::new(Self(provider))
380    }
381}
382
383impl rustls::client::danger::ServerCertVerifier for SkipServerVerification {
384    fn verify_server_cert(
385        &self,
386        _end_entity: &CertificateDer<'_>,
387        _intermediates: &[CertificateDer<'_>],
388        _server_name: &ServerName<'_>,
389        _ocsp: &[u8],
390        _now: UnixTime,
391    ) -> Result<rustls::client::danger::ServerCertVerified, rustls::Error> {
392        Ok(rustls::client::danger::ServerCertVerified::assertion())
393    }
394
395    fn verify_tls12_signature(
396        &self,
397        message: &[u8],
398        cert: &CertificateDer<'_>,
399        dss: &rustls::DigitallySignedStruct,
400    ) -> Result<rustls::client::danger::HandshakeSignatureValid, rustls::Error> {
401        rustls::crypto::verify_tls12_signature(
402            message,
403            cert,
404            dss,
405            &self.0.signature_verification_algorithms,
406        )
407    }
408
409    fn verify_tls13_signature(
410        &self,
411        message: &[u8],
412        cert: &CertificateDer<'_>,
413        dss: &rustls::DigitallySignedStruct,
414    ) -> Result<rustls::client::danger::HandshakeSignatureValid, rustls::Error> {
415        rustls::crypto::verify_tls13_signature(
416            message,
417            cert,
418            dss,
419            &self.0.signature_verification_algorithms,
420        )
421    }
422
423    fn supported_verify_schemes(&self) -> Vec<rustls::SignatureScheme> {
424        self.0.signature_verification_algorithms.supported_schemes()
425    }
426}