perf/
client.rs

1#[cfg(feature = "json-output")]
2use std::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 quinn::{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 = quinn::EndpointConfig::default();
102    endpoint_cfg.max_udp_payload_size(opt.common.max_udp_payload_size)?;
103
104    let endpoint = quinn::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 = Arc::new(QuicClientConfig::try_from(crypto)?);
130    let mut config = quinn::ClientConfig::new(match opt.common.no_protection {
131        true => Arc::new(NoProtectionClientConfig::new(crypto)),
132        false => 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        loop {
170            let start = Instant::now();
171            tokio::time::sleep(interval_duration).await;
172            {
173                stats.on_interval(start, &stream_stats);
174
175                stats.print();
176                if opt.common.conn_stats {
177                    println!("{:?}\n", connection.stats());
178                }
179            }
180        }
181    };
182
183    tokio::select! {
184        _ = drive_fut => {}
185        _ = stats_fut => {}
186        _ = tokio::signal::ctrl_c() => {
187            info!("shutting down");
188            connection.close(0u32.into(), b"interrupted");
189        }
190        // Add a small duration so the final interval can be reported
191        _ = tokio::time::sleep(Duration::from_secs(opt.duration) + Duration::from_millis(200)) => {
192            info!("shutting down");
193            connection.close(0u32.into(), b"done");
194        }
195    }
196
197    endpoint.wait_idle().await;
198
199    #[cfg(feature = "json-output")]
200    if let Some(path) = opt.json {
201        stats.print_json(path.as_path())?;
202    }
203
204    Ok(())
205}
206
207async fn drain_stream(
208    mut stream: quinn::RecvStream,
209    download: u64,
210    stream_stats: OpenStreamStats,
211) -> Result<()> {
212    if download == 0 {
213        return Ok(());
214    }
215
216    #[rustfmt::skip]
217    let mut bufs = [
218        Bytes::new(), Bytes::new(), Bytes::new(), Bytes::new(),
219        Bytes::new(), Bytes::new(), Bytes::new(), Bytes::new(),
220        Bytes::new(), Bytes::new(), Bytes::new(), Bytes::new(),
221        Bytes::new(), Bytes::new(), Bytes::new(), Bytes::new(),
222        Bytes::new(), Bytes::new(), Bytes::new(), Bytes::new(),
223        Bytes::new(), Bytes::new(), Bytes::new(), Bytes::new(),
224        Bytes::new(), Bytes::new(), Bytes::new(), Bytes::new(),
225        Bytes::new(), Bytes::new(), Bytes::new(), Bytes::new(),
226    ];
227    let download_start = Instant::now();
228    let recv_stream_stats = stream_stats.new_receiver(&stream, download);
229
230    let mut first_byte = true;
231
232    while let Some(size) = stream.read_chunks(&mut bufs[..]).await? {
233        if first_byte {
234            recv_stream_stats.on_first_byte(download_start.elapsed());
235            first_byte = false;
236        }
237        let bytes_received = bufs[..size].iter().map(|b| b.len()).sum();
238        recv_stream_stats.on_bytes(bytes_received);
239    }
240
241    if first_byte {
242        recv_stream_stats.on_first_byte(download_start.elapsed());
243    }
244    recv_stream_stats.finish(download_start.elapsed());
245
246    debug!("response finished on {}", stream.id());
247    Ok(())
248}
249
250async fn drive_uni(
251    connection: quinn::Connection,
252    stream_stats: OpenStreamStats,
253    concurrency: u64,
254    upload: u64,
255    download: u64,
256) -> Result<()> {
257    if concurrency == 0 {
258        return Ok(());
259    }
260
261    let sem = Arc::new(Semaphore::new(concurrency as usize));
262
263    loop {
264        let permit = sem.clone().acquire_owned().await.unwrap();
265        let send = connection.open_uni().await?;
266        let stream_stats = stream_stats.clone();
267
268        debug!("sending request on {}", send.id());
269        let connection = connection.clone();
270        tokio::spawn(async move {
271            if let Err(e) = request_uni(send, connection, upload, download, stream_stats).await {
272                error!("sending request failed: {:#}", e);
273            }
274
275            drop(permit);
276        });
277    }
278}
279
280async fn request_uni(
281    send: quinn::SendStream,
282    conn: quinn::Connection,
283    upload: u64,
284    download: u64,
285    stream_stats: OpenStreamStats,
286) -> Result<()> {
287    request(send, upload, download, stream_stats.clone()).await?;
288    let recv = conn.accept_uni().await?;
289    drain_stream(recv, download, stream_stats).await?;
290    Ok(())
291}
292
293async fn request(
294    mut send: quinn::SendStream,
295    mut upload: u64,
296    download: u64,
297    stream_stats: OpenStreamStats,
298) -> Result<()> {
299    let upload_start = Instant::now();
300    send.write_all(&download.to_be_bytes()).await?;
301    if upload == 0 {
302        send.finish().unwrap();
303        return Ok(());
304    }
305
306    let send_stream_stats = stream_stats.new_sender(&send, upload);
307
308    static DATA: [u8; 1024 * 1024] = [42; 1024 * 1024];
309    while upload > 0 {
310        let chunk_len = upload.min(DATA.len() as u64);
311        send.write_chunk(Bytes::from_static(&DATA[..chunk_len as usize]))
312            .await
313            .context("sending response")?;
314        send_stream_stats.on_bytes(chunk_len as usize);
315        upload -= chunk_len;
316    }
317    send.finish().unwrap();
318    // Wait for stream to close
319    _ = send.stopped().await;
320    send_stream_stats.finish(upload_start.elapsed());
321
322    debug!("upload finished on {}", send.id());
323    Ok(())
324}
325
326async fn drive_bi(
327    connection: quinn::Connection,
328    stream_stats: OpenStreamStats,
329    concurrency: u64,
330    upload: u64,
331    download: u64,
332) -> Result<()> {
333    if concurrency == 0 {
334        return Ok(());
335    }
336
337    let sem = Arc::new(Semaphore::new(concurrency as usize));
338
339    loop {
340        let permit = sem.clone().acquire_owned().await.unwrap();
341        let (send, recv) = connection.open_bi().await?;
342        let stream_stats = stream_stats.clone();
343
344        debug!("sending request on {}", send.id());
345        tokio::spawn(async move {
346            if let Err(e) = request_bi(send, recv, upload, download, stream_stats).await {
347                error!("request failed: {:#}", e);
348            }
349
350            drop(permit);
351        });
352    }
353}
354
355async fn request_bi(
356    send: quinn::SendStream,
357    recv: quinn::RecvStream,
358    upload: u64,
359    download: u64,
360    stream_stats: OpenStreamStats,
361) -> Result<()> {
362    request(send, upload, download, stream_stats.clone()).await?;
363    drain_stream(recv, download, stream_stats).await?;
364    Ok(())
365}
366
367#[derive(Debug)]
368struct SkipServerVerification(Arc<rustls::crypto::CryptoProvider>);
369
370impl SkipServerVerification {
371    fn new(provider: Arc<rustls::crypto::CryptoProvider>) -> Arc<Self> {
372        Arc::new(Self(provider))
373    }
374}
375
376impl rustls::client::danger::ServerCertVerifier for SkipServerVerification {
377    fn verify_server_cert(
378        &self,
379        _end_entity: &CertificateDer<'_>,
380        _intermediates: &[CertificateDer<'_>],
381        _server_name: &ServerName<'_>,
382        _ocsp: &[u8],
383        _now: UnixTime,
384    ) -> Result<rustls::client::danger::ServerCertVerified, rustls::Error> {
385        Ok(rustls::client::danger::ServerCertVerified::assertion())
386    }
387
388    fn verify_tls12_signature(
389        &self,
390        message: &[u8],
391        cert: &CertificateDer<'_>,
392        dss: &rustls::DigitallySignedStruct,
393    ) -> Result<rustls::client::danger::HandshakeSignatureValid, rustls::Error> {
394        rustls::crypto::verify_tls12_signature(
395            message,
396            cert,
397            dss,
398            &self.0.signature_verification_algorithms,
399        )
400    }
401
402    fn verify_tls13_signature(
403        &self,
404        message: &[u8],
405        cert: &CertificateDer<'_>,
406        dss: &rustls::DigitallySignedStruct,
407    ) -> Result<rustls::client::danger::HandshakeSignatureValid, rustls::Error> {
408        rustls::crypto::verify_tls13_signature(
409            message,
410            cert,
411            dss,
412            &self.0.signature_verification_algorithms,
413        )
414    }
415
416    fn supported_verify_schemes(&self) -> Vec<rustls::SignatureScheme> {
417        self.0.signature_verification_algorithms.supported_schemes()
418    }
419}